Spring Boot 统一修改 Controller 返回值

需求

在Controller返回对象时，包装一层统一的返回值对象。

如：返回了String对象，需包装成{"code":0, "data":"string"}。

HandlerMethodReturnValueHandler

自定义 HandlerMethodReturnValueHandler

public class MyReturnValueHandler implements HandlerMethodReturnValueHandler {

    private final HandlerMethodReturnValueHandler handler;

    public MyReturnValueHandler(HandlerMethodReturnValueHandler handler) {
        this.handler = handler;
    }

    @Override
    public boolean supportsReturnType(MethodParameter returnType) {
        return (AnnotatedElementUtils.hasAnnotation(returnType.getContainingClass(), ResponseBody.class) ||
            returnType.hasMethodAnnotation(ResponseBody.class));
    }

    @Override
    public void handleReturnValue(Object returnValue, MethodParameter returnType,
        ModelAndViewContainer mavContainer, NativeWebRequest webRequest) throws Exception {
        Map<String, Object> map = new HashMap<>();
        map.put("code", 0);
        map.put("data", returnValue);
        handler.handleReturnValue(map, returnType, mavContainer, webRequest);
    }
}

注入 HandlerMethodReturnValueHandler

@Configuration
public class ReturnValueConfig implements InitializingBean {

    private final RequestMappingHandlerAdapter requestMappingHandlerAdapter;

    public ReturnValueConfig(RequestMappingHandlerAdapter requestMappingHandlerAdapter) {
        this.requestMappingHandlerAdapter = requestMappingHandlerAdapter;
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        List<HandlerMethodReturnValueHandler> originHandlers = requestMappingHandlerAdapter.getReturnValueHandlers();
        List<HandlerMethodReturnValueHandler> newHandlers = new ArrayList<>(originHandlers.size());
        for (HandlerMethodReturnValueHandler originHandler : originHandlers) {
            if (originHandler instanceof RequestResponseBodyMethodProcessor) {
                newHandlers.add(new MyReturnValueHandler(originHandler));
            }else{
                newHandlers.add(originHandler);
            }
        }
        requestMappingHandlerAdapter.setReturnValueHandlers(newHandlers);
    }
}

ResponseBodyAdvice

使用@RestControllerAdvice或@ControllerAdvice注入自定义ResponseBodyAdvice。

@RestControllerAdvice
public class ControllerAdvice implements ResponseBodyAdvice {

    @Override
    public boolean supports(MethodParameter returnType, Class converterType) {
        return true;
    }

    @Override
    public Object beforeBodyWrite(Object body, MethodParameter returnType,
                                  MediaType selectedContentType, Class selectedConverterType,
                                  ServerHttpRequest request, ServerHttpResponse response) {
        if (body instanceof Map) {
            ((Map) body).put("code", "0");
        }
        return body;
    }
}

说明

HandlerMethodReturnValueHandler都会先于ResponseBodyAdvice执行，所有在HandlerMethodReturnValueHandler的改动都会带到ResponseBodyAdvice。