Limit output payload response in CXF JAX-RS based service

Question

I have multiple jax-rs services built using cxf/spring. I want to control the output payload response size of all services. For simplicity sake, let's say none of api's in any of the services should ever return a JSON response payload more than 500 characters and I want to control this in one place instead of relying on individual services to adhere to this requirement. (We already have other features built into the custom framework/base component that all services depend on).

I have tried implementing this using JAX-RS's WriterInterceptor , ContainerResponseFilter and CXF's Phase Interceptor , but none of the approaches seem to be completely satisfy my requirement. More details on what I've done so far:

Option 1: (WriterInteceptor) In the overridden method, I get the ouputstream and set the max size of the cache to 500. When I invoke an api that returns more than 500 characters in the response payload, I get an HTTP 400 Bad Request status, but the response body contains the entire JSON payload.

@Provider
public class ResponsePayloadInterceptor implements WriterInterceptor {
    private static final Logger LOGGER = LoggerFactory.getLogger(ResponsePayloadInterceptor.class);

    @Override
    public void aroundWriteTo(WriterInterceptorContext context) throws IOException, WebApplicationException {
        final OutputStream outputStream = context.getOutputStream();

        CacheAndWriteOutputStream cacheAndWriteOutputStream = new CacheAndWriteOutputStream(outputStream);
        cacheAndWriteOutputStream.setMaxSize(500);
        context.setOutputStream(cacheAndWriteOutputStream);

        context.proceed();
    }
}

Option 2a: (CXF Phase Inteceptor) In the overridden method, I get the response as String from the ouputstream and check it's size. If it's greater than 500, I create a new Response object with only the data Too much data and set it in the message. Even if the response is > 500 characters, I get an HTTP 200 OK status with the entire JSON. Only when I use the phase as POST_MARSHAL or a later phase, I'm able to get hold of the JSON response and check it's length, but by that time the response has already been streamed to the client.

@Provider
public class ResponsePayloadInterceptor extends AbstractPhaseInterceptor<Message> {
    private static final Logger LOGGER = LoggerFactory.getLogger(ResponsePayloadInterceptor.class);

    public ResponsePayloadInterceptor() {
        super(Phase.POST_MARSHAL);
    }

    @Override
    public void handleMessage(Message message) throws Fault {
        LOGGER.info("handleMessage() - Response intercepted");
        try {
            OutputStream outputStream = message.getContent(OutputStream.class);
...
            CachedOutputStream cachedOutputStream = (CachedOutputStream) outputStream;
            String responseBody = IOUtils.toString(cachedOutputStream.getInputStream(), "UTF-8");
...
            LOGGER.info("handleMessage() - Response: {}", responseBody);
            LOGGER.info("handleMessage() - Response Length: {}", responseBody.length());
            if (responseBody.length() > 500) {
                Response response = Response.status(Response.Status.BAD_REQUEST)
                                            .entity("Too much data").build();
                message.getExchange().put(Response.class, response);
            }
        } catch (IOException e) {
            LOGGER.error("handleMessage() - Error");
            e.printStackTrace();
        }
    }
}

Option 2b: (CXF Phase Inteceptor) Same as above, but only the contents of if block is changed. If response length is greater than 500, I create a new output stream with the string Too much data and set it in message. But if the response payload is > 500 characters, I still get an HTTP 200 OK status with an invalid JSON response (entire JSON + additional text) ie, the response looks like this: [{"data":"", ...}, {...}]Too much data (the text 'Too much data' is appended to the JSON)

        if (responseBody.length() > 500) {
            InputStream inputStream = new ByteArrayInputStream("Too much data".getBytes("UTF-8"));
            outputStream.flush();
            IOUtils.copy(inputStream, outputStream);

            OutputStream out = new CachedOutputStream();
            out.write("Too much data".getBytes("UTF-8"));
            message.setContent(OutputStream.class, out);
        }

