From 32a2485aa4071e4faf67da49716ec6539f0c29f6 Mon Sep 17 00:00:00 2001
From: Andre Kurait <akurait@amazon.com>
Date: Thu, 25 Apr 2024 20:05:19 -0500
Subject: [PATCH] Harden NettyJsonContentAuthSigner

Signed-off-by: Andre Kurait <akurait@amazon.com>
---
 .../replay/ParsedHttpMessagesAsDicts.java     |  2 +-
 .../http/NettyJsonContentAuthSigner.java      | 36 +++++++++++++------
 2 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/ParsedHttpMessagesAsDicts.java b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/ParsedHttpMessagesAsDicts.java
index 93df879de3..56179f694e 100644
--- a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/ParsedHttpMessagesAsDicts.java
+++ b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/ParsedHttpMessagesAsDicts.java
@@ -104,7 +104,7 @@ public static void fillStatusCodeMetrics(@NonNull IReplayContexts.ITupleHandling
 
     private static Map<String, Object> fillMap(LinkedHashMap<String, Object> map,
                                                HttpHeaders headers, ByteBuf content) {
-        try (var bufHolder = RefSafeHolder.create(Base64.encode(content, Base64Dialect.URL_SAFE))) {
+        try (var bufHolder = RefSafeHolder.create(Base64.encode(content, false, Base64Dialect.STANDARD))) {
             var buf = bufHolder.get();
             assert buf != null : "Base64.encode should not return null";
             String base64body = buf.toString(StandardCharsets.UTF_8);
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..cc5b425a14 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.channelUnregistered(ctx);
+    }
+
 }
\ No newline at end of file