Java Web Flux : How to fetch body from ServerHttpResponse?

Question

I have implemented custom ServerHttpRequestDecorator to fetch request body But unable to do the same for response.

@Override
    public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain)
    {
        long start = System.currentTimeMillis();
        LogServerWebExchangeDecorator exchangeDecorator = new LogServerWebExchangeDecorator(
                exchange);
        exchangeDecorator.getResponse().beforeCommit(() ->
        {
            logRequest(exchangeDecorator.getRequest(), exchangeDecorator.getRequestBody());
            logResponse(exchangeDecorator.getRequest(), exchangeDecorator.getResponse(), start);
            return Mono.empty();
        });
        return chain.filter(exchangeDecorator);
    }

How can achieve the same?

Answer 1

private class LogServerHttpResponseDecorator extends ServerHttpResponseDecorator
    {
        private final FastByteArrayOutputStream fastByteArrayOutputStream;
        public LogServerHttpResponseDecorator(ServerHttpResponse response, LoggingContext loggingContext)
        {
            super(response);
            fastByteArrayOutputStream = new FastByteArrayOutputStream();
            response.beforeCommit(LoggingContextUtils.withContext(loggingContext, (Supplier<Mono<Void>>) () ->
                {
                    logResponse(byteArrayToString(fastByteArrayOutputStream.toByteArray(),
                            getHeaders()));
                    return Mono.empty();
                }));
        }

        @Override
        public Mono<Void> writeWith(Publisher<? extends DataBuffer> body)
        {
            return super.writeWith(Flux.from(body).map(buffer ->
            {
                readBuffer(fastByteArrayOutputStream, buffer);
                return buffer;
            }));
        }

        @Override
        public Mono<Void> writeAndFlushWith(
                Publisher<? extends Publisher<? extends DataBuffer>> body)
        {
            return super.writeAndFlushWith(Flux.from(body).map(publisher ->
                    Flux.from(publisher).map(buffer ->
                    {
                        readBuffer(fastByteArrayOutputStream, buffer);
                        return buffer;
                    })));
        }
    }
private void readBuffer(FastByteArrayOutputStream fastByteArrayOutputStream, DataBuffer buffer)
    {
        try
        {
            Channels.newChannel(fastByteArrayOutputStream)
                    .write(buffer.asByteBuffer().asReadOnlyBuffer());
        }
        catch (IOException e)
        {
            if (LoggingUtils.isLoggable(LOGGER, Level.ERROR))
            {
                LoggingUtils.log(LOGGER, Level.ERROR, null, "Unable to log", e);
            }
        }
    }