JoyLau's Blog

/**
 * 根据 spring 源码优化
 * 优化了文件下载请求解析报文带来的内存使用
 * 添加了如下代码判断，当请求返回的头信息非 json 响应时， 将不再解析报文
 * public Mono<Void> writeWith(Publisher<? extends DataBuffer> body)
 *    if (!Utils.isJsonResponse(exchange)) {
 *        config.getRewriteFunction().apply(exchange, null);
 *        return getDelegate().writeWith(body);
 *    }
 *
 */
public class ModifyResponseBodyGatewayFilterFactoryCopy extends AbstractGatewayFilterFactory<ModifyResponseBodyGatewayFilterFactoryCopy.Config> {

    private final Map<String, MessageBodyDecoder> messageBodyDecoders;

    private final Map<String, MessageBodyEncoder> messageBodyEncoders;

    private final List<HttpMessageReader<?>> messageReaders;

    public ModifyResponseBodyGatewayFilterFactoryCopy(List<HttpMessageReader<?>> messageReaders,
                                                      Set<MessageBodyDecoder> messageBodyDecoders, Set<MessageBodyEncoder> messageBodyEncoders) {
        super(Config.class);
        this.messageReaders = messageReaders;
        this.messageBodyDecoders = messageBodyDecoders.stream()
                .collect(Collectors.toMap(MessageBodyDecoder::encodingType, identity()));
        this.messageBodyEncoders = messageBodyEncoders.stream()
                .collect(Collectors.toMap(MessageBodyEncoder::encodingType, identity()));
    }

    @Override
    public GatewayFilter apply(Config config) {
        ModifyResponseGatewayFilter gatewayFilter = new ModifyResponseGatewayFilter(config);
        gatewayFilter.setFactory(this);
        return gatewayFilter;
    }

    public static class Config {

        private Class inClass;

        private Class outClass;

        private Map<String, Object> inHints;

        private Map<String, Object> outHints;

        private String newContentType;

        private RewriteFunction rewriteFunction;

        public Class getInClass() {
            return inClass;
        }

        public Config setInClass(Class inClass) {
            this.inClass = inClass;
            return this;
        }

        public Class getOutClass() {
            return outClass;
        }

        public Config setOutClass(Class outClass) {
            this.outClass = outClass;
            return this;
        }

        public Map<String, Object> getInHints() {
            return inHints;
        }

        public Config setInHints(Map<String, Object> inHints) {
            this.inHints = inHints;
            return this;
        }

        public Map<String, Object> getOutHints() {
            return outHints;
        }

        public Config setOutHints(Map<String, Object> outHints) {
            this.outHints = outHints;
            return this;
        }

        public String getNewContentType() {
            return newContentType;
        }

        public Config setNewContentType(String newContentType) {
            this.newContentType = newContentType;
            return this;
        }

        public RewriteFunction getRewriteFunction() {
            return rewriteFunction;
        }

        public Config setRewriteFunction(RewriteFunction rewriteFunction) {
            this.rewriteFunction = rewriteFunction;
            return this;
        }

        public <T, R> Config setRewriteFunction(Class<T> inClass, Class<R> outClass,
                                                RewriteFunction<T, R> rewriteFunction) {
            setInClass(inClass);
            setOutClass(outClass);
            setRewriteFunction(rewriteFunction);
            return this;
        }

    }

    public class ModifyResponseGatewayFilter implements GatewayFilter, Ordered {

        private final Config config;

        private GatewayFilterFactory<Config> gatewayFilterFactory;

        public ModifyResponseGatewayFilter(Config config) {
            this.config = config;
        }

        @Override
        public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
            return chain.filter(exchange.mutate().response(new ModifiedServerHttpResponse(exchange, config)).build());
        }

        @Override
        public int getOrder() {
            return NettyWriteResponseFilter.WRITE_RESPONSE_FILTER_ORDER - 1;
        }

        @Override
        public String toString() {
            Object obj = (this.gatewayFilterFactory != null) ? this.gatewayFilterFactory : this;
            return filterToStringCreator(obj).append("New content type", config.getNewContentType())
                    .append("In class", config.getInClass()).append("Out class", config.getOutClass()).toString();
        }

