Support modifying inlineBinaryBody in request transform

Signed-off-by: Andre Kurait <[email protected]>
opensearch-project · Oct 6, 2024 · 9a5d323 · 9a5d323
1 parent e072e95
commit 9a5d323
Show file tree

Hide file tree

Showing 5 changed files with 96 additions and 6 deletions.
diff --git a/...opensearch/migrations/replay/datahandlers/http/NettyDecodedHttpRequestConvertHandler.java b/...opensearch/migrations/replay/datahandlers/http/NettyDecodedHttpRequestConvertHandler.java
@@ -1,5 +1,6 @@
 package org.opensearch.migrations.replay.datahandlers.http;
 
+import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -102,7 +103,7 @@ public void channelRead(@NonNull ChannelHandlerContext ctx, @NonNull Object msg)
         }
     }
 
-    private static HttpJsonRequestWithFaultingPayload transform(
+    public static HttpJsonRequestWithFaultingPayload transform(
         IJsonTransformer transformer,
         HttpJsonRequestWithFaultingPayload httpJsonMessage
     ) {
@@ -130,10 +131,19 @@ private static HttpJsonRequestWithFaultingPayload transform(
 
         if (innerPayloadByteBuf != null) {
             // replace protected byteBuf if was hidden and still there
-            var transformedByteBuf = httpJsonMessage.payload().get(JsonKeysForHttpMessage.INLINED_BINARY_BODY_DOCUMENT_KEY);
+            var transformedValue = httpJsonMessage.payload().get(JsonKeysForHttpMessage.INLINED_BINARY_BODY_DOCUMENT_KEY);
             assert protectedByteBuf != null : "Expected protectedByteBuf to be defined if innerPayloadByteBuf is";
-            if (protectedByteBuf.equals(transformedByteBuf)) {
+            if (protectedByteBuf.equals(transformedValue)) {
                 httpJsonMessage.payload().put(JsonKeysForHttpMessage.INLINED_BINARY_BODY_DOCUMENT_KEY, innerPayloadByteBuf);
+            } else {
+                innerPayloadByteBuf.release();
+                if (transformedValue instanceof String) {
+                    httpJsonMessage.payload().put(JsonKeysForHttpMessage.INLINED_BINARY_BODY_DOCUMENT_KEY,
+                        Unpooled.wrappedBuffer(((String) transformedValue).getBytes(StandardCharsets.UTF_8)));
+                } else if (!(transformedValue instanceof ByteBuf)) {
+                    throw new UnsupportedOperationException("Type of " + JsonKeysForHttpMessage.INLINED_BINARY_BODY_DOCUMENT_KEY
+                 + " not supported.");
+                }
             }
         }
 

diff --git a/.../java/org/opensearch/migrations/replay/datahandlers/http/NettyJsonBodyConvertHandler.java b/.../java/org/opensearch/migrations/replay/datahandlers/http/NettyJsonBodyConvertHandler.java
@@ -27,8 +27,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
             }
             HttpJsonRequestWithFaultingPayload newHttpJson;
             try {
-                var output = transformer.transformJson(httpMsg);
-                newHttpJson = new HttpJsonRequestWithFaultingPayload(output);
+                newHttpJson = NettyDecodedHttpRequestConvertHandler.transform(transformer, httpMsg);
             } catch (Exception e) {
                 var remainingBytes = httpMsg.payload().get(JsonKeysForHttpMessage.INLINED_BINARY_BODY_DOCUMENT_KEY);
                 ReferenceCountUtil.release(remainingBytes); // release because we're not passing it along for cleanup

diff --git a/.../org/opensearch/migrations/replay/datahandlers/http/HttpJsonTransformingConsumerTest.java b/.../org/opensearch/migrations/replay/datahandlers/http/HttpJsonTransformingConsumerTest.java
@@ -57,7 +57,8 @@ private static Stream<Arguments> provideTestParameters() {
             "/requests/raw/post_formUrlEncoded_withLargeHeader.txt",
             "/requests/raw/post_formUrlEncoded_withDuplicateHeaders.txt",
             "/requests/raw/get_withAuthHeader.txt",
-            "/requests/raw/post_json_gzip.gz"
+            "/requests/raw/post_json_gzip.gz",
+            "/requests/raw/post_withPlainText.txt",
         };
 
         return Stream.of(attemptedChunks)
@@ -160,6 +161,49 @@ public void testRemoveAuthHeadersWorks() throws Exception {
         Assertions.assertArrayEquals(testBytes, testPacketCapture.getBytesCaptured());
     }
 
+    @Test
+    @Tag("longTest")
+    public void testRemoveBinaryPayloadWorks() throws Exception {
+        final var dummyAggregatedResponse = new AggregatedRawResponse(null, 17, Duration.ZERO, List.of(), null);
+        var testPacketCapture = new TestCapturePacketToHttpHandler(Duration.ofMillis(100), dummyAggregatedResponse);
+        String redactBody = "{ " +
+            "    \"operation\": \"modify-overwrite-beta\", " +
+            "    \"spec\": { " +
+            "       \"payload\": { " +
+            "         \"inlinedBinaryBody\": \"REDACTED\" " +
+            "       } " +
+            "   } " +
+            "}";
+        String fullConfig = "[{\"JsonJoltTransformerProvider\": { \"script\": " + redactBody + "}}]";
+        IJsonTransformer jsonJoltTransformer = new TransformationLoader().getTransformerFactoryLoader(null, null, fullConfig);
+
+        var transformingHandler = new HttpJsonTransformingConsumer<>(
+            jsonJoltTransformer,
+            null,
+            testPacketCapture,
+            rootContext.getTestConnectionRequestContext(0)
+        );
+        byte[] testBytes;
+        try (
+            var sampleStream = HttpJsonTransformingConsumer.class.getResourceAsStream(
+                "/requests/raw/post_withPlainText.txt"
+            )
+        ) {
+            assert sampleStream != null;
+            testBytes = sampleStream.readAllBytes();
+        }
+        transformingHandler.consumeBytes(testBytes);
+        var returnedResponse = transformingHandler.finalizeRequest().get();
+        var expectedString = new String(testBytes, StandardCharsets.UTF_8)
+            .replace("This is a test\r\n","REDACTED")
+            .replace("Content-Length: 15", "Content-Length: 8");
+        Assertions.assertEquals(expectedString, testPacketCapture.getCapturedAsString());
+        Assertions.assertArrayEquals(expectedString.getBytes(StandardCharsets.UTF_8),
+            testPacketCapture.getBytesCaptured());
+        Assertions.assertEquals(HttpRequestTransformationStatus.completed(), returnedResponse.transformationStatus);
+        Assertions.assertNull(returnedResponse.transformationStatus.getException());
+    }
+
     @Test
     public void testPartialBodyIsPassedThrough() throws Exception {
         final var dummyAggregatedResponse = new AggregatedRawResponse(null, 17, Duration.ZERO, List.of(), null);

diff --git a/...rafficReplayer/src/test/java/org/opensearch/migrations/replay/util/RefSafeHolderTest.java b/...rafficReplayer/src/test/java/org/opensearch/migrations/replay/util/RefSafeHolderTest.java
@@ -0,0 +1,30 @@
+package org.opensearch.migrations.replay.util;
+
+import org.junit.jupiter.api.Test;
+
+import io.netty.util.ReferenceCountUtil;
+import org.mockito.MockedStatic;
+
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.Mockito.*;
+
+public class RefSafeHolderTest {
+    @Test
+    void testCloseWithNonNullResource() {
+        try (MockedStatic<ReferenceCountUtil> mockedStatic = mockStatic(ReferenceCountUtil.class)) {
+            Object resource = new Object();
+            try (var ignored = RefSafeHolder.create(resource)) {
+                mockedStatic.verify(() -> ReferenceCountUtil.release(resource), times(0));
+            }
+            mockedStatic.verify(() -> ReferenceCountUtil.release(resource), times(1));
+        }
+    }
+
+    @Test
+    void testCloseWithNullResource() {
+        var holder = RefSafeHolder.create(null);
+        assertNull(holder.get());
+        assertEquals("RefSafeHolder{null}", holder.toString());
+        holder.close();
+    }
+}
diff --git a/TrafficCapture/trafficReplayer/src/test/resources/requests/raw/post_withPlainText.txt b/TrafficCapture/trafficReplayer/src/test/resources/requests/raw/post_withPlainText.txt
@@ -0,0 +1,7 @@
+POST /test HTTP/1.1
+Host: foo.example
+Authorization: Basic YWRtaW46YWRtaW4=
+Content-Type: text/plain
+Content-Length: 15
+
+This is a test