From eb0c6fb46ca70c307600bce0ad68f0749cf9b8ff Mon Sep 17 00:00:00 2001
From: pangdayuan <xinyuan_wang@trip.com>
Date: Wed, 10 Jan 2024 14:00:34 +0800
Subject: [PATCH] fix: the interface request fails

---
 .../interceptor/MetricInterceptor.java        | 53 +++++++++++++++++--
 1 file changed, 48 insertions(+), 5 deletions(-)

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..03ef3100 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
@@ -4,8 +4,12 @@
 import com.google.common.collect.Maps;
 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;
@@ -46,18 +50,18 @@ 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)) {
+    ContentCachingRequestWrapper cachingRequest= toCachingRequest(request);
+    ContentCachingResponseWrapper cachingResponse = toCachingResponse(response);
+    if (cachingRequest == null || cachingResponse == null) {
       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();
 
     long startTime = (Long) request.getAttribute(START_TIME);
 
@@ -66,6 +70,45 @@ public void postHandle(HttpServletRequest request, HttpServletResponse response,
     recordPayloadInfo(clientApp, category, request.getRequestURI(), requestLength, responseLength, endTime - startTime);
   }
 
+  @Override
+  public void afterCompletion(HttpServletRequest request, HttpServletResponse response,
+      Object handler, Exception ex) throws Exception {
+    ContentCachingResponseWrapper cachingResponse = toCachingResponse(response);
+    if (cachingResponse != null) {
+      cachingResponse.copyBodyToResponse();
+    }
+  }
+
+  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)) {