refactor: implement lazy request body loading with dedicated HttpRequestBody class

• Created new HttpRequestBody class to encapsulate lazy loading logic and state
• Refactored PluginHttpContext to use HttpRequestBody instead of managing state directly
• Updated ProxyWasmFilter to use shared HttpRequestBody instance across plugins
• Improved request body processing loop with proper action handling
• Enhanced type safety by removing Object-based field types
• Moved body state management from context to dedicated supplier class

This refactoring addresses memory efficiency concerns by only loading request bodies
when WASM modules actually access them via proxy_get_buffer_bytes.

Signed-off-by: Hiram Chirino <hiram@hiramchirino.com>

Author: chirino

proxy-wasm-java-host/src/main/java/io/roastedroot/proxywasm/internal/HttpRequestBody.java

@@ -0,0 +1,55 @@
+package io.roastedroot.proxywasm.internal;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.function.Supplier;
+
+/**
+ * Holds the HTTP request body, loading it from a stream when first accessed.
+ */
+public class HttpRequestBody {
+
+    private byte[] body;
+    private boolean loaded = false;
+    private final Supplier<InputStream> streamSupplier;
+
+    public HttpRequestBody(Supplier<InputStream> streamSupplier) {
+        this.streamSupplier = streamSupplier;
+    }
+
+    public byte[] get() {
+        if (!loaded) {
+            try {
+                body = streamSupplier.get().readAllBytes();
+            } catch (IOException e) {
+                throw new RuntimeException("Failed to read request body", e);
+            }
+            loaded = true;
+        }
+        return body;
+    }
+
+    /**
+     * Returns true if the request body has been loaded from the stream.
+     */
+    public boolean isLoaded() {
+        return loaded;
+    }
+
+    /**
+     * Sets the request body directly, marking it as loaded.
+     * This is used when the body is modified by WASM plugins.
+     */
+    public void setBody(byte[] body) {
+        this.body = body;
+        this.loaded = true;
+    }
+
+    /**
+     * Returns the request body if it has been loaded, null otherwise.
+     * This allows checking if the body was accessed without triggering a load.
+     */
+    public byte[] getBodyIfLoaded() {
+        return loaded ? body : null;
+    }
+}

proxy-wasm-java-host/src/main/java/io/roastedroot/proxywasm/internal/PluginHttpContext.java

