How can I safely intercept the Response stream in a custom Owin Middleware

Question

I'm trying to write a simple OWIN Middleware, in order to intercept the response stream. What I'm trying to do is replace the original stream with custom Stream-based class, where I will be able to intercept writes to the response stream.

However, I'm facing some issues because I cannot know when the response has been completely written to by inner middleware components in the chain. The Dispose override of the Stream is never called. So I don't know when it's time to perform my processing, which should happen at the end of the response Stream.

Here is a sample code:

public sealed class CustomMiddleware: OwinMiddleware
{
    public CustomMiddleware(OwinMiddleware next)
        : base(next)
    {
    }

    public override async Task Invoke(IOwinContext context)
    {
        var request = context.Request;
        var response = context.Response;

        // capture response stream

        var vr = new MemoryStream();
        var responseStream = new ResponseStream(vr, response.Body);

        response.OnSendingHeaders(state =>
        {
            var resp = (state as IOwinContext).Response;
            var contentLength = resp.Headers.ContentLength;

            // contentLength == null for Chunked responses

        }, context);

        // invoke the next middleware in the pipeline

        await Next.Invoke(context);
    }
}

public sealed class ResponseStream : Stream
{
    private readonly Stream stream_; // MemoryStream
    private readonly Stream output_; // Owin response
    private long writtenBytes_ = 0L;

    public ResponseStream(Stream stream, Stream output)
    {
        stream_ = stream;
        output_ = output;
    }

    ... // System.IO.Stream implementation

    public override void Write(byte[] buffer, int offset, int count)
    {
        // capture writes to the response stream in our local stream
        stream_.Write(buffer, offset, count);

        // write to the real output stream
        output_.Write(buffer, offset, count);

        // update the number of bytes written

        writtenBytes_ += count;

        // how do we know the response is complete ?
        // we could check that the number of bytes written
        // is equal to the content length, but content length
        // is not available for Chunked responses.
    }

    protected override void Dispose(bool disposing)
    {
        // we could perform our processing
        // when the stream is disposed of.
        // however, this method is never called by
        // the OWIN/Katana infrastructure.
    }
}

As I've alluded to in the comments from the code above, there are two strategies that I can think of in order to detect whether the response is complete.

a) I can record the number of bytes written to the response stream and correlate that to the expected response length. However, in the case of responses which use the Chunked Transfer Encoding, the length is not known.

b) I can decide that the response stream is complete when Dispose is called on the response stream. However, the OWIN/Katana infrastructure never calls Dispose on the replaced stream.

I have been investigating Opaque Streaming in order to see whether manipulating the underlying HTTP protocol would be a feasible approach, but I don't seem to find whether Katana supports Opaque Streaming or not.

Is there a way to achieve what I want ?

Answer 1

I do not think you will need a sub-classed stream but then here is how you can read the response. Just ensure this middleware is the first one in the OWIN pipeline so that it will be the last one to inspect the response.

using AppFunc = Func<IDictionary<string, object>, Task>;

public class CustomMiddleware
{
    private readonly AppFunc next;

    public CustomMiddleware(AppFunc next)
    {
        this.next = next;
    }

    public async Task Invoke(IDictionary<string, object> env)
    {
        IOwinContext context = new OwinContext(env);

        // Buffer the response
        var stream = context.Response.Body;
        var buffer = new MemoryStream();
        context.Response.Body = buffer;

        await this.next(env);

        buffer.Seek(0, SeekOrigin.Begin);
        var reader = new StreamReader(buffer);
        string responseBody = await reader.ReadToEndAsync();

        // Now, you can access response body.
        Debug.WriteLine(responseBody);

        // You need to do this so that the response we buffered
        // is flushed out to the client application.
        buffer.Seek(0, SeekOrigin.Begin);
        await buffer.CopyToAsync(stream);
    }
}

BTW, as far as I know, deriving from OwinMiddleware is not considered a good practice because OwinMiddleware is specific to Katana. It is however nothing to do with your problem though.