Fix a bug that would overcount bytes read from the client due to the PushbackInputStream.

Author: robotdan

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

@@ -18,6 +18,8 @@
 import java.io.IOException;
 import java.io.InputStream;
 
+import io.fusionauth.http.server.Instrumenter;
+
 /**
  * An input stream that allows a portion of bytes read into a buffer to be pushed back to be read again.
  *
@@ -28,14 +30,18 @@ public class PushbackInputStream extends InputStream {
 
   private final InputStream delegate;
 
+  private final Instrumenter instrumenter;
+
   private byte[] buffer;
 
   private int bufferEndPosition;
 
   private int bufferPosition;
 
-  public PushbackInputStream(InputStream delegate) {
+
+  public PushbackInputStream(InputStream delegate, Instrumenter instrumenter) {
     this.delegate = delegate;
+    this.instrumenter = instrumenter;
   }
 
   public int getAvailableBufferedBytesRemaining() {
@@ -78,10 +84,17 @@ public int read(byte[] b, int off, int len) throws IOException {
       //   complete processing of the bytes we just read from the buffer in order to send the HTTP response.
       //   The end result is the client will block while waiting for us to send a response until we take an exception waiting
       //   for the read timeout.
+      // - Do not count this as a read from the client, that would double count.
+      //   We should only be counting bytes read from the delegate.
       return read;
     }
 
-    return delegate.read(b, off, len);
+    var read = delegate.read(b, off, len);
+    if (read > 0 && instrumenter != null) {
+      instrumenter.readFromClient(read);
+    }
+
+    return read;
   }
 
   @Override

src/main/java/io/fusionauth/http/server/internal/HTTPWorker.java

@@ -21,13 +21,11 @@
 import java.net.SocketException;
 import java.net.SocketTimeoutException;
 
-import io.fusionauth.http.server.io.ConnectionClosedException;
 import io.fusionauth.http.HTTPValues;
 import io.fusionauth.http.HTTPValues.Connections;
 import io.fusionauth.http.HTTPValues.Headers;
 import io.fusionauth.http.HTTPValues.Protocols;
 import io.fusionauth.http.ParseException;
-import io.fusionauth.http.server.io.TooManyBytesToDrainException;
 import io.fusionauth.http.io.PushbackInputStream;
 import io.fusionauth.http.log.Logger;
 import io.fusionauth.http.server.HTTPHandler;
@@ -36,11 +34,13 @@
 import io.fusionauth.http.server.HTTPResponse;
 import io.fusionauth.http.server.HTTPServerConfiguration;
 import io.fusionauth.http.server.Instrumenter;
+import io.fusionauth.http.server.io.ConnectionClosedException;
 import io.fusionauth.http.server.io.HTTPInputStream;
 import io.fusionauth.http.server.io.HTTPOutputStream;
 import io.fusionauth.http.server.io.Throughput;
 import io.fusionauth.http.server.io.ThroughputInputStream;
 import io.fusionauth.http.server.io.ThroughputOutputStream;
+import io.fusionauth.http.server.io.TooManyBytesToDrainException;
 import io.fusionauth.http.util.HTTPTools;
 
 /**
@@ -80,7 +80,7 @@ public HTTPWorker(Socket socket, HTTPServerConfiguration configuration, Instrume
     this.throughput = throughput;
     this.buffers = new HTTPBuffers(configuration);
     this.logger = configuration.getLoggerFactory().getLogger(HTTPWorker.class);
-    this.inputStream = new PushbackInputStream(new ThroughputInputStream(socket.getInputStream(), throughput));
+    this.inputStream = new PushbackInputStream(new ThroughputInputStream(socket.getInputStream(), throughput), instrumenter);
     this.state = State.Read;
     this.startInstant = System.currentTimeMillis();
     logger.trace("[{}] Starting HTTP worker.", Thread.currentThread().threadId());
@@ -123,7 +123,7 @@ public void run() {
         // Not this line of code will block
         // - When a client is using Keep-Alive - we will loop and block here while we wait for the client to send us bytes.
         byte[] requestBuffer = buffers.requestBuffer();
-        HTTPTools.parseRequestPreamble(inputStream, request, requestBuffer, instrumenter, () -> state = State.Read);
+        HTTPTools.parseRequestPreamble(inputStream, request, requestBuffer, () -> state = State.Read);
         if (logger.isTraceEnabled()) {
           int availableBufferedBytes = inputStream.getAvailableBufferedBytesRemaining();
           if (availableBufferedBytes != 0) {

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

@@ -143,8 +143,7 @@ public int read(byte[] b, int off, int len) 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.
+    // - If we have read past the end of the current request, push those bytes back onto the InputStream.
     int read = delegate.read(b, off, len);
     if (fixedLength && read > 0) {
       int extraBytes = (int) (read - bytesRemaining);
@@ -154,10 +153,6 @@ public int read(byte[] b, int off, int len) throws IOException {
     }
 
     if (read > 0) {
-      if (instrumenter != null) {
-        instrumenter.readFromClient(read);
-      }
-
       if (fixedLength) {
         bytesRemaining -= read;
       }

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

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, FusionAuth, All Rights Reserved
+ * Copyright (c) 2024-2025, FusionAuth, All Rights Reserved
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -44,7 +44,7 @@ public void close() throws IOException {
   }
 
   @Override
-  public void mark(int readlimit) {
+  public void mark(@SuppressWarnings("SpellCheckingInspection") int readlimit) {
     delegate.mark(readlimit);
   }
 

src/main/java/io/fusionauth/http/util/HTTPTools.java

@@ -39,7 +39,6 @@
 import io.fusionauth.http.log.LoggerFactory;
 import io.fusionauth.http.server.HTTPRequest;
 import io.fusionauth.http.server.HTTPResponse;
-import io.fusionauth.http.server.Instrumenter;
 import io.fusionauth.http.server.io.ConnectionClosedException;
 
 public final class HTTPTools {
@@ -93,8 +92,8 @@ public static boolean isHexadecimalCharacter(byte ch) {
    */
   public static boolean isTokenCharacter(byte ch) {
     return ch == '!' || ch == '#' || ch == '$' || ch == '%' || ch == '&' || ch == '\'' || ch == '*' || ch == '+' || ch == '-' || ch == '.' ||
-           ch == '^' || ch == '_' || ch == '`' || ch == '|' || ch == '~' || (ch >= 'A' && ch <= 'Z') || (ch >= 'a' && ch <= 'z') ||
-           (ch >= '0' && ch <= '9');
+        ch == '^' || ch == '_' || ch == '`' || ch == '|' || ch == '~' || (ch >= 'A' && ch <= 'Z') || (ch >= 'a' && ch <= 'z') ||
+        (ch >= '0' && ch <= '9');
   }
 
   /**
@@ -262,13 +261,10 @@ public static HeaderValue parseHeaderValue(String value) {
    * @param inputStream   The input stream to read the preamble from.
    * @param request       The HTTP request to populate.
    * @param requestBuffer A buffer used for reading to help reduce memory thrashing.
-   * @param instrumenter  The Instrumenter that is informed of bytes read.
    * @param readObserver  An observer that is called once one byte has been read.
    * @throws IOException If the read fails.
    */