        public void setFactory(GatewayFilterFactory<Config> gatewayFilterFactory) {
            this.gatewayFilterFactory = gatewayFilterFactory;
        }

    }

    protected class ModifiedServerHttpResponse extends ServerHttpResponseDecorator {

        private final ServerWebExchange exchange;

        private final Config config;

        public ModifiedServerHttpResponse(ServerWebExchange exchange, Config config) {
            super(exchange.getResponse());
            this.exchange = exchange;
            this.config = config;
        }

        @SuppressWarnings("unchecked")
        @Override
        public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
            if (!Utils.isJsonResponse(exchange)) {
                config.getRewriteFunction().apply(exchange, null);
                return getDelegate().writeWith(body);
            }

            Class inClass = config.getInClass();
            Class outClass = config.getOutClass();

            String originalResponseContentType = exchange.getAttribute(ORIGINAL_RESPONSE_CONTENT_TYPE_ATTR);

            HttpHeaders httpHeaders = new HttpHeaders();
            // explicitly add it in this way instead of
            // 'httpHeaders.setContentType(originalResponseContentType)'
            // this will prevent exception in case of using non-standard media
            // types like "Content-Type: image"
            httpHeaders.add(HttpHeaders.CONTENT_TYPE, originalResponseContentType);

            ClientResponse clientResponse = prepareClientResponse(body, httpHeaders);

            // TODO: flux or mono
            Mono modifiedBody = extractBody(exchange, clientResponse, inClass)
                    .flatMap(originalBody -> config.getRewriteFunction().apply(exchange, originalBody))
                    .switchIfEmpty(Mono.defer(() -> (Mono) config.getRewriteFunction().apply(exchange, null)));

            BodyInserter bodyInserter = BodyInserters.fromPublisher(modifiedBody, outClass);
            CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(exchange,
                    exchange.getResponse().getHeaders());
            return bodyInserter.insert(outputMessage, new BodyInserterContext()).then(Mono.defer(() -> {
                Mono<DataBuffer> messageBody = writeBody(getDelegate(), outputMessage, outClass);
                HttpHeaders headers = getDelegate().getHeaders();
                if (!headers.containsKey(HttpHeaders.TRANSFER_ENCODING)
                        || headers.containsKey(HttpHeaders.CONTENT_LENGTH)) {
                    messageBody = messageBody.doOnNext(data -> headers.setContentLength(data.readableByteCount()));
                }
                // TODO: fail if isStreamingMediaType?
                return getDelegate().writeWith(messageBody);
            }));
        }

        @Override
        public Mono<Void> writeAndFlushWith(Publisher<? extends Publisher<? extends DataBuffer>> body) {
            return writeWith(Flux.from(body).flatMapSequential(p -> p));
        }

        private ClientResponse prepareClientResponse(Publisher<? extends DataBuffer> body, HttpHeaders httpHeaders) {
            ClientResponse.Builder builder;
            builder = ClientResponse.create(exchange.getResponse().getStatusCode(), messageReaders);
            return builder.headers(headers -> headers.putAll(httpHeaders)).body(Flux.from(body)).build();
        }

        private <T> Mono<T> extractBody(ServerWebExchange exchange, ClientResponse clientResponse, Class<T> inClass) {
            // if inClass is byte[] then just return body, otherwise check if
            // decoding required
            if (byte[].class.isAssignableFrom(inClass)) {
                return clientResponse.bodyToMono(inClass);
            }

            List<String> encodingHeaders = exchange.getResponse().getHeaders().getOrEmpty(HttpHeaders.CONTENT_ENCODING);
            for (String encoding : encodingHeaders) {
                MessageBodyDecoder decoder = messageBodyDecoders.get(encoding);
                if (decoder != null) {
                    return clientResponse.bodyToMono(byte[].class).publishOn(Schedulers.parallel()).map(decoder::decode)
                            .map(bytes -> exchange.getResponse().bufferFactory().wrap(bytes))
                            .map(buffer -> prepareClientResponse(Mono.just(buffer),
                                    exchange.getResponse().getHeaders()))
                            .flatMap(response -> response.bodyToMono(inClass));
                }
            }

            return clientResponse.bodyToMono(inClass);
        }

        private Mono<DataBuffer> writeBody(ServerHttpResponse httpResponse, CachedBodyOutputMessage message,
                                           Class<?> outClass) {
            Mono<DataBuffer> response = DataBufferUtils.join(message.getBody());
            if (byte[].class.isAssignableFrom(outClass)) {
                return response;
            }

            List<String> encodingHeaders = httpResponse.getHeaders().getOrEmpty(HttpHeaders.CONTENT_ENCODING);
            for (String encoding : encodingHeaders) {
                MessageBodyEncoder encoder = messageBodyEncoders.get(encoding);
                if (encoder != null) {
                    DataBufferFactory dataBufferFactory = httpResponse.bufferFactory();
                    response = response.publishOn(Schedulers.parallel()).map(buffer -> {
                        byte[] encodedResponse = encoder.encode(buffer);
                        DataBufferUtils.release(buffer);
                        return encodedResponse;
                    }).map(dataBufferFactory::wrap);
                    break;
                }
            }

            return response;
        }

    }

}

