fix: metric interceptor output response (#148)

* fix: metric interceptor output response

arex-storage-config/pom.xml

@@ -45,7 +45,7 @@
   <parent>
     <artifactId>arex-storage-service</artifactId>
     <groupId>com.arextest</groupId>
-    <version>1.0.58</version>
+    <version>1.0.59</version>
   </parent>
 
   <properties>

arex-storage-model/pom.xml

@@ -7,7 +7,7 @@
   <parent>
     <artifactId>arex-storage-service</artifactId>
     <groupId>com.arextest</groupId>
-    <version>1.0.58</version>
+    <version>1.0.59</version>
   </parent>
 
   <profiles>

arex-storage-web-api/pom.xml

@@ -123,7 +123,7 @@
   <parent>
     <artifactId>arex-storage-service</artifactId>
     <groupId>com.arextest</groupId>
-    <version>1.0.58</version>
+    <version>1.0.59</version>
   </parent>
 
   <profiles>

arex-storage-web-api/src/main/java/com/arextest/storage/interceptor/MetricInterceptor.java

@@ -2,16 +2,19 @@
 
 import com.arextest.storage.metric.MetricListener;
 import com.google.common.collect.Maps;
+import java.io.IOException;
 import java.util.List;
 import java.util.Map;
+import javax.servlet.ServletRequest;
+import javax.servlet.ServletResponse;
 import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletRequestWrapper;
 import javax.servlet.http.HttpServletResponse;
+import javax.servlet.http.HttpServletResponseWrapper;
 import org.apache.commons.collections4.CollectionUtils;
 import org.apache.commons.lang3.StringUtils;
-import org.springframework.lang.Nullable;
 import org.springframework.stereotype.Component;
 import org.springframework.web.servlet.HandlerInterceptor;
-import org.springframework.web.servlet.ModelAndView;
 import org.springframework.web.util.ContentCachingRequestWrapper;
 import org.springframework.web.util.ContentCachingResponseWrapper;
 
@@ -44,20 +47,24 @@ public boolean preHandle(HttpServletRequest request, HttpServletResponse respons
   }
 
   @Override
-  public void postHandle(HttpServletRequest request, HttpServletResponse response, Object handler,
-      @Nullable ModelAndView modelAndView) throws Exception {
-    if (!(request instanceof ContentCachingRequestWrapper && response instanceof ContentCachingResponseWrapper)) {
+  public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex)
+      throws IOException {
+    ContentCachingRequestWrapper cachingRequest= toCachingRequest(request);
+    ContentCachingResponseWrapper cachingResponse = toCachingResponse(response);
+    if (cachingRequest == null || cachingResponse == null) {
+      if (cachingResponse != null) {
+        cachingResponse.copyBodyToResponse();
+      }
       return;
     }
 
     long endTime = System.currentTimeMillis();
 
-    byte[] requestBody = ((ContentCachingRequestWrapper) request).getContentAsByteArray();
+    byte[] requestBody = cachingRequest.getContentAsByteArray();
     int requestLength = requestBody.length;
 
-    ContentCachingResponseWrapper responseWrapper = (ContentCachingResponseWrapper) response;
-    int responseLength = responseWrapper.getContentSize();
-    responseWrapper.copyBodyToResponse();
+    int responseLength = cachingResponse.getContentSize();
+    cachingResponse.copyBodyToResponse();
 
     long startTime = (Long) request.getAttribute(START_TIME);
 
@@ -66,6 +73,36 @@ public void postHandle(HttpServletRequest request, HttpServletResponse response,
     recordPayloadInfo(clientApp, category, request.getRequestURI(), requestLength, responseLength, endTime - startTime);
   }
 
+  private ContentCachingRequestWrapper toCachingRequest(HttpServletRequest request) {
+    if (request instanceof ContentCachingRequestWrapper) {
+      return (ContentCachingRequestWrapper) request;
+    }
+
+    if (request instanceof HttpServletRequestWrapper) {
+      ServletRequest cachingRequest = ((HttpServletRequestWrapper) request).getRequest();
+      if (cachingRequest instanceof ContentCachingRequestWrapper) {
+        return (ContentCachingRequestWrapper) cachingRequest;
+      }
+    }
+
+    return null;
+  }
+
+  private ContentCachingResponseWrapper toCachingResponse(HttpServletResponse response) {
+    if (response instanceof ContentCachingResponseWrapper) {
+      return (ContentCachingResponseWrapper) response;
+    }
+
+    if (response instanceof HttpServletResponseWrapper) {
+      ServletResponse cachingResponse = ((HttpServletResponseWrapper) response).getResponse();
+      if (cachingResponse instanceof ContentCachingResponseWrapper) {
+        return (ContentCachingResponseWrapper) cachingResponse;
+      }
+    }
+
+    return null;
+  }
+
   public void recordPayloadInfo(String clientApp, String category,
       String path, int requestLength, int responseLength, long executeMillis) {
     if (CollectionUtils.isEmpty(metricListeners)) {

pom.xml

@@ -408,5 +408,5 @@
   <url>https://github.com/arextest/arex-storage</url>
 
 
-  <version>1.0.58</version>
+  <version>1.0.59</version>
 </project>