From 105e1d5436f007acfadfd669c2b4550ae8faf32d Mon Sep 17 00:00:00 2001
From: Andrew Godwin <andrew@aeracode.org>
Date: Sun, 4 Feb 2018 12:08:57 -0800
Subject: [PATCH] Don't apply HTTP timeout to WebSocket connections!

---
 daphne/http_protocol.py |  1 +
 tests/http_base.py      |  5 ++++-
 tests/test_websocket.py | 29 +++++++++++++++++++++++++++++
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/daphne/http_protocol.py b/daphne/http_protocol.py
index 40694c0..37e43eb 100755
--- a/daphne/http_protocol.py
+++ b/daphne/http_protocol.py
@@ -118,6 +118,7 @@ class WebRequest(http.Request):
                 protocol.dataReceived(data)
                 # Remove our HTTP reply channel association
                 logger.debug("Upgraded connection %s to WebSocket", self.client_addr)
+                self.server.discard_protocol(self)
                 # Resume the producer so we keep getting data, if it's available as a method
                 self.channel._networkProducer.resumeProducing()
 
diff --git a/tests/http_base.py b/tests/http_base.py
index fc1f7fe..afb67f8 100644
--- a/tests/http_base.py
+++ b/tests/http_base.py
@@ -18,8 +18,9 @@ class DaphneTestingInstance:
     Works as a context manager.
     """
 
-    def __init__(self, xff=False):
+    def __init__(self, xff=False, http_timeout=60):
         self.xff = xff
+        self.http_timeout = http_timeout
         self.host = "127.0.0.1"
 
     def port_in_use(self, port):
@@ -59,6 +60,8 @@ class DaphneTestingInstance:
         # Optionally enable X-Forwarded-For support.
         if self.xff:
             daphne_args += ["--proxy-headers"]
+        if self.http_timeout:
+            daphne_args += ["--http-timeout=%i" % self.http_timeout]
         # Start up process and make sure it begins listening. Try this 3 times.
         for _ in range(3):
             self.process = subprocess.Popen(daphne_args + ["daphne.test_application:TestApplication"])
diff --git a/tests/test_websocket.py b/tests/test_websocket.py
index 409aac4..a55d889 100644
--- a/tests/test_websocket.py
+++ b/tests/test_websocket.py
@@ -1,6 +1,7 @@
 # coding: utf8
 
 import collections
+import time
 from urllib import parse
 
 from hypothesis import given, settings
@@ -242,3 +243,31 @@ class TestWebsocket(DaphneTestCase):
             # Make sure it got our frame
             _, messages = test_app.get_received()
             assert messages[1] == {"type": "websocket.receive", "bytes": b"what is here? \xe2"}
+
+    def test_http_timeout(self):
+        """
+        Tests that the HTTP timeout doesn't kick in for WebSockets
+        """
+        with DaphneTestingInstance(http_timeout=1) as test_app:
+            # Connect
+            test_app.add_send_messages([
+                {
+                    "type": "websocket.accept",
+                }
+            ])
+            sock, _ = self.websocket_handshake(test_app)
+            _, messages = test_app.get_received()
+            self.assert_valid_websocket_connect_message(messages[0])
+            # Wait 2 seconds
+            time.sleep(2)
+            # Prep frame for it to send
+            test_app.add_send_messages([
+                {
+                    "type": "websocket.send",
+                    "text": "cake",
+                }
+            ])
+            # Send it a frame
+            self.websocket_send_frame(sock, "still alive?")
+            # Receive a frame and make sure it's correct
+            assert self.websocket_receive_frame(sock) == "cake"