Is there a way to defer the closing of a WebSocket connection?

[lawrence-forooghian]

I've written an addon which allows an external process to control mitmproxy’s interception of WebSocket messages. It works as follows:

1. Upon receiving a websocket_message hook, the addon calls drop() on the message, and invokes an RPC call on the external process to ask it what to do with the message. (In order to preserve message ordering, subsequent messages intercepted on that flow, in the direction described by from_client, will be queued by the addon pending the external process’s reply.)
2. The external process replies, telling the addon to either drop the message, or to replace it with a message with some given data.
3. If told to replace the message, the addon injects this replacement into the flow, per this example.

This works well, except in the case where the server sends a final WebSocket message and then immediately closes the connection. In this case, mitmproxy will close its WebSocket connection with the client before receiving the external process’s reply about what to do with the final WebSocket message. This means that the client does not receive this final message.

The behaviour that I would like is the following:

1. My addon finds out that the server closed the WebSocket connection (I guess this would be either just a TCP close, or a CLOSE frame followed by a TCP close)
2. At this point, my addon is able to instruct mitmproxy to keep the WebSocket connection with the client open for the time being.
3. The addon is now able to, at its leisure, inject further messages towards the client.
4. Finally, at its leisure, the addon tells mitmproxy to propagate the server-initiated close to the client (i.e. forward any CLOSE frame and close the TCP connection).

It seems to me that this is not currently possible through the addons API. Is that correct? If so, a couple of questions:

• is it something that you'd consider adding support for, or would accept a PR for (obviously, pending some conversations about what the API would look like)?
• is it something that I could achieve by dropping down to a lower level flow (e.g. TCPFlow) — would that allow me to defer the closing of mitmproxy’s TCP connection with the client?

Thank you!

[mhils]

It seems to me that this is not currently possible through the addons API. Is that correct?

Yes.

is it something that you'd consider adding support for, or would accept a PR for (obviously, pending some conversations about what the API would look like)?

I would be really curious to hear API suggestions and then evaluate how much time we want to spend on it now. :) A clean API would be nice here, but it's not super easy to design.

The immediate thought that comes to my mind that it would be probably be most useful to extend our ASGI support to WebSockets (see wsgi-flask-app.py). If we could get this to support WebSockets, users would have a really powerful API to handle WebSocket connections themselves (with optional upstream communication).

is it something that I could achieve by dropping down to a lower level flow (e.g. TCPFlow) — would that allow me to defer the closing of mitmproxy’s TCP connection with the client?

I don't think so, no. The injection mechanism is really quite limited at the moment.

[lawrence-forooghian]

Thank you for your quick reply, @mhils!

Whilst I am keen for this functionality to exist in mitmproxy, and would love to help implement it, it’s unfortunately unlikely that it’s something I’ll be able to get approval to put much time into unless there’s no obvious alternative. I think that if I can find a workaround that addresses my immediate use case, then I might end up doing that. I have an alternative in mind which I’m not quite sure how to implement; I will open a separate question about it.