From 9a5d323c893fa3fd0f415e815d945fc606218d63 Mon Sep 17 00:00:00 2001 From: Andre Kurait Date: Sun, 6 Oct 2024 00:18:05 -0500 Subject: [PATCH] Support modifying inlineBinaryBody in request transform Signed-off-by: Andre Kurait --- ...NettyDecodedHttpRequestConvertHandler.java | 16 +++++-- .../http/NettyJsonBodyConvertHandler.java | 3 +- .../HttpJsonTransformingConsumerTest.java | 46 ++++++++++++++++++- .../replay/util/RefSafeHolderTest.java | 30 ++++++++++++ .../requests/raw/post_withPlainText.txt | 7 +++ 5 files changed, 96 insertions(+), 6 deletions(-) create mode 100644 TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/util/RefSafeHolderTest.java create mode 100644 TrafficCapture/trafficReplayer/src/test/resources/requests/raw/post_withPlainText.txt diff --git a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/datahandlers/http/NettyDecodedHttpRequestConvertHandler.java b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/datahandlers/http/NettyDecodedHttpRequestConvertHandler.java index 6d48f58ce..4063c8ec7 100644 --- a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/datahandlers/http/NettyDecodedHttpRequestConvertHandler.java +++ b/TrafficCapture/trafficReplayer/src/main/java/org/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/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/datahandlers/http/NettyJsonBodyConvertHandler.java b/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/datahandlers/http/NettyJsonBodyConvertHandler.java index f8f532fa3..94c0ca2eb 100644 --- a/TrafficCapture/trafficReplayer/src/main/java/org/opensearch/migrations/replay/datahandlers/http/NettyJsonBodyConvertHandler.java +++ b/TrafficCapture/trafficReplayer/src/main/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/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/datahandlers/http/HttpJsonTransformingConsumerTest.java b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/datahandlers/http/HttpJsonTransformingConsumerTest.java index 995239005..407c3e846 100644 --- a/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/datahandlers/http/HttpJsonTransformingConsumerTest.java +++ b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/datahandlers/http/HttpJsonTransformingConsumerTest.java @@ -57,7 +57,8 @@ private static Stream 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/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/util/RefSafeHolderTest.java b/TrafficCapture/trafficReplayer/src/test/java/org/opensearch/migrations/replay/util/RefSafeHolderTest.java new file mode 100644 index 000000000..f6ea1ef1d --- /dev/null +++ b/TrafficCapture/trafficReplayer/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 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 new file mode 100644 index 000000000..4cbc86e05 --- /dev/null +++ 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