From ecde7138da6a0e25a92b2b73d5938573d2dcd6bb Mon Sep 17 00:00:00 2001
From: Andre Kurait <akurait@amazon.com>
Date: Thu, 25 Apr 2024 19:26:16 -0500
Subject: [PATCH] NettyJsonContentAuthSigner Hardening

Signed-off-by: Andre Kurait <akurait@amazon.com>
---
 .../http/NettyJsonContentAuthSigner.java      | 36 +++++++++++++------
 1 file changed, 25 insertions(+), 11 deletions(-)

diff --git a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/datahandlers/http/NettyJsonContentAuthSigner.java b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/datahandlers/http/NettyJsonContentAuthSigner.java
index 31ceebc530..dfbf10e19e 100644
--- a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/datahandlers/http/NettyJsonContentAuthSigner.java
+++ b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/datahandlers/http/NettyJsonContentAuthSigner.java
@@ -12,11 +12,11 @@
 public class NettyJsonContentAuthSigner extends ChannelInboundHandlerAdapter {
     IAuthTransformer.StreamingFullMessageTransformer signer;
     HttpJsonMessageWithFaultingPayload httpMessage;
-    List<HttpContent> receivedHttpContents;
+    List<HttpContent> httpContentsBuffer;
 
     public NettyJsonContentAuthSigner(IAuthTransformer.StreamingFullMessageTransformer signer) {
         this.signer = signer;
-        this.receivedHttpContents = new ArrayList<>();
+        this.httpContentsBuffer = new ArrayList<>();
     }
 
     @Override
@@ -24,23 +24,37 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
         if (msg instanceof HttpJsonMessageWithFaultingPayload) {
             httpMessage = (HttpJsonMessageWithFaultingPayload) msg;
         } else if (msg instanceof HttpContent) {
-            receivedHttpContents.add(((HttpContent) msg).retainedDuplicate());
             var httpContent = (HttpContent) msg;
+            httpContentsBuffer.add(httpContent);
             signer.consumeNextPayloadPart(httpContent.content().nioBuffer());
             if  (msg instanceof LastHttpContent) {
-                finalizeSignature(ctx);
+                signer.finalizeSignature(httpMessage);
+                flushDownstream(ctx);
             }
         } else {
             super.channelRead(ctx, msg);
         }
     }
 
-    private void finalizeSignature(ChannelHandlerContext ctx) {
-        signer.finalizeSignature(httpMessage);
-        ctx.fireChannelRead(httpMessage);
-        receivedHttpContents.stream().forEach(content->{
-            ctx.fireChannelRead(content);
-            content.content().release();
-        });
+    private void flushDownstream(ChannelHandlerContext ctx) {
+        if(httpMessage != null) {
+            ctx.fireChannelRead(httpMessage);
+            httpMessage = null;
+        }
+        httpContentsBuffer.forEach(ctx::fireChannelRead);
+        httpContentsBuffer.clear();
+    }
+
+    @Override
+    public void handlerRemoved(ChannelHandlerContext ctx) throws Exception {
+        flushDownstream(ctx);
+        super.handlerRemoved(ctx);
     }
+
+    @Override
+    public void channelUnregistered(ChannelHandlerContext ctx) throws Exception {
+        flushDownstream(ctx);
+        super.channelInactive(ctx);
+    }
+
 }
\ No newline at end of file