wsgi: websocket ALREADY_HANDLED flag on corolocal

Hardtack · web-flow · commit df0bc00c3b75 · 2021-05-02T16:10:52.000+03:00
#543
diff --git a/eventlet/websocket.py b/eventlet/websocket.py
@@ -135,7 +135,8 @@ def __call__(self, environ, start_response):
         ws._send_closing_frame(True)
         # use this undocumented feature of eventlet.wsgi to ensure that it
         # doesn't barf on the fact that we didn't call start_response
-        return wsgi.ALREADY_HANDLED
+        wsgi.WSGI_LOCAL.already_handled = True
+        return []
 
     def _handle_legacy_request(self, environ):
         if 'eventlet.input' in environ:
diff --git a/eventlet/wsgi.py b/eventlet/wsgi.py
@@ -9,6 +9,7 @@
 import eventlet
 from eventlet import greenio
 from eventlet import support
+from eventlet.corolocal import local
 from eventlet.green import BaseHTTPServer
 from eventlet.green import socket
 import six
@@ -69,19 +70,7 @@ class ChunkReadError(ValueError):
     pass
 
 
-# special flag return value for apps
-class _AlreadyHandled(object):
-
-    def __iter__(self):
-        return self
-
-    def next(self):
-        raise StopIteration
-
-    __next__ = next
-
-
-ALREADY_HANDLED = _AlreadyHandled()
+WSGI_LOCAL = local()
 
 
 class Input(object):
@@ -570,14 +559,12 @@ def cap(x):
 
         try:
             try:
+                WSGI_LOCAL.already_handled = False
                 result = self.application(self.environ, start_response)
-                if (isinstance(result, _AlreadyHandled)
-                        or isinstance(getattr(result, '_obj', None), _AlreadyHandled)):
-                    self.close_connection = 1
-                    return
 
                 # Set content-length if possible
-                if not headers_sent and hasattr(result, '__len__') and \
+                if headers_set and \
+                        not headers_sent and hasattr(result, '__len__') and \
                         'Content-Length' not in [h for h, _v in headers_set[1]]:
                     headers_set[1].append(('Content-Length', str(sum(map(len, result)))))
 
@@ -599,6 +586,9 @@ def cap(x):
                         towrite = []
                         just_written_size = towrite_size
                         towrite_size = 0
+                if WSGI_LOCAL.already_handled:
+                    self.close_connection = 1
+                    return
                 if towrite:
                     just_written_size = towrite_size
                     write(b''.join(towrite))
diff --git a/tests/websocket_test.py b/tests/websocket_test.py
@@ -528,6 +528,38 @@ def test_close_idle(self):
         with eventlet.Timeout(1):
             pool.waitall()
 
+    def test_wrapped_wsgi(self):
+        site = self.site
+
+        def wrapper(environ, start_response):
+            for chunk in site(environ, start_response):
+                yield chunk
+
+        self.site = wrapper
+        self.spawn_server()
+        connect = [
+            "GET /range HTTP/1.1",
+            "Upgrade: WebSocket",
+            "Connection: Upgrade",
+            "Host: {}:{}".format(*self.server_addr),
+            "Origin: http://{}:{}".format(*self.server_addr),
+            "WebSocket-Protocol: ws",
+        ]
+        sock = eventlet.connect(self.server_addr)
+
+        sock.sendall(six.b("\r\n".join(connect) + "\r\n\r\n"))
+        resp = sock.recv(1024)
+        headers, result = resp.split(b"\r\n\r\n")
+        msgs = [result.strip(b"\x00\xff")]
+        msgs.extend(sock.recv(20).strip(b"\x00\xff") for _ in range(10))
+        expect = [six.b("msg {}".format(i)) for i in range(10)] + [b""]
+        assert msgs == expect
+        # In case of server error, server will write HTTP 500 response to the socket
+        msg = sock.recv(20)
+        assert not msg
+        sock.close()
+        eventlet.sleep(0.01)
+
 
 class TestWebSocketSSL(tests.wsgi_test._TestBase):
     def set_site(self):
diff --git a/tests/wsgi_test.py b/tests/wsgi_test.py
@@ -103,7 +103,8 @@ def chunked_post(env, start_response):
 
 def already_handled(env, start_response):
     start_response('200 OK', [('Content-type', 'text/plain')])
-    return wsgi.ALREADY_HANDLED
+    wsgi.WSGI_LOCAL.already_handled = True
+    return []
 
 
 class Site(object):
@@ -115,8 +116,7 @@ def __call__(self, env, start_response):
 
 
 class IterableApp(object):
-
-    def __init__(self, send_start_response=False, return_val=wsgi.ALREADY_HANDLED):
+    def __init__(self, send_start_response=False, return_val=()):
         self.send_start_response = send_start_response
         self.return_val = return_val
         self.env = {}
@@ -125,6 +125,8 @@ def __call__(self, env, start_response):
         self.env = env
         if self.send_start_response:
             start_response('200 OK', [('Content-type', 'text/plain')])
+        else:
+            wsgi.WSGI_LOCAL.already_handled = True
         return self.return_val
 
 
