如何修改response的header和body

如何修改response的header和body

shenyu是基于 webflux 的响应式编程项目，由于之前没有接触过，我在开发 modifyResponse插件的时候，开始有点摸不着头脑，然后大佬推荐了一个链接：
https://htmlpreview.github.io/?https://github.com/get-set/reactor-core/blob/master-zh/src/docs/index.html

看完我对响应性编程有了初步的了解。

修改header

1.获取HttpHeaders

ServerHttpResponse response = exchange.getResponse();
HttpHeaders httpHeaders = response.getHeaders();

2.HttpHeaders提供了 add、set、addAll、remove 等函数来操作数据。

if (Objects.nonNull(modifyResponseRuleHandle.getAddHeaders()) && MapUtils.isNotEmpty(modifyResponseRuleHandle.getAddHeaders())) {
 	Map<String, String> addHeaderMap = modifyResponseRuleHandle.getAddHeaders();
    addHeaderMap.entrySet().stream().forEach(a -> httpHeaders.add(a.getKey(), a.getValue()));
}

if (Objects.nonNull(modifyResponseRuleHandle.getSetHeaders()) && MapUtils.isNotEmpty(modifyResponseRuleHandle.getSetHeaders())) {
    Map<String, String> setHeaderMap = modifyResponseRuleHandle.getSetHeaders();
    setHeaderMap.entrySet().stream().forEach(a -> httpHeaders.set(a.getKey(), a.getValue()));
}

if (Objects.nonNull(modifyResponseRuleHandle.getReplaceHeaderKeys()) && MapUtils.isNotEmpty(modifyResponseRuleHandle.getReplaceHeaderKeys())) {
    Map<String, String> replaceHeaderMap = modifyResponseRuleHandle.getReplaceHeaderKeys();
    replaceHeaderMap.entrySet().stream().forEach(a -> {
        httpHeaders.addAll(a.getValue(), httpHeaders.get(a.getKey()));
        httpHeaders.remove(a.getKey());
    });
}

if (Objects.nonNull(modifyResponseRuleHandle.getRemoveHeaderKeys()) && !CollectionUtils.isEmpty(modifyResponseRuleHandle.getRemoveHeaderKeys())) {
    Set<String> removeHeaderList = modifyResponseRuleHandle.getRemoveHeaderKeys();
    removeHeaderList.stream().forEach(a -> httpHeaders.remove(a));
}

修改statusCode

1.获取ClientResponse

ClientResponse clientResponse = exchange.getAttribute(Constants.CLIENT_RESPONSE_ATTR);

2.设置code

response.setStatusCode(HttpStatus.valueOf(modifyResponseRuleHandle.getStatusCode()));

修改body

这一步是最难的，我也是请教了大佬才解决的问题。

关键点是自定义一个类 ModifyServerHttpResponse 继承 ServerHttpResponseDecorator ，然后重写 writeWith 函数。然后从 DataBuffer 里读出 body 的内容，修改完再写回去。

1.执行入口
doExecute（）最后调用 ModifyServerHttpResponse 的初始化

 return chain.execute(exchange.mutate()
                .response(new ModifyServerHttpResponse(exchange, modifyResponseRuleHandle)).build());

2.WebClientMessageWriter 在 writeWith 的时候，会走到子类的重写函数里

 return response.writeWith(clientResponse.body(BodyExtractors.toDataBuffers())).doOnCancel(() -> clean(exchange));

3.重写 body ，这一块有参考 SpringCloudGateway

public Mono<Void> writeWith(@NonNull final Publisher<? extends DataBuffer> body) {
   HttpHeaders httpHeaders = new HttpHeaders();
   httpHeaders.setContentType(MediaType.APPLICATION_JSON);

 	//重新生成response
    ClientResponse clientResponse = prepareClientResponse(body, httpHeaders);
    //修改body里的内容，并返回Mono字节流
    Mono<byte[]> modifiedBody = clientResponse.bodyToMono(byte[].class)
            .flatMap(originalBody -> Mono.just(updateResponse(originalBody)));

	//生成BodyInserter
    BodyInserter<Mono<byte[]>, ReactiveHttpOutputMessage> bodyInserter =
            BodyInserters.fromPublisher(modifiedBody, byte[].class);
    //保存body的信息
    CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(exchange,
            exchange.getResponse().getHeaders());
            
    //把body写回resposne        
    return bodyInserter.insert(outputMessage, new BodyInserterContext()).then(Mono.defer(() -> {
        Mono<DataBuffer> messageBody = DataBufferUtils.join(outputMessage.getBody());
        HttpHeaders headers = getDelegate().getHeaders();
        if (!headers.containsKey(HttpHeaders.TRANSFER_ENCODING)
                || headers.containsKey(HttpHeaders.CONTENT_LENGTH)) {
            //重写body的长度
            messageBody = messageBody.doOnNext(data -> headers.setContentLength(data.readableByteCount()));
        }
        return getDelegate().writeWith(messageBody);
    }));

}

因为body的内容可能是多层级的，例如

{
    "code":0,
    "msg":"sucess",
    "data":{
        "nameInfo":{
            "name":"test"
        },
        "body":{
            "age":"18"
        }
    }
}

所以修改 body 内容，用的是 JsonPath，需要在插件的 pom 文件里加上下面的依赖

<dependency>
            <groupId>com.jayway.jsonpath</groupId>
            <artifactId>json-path</artifactId>
            <version>2.4.0</version>
        </dependency>

具体代码，可以参考 shenyu 的代码仓库 https://github.com/apache/incubator-shenyu
里面的 shenyu-plugin-modify-reponse 模块