From f92417457c3df92fa7f6a3298d662a16e899638e Mon Sep 17 00:00:00 2001
From: pangdayuan1 <116159079+pangdayuan1@users.noreply.github.com>
Date: Thu, 11 Jan 2024 11:58:41 +0800
Subject: [PATCH] fix: metric interceptor output response (#148)

* fix: metric interceptor output response
---
 arex-storage-config/pom.xml                   |  2 +-
 arex-storage-model/pom.xml                    |  2 +-
 arex-storage-web-api/pom.xml                  |  2 +-
 .../interceptor/MetricInterceptor.java        | 55 ++++++++++++++++---
 pom.xml                                       |  2 +-
 5 files changed, 50 insertions(+), 13 deletions(-)

diff --git a/arex-storage-config/pom.xml b/arex-storage-config/pom.xml
index ef9ee108..f941e175 100644
--- a/arex-storage-config/pom.xml
+++ b/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>
diff --git a/arex-storage-model/pom.xml b/arex-storage-model/pom.xml
index bc7ccc92..1e6619bc 100644
--- a/arex-storage-model/pom.xml
+++ b/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>
diff --git a/arex-storage-web-api/pom.xml b/arex-storage-web-api/pom.xml
index e556cd8f..b5fc17c2 100644
--- a/arex-storage-web-api/pom.xml
+++ b/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>
diff --git a/arex-storage-web-api/src/main/java/com/arextest/storage/interceptor/MetricInterceptor.java b/arex-storage-web-api/src/main/java/com/arextest/storage/interceptor/MetricInterceptor.java
index 7a21a5b6..b4508156 100644
--- a/arex-storage-web-api/src/main/java/com/arextest/storage/interceptor/MetricInterceptor.java
+++ b/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)) {
diff --git a/pom.xml b/pom.xml
index 15049b45..e31e5133 100644
--- a/pom.xml
+++ b/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>
\ No newline at end of file