Spring Gateway mutate response body on isError

Question

I have requirement to mutate the response body for 4xx and 5xx responses.

incoming response would look like

{
  "reason": "sample reason"
}

I need to change that to

{
  "userMessage": "sample reason"
}

Here is my attempt

public class ErrorResponseFilter implements GlobalFilter, Ordered {
    private static final String REASON = "reason";
    private static final String USER_MESSAGE = "userMessage";
    private final ModifyResponseBodyGatewayFilterFactory modifyResponseBodyGatewayFilterFactory;

    public ErrorResponseFilter(ModifyResponseBodyGatewayFilterFactory modifyResponseBodyGatewayFilterFactory) {
        this.modifyResponseBodyGatewayFilterFactory = modifyResponseBodyGatewayFilterFactory;
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        return modifyResponseBodyGatewayFilterFactory
                .apply(c -> c.setInClass(JsonNode.class)
                        .setOutClass(JsonNode.class)
                        .setRewriteFunction((RewriteFunction<JsonNode, JsonNode>) (exchange1, body) -> {
                            if (null == body) return Mono.empty();

                            var originalResponse = exchange1.getResponse();
                            log.info("COOOODEDDDE: " + originalResponse.getStatusCode());
                            HttpStatusCode statusCode = originalResponse.getStatusCode();

                            if (statusCode != null && !statusCode.isError()) {
                                if (body.isObject()) {
                                    ObjectNode node = (ObjectNode) body;
                                    node.set(USER_MESSAGE, node.get(REASON));
                                    node.remove(REASON);
                                }
                            }
                            return Mono.just(body);
                        })
                )
                .filter(exchange, chain);
    }

    @Override
    public int getOrder() {
        return Ordered.HIGHEST_PRECEDENCE;
    }
}

However, this doesn't work, HTTP code is 200 OK all the time

Update 1:

I had attempted to use a bean

 @Bean
    public RouteLocator routes(RouteLocatorBuilder builder) {
        return builder.routes()
                .route("rewrite_response_user_message", r -> r.alwaysTrue()
                        .filters(f -> f.modifyResponseBody(JsonNode.class, JsonNode.class, (exchange, u) -> {
                            var originalResponse = exchange.getResponse();
                            HttpStatusCode statusCode = originalResponse.getStatusCode();
                            log.info("HTTP CODE: " + originalResponse.getStatusCode());
                            if(u == null) return Mono.empty();
                            if (statusCode != null && statusCode.isError()) {
                                if (u.isObject()) {
                                    ObjectNode node = (ObjectNode) u;
                                    node.set("userMessage", node.get("response"));
                                    node.remove("response");
                                }
                            }
                            return Mono.just(u);
                        }))
                        .uri(uri))
                .build();
    }

Without the bean a route returns 204 NO CONTENT but when the bean is enabled the same prints 404 in above code.

Update 2:

Looks like, a filter that applies to all routes cannot be added via Java DSL.

Update 3: I got it to work for some scenarios with following code and using it as a default-filter

public class ErrorReasonToUserMessageGatewayFilterFactory
        extends AbstractGatewayFilterFactory<ErrorReasonToUserMessageGatewayFilterFactory.Config> implements Ordered {
    private static final String REASON = "reason";
    private static final String USER_MESSAGE = "userMessage";
    private final ModifyResponseBodyGatewayFilterFactory modifyResponseBodyFilterFactory;

    public ErrorReasonToUserMessageGatewayFilterFactory(
            ModifyResponseBodyGatewayFilterFactory modifyResponseBodyFilterFactory) {
        super(Config.class);
        this.modifyResponseBodyFilterFactory = modifyResponseBodyFilterFactory;
    }

    @Override
    public GatewayFilter apply(Config config) {
        return modifyResponseBodyFilterFactory.apply(
                c -> c.setRewriteFunction(JsonNode.class, JsonNode.class, (swe, body) -> Optional.ofNullable(body)
                        .map(rewriteReasonToUserMessageFunction(swe, body))
                        .orElseGet(Mono::empty)));
    }

    private Function<JsonNode, Mono<JsonNode>> rewriteReasonToUserMessageFunction(
            ServerWebExchange swe, JsonNode body) {
        return s -> {
            HttpStatusCode statusCode = swe.getResponse().getStatusCode();
            if (statusCode != null && statusCode.isError() && body instanceof ObjectNode obj && body.has(REASON)) {
                obj.set(USER_MESSAGE, obj.get(REASON));
                obj.remove(REASON);
                return Mono.just(obj);
            }
            return Mono.just(body);
        };
    }

    @Override
    public int getOrder() {
        return Ordered.HIGHEST_PRECEDENCE;
    }

    public static class Config {}
}

Above work for scenarios where either the Content-Type: application/json or missing. it fails for application/octet-stream as the payload cannot be converted to JsonNode.

Question is how can I apply the above filter conditionally to 4xx/5xx and/or payload cannot be converted to JsonNode?

Answer 1

I got it to work thanks to @spencergibb comment and an example found here Here is the my updated working code that satisfies my needs

@Component
public class ErrorResponseGlobalPostFilter implements GlobalFilter, Ordered {
    private static final String REASON = "reason";
    private static final String USER_MESSAGE = "userMessage";
    private final ModifyResponseBodyGatewayFilterFactory modifyResponseBodyGatewayFilterFactory;
    private final ModifyResponseBodyGatewayFilterFactory.Config config;

    public ErrorResponseGlobalPostFilter(
            ModifyResponseBodyGatewayFilterFactory modifyResponseBodyGatewayFilterFactory) {
        this.modifyResponseBodyGatewayFilterFactory = modifyResponseBodyGatewayFilterFactory;
        config = new ModifyResponseBodyGatewayFilterFactory.Config();
        config.setInClass(JsonNode.class);
        config.setOutClass(JsonNode.class);
        config.setRewriteFunction(JsonNode.class, JsonNode.class, (swe, body) -> Optional.ofNullable(body)
                .map(rewriteReasonToUserMessageFunction(swe, body))
                .orElseGet(Mono::empty));
    }

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

    private Function<JsonNode, Mono<JsonNode>> rewriteReasonToUserMessageFunction(
            ServerWebExchange swe, JsonNode body) {
        return s -> {
            HttpStatusCode statusCode = swe.getResponse().getStatusCode();
            if (statusCode != null && statusCode.isError() && body instanceof ObjectNode obj && body.has(REASON)) {
                obj.set(USER_MESSAGE, obj.get(REASON));
                obj.remove(REASON);
                return Mono.just(obj);
            }
            return Mono.just(body);
        };
    }

    protected class ErrorResponseBuilder extends ServerHttpResponseDecorator {

        private final ServerWebExchange exchange;

        public ErrorResponseBuilder(ServerWebExchange exchange) {
            super(exchange.getResponse());
            this.exchange = exchange;
        }

        @Override
        @NonNull
        public Mono<Void> writeWith(@NonNull Publisher<? extends DataBuffer> publishedBody) {
            HttpStatusCode statusCode = exchange.getResponse().getStatusCode();
            if (statusCode != null
                    && statusCode.isError()
                    && !Objects.equals(
                            exchange.getResponse().getHeaders().getContentType(), MediaType.APPLICATION_OCTET_STREAM)) {
                ModifyResponseBodyGatewayFilterFactory.ModifiedServerHttpResponse modifiedServerHttpResponse =
                        modifyResponseBodyGatewayFilterFactory.new ModifiedServerHttpResponse(exchange, config);

                return modifiedServerHttpResponse.writeWith(publishedBody);
            }
            return super.writeWith(publishedBody);
        }
    }

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