@@ -13,7 +13,9 @@ public class PluginHttpContext {
     private final long startedAt = System.currentTimeMillis();
 
     final HashMap<List<String>, byte[]> properties = new HashMap<>();
-    private byte[] httpRequestBody = new byte[0];
+    private HttpRequestBody httpRequestBodyState;
+
+    // Other body buffers and state fields (not lazy)
     private byte[] grpcReceiveBuffer = new byte[0];
     private byte[] upstreamData = new byte[0];
     private byte[] downStreamData = new byte[0];
@@ -22,12 +24,45 @@ public class PluginHttpContext {
     private Action action;
     private CountDownLatch resumeLatch;
 
-    PluginHttpContext(Plugin plugin, HttpRequestAdaptor requestAdaptor) {
+    public PluginHttpContext(Plugin plugin, HttpRequestAdaptor requestAdaptor) {
         this.plugin = plugin;
         this.requestAdaptor = requestAdaptor;
         this.context = plugin.wasm.createHttpContext(new HandlerImpl());
     }
 
+    /**
+     * Sets the lazy request body supplier.
+     */
+    public void setHttpRequestBodyState(HttpRequestBody supplier) {
+        this.httpRequestBodyState = supplier;
+    }
+
+    /**
+     * Gets the lazy request body supplier.
+     */
+    public HttpRequestBody getHttpRequestBodyState() {
+        return httpRequestBodyState;
+    }
+
+    /**
+     * Gets the HTTP request body, triggering lazy loading if needed.
+     */
+    public byte[] getHttpRequestBody() {
+        if (httpRequestBodyState != null) {
+            return httpRequestBodyState.get();
+        }
+        return new byte[0];
+    }
+
+    /**
+     * Sets the HTTP request body, updating the supplier if present.
+     */
+    public void setHttpRequestBody(byte[] body) {
+        if (httpRequestBodyState != null && body != null) {
+            httpRequestBodyState.setBody(body);
+        }
+    }
+
     public Plugin plugin() {
         return plugin;
     }
@@ -68,14 +103,6 @@ public void maybePause() {
         }
     }
 
-    public byte[] getHttpRequestBody() {
-        return httpRequestBody;
-    }
-
-    public void setHttpRequestBody(byte[] httpRequestBody) {
-        this.httpRequestBody = httpRequestBody;
-    }
-
     public byte[] getHttpResponseBody() {
         return httpResponseBody;
     }
@@ -156,17 +183,21 @@ public ProxyMap getGrpcReceiveTrailerMetaData() {
 
         @Override
         public byte[] getHttpRequestBody() {
-            return httpRequestBody;
+            return PluginHttpContext.this
+                    .getHttpRequestBody(); // Delegate to outer class for lazy loading
         }
 
         @Override
         public WasmResult setHttpRequestBody(byte[] body) {
-            httpRequestBody = body;
+            PluginHttpContext.this.setHttpRequestBody(body); // Delegate to outer class
             return WasmResult.OK;
         }
 
         public void appendHttpRequestBody(byte[] body) {
-            httpRequestBody = Helpers.append(httpRequestBody, body);
+            byte[] currentBody =
+                    PluginHttpContext.this
+                            .getHttpRequestBody(); // This will trigger lazy loading if needed
+            PluginHttpContext.this.setHttpRequestBody(Helpers.append(currentBody, body));
         }
 
         @Override

proxy-wasm-jaxrs/src/main/java/io/roastedroot/proxywasm/jaxrs/ProxyWasmFilter.java

@@ -4,6 +4,7 @@
 
 import io.roastedroot.proxywasm.StartException;
 import io.roastedroot.proxywasm.internal.Action;
+import io.roastedroot.proxywasm.internal.HttpRequestBody;
 import io.roastedroot.proxywasm.internal.Plugin;
 import io.roastedroot.proxywasm.internal.PluginHttpContext;
 import io.roastedroot.proxywasm.internal.Pool;
@@ -25,14 +26,19 @@
 import java.util.logging.Logger;
 
 /**
- * Implements the JAX-RS {@link ContainerRequestFilter}, {@link ContainerResponseFilter},
- * and {@link WriterInterceptor} interfaces to intercept HTTP requests and responses,
+ * Implements the JAX-RS {@link ContainerRequestFilter},
+ * {@link ContainerResponseFilter},
+ * and {@link WriterInterceptor} interfaces to intercept HTTP requests and
+ * responses,
  * allowing Proxy-Wasm plugins to process them.
  *
- * <p>This filter is registered by the {@link ProxyWasmFeature}. It interacts with
- * {@link Plugin} instances obtained from configured {@link Pool}s to execute the
+ * <p>
+ * This filter is registered by the {@link ProxyWasmFeature}. It interacts with
+ * {@link Plugin} instances obtained from configured {@link Pool}s to execute
+ * the
  * appropriate Proxy-Wasm ABI functions (e.g., {@code on_http_request_headers},
- * {@code on_http_response_body}) at different stages of the JAX-RS request/response lifecycle.
+ * {@code on_http_response_body}) at different stages of the JAX-RS
+ * request/response lifecycle.
  *
  * @see ProxyWasmFeature
  * @see ProxyWasm
@@ -50,7 +56,8 @@ public class ProxyWasmFilter
     /**
      * Constructs a ProxyWasmFilter.
      *
-     * @param pluginPools A list of {@link Pool} instances, each managing a pool of {@link Plugin}
+     * @param pluginPools A list of {@link Pool} instances, each managing a pool of
+     *                    {@link Plugin}
      *                    instances for a specific Wasm module.
      */
     public ProxyWasmFilter(List<Pool> pluginPools) {
@@ -76,10 +83,15 @@ public void release() {
     /**
      * Intercepts incoming JAX-RS requests before they reach the resource method.
      *
-     * <p>This method iterates through the configured plugin pools, borrows a {@link Plugin}
-     * instance from each, creates a {@link PluginHttpContext}, and calls the plugin's
-     * {@code on_http_request_headers} and potentially {@code on_http_request_body} functions.
-     * It handles potential early responses or modifications dictated by the plugins.
+     * <p>
+     * This method iterates through the configured plugin pools, borrows a
+     * {@link Plugin}
+     * instance from each, creates a {@link PluginHttpContext}, and calls the
+     * plugin's
+     * {@code on_http_request_headers} and potentially {@code on_http_request_body}
+     * functions.
+     * It handles potential early responses or modifications dictated by the
+     * plugins.
      *
      * @param requestContext The JAX-RS request context.
      * @throws IOException If an I/O error occurs, typically during body processing.
@@ -89,6 +101,7 @@ public void filter(ContainerRequestContext requestContext) throws IOException {
 
         ArrayList<FilterContext> filterContexts = new ArrayList<>();
         requestContext.setProperty(FILTER_CONTEXT, filterContexts);
+
         for (var pluginPool : pluginPools) {
             try {
                 Plugin plugin = pluginPool.borrow();
@@ -100,6 +113,7 @@ public void filter(ContainerRequestContext requestContext) throws IOException {
                                     serverAdaptor.httpRequestAdaptor(requestContext);
                     requestAdaptor.setRequestContext(requestContext);
                     var httpContext = plugin.createHttpContext(requestAdaptor);
+
                     filterContexts.add(new FilterContext(pluginPool, plugin, httpContext));
                 } finally {
                     plugin.unlock();
@@ -116,6 +130,17 @@ public void filter(ContainerRequestContext requestContext) throws IOException {
                 return;
             }
         }
+
+        // Create a shared lazy body supplier for all plugins
+        HttpRequestBody bodySupplier = new HttpRequestBody(() -> requestContext.getEntityStream());
+
+        // Set up lazy providers for all plugins that need the body
+        for (var filterContext : filterContexts) {
+            if (filterContext.httpContext.context().hasOnRequestBody()) {
+                filterContext.httpContext.setHttpRequestBodyState(bodySupplier);
+            }
+        }
+
         for (var filterContext : filterContexts) {
             filter(requestContext, filterContext);
         }
@@ -147,29 +172,30 @@ private void filter(ContainerRequestContext requestContext, FilterContext filter
             // the plugin may not be interested in the request body.
             if (httpContext.context().hasOnRequestBody()) {
 
-                // TODO: find out if it's more efficient to read the body in chunks and do multiple
-                // callOnRequestBody calls.
-                byte[] bytes = requestContext.getEntityStream().readAllBytes();
+                HttpRequestBody httpRequestBodyState = httpContext.getHttpRequestBodyState();
 
-                httpContext.setHttpRequestBody(bytes);
-                var action = httpContext.context().callOnRequestBody(true);
-                bytes = httpContext.getHttpRequestBody();
-                if (action == Action.CONTINUE) {
-                    // continue means plugin is done reading the body.
-                    httpContext.setHttpRequestBody(null);
-                } else {
+                while (true) {
+                    // if we streamed body updates, then endOfStream would be initially false
+                    var action = httpContext.context().callOnRequestBody(true);
+
+                    // does the plugin want to respond early?
+                    var sendResponse = httpContext.consumeSentHttpResponse();
+                    if (sendResponse != null) {
+                        requestContext.abortWith(toResponse(sendResponse));
+                        return;
+                    }
+
+                    if (action == Action.CONTINUE) {
+                        break;
+                    }
                     httpContext.maybePause();
                 }
 
-                // does the plugin want to respond early?
-                var sendResponse = httpContext.consumeSentHttpResponse();
-                if (sendResponse != null) {
-                    requestContext.abortWith(toResponse(sendResponse));
-                    return;
+                // Body was accessed and potentially modified, update the request stream
+                if (httpRequestBodyState.isLoaded()) {
+                    byte[] bytes = httpRequestBodyState.getBodyIfLoaded();
+                    requestContext.setEntityStream(new java.io.ByteArrayInputStream(bytes));
                 }
-
-                // plugin may have modified the body
-                requestContext.setEntityStream(new java.io.ByteArrayInputStream(bytes));
             }
 
         } finally {
@@ -184,10 +210,14 @@ private static Response internalServerError() {
     /**
      * Intercepts outgoing JAX-RS responses before the entity is written.
      *
-     * <p>This method iterates through the configured plugin pools, retrieves the
-     * {@link PluginHttpContext} created during the request phase, and calls the plugin's
-     * {@code on_http_response_headers} function. It handles potential modifications to the
-     * response headers dictated by the plugins. If the response has no entity but the plugin
+     * <p>
+     * This method iterates through the configured plugin pools, retrieves the
+     * {@link PluginHttpContext} created during the request phase, and calls the
+     * plugin's
+     * {@code on_http_response_headers} function. It handles potential modifications
+     * to the
+     * response headers dictated by the plugins. If the response has no entity but
+     * the plugin
      * implements {@code on_http_response_body}, it invokes that callback as well.
      *
      * @param requestContext  The JAX-RS request context.
@@ -274,15 +304,20 @@ private void filter(
     /**
      * Intercepts the response body writing process.
      *
-     * <p>This method is called when the JAX-RS framework is about to serialize and write
+     * <p>
+     * This method is called when the JAX-RS framework is about to serialize and
+     * write
      * the response entity. It captures the original response body, allows plugins
-     * (via {@code on_http_response_body}) to inspect or modify it, and then writes the
+     * (via {@code on_http_response_body}) to inspect or modify it, and then writes
+     * the
      * potentially modified body to the original output stream. It handles potential
      * early responses dictated by the plugins during body processing.
      *
      * @param ctx The JAX-RS writer interceptor context.
-     * @throws IOException             If an I/O error occurs during stream processing.
-     * @throws WebApplicationException If a plugin decides to abort processing and send an
+     * @throws IOException             If an I/O error occurs during stream
+     *                                 processing.
+     * @throws WebApplicationException If a plugin decides to abort processing and
+     *                                 send an
      *                                 alternative response during body filtering.
      */
     @Override