Protect readLine() against DoS

pom.xml

@@ -89,7 +89,7 @@
         <grpc.version>1.35.0</grpc.version>
 
         <air.javadoc.lint>-missing</air.javadoc.lint>
-        <versions.java-security-toolkit>1.2.0</versions.java-security-toolkit>
+        <versions.java-security-toolkit>1.2.1</versions.java-security-toolkit>
     </properties>
 
     <modules>
@@ -2214,6 +2214,7 @@
             <dependency>
                 <groupId>io.github.pixee</groupId>
                 <artifactId>java-security-toolkit</artifactId>
+
                 <version>${versions.java-security-toolkit}</version>
             </dependency>
         </dependencies>

presto-atop/src/test/java/com/facebook/presto/atop/TestingAtopFactory.java

@@ -15,6 +15,7 @@
 
 import com.google.common.base.Joiner;
 import com.google.common.base.Splitter;
+import io.github.pixee.security.BoundedLineReader;
 
 import java.io.BufferedReader;
 import java.io.IOException;
@@ -53,7 +54,7 @@ private TestingAtop(InputStream dataStream, ZonedDateTime date)
             this.date = date;
             this.reader = new BufferedReader(new InputStreamReader(dataStream));
             try {
-                line = reader.readLine();
+                line = BoundedLineReader.readLine(reader, 5_000_000);
             }
             catch (IOException e) {
                 line = null;
@@ -74,7 +75,7 @@ public String next()
             }
             String currentLine = line;
             try {
-                line = reader.readLine();
+                line = BoundedLineReader.readLine(reader, 5_000_000);
             }
             catch (IOException e) {
                 line = null;

presto-druid/src/main/java/com/facebook/presto/druid/DruidBrokerPageSource.java

@@ -31,6 +31,7 @@
 import com.google.common.collect.ImmutableList;
 import io.airlift.slice.Slice;
 import io.airlift.slice.Slices;
+import io.github.pixee.security.BoundedLineReader;
 import org.joda.time.DateTime;
 import org.joda.time.chrono.ISOChronology;
 import org.joda.time.format.DateTimeFormatter;
@@ -115,7 +116,7 @@ public Page getNextPage()
                 handle -> ((DruidColumnHandle) handle).getColumnName().equals("errorMessage"));
         try {
             String readLine;
-            while ((readLine = responseStream.readLine()) != null) {
+            while ((readLine = BoundedLineReader.readLine(responseStream, 5_000_000)) != null) {
                 // if read a blank line,it means read finish
                 if (readLine.isEmpty()) {
                     finished = true;

presto-druid/src/main/java/com/facebook/presto/druid/segment/SmooshedColumnSource.java

@@ -14,6 +14,7 @@
 package com.facebook.presto.druid.segment;
 
 import com.facebook.presto.spi.PrestoException;
+import io.github.pixee.security.BoundedLineReader;
 
 import java.io.BufferedReader;
 import java.io.ByteArrayInputStream;
@@ -80,7 +81,7 @@ private void loadSmooshFileMetadata()
         try {
             byte[] metadata = indexFileSource.readFile(SMOOSH_METADATA_FILE_NAME);
             BufferedReader in = new BufferedReader(new InputStreamReader(new ByteArrayInputStream(metadata)));
-            String line = in.readLine();
+            String line = BoundedLineReader.readLine(in, 5_000_000);
             if (line == null) {
                 throw new PrestoException(DRUID_SEGMENT_LOAD_ERROR, format("Malformed metadata file: first line should be version,maxChunkSize,numChunks, got null."));
             }
@@ -94,7 +95,7 @@ private void loadSmooshFileMetadata()
             }
 
             while (true) {
-                line = in.readLine();
+                line = BoundedLineReader.readLine(in, 5_000_000);
                 if (line == null) {
                     break;
                 }

presto-local-file/src/main/java/com/facebook/presto/localfile/LocalFileRecordCursor.java

@@ -26,6 +26,7 @@
 import com.google.common.collect.Iterables;
 import io.airlift.slice.Slice;
 import io.airlift.slice.Slices;
+import io.github.pixee.security.BoundedLineReader;
 
 import java.io.BufferedReader;
 import java.io.File;
@@ -319,7 +320,7 @@ public List<String> readFields()
                 if (reader == null) {
                     return null;
                 }
-                String line = reader.readLine();
+                String line = BoundedLineReader.readLine(reader, 5_000_000);
                 if (line != null) {
                     fields = LINE_SPLITTER.splitToList(line);
                     if (!newReader || meetsPredicate(fields)) {

presto-main/pom.xml

@@ -489,6 +489,10 @@
             <artifactId>ratis-common</artifactId>
             <optional>true</optional>
         </dependency>
+        <dependency>
+            <groupId>io.github.pixee</groupId>
+            <artifactId>java-security-toolkit</artifactId>
+        </dependency>
     </dependencies>
 
     <build>

presto-main/src/main/java/com/facebook/presto/operator/HttpRpcShuffleClient.java

@@ -29,6 +29,7 @@
 import io.airlift.slice.InputStreamSliceInput;
 import io.airlift.slice.SliceInput;
 import io.airlift.units.DataSize;
+import io.github.pixee.security.BoundedLineReader;
 
 import javax.annotation.concurrent.ThreadSafe;
 
@@ -161,7 +162,7 @@ public PagesResponse handle(Request request, Response response)
                     try (BufferedReader reader = new BufferedReader(new InputStreamReader(response.getInputStream(), UTF_8))) {
                         // Get up to 1000 lines for debugging
                         for (int i = 0; i < 1000; i++) {
-                            String line = reader.readLine();
+                            String line = BoundedLineReader.readLine(reader, 5_000_000);
                             // Don't output more than 100KB
                             if (line == null || body.length() + line.length() > 100 * 1024) {
                                 break;

presto-spark-base/src/main/java/com/facebook/presto/spark/execution/nativeprocess/NativeExecutionProcess.java

@@ -31,6 +31,7 @@
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.SettableFuture;
 import io.airlift.units.Duration;
+import io.github.pixee.security.BoundedLineReader;
 import io.github.pixee.security.SystemCommand;
 import org.apache.spark.SparkEnv$;
 import org.apache.spark.SparkFiles;
@@ -443,7 +444,7 @@ public void run()
                     BufferedWriter writer = new BufferedWriter(new OutputStreamWriter(outputStream, UTF_8))) {
                 String line;
                 boolean aborted = false;
-                while ((line = reader.readLine()) != null) {
+                while ((line = BoundedLineReader.readLine(reader, 5_000_000)) != null) {
                     if (!aborted && line.startsWith("*** Aborted")) {
                         aborted = true;
                     }