@Configuration
public class ModifyResponseBodyGatewayFilterAutoConfiguration {

    @Bean
    @ConditionalOnEnabledFilter
    public ModifyResponseBodyGatewayFilterFactoryCopy modifyResponseBodyGatewayFilterFactoryCopy(
            ServerCodecConfigurer codecConfigurer, Set<MessageBodyDecoder> bodyDecoders,
            Set<MessageBodyEncoder> bodyEncoders) {
        return new ModifyResponseBodyGatewayFilterFactoryCopy(codecConfigurer.getReaders(), bodyDecoders, bodyEncoders);
    }
}

@Slf4j
@Component
@AllArgsConstructor
public class ResponseLogFilter implements GlobalFilter, Ordered {
    private final ModifyResponseBodyGatewayFilterFactoryCopy modifyResponseBodyGatewayFilterFactoryCopy;

    private final AsyncResponseHandler asyncResponseHandler;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ModifyResponseBodyGatewayFilterFactoryCopy.Config config = new ModifyResponseBodyGatewayFilterFactoryCopy.Config()
                .setRewriteFunction(byte[].class, byte[].class, (e, bytes) -> {
                    asyncResponseHandler.handle(e, bytes);
                    return Mono.justOrEmpty(bytes);
                });
        return modifyResponseBodyGatewayFilterFactoryCopy.apply(config).filter(exchange, chain);
    }

    @Override
    public int getOrder() {
        return NettyWriteResponseFilter.WRITE_RESPONSE_FILTER_ORDER - 98;
    }
}

@Slf4j
@Component
@AllArgsConstructor
public class AsyncResponseHandler {

    private final Jackson2HashMapper jackson2HashMapper;

    private final RedisTemplate<String, AbstractLog> logRedisTemplate;

    private final ObjectMapper objectMapper;

    @Async("logTaskPool")
    @SuppressWarnings("unchecked")
    public void handle(ServerWebExchange exchange, byte[] bytes) {
        ServerHttpResponse response = exchange.getResponse();
        try {
            Future<AccessLog> future =
                    (Future<AccessLog>) exchange.getAttributes().get(Constant.ACCESS_LOG_REQUEST_FUTURE_ATTR);
            AccessLog accessLog = future.get(10, TimeUnit.SECONDS);
            accessLog.setResponseHeaders(response.getHeaders().toSingleValueMap().toString());
            try {
                String aud = JwtUtils.verifyTokenSubject(accessLog.getToken());
                accessLog.setOriginType(OriginType.getByValue(aud));
            } catch (Exception e) {
                accessLog.setOriginType(OriginType.OTHER);
            }
            accessLog.setTakenTime(System.currentTimeMillis() - accessLog.getTakenTime());
            accessLog.setHttpCode(response.getRawStatusCode());
            if (!Utils.isJsonResponse(exchange)) {
                accessLog.setResponse("非 json 报文");
            }
            if (Utils.isDownloadResponse(exchange)) {
                accessLog.setResponse("二进制文件");
            }
            Optional.ofNullable(bytes).ifPresent(bs -> {
                if (bytes.length <= DataSize.ofKilobytes(256).toBytes()) {
                    // 小于指定大小报文进行转化(考虑到文件下载的响应报文)
                    accessLog.setResponse(new String(bytes, StandardCharsets.UTF_8));
                } else {
                    accessLog.setResponse("报文过长");
                }
            });
            logRedisTemplate.opsForStream().add(LogConstant.ACCESS_LOG_KEY_NAME, jackson2HashMapper.toHash(accessLog));
            // 进行修剪，限制其最大长度, 防止内存过高
            logRedisTemplate.opsForStream().trim(LogConstant.ACCESS_LOG_KEY_NAME, LogConstant.ACCESS_LOG_MAX_LENGTH,
                    true);
            if (log.isDebugEnabled()) {
                String logger = objectMapper.writerWithDefaultPrettyPrinter().writeValueAsString(accessLog);
                log.debug("log: \n{}", logger);
            }
        } catch (Exception e) {
            log.warn("access log save error: ", e);
        }
    }
}

