Working

Author: robotdan

src/main/java/io/fusionauth/http/io/ChunkedInputStream.java

@@ -40,6 +40,8 @@ public class ChunkedInputStream extends InputStream {
 
   private int bufferLength;
 
+  private int chunkBytesRead;
+
   private int chunkBytesRemaining;
 
   private int chunkSize;
@@ -52,119 +54,127 @@ public ChunkedInputStream(PushbackInputStream delegate, int bufferSize) {
   }
 
   @Override
-  public int read(byte[] b, int off, int len) throws IOException {
-    return processChunk(b, off, len);
-  }
-
-  @Override
-  public int read() throws IOException {
-    var read = read(b1);
-    if (read <= 0) {
-      return read;
-    }
-
-    return b1[0] & 0xFF;
-  }
-
-  private int processChunk(byte[] destination, int offset, int length) throws IOException {
-    int totalRead = 0;
-    while (totalRead < length) {
-      // Bail early if the state machine is done
-      if (state == ChunkedBodyState.Complete) {
-        // We need to push back any remaining bytes to the InputStream since we may have read more bytes than we needed.
-        int leftOver = bufferLength - bufferIndex;
-        if (leftOver > 0) {
-          delegate.push(buffer, bufferIndex, leftOver);
-        }
-
-        return -1;
+  public int read(byte[] destination, int dOff, int dLen) throws IOException {
+    int dIndex = dOff;
+    while (dIndex < dLen) {
+      if (state == ChunkedInputStream.ChunkedBodyState.Complete) {
+        pushBackOverReadBytes();
+        break;
       }
 
       // Read some more if we are out of bytes
       if (bufferIndex >= bufferLength) {
         bufferIndex = 0;
         bufferLength = delegate.read(buffer);
-        if (bufferLength > 0) {
-          totalRead += bufferLength;
-        }
       }
 
-      for (; bufferIndex < bufferLength; bufferIndex++) {
+      // Process the buffer
+      while (bufferIndex < bufferLength) {
         ChunkedBodyState nextState;
         try {
-          nextState = state.next(buffer[bufferIndex], chunkSize, chunkSize - chunkBytesRemaining);
+          nextState = state.next(buffer[bufferIndex], chunkSize, chunkBytesRead);
         } catch (ParseException e) {
           // This allows us to add the index to the exception. Useful for debugging.
           e.setIndex(bufferIndex);
           throw e;
         }
 
-        // We are DONE!
-        if (nextState == ChunkedBodyState.Complete) {
+        // We have reached the end of the encoded payload. Push back any additional bytes read.
+        if (state == ChunkedBodyState.Complete) {
           state = nextState;
-          bufferIndex++;
-          int leftOver = bufferLength - bufferIndex;
-          if (leftOver > 0) {
-            delegate.push(buffer, bufferIndex, leftOver);
-          }
-
-          return -1;
+          pushBackOverReadBytes();
+          break;
         }
 
-        // Record the size hex digit
+        // Capture the character to calculate the next chunk size
         if (nextState == ChunkedBodyState.ChunkSize) {
           headerSizeHex.appendCodePoint(buffer[bufferIndex]);
           state = nextState;
+          bufferIndex++;
           continue;
         }
 
-        // Capture the chunk size
+        // We have found the chunk, this means we can now convert the captured chunk size bytes and then try and read the chunk.
         if (state != ChunkedBodyState.Chunk && nextState == ChunkedBodyState.Chunk) {
-          // This means we finished reading the size and are ready to start processing
           if (headerSizeHex.isEmpty()) {
             throw new ChunkException("Chunk size is missing");
           }
 
           // This is the start of a chunk, so set the size and counter and reset the size hex string
           chunkSize = (int) Long.parseLong(headerSizeHex, 0, headerSizeHex.length(), 16);
 
+          chunkBytesRead = 0;
           chunkBytesRemaining = chunkSize;
           headerSizeHex.delete(0, headerSizeHex.length());
 
-          // AF\r1234 i=3 length=42 read=7 chunkSize=175(AF)
+          // A chunk size of 0 indicates this is the terminating chunk. Continue and we will expect the state machine
+          // to process the final CRLF and hit the Complete state.
           if (chunkSize == 0) {
             state = nextState;
-            return 0;
+            bufferIndex++;
+            continue;
           }
         }
 
-        int remainingInBuffer = bufferLength - bufferIndex;
-        int lengthToCopy = Math.min(Math.min(chunkBytesRemaining, remainingInBuffer), length); // That's an ugly baby!
+        int lengthToCopy;
+        if (chunkBytesRemaining > 0) {
+          int remainingInBuffer = bufferLength - bufferIndex;
+          lengthToCopy = Math.min(Math.min(chunkBytesRemaining, remainingInBuffer), dLen - dIndex); // That's an ugly baby!
 
-        // If we don't have anything to copy, continue.
-        if (lengthToCopy == 0) {
+          // This means we don't have room in the destination buffer
+          if (lengthToCopy == 0) {
+            state = nextState;
+            bufferIndex++;
+            return dIndex - dOff;
+          }
+        } else {
+          // Nothing to do, continue to the next state.
           state = nextState;
+          bufferIndex++;
           continue;
         }
 
-        // TODO : Daniel : If we wanted to handle more than one chunk at a time, I think we would potentially continue here
-        //        and see if we can get more than one chunk completed before copying back to the destination.
-        //        Discuss with Brian.
-        //        Should we try this right now, or is this ok? Load testing with a small body,
-        //        Chunked is slower than fixed length, I suppose this makes sense. In practice I don't think browsers
-        //        and such use chunked for small requests.
-
-        System.arraycopy(buffer, bufferIndex, destination, offset, lengthToCopy);
+        // Copy 'lengthToCopy' to the destination buffer
+        System.arraycopy(buffer, bufferIndex, destination, dIndex, lengthToCopy);
         bufferIndex += lengthToCopy;
+        chunkBytesRead += lengthToCopy;
         chunkBytesRemaining -= lengthToCopy;
+        dIndex += lengthToCopy;
         state = nextState;
-        return lengthToCopy;
+
+        // If we have bytes to copy, we are at the correct location in the state machine, If we don't have room, break.
+        // - This will break the while loop, and return at the end of this method the total bytes we have written to the destination buffer.
+        if (dIndex == dLen) {
+          break;
+        }
       }
     }
 
-    return 0;
+    int total = dIndex - dOff;
+    return total == 0 ? -1 : total;
+  }
+
+  @Override
+  public int read() throws IOException {
+    var read = read(b1);
+    if (read <= 0) {
+      return read;
+    }
+
+    return b1[0] & 0xFF;
+  }
+
+  private void pushBackOverReadBytes() {
+    int leftOver = bufferLength - bufferIndex;
+    if (leftOver > 0) {
+      delegate.push(buffer, bufferIndex, leftOver);
+
+      // Move the pointer to the end of the buffer, We have used up the bytes by pushing them back.
+      bufferIndex = bufferLength;
+    }
   }
 
+
   public enum ChunkedBodyState {
     ChunkExtensionStart {
       @Override
@@ -266,17 +276,18 @@ public ChunkedBodyState next(byte ch, long length, long bytesRead) {
         return Chunk;
       }
     },
-
     Chunk {
       @Override
       public ChunkedBodyState next(byte ch, long length, long bytesRead) {
-        if (bytesRead < length) {
-          return Chunk;
+        if (length == 0) {
+          return Complete;
         } else if (bytesRead == length && ch == '\r') {
           return ChunkCR;
+        } else if (bytesRead < length) {
+          return Chunk;
+        } else {
+          throw makeParseException(ch, this);
         }
-
-        throw makeParseException(ch, this);
       }
     },
 
@@ -311,6 +322,6 @@ public ChunkedBodyState next(byte ch, long length, long bytesRead) {
       }
     };
 
-    public abstract ChunkedBodyState next(byte ch, long length, long bytesRead);
+    public abstract ChunkedInputStream.ChunkedBodyState next(byte ch, long length, long bytesRead);
   }
 }

src/main/java/io/fusionauth/http/server/Configurable.java

@@ -47,7 +47,7 @@ default T withBaseDir(Path baseDir) {
   }
 
   /**
-   * Sets the buffer size for the chunked input stream. Defaults to 2k.
+   * Sets the buffer size for the chunked input stream. Defaults to 4k.
    *
    * @param chunkedBufferSize the buffer size used to read a request body that was encoded using 'chunked' transfer-encoding.
    * @return This.

src/main/java/io/fusionauth/http/server/HTTPServerConfiguration.java

@@ -34,7 +34,7 @@ public class HTTPServerConfiguration implements Configurable<HTTPServerConfigura
 
   private Path baseDir = Path.of("");
 
-  private int chunkedBufferSize = 2 * 1024; // 2k bytes
+  private int chunkedBufferSize = 4 * 1024; // 4k bytes
 
   private boolean compressByDefault = true;
 

src/main/java/io/fusionauth/http/server/io/HTTPInputStream.java

@@ -122,14 +122,19 @@ public int read() throws IOException {
   }
 
   @Override
-  public int read(byte[] buffer) throws IOException {
-    return read(buffer, 0, buffer.length);
+  public int read(byte[] b) throws IOException {
+    return read(b, 0, b.length);
   }
 
   @Override
-  public int read(byte[] buffer, int offset, int length) throws IOException {
+  public int read(byte[] b, int off, int len) throws IOException {
+    if (len == 0) {
+      return 0;
+    }
+
     // If this is a fixed length request, and we have less than or equal to 0 bytes remaining, return -1
-    if (!request.isChunked() && bytesRemaining <= 0) {
+    boolean fixedLength = !request.isChunked();
+    if (fixedLength && bytesRemaining <= 0) {
       return -1;
     }
 
@@ -140,19 +145,24 @@ public int read(byte[] buffer, int offset, int length) throws IOException {
     // When we have a fixed length request, read beyond the remainingBytes if possible.
     // - Under heavy load we may be able to start reading the next request. Just push those bytes
     //   back onto the InputStream and we will read them later.
-    int read = delegate.read(buffer, offset, length);
-    if (!request.isChunked()) {
+    int read = delegate.read(b, off, len);
+    if (fixedLength && read > 0) {
       int extraBytes = (int) (read - bytesRemaining);
       if (extraBytes > 0) {
-        pushbackInputStream.push(buffer, (int) bytesRemaining, extraBytes);
+        pushbackInputStream.push(b, (int) bytesRemaining, extraBytes);
       }
     }
 
-    if (instrumenter != null) {
-      instrumenter.readFromClient(read);
+    if (read > 0) {
+      if (instrumenter != null) {
+        instrumenter.readFromClient(read);
+      }
+
+      if (fixedLength) {
+        bytesRemaining -= read;
+      }
     }
 
-    bytesRemaining -= read;
     return read;
   }
 

src/test/java/io/fusionauth/http/BaseSocketTest.java

@@ -36,9 +36,7 @@ protected Builder withRequest(String request) {
   private void assertResponse(String request, String chunkedExtension, String response) throws Exception {
     HTTPHandler handler = (req, res) -> {
       // Read the request body
-      byte[] bodyBytes = req.getInputStream().readAllBytes();
-      System.out.println(bodyBytes.length);
-
+      req.getInputStream().readAllBytes();
       res.setStatus(200);
     };
 

src/test/java/io/fusionauth/http/io/ChunkedInputStreamTest.java

@@ -120,22 +120,10 @@ public void partialChunks() throws IOException {
             0\r
             \r
             """), 1024);
-    assertEquals(inputStream.read(buf), 8);
-    var result = new String(buf, 0, 8);
-    assertEquals(result, "12345678");
-
-    assertEquals(inputStream.read(buf), 2);
-    result = new String(buf, 0, 2);
-    assertEquals(result, "90");
-
-    assertEquals(inputStream.read(buf), 20);
-    result = new String(buf, 0, 20);
-    assertEquals(result, "12345678901234567890");
-
-    assertEquals(inputStream.read(buf), 30);
-    result = new String(buf, 0, 30);
-    assertEquals(result, "123456789012345678901234567890");
-    assertEquals(inputStream.read(), 0);
+    // All chunks will be read on the first attempt because the buffer is large enough
+    assertEquals(inputStream.read(buf), 60);
+    var result = new String(buf, 0, 60);
+    assertEquals(result, "123456789012345678901234567890123456789012345678901234567890");
     assertEquals(inputStream.read(), -1);
   }
 
@@ -153,15 +141,12 @@ public void partialHeader() throws IOException {
             0\r
             \r
             """), 1024);
-    assertEquals(inputStream.read(buf), 10);
-    var result = new String(buf, 0, 10);
-    assertEquals(result, "1234567890");
-
-    assertEquals(inputStream.read(buf), 20);
-    result = new String(buf, 0, 20);
-    assertEquals(result, "12345678901234567890");
-    assertEquals(inputStream.read(), 0);
-    assertEquals(inputStream.read(), -1);
+
+    // All chunks will be read on the first attempt because the buffer is large enough
+    assertEquals(inputStream.read(buf), 30);
+    var result = new String(buf, 0, 30);
+    assertEquals(result, "123456789012345678901234567890");
+    assertEquals(inputStream.read(buf), -1);
   }
 
   private Builder withBody(String body) {
@@ -224,7 +209,8 @@ public int read(byte[] b, int off, int len) {
         return -1;
       }
 
-      // We may only read part way through one of the parts. If we didn't read all the way through, use the subPartIndex
+      // We may only read part way through one of the parts.
+      // If we didn't read all the way through, use the subPartIndex
       int read = Math.min(parts[partsIndex].length - subPartIndex, b.length);
       System.arraycopy(parts[partsIndex], 0, b, 0, read);
       if (read < parts[partsIndex].length - subPartIndex) {