From ac9852122fd70c1f6bc52857f2cc0fc36bd10adb Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Tue, 16 Mar 2021 14:52:01 +0100
Subject: [PATCH] add websocket.timestamp_end

---
 examples/addons/websocket-inject-message.py | 2 +-
 mitmproxy/io/compat.py                      | 1 +
 mitmproxy/proxy/layers/websocket.py         | 2 ++
 mitmproxy/websocket.py                      | 4 ++++
 4 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/examples/addons/websocket-inject-message.py b/examples/addons/websocket-inject-message.py
index 012d4b27e..916e2c0a9 100644
--- a/examples/addons/websocket-inject-message.py
+++ b/examples/addons/websocket-inject-message.py
@@ -21,7 +21,7 @@ def websocket_message(flow):
 
 async def inject_async(flow: http.HTTPFlow):
     msg = "hello from mitmproxy! "
-    while flow.server_conn.connected:
+    while flow.websocket.timestamp_end is None:
         ctx.master.commands.call("inject", [flow], False, msg)
         await asyncio.sleep(1)
         msg = msg[1:] + msg[:1]
diff --git a/mitmproxy/io/compat.py b/mitmproxy/io/compat.py
index 16887dd1a..795dda8af 100644
--- a/mitmproxy/io/compat.py
+++ b/mitmproxy/io/compat.py
@@ -291,6 +291,7 @@ def convert_11_12(data):
             "closed_by_client": ws_flow["close_sender"] == "client",
             "close_code": ws_flow["close_code"],
             "close_reason": ws_flow["close_reason"],
+            "timestamp_end": data.get("server_conn", {}).get("timestamp_end", None),
         }
 
     else:
diff --git a/mitmproxy/proxy/layers/websocket.py b/mitmproxy/proxy/layers/websocket.py
index ca15ed9ff..cdf10ab1d 100644
--- a/mitmproxy/proxy/layers/websocket.py
+++ b/mitmproxy/proxy/layers/websocket.py
@@ -1,3 +1,4 @@
+import time
 from dataclasses import dataclass
 from typing import Iterator, List
 
@@ -189,6 +190,7 @@ class WebsocketLayer(layer.Layer):
                 )
                 yield dst_ws.send2(ws_event)
             elif isinstance(ws_event, wsproto.events.CloseConnection):
+                self.flow.websocket.timestamp_end = time.time()
                 self.flow.websocket.closed_by_client = from_client
                 self.flow.websocket.close_code = ws_event.code
                 self.flow.websocket.close_reason = ws_event.reason
diff --git a/mitmproxy/websocket.py b/mitmproxy/websocket.py
index 841944058..7fd0fcb02 100644
--- a/mitmproxy/websocket.py
+++ b/mitmproxy/websocket.py
@@ -106,11 +106,15 @@ class WebSocketData(stateobject.StateObject):
     close_reason: Optional[str] = None
     """[Close Reason](https://tools.ietf.org/html/rfc6455#section-7.1.6)"""
 
+    timestamp_end: Optional[float] = None
+    """*Timestamp:* WebSocket connection closed."""
+
     _stateobject_attributes = dict(
         messages=List[WebSocketMessage],
         closed_by_client=bool,
         close_code=int,
         close_reason=str,
+        timestamp_end=float,
     )
 
     def __init__(self):