@Slf4j
@Component
@AllArgsConstructor
public class RequestLogFilter implements GlobalFilter, Ordered {

   private final ModifyRequestBodyGatewayFilterFactory modifyRequestBodyGatewayFilterFactory;

   private final AsyncRequestHandler asyncRequestHandle;

   @Override
   public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
      // 不解析 body
      if (Utils.isUploadRequest(exchange)) {
         process(exchange, null);
         return chain.filter(exchange);
      }
      ModifyRequestBodyGatewayFilterFactory.Config modifyRequestConfig =
              new ModifyRequestBodyGatewayFilterFactory.Config()
                      .setRewriteFunction(byte[].class, byte[].class, (e, bytes) -> {
                         process(e, bytes);
                         return Mono.justOrEmpty(bytes);
                      });
      return modifyRequestBodyGatewayFilterFactory.apply(modifyRequestConfig).filter(exchange, chain);
   }

   @Override
   public int getOrder() {
      return -100;
   }

   private void process(ServerWebExchange exchange, byte[] bytes) {
      // 设置当前请求时间
      exchange.getAttributes().put(Constant.REQUEST_START_TIME_ATTR, System.currentTimeMillis());
      exchange.getAttributes().put(Constant.ACCESS_LOG_REQUEST_FUTURE_ATTR,
              asyncRequestHandle.handle(exchange, bytes));
   }
}

  /**
  * 判断是否文件上传请求
  *
  * @param exchange ServerWebExchange
  * @return boolean
  */
 public static boolean isUploadRequest(ServerWebExchange exchange){
   return MediaType.MULTIPART_FORM_DATA.isCompatibleWith(requestContentType(exchange));
}

@Component
@Slf4j
@AllArgsConstructor
public class AsyncRequestHandler {
   @Async("logTaskPool")
   public Future<AccessLog> handle(ServerWebExchange exchange, byte[] bytes) {
      return new AsyncResult<>(wrapperAccessLog(exchange, bytes));
   }

   /**
    * 保证访问日志请求体
    *
    * @param exchange ServerWebExchange
    * @param bytes    data to send
    * @return log
    */
   private AccessLog wrapperAccessLog(ServerWebExchange exchange, byte[] bytes) {
      ServerHttpRequest request = exchange.getRequest();
      AccessLog accessLog = new AccessLog();
      accessLog.setToken(Utils.getToken(exchange));
      accessLog.setTime(LocalDateTime.now());
      accessLog.setApplication(getApplicationName(exchange));
      accessLog.setIp(Utils.getIp(exchange));
      accessLog.setUri(request.getURI().toString());
      accessLog.setHttpMethod(HttpMethod.valueOf(request.getMethodValue()));
      // 临时设置当前时间，后面替换成真正耗时
      accessLog.setTakenTime((long) exchange.getAttributes().get(Constant.REQUEST_START_TIME_ATTR));
      accessLog.setRequestHeaders(request.getHeaders().toSingleValueMap().toString());
      if(Utils.isUploadRequest(exchange)) {
         accessLog.setRequest("二进制文件");
      }
      Optional.ofNullable(bytes).ifPresent(bs -> {
         if (bytes.length <= DataSize.ofKilobytes(256).toBytes()) {
            // 小于指定大小的报文进行转化(考虑到文件上传的请求报文)
            accessLog.setRequest(new String(bytes, StandardCharsets.UTF_8));
         } else {
            accessLog.setRequest("报文过长");
         }
      });
      return accessLog;
   }

