From 2f15a34520fccb0668f69ddf2c95a9eddce1068a Mon Sep 17 00:00:00 2001
From: johnniang <johnniang@riseup.net>
Date: Thu, 21 Mar 2019 10:55:11 +0800
Subject: [PATCH] Refactor CommonResultControllerAdvice

---
 .../base/CommonResultControllerAdvice.java    | 42 +++++++++++++++----
 1 file changed, 35 insertions(+), 7 deletions(-)

diff --git a/src/main/java/cc/ryanc/halo/web/controller/base/CommonResultControllerAdvice.java b/src/main/java/cc/ryanc/halo/web/controller/base/CommonResultControllerAdvice.java
index 26116c651..0dc1cf57e 100644
--- a/src/main/java/cc/ryanc/halo/web/controller/base/CommonResultControllerAdvice.java
+++ b/src/main/java/cc/ryanc/halo/web/controller/base/CommonResultControllerAdvice.java
@@ -4,11 +4,15 @@ import cc.ryanc.halo.model.support.BaseResponse;
 import org.springframework.core.MethodParameter;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.MediaType;
+import org.springframework.http.converter.HttpMessageConverter;
+import org.springframework.http.converter.json.AbstractJackson2HttpMessageConverter;
 import org.springframework.http.converter.json.MappingJacksonValue;
 import org.springframework.http.server.ServerHttpRequest;
 import org.springframework.http.server.ServerHttpResponse;
+import org.springframework.lang.NonNull;
+import org.springframework.lang.Nullable;
 import org.springframework.web.bind.annotation.ControllerAdvice;
-import org.springframework.web.servlet.mvc.method.annotation.AbstractMappingJacksonResponseBodyAdvice;
+import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;
 
 /**
  * Controller advice for comment result.
@@ -16,14 +20,37 @@ import org.springframework.web.servlet.mvc.method.annotation.AbstractMappingJack
  * @author johnniang
  */
 @ControllerAdvice("cc.ryanc.halo.web.controller")
-public class CommonResultControllerAdvice extends AbstractMappingJacksonResponseBodyAdvice {
+public class CommonResultControllerAdvice implements ResponseBodyAdvice<Object> {
 
     @Override
-    protected void beforeBodyWriteInternal(MappingJacksonValue bodyContainer,
-                                           MediaType contentType,
-                                           MethodParameter returnType,
-                                           ServerHttpRequest request,
-                                           ServerHttpResponse response) {
+    public boolean supports(MethodParameter returnType, Class<? extends HttpMessageConverter<?>> converterType) {
+        return AbstractJackson2HttpMessageConverter.class.isAssignableFrom(converterType);
+    }
+
+    @Override
+    @NonNull
+    public final Object beforeBodyWrite(@Nullable Object body, MethodParameter returnType,
+                                        MediaType contentType, Class<? extends HttpMessageConverter<?>> converterType,
+                                        ServerHttpRequest request, ServerHttpResponse response) {
+        MappingJacksonValue container = getOrCreateContainer(body);
+        // The contain body will never be null
+        beforeBodyWriteInternal(container, contentType, returnType, request, response);
+        return container;
+    }
+
+    /**
+     * Wrap the body in a {@link MappingJacksonValue} value container (for providing
+     * additional serialization instructions) or simply cast it if already wrapped.
+     */
+    private MappingJacksonValue getOrCreateContainer(Object body) {
+        return (body instanceof MappingJacksonValue ? (MappingJacksonValue) body : new MappingJacksonValue(body));
+    }
+
+    private void beforeBodyWriteInternal(MappingJacksonValue bodyContainer,
+                                         MediaType contentType,
+                                         MethodParameter returnType,
+                                         ServerHttpRequest request,
+                                         ServerHttpResponse response) {
         // Get return body
         Object returnBody = bodyContainer.getValue();
 
@@ -39,4 +66,5 @@ public class CommonResultControllerAdvice extends AbstractMappingJacksonResponse
         bodyContainer.setValue(baseResponse);
         response.setStatusCode(HttpStatus.valueOf(baseResponse.getStatus()));
     }
+
 }