Option 3: (ContainerResponseFilter) Using the ContainerResponseFilter, I added a Content-Length response header with value as 500. If response length is > 500, I get an HTTP 200 OK status with an invalid JSON response (truncated to 500 characters). If the response length is < 500, still get an HTTP 200 OK status, but the client waits for more data to be returned by the server (as expected) and times out, which isn't a desirable solution.

@Provider
public class ResponsePayloadFilter implements ContainerResponseFilter {
    private static final Logger LOGGER = LoggerFactory.getLogger(ResponsePayloadFilter.class);

    @Override
    public void filter(ContainerRequestContext requestContext, ContainerResponseContext responseContext) throws IOException {
        LOGGER.info("filter() - Response intercepted");
        CachedOutputStream cos = (CachedOutputStream) responseContext.getEntityStream();
        StringBuilder responsePayload = new StringBuilder();
        ByteArrayOutputStream out = new ByteArrayOutputStream();

        if (cos.getInputStream().available() > 0) {
            IOUtils.copy(cos.getInputStream(), out);
            byte[] responseEntity = out.toByteArray();
            responsePayload.append(new String(responseEntity));
        }

        LOGGER.info("filter() - Content: {}", responsePayload.toString());
        responseContext.getHeaders().add("Content-Length", "500");
    }
}

Any suggestions on how I can tweak the above approaches to get what I want or any other different pointers?

Answer 1

I resolved this partially using help from this answer . I say partially because I'm successfully able to control the payload, but the not the response status code. Ideally, if the response length is greater than 500 and I modify the message content, I would like to send a different response status code (other than 200 OK). But this is a good enough solution for me to proceed at this point. If I figure out how to update the status code as well, I'll come back and update this answer.

import org.apache.commons.io.IOUtils;
import org.apache.cxf.interceptor.Fault;
import org.apache.cxf.io.CachedOutputStream;
import org.apache.cxf.message.Message;
import org.apache.cxf.phase.AbstractPhaseInterceptor;
import org.apache.cxf.phase.Phase;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;

public class ResponsePayloadInterceptor extends AbstractPhaseInterceptor<Message> {
    private static final Logger LOGGER = LoggerFactory.getLogger(ResponsePayloadInterceptor.class);

    public ResponsePayloadInterceptor() {
        super(Phase.PRE_STREAM);
    }

    @Override
    public void handleMessage(Message message) throws Fault {
        LOGGER.info("handleMessage() - Response intercepted");
        try {
            OutputStream outputStream = message.getContent(OutputStream.class);
            CachedOutputStream cachedOutputStream = new CachedOutputStream();
            message.setContent(OutputStream.class, cachedOutputStream);

            message.getInterceptorChain().doIntercept(message);

            cachedOutputStream.flush();
            cachedOutputStream.close();

            CachedOutputStream newCachedOutputStream = (CachedOutputStream) message.getContent(OutputStream.class);
            String currentResponse = IOUtils.toString(newCachedOutputStream.getInputStream(), "UTF-8");
            newCachedOutputStream.flush();
            newCachedOutputStream.close();

            if (currentResponse != null) {
                LOGGER.info("handleMessage() - Response: {}", currentResponse);
                LOGGER.info("handleMessage() - Response Length: {}", currentResponse.length());

                if (currentResponse.length() > 500) {
                    InputStream replaceInputStream = IOUtils.toInputStream("{\"message\":\"Too much data\"}", "UTF-8");

                    IOUtils.copy(replaceInputStream, outputStream);
                    replaceInputStream.close();

                    message.setContent(OutputStream.class, outputStream);
                    outputStream.flush();
                    outputStream.close();
                } else {
                    InputStream replaceInputStream = IOUtils.toInputStream(currentResponse, "UTF-8");

                    IOUtils.copy(replaceInputStream, outputStream);
                    replaceInputStream.close();

                    message.setContent(OutputStream.class, outputStream);
                    outputStream.flush();
                    outputStream.close();
                }
            }
        } catch (IOException e) {
            LOGGER.error("handleMessage() - Error", e);
            throw new RuntimeException(e);
        }
    }