   /**
    * 获取服务名称
    *
    * @param exchange ServerWebExchange
    * @return name of the application
    */
   private String getApplicationName(ServerWebExchange exchange) {
      String routingId =
              (String) exchange.getAttributes().get(ServerWebExchangeUtils.GATEWAY_PREDICATE_MATCHED_PATH_ROUTE_ID_ATTR);
      // 自动注册的微服务
      if (routingId.startsWith(Constant.MODULE_SUB_PREFIX)) {
         return routingId.substring(Constant.MODULE_SUB_PREFIX.length());
      } else {
         return routingId;
      }
   }

}

public class AccessLog extends AbstractLog {

   /**
    * 所属服务
    */
   private String application;

   /**
    * ip
    */
   private String ip;

   /**
    * uri
    */
   private String uri;

   /**
    * 请求方法
    */
   private HttpMethod httpMethod;

   /**
    * 消耗时间，单位毫秒
    */
   private Long takenTime;

   /**
    * http 状态码
    */
   private Integer httpCode;

   /**
    * 请求报文
    */
   private String request;

   /**
    * 响应报文
    */
   private String response;

   /**
    * 请求头信息
    */
   private String requestHeaders;

   /**
    * 响应头信息
    */
   private String responseHeaders;


   public String getApplication() {
      return application;
   }

   public void setApplication(String application) {
      this.application = application;
   }

   public String getIp() {
      return ip;
   }

   public void setIp(String ip) {
      this.ip = ip;
   }

   public String getUri() {
      return uri;
   }

   public void setUri(String uri) {
      this.uri = uri;
   }

   public HttpMethod getHttpMethod() {
      return httpMethod;
   }

   public void setHttpMethod(HttpMethod httpMethod) {
      this.httpMethod = httpMethod;
   }

   public Long getTakenTime() {
      return takenTime;
   }

   public void setTakenTime(Long takenTime) {
      this.takenTime = takenTime;
   }

   public Integer getHttpCode() {
      return httpCode;
   }

   public void setHttpCode(Integer httpCode) {
      this.httpCode = httpCode;
   }

   public String getRequest() {
      return request;
   }

   public void setRequest(String request) {
      this.request = request;
   }

   public String getResponse() {
      return response;
   }

   public void setResponse(String response) {
      this.response = response;
   }

   public String getRequestHeaders() {
      return requestHeaders;
   }

   public void setRequestHeaders(String requestHeaders) {
      this.requestHeaders = requestHeaders;
   }

   public String getResponseHeaders() {
      return responseHeaders;
   }

   public void setResponseHeaders(String responseHeaders) {
      this.responseHeaders = responseHeaders;
   }
}

public abstract class AbstractLog {

   /**
    * 日志 id
    */
   private String id;

   /**
    * 访问 token
    */
   private String token;

   /**
    * 用户 id
    */
   private String userid;

   /**
    * 用户名（用户账号)
    */
   private String username;

   /**
    * 用户姓名
    */
   private String name;

   /**
    * 操作时间
    */
   private LocalDateTime time;

   /**
    * 操作来源
    */
   private OriginType originType;


   public String getId() {
      return id;
   }

   public void setId(String id) {
      this.id = id;
   }

   public String getToken() {
      return token;
   }

   public void setToken(String token) {
      this.token = token;
   }

   public String getUserid() {
      return userid;
   }

   public void setUserid(String userid) {
      this.userid = userid;
   }

   public String getUsername() {
      return username;
   }

   public void setUsername(String username) {
      this.username = username;
   }

   public String getName() {
      return name;
   }

   public void setName(String name) {
      this.name = name;
   }

   public LocalDateTime getTime() {
      return time;
   }

   public void setTime(LocalDateTime time) {
      this.time = time;
   }

   public OriginType getOriginType() {
      return originType;
   }

   public void setOriginType(OriginType originType) {
      this.originType = originType;
   }
}