Merge pull request #23 from SmartBear/explore-646-limit-response-size

feat: limit response size
SmartBear · Oct 13, 2022 · c015749 · c015749
2 parents f192e42 + 4a24b0b
commit c015749
Show file tree

Hide file tree

Showing 8 changed files with 121 additions and 17 deletions.
diff --git a/logbook-httpclient5/src/main/java/org/zalando/logbook/httpclient5/HttpEntities.java b/logbook-httpclient5/src/main/java/org/zalando/logbook/httpclient5/HttpEntities.java
@@ -19,8 +19,8 @@ interface Copy extends HttpEntity {
         byte[] getBody();
     }
 
-    Copy copy(final HttpEntity entity) throws IOException {
-        final byte[] body = toByteArray(entity);
+    Copy copy(final HttpEntity entity, int sizeLimit) throws IOException {
+        final byte[] body = toByteArray(entity, sizeLimit);
         ContentType contentType = ContentType.parse(entity.getContentType());
         boolean chunked = entity.isChunked();
         String contentEncoding = entity.getContentEncoding();

diff --git a/logbook-httpclient5/src/main/java/org/zalando/logbook/httpclient5/LocalRequest.java b/logbook-httpclient5/src/main/java/org/zalando/logbook/httpclient5/LocalRequest.java
@@ -76,7 +76,7 @@ public State buffer(final HttpRequest request, final EntityDetails entity) throw
                 if (httpEntity == null) {
                     return new Passing();
                 } else {
-                    final HttpEntities.Copy copy = HttpEntities.copy(httpEntity);
+                    final HttpEntities.Copy copy = HttpEntities.copy(httpEntity, Integer.MAX_VALUE);
                     original.setEntity(copy);
                     return new Buffering(copy.getBody());
                 }

diff --git a/...ient5/src/main/java/org/zalando/logbook/httpclient5/LogbookHttpAsyncResponseConsumer.java b/...ient5/src/main/java/org/zalando/logbook/httpclient5/LogbookHttpAsyncResponseConsumer.java
@@ -46,7 +46,7 @@ public void consumeResponse(HttpResponse response, EntityDetails entityDetails,
 
     @Override
     public void consume(ByteBuffer src) throws IOException {
-        stage.process(new RemoteResponse(this.response, this.entityDetails, src)).write();
+        stage.process(new RemoteResponse(this.response, this.entityDetails, src, Integer.MAX_VALUE)).write();
         delegate().consume(src);
     }
 

diff --git a/...client5/src/main/java/org/zalando/logbook/httpclient5/LogbookHttpResponseInterceptor.java b/...client5/src/main/java/org/zalando/logbook/httpclient5/LogbookHttpResponseInterceptor.java
@@ -21,11 +21,19 @@
  */
 @API(status = EXPERIMENTAL)
 public final class LogbookHttpResponseInterceptor implements HttpResponseInterceptor {
+    final int maxResponseSize;
+    public LogbookHttpResponseInterceptor() {
+        maxResponseSize = Integer.MAX_VALUE;
+    }
+
+    public LogbookHttpResponseInterceptor(int maxResponseSize) {
+        this.maxResponseSize = maxResponseSize;
+    }
 
     @Override
     public void process(HttpResponse original, EntityDetails entity, HttpContext context) throws IOException {
         final ResponseProcessingStage stage = find(context);
-        stage.process(new RemoteResponse(original)).write();
+        stage.process(new RemoteResponse(original, maxResponseSize)).write();
     }
 
     private ResponseProcessingStage find(final HttpContext context) {

diff --git a/logbook-httpclient5/src/main/java/org/zalando/logbook/httpclient5/RemoteResponse.java b/logbook-httpclient5/src/main/java/org/zalando/logbook/httpclient5/RemoteResponse.java
@@ -25,11 +25,20 @@ final class RemoteResponse implements org.zalando.logbook.HttpResponse {
     private final HttpResponse response;
     private final EntityDetails entityDetails;
     private final ByteBuffer body;
+    private final int sizeLimit;
 
     RemoteResponse(HttpResponse response) {
         this.response = response;
         this.body = null;
         this.entityDetails = null;
+        this.sizeLimit = Integer.MAX_VALUE;
+    }
+
+    RemoteResponse(HttpResponse response, int sizeLimit) {
+        this.response = response;
+        this.body = null;
+        this.entityDetails = null;
+        this.sizeLimit = sizeLimit;
     }
 
     private interface State {
@@ -42,7 +51,7 @@ default State without() {
             return this;
         }
 
-        default State buffer(final HttpResponse response) throws IOException {
+        default State buffer(final HttpResponse response, int sizeLimit) throws IOException {
             return this;
         }
 
@@ -73,14 +82,14 @@ public State without() {
         }
 
         @Override
-        public State buffer(final HttpResponse response) throws IOException {
+        public State buffer(final HttpResponse response, int sizeLimit) throws IOException {
             if (response instanceof ClassicHttpResponse) {
                 ClassicHttpResponse classicResponse = (ClassicHttpResponse) response;
-                HttpEntity entity = (classicResponse).getEntity();
+                HttpEntity entity = classicResponse.getEntity();
                 if (entity == null) {
                     return new Passing();
                 } else {
-                    final HttpEntities.Copy copy = HttpEntities.copy(entity);
+                    final HttpEntities.Copy copy = HttpEntities.copy(entity, sizeLimit);
                     classicResponse.setEntity(copy);
                     return new Buffering(copy.getBody());
                 }
@@ -212,7 +221,7 @@ public RemoteResponse withoutBody() {
 
     @Override
     public byte[] getBody() {
-        return state.updateAndGet(throwingUnaryOperator(state -> (body != null) ? state.buffer(Objects.requireNonNull(entityDetails), body) : state.buffer(response))).getBody();
+        return state.updateAndGet(throwingUnaryOperator(state -> (body != null) ? state.buffer(Objects.requireNonNull(entityDetails), body) : state.buffer(response, sizeLimit))).getBody();
     }
 
 }
diff --git a/...ok-httpclient5/src/test/java/org/zalando/logbook/httpclient5/LimitedResponseSizeTest.java b/...ok-httpclient5/src/test/java/org/zalando/logbook/httpclient5/LimitedResponseSizeTest.java
@@ -0,0 +1,86 @@
+package org.zalando.logbook.httpclient5;
+
+import com.github.restdriver.clientdriver.ClientDriver;
+import com.github.restdriver.clientdriver.ClientDriverFactory;
+import org.apache.hc.client5.http.classic.methods.HttpGet;
+import org.apache.hc.client5.http.impl.classic.CloseableHttpClient;
+import org.apache.hc.client5.http.impl.classic.CloseableHttpResponse;
+import org.apache.hc.client5.http.impl.classic.HttpClientBuilder;
+import org.apache.hc.core5.http.ParseException;
+import org.apache.hc.core5.http.io.entity.EntityUtils;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+import org.zalando.logbook.Correlation;
+import org.zalando.logbook.DefaultHttpLogFormatter;
+import org.zalando.logbook.DefaultSink;
+import org.zalando.logbook.HttpLogWriter;
+import org.zalando.logbook.Logbook;
+import org.zalando.logbook.TestStrategy;
+
+import java.io.IOException;
+
+import static com.github.restdriver.clientdriver.ClientDriverRequest.Method.GET;
+import static com.github.restdriver.clientdriver.RestClientDriver.giveResponse;
+import static com.github.restdriver.clientdriver.RestClientDriver.onRequestTo;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class LimitedResponseSizeTest {
+
+    private static final int MAX_RESPONSE_SIZE = 10;
+    final ClientDriver driver = new ClientDriverFactory().createClientDriver();
+
+    final HttpLogWriter writer = mock(HttpLogWriter.class);
+
+    protected final Logbook logbook = Logbook.builder()
+            .strategy(new TestStrategy())
+            .sink(new DefaultSink(new DefaultHttpLogFormatter(), writer))
+            .build();
+
+    private final CloseableHttpClient client = HttpClientBuilder.create()
+            .addRequestInterceptorFirst(new LogbookHttpRequestInterceptor(logbook))
+            .addResponseInterceptorFirst(new LogbookHttpResponseInterceptor(MAX_RESPONSE_SIZE))
+            .build();
+
+    @BeforeEach
+    void defaultBehaviour() {
+        when(writer.isActive()).thenReturn(true);
+    }
+
+    @AfterEach
+    void stop() throws IOException {
+        client.close();
+    }
+
+    @Test
+    void shouldLogRequestWithoutBody() throws IOException, ParseException {
+        driver.addExpectation(onRequestTo("/").withMethod(GET),
+                giveResponse("Hello, world!", "text/plain"));
+
+        final HttpGet request = new HttpGet(driver.getBaseUrl());
+
+        final CloseableHttpResponse response = client.execute(request);
+
+        assertThat(response.getCode()).isEqualTo(200);
+        assertThat(response.getEntity()).isNotNull();
+        final String cappedResponse = "Hello, world!".substring(0, 10);
+        assertThat(EntityUtils.toString(response.getEntity())).isEqualTo(cappedResponse);
+
+        final String message = captureResponse();
+
+        assertThat(message)
+                .startsWith("Incoming Response:")
+                .contains("HTTP/1.1 200 OK", "Content-Type: text/plain", cappedResponse);
+    }
+
+    private String captureResponse() throws IOException {
+        final ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
+        verify(writer).write(any(Correlation.class), captor.capture());
+        return captor.getValue();
+    }
+}
diff --git a/logbook-parent/pom.xml b/logbook-parent/pom.xml
@@ -25,21 +25,21 @@
         <json-path.version>2.7.0</json-path.version>
         <kotlin.version>1.6.10</kotlin.version>
         <ktor.version>1.6.8</ktor.version>
-        <netty.version>4.1.81.Final</netty.version>
+        <netty.version>4.1.82.Final</netty.version>
         <reactor-netty.version>1.0.22</reactor-netty.version>
         <slf4j.version>1.7.36</slf4j.version>
 
         <spring4.version>4.3.30.RELEASE</spring4.version>
         <spring-boot1.version>1.5.22.RELEASE</spring-boot1.version>
         <spring-security4.version>4.2.20.RELEASE</spring-security4.version>
 
-        <spring5.version>5.3.16</spring5.version>
+        <spring5.version>5.3.23</spring5.version>
         <spring-boot2.version>2.6.4</spring-boot2.version>
-        <spring-security5.version>5.6.0</spring-security5.version>
+        <spring-security5.version>5.6.7</spring-security5.version>
 
-        <spring.version>5.3.16</spring.version>
+        <spring.version>5.3.23</spring.version>
         <spring-boot.version>2.6.4</spring-boot.version>
-        <spring-security.version>5.6.2</spring-security.version>
+        <spring-security.version>5.6.7</spring-security.version>
 
         <junit.version>5.8.2</junit.version>
         <mockito.version>4.3.1</mockito.version>

diff --git a/logbook-spring/src/test/java/org/zalando/logbook/spring/RemoteResponseTest.java b/logbook-spring/src/test/java/org/zalando/logbook/spring/RemoteResponseTest.java
@@ -2,6 +2,7 @@
 
 import org.junit.jupiter.api.Test;
 import org.springframework.http.HttpStatus;
+import org.springframework.http.client.ClientHttpResponse;
 import org.springframework.mock.http.client.MockClientHttpResponse;
 
 import java.io.IOException;
@@ -15,7 +16,7 @@ class RemoteResponseTest {
 
     @Test
     void statusCanThrow() throws IOException {
-        MockClientHttpResponse response = mock(MockClientHttpResponse.class);
+        ClientHttpResponse response = mock(ClientHttpResponse.class);
         when(response.getRawStatusCode()).thenThrow(new IOException("io exception"));
         assertThatThrownBy(() -> unit(response).getStatus()).hasMessageContaining("io exception");
     }
@@ -34,7 +35,7 @@ private MockClientHttpResponse helloWorld() {
         return new MockClientHttpResponse("hello world".getBytes(), HttpStatus.OK);
     }
 
-    private org.zalando.logbook.HttpResponse unit(MockClientHttpResponse response) {
+    private org.zalando.logbook.HttpResponse unit(ClientHttpResponse response) {
         return new RemoteResponse(response);
     }
 }