-  public static void parseRequestPreamble(PushbackInputStream inputStream, HTTPRequest request, byte[] requestBuffer,
-                                          Instrumenter instrumenter,
-                                          Runnable readObserver)
+  public static void parseRequestPreamble(PushbackInputStream inputStream, HTTPRequest request, byte[] requestBuffer, Runnable readObserver)
       throws IOException {
     RequestPreambleState state = RequestPreambleState.RequestMethod;
     var valueBuffer = new ByteArrayOutputStream(512);
@@ -288,10 +284,6 @@ public static void parseRequestPreamble(PushbackInputStream inputStream, HTTPReq
 
       logger.trace("Read [{}] from client for preamble.", read);
 
-      if (instrumenter != null) {
-        instrumenter.readFromClient(read);
-      }
-
       // Tell the callback that we've read at least one byte
       readObserver.run();
 

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

@@ -435,6 +435,15 @@ public void large_body(String scheme) throws Exception {
 
       assertEquals(response.statusCode(), 200);
       assertEquals(response.body(), bytes);
+
+      // This assertion is always true unless we change what we are writing above. But it is here for reference.
+      assertEquals(bytes.length, 16_804);
+      // This should prove the PushbackInputStream isn't double counting.
+      // - We should expect the bytes read to be roughly equivalent to the payload length. We can add some bytes to account
+      //   for the HTTP preamble, and for HTTPs additional overhead is incurred due to encryption. But if we are counting incorrectly
+      //   due to bytes being pushed back and counted again, the numbers would be almost double.
+      // - So we should expect the bytes read to be within the ball bark of the payload length.
+      assertEquals(instrumenter.getBytesRead(), scheme.equals("http") ? 17_032: 16_933);
     }
   }
 

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

@@ -182,7 +182,7 @@ private Builder withBody(String body) {
   }
 
   private PushbackInputStream withParts(String... parts) {
-    return new PushbackInputStream(new PieceMealInputStream(parts));
+    return new PushbackInputStream(new PieceMealInputStream(parts), null);
   }
 
   @SuppressWarnings("UnusedReturnValue")
@@ -221,7 +221,7 @@ public Builder assertNextRead(ThrowingFunction<ChunkedInputStream, Integer> func
     }
 
     public Builder assertResult(String expected) throws IOException {
-      pushbackInputStream = new PushbackInputStream(new ByteArrayInputStream(body.getBytes(StandardCharsets.UTF_8)));
+      pushbackInputStream = new PushbackInputStream(new ByteArrayInputStream(body.getBytes(StandardCharsets.UTF_8)), null);
       chunkedInputStream = new ChunkedInputStream(pushbackInputStream, 2048);
 
       String actual = new String(chunkedInputStream.readAllBytes(), StandardCharsets.UTF_8);