From a73b073d175efec27eb14a24cef00ef3e454b0fb Mon Sep 17 00:00:00 2001
From: Randall Leeds <randall@bleeds.info>
Date: Wed, 27 Dec 2023 17:47:06 -0800
Subject: [PATCH] Force connection close after protocol upgrade

A new protocol, such as WebSockets or HTTP/2, may manage the framing of
multiple messages or requests of that protocol over a single connection,
but the connection should close once the WSGI callable returns.

- Close connections after completion of any upgraded response.

- Close connections after a response with status code less than 200. Any
  such response should be only for a protocol upgrade since applications
  never see any `Expect` header.

- Ensure that the `status_code` member always exists on instances of the
  `Response` class (close #1210).
---
 gunicorn/http/wsgi.py | 22 ++++++++++++++--------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/gunicorn/http/wsgi.py b/gunicorn/http/wsgi.py
index 6f3d9b68f..961765ae6 100644
--- a/gunicorn/http/wsgi.py
+++ b/gunicorn/http/wsgi.py
@@ -205,6 +205,7 @@ def __init__(self, req, sock, cfg):
         self.sock = sock
         self.version = SERVER
         self.status = None
+        self.status_code = None
         self.chunked = False
         self.must_close = False
         self.headers = []
@@ -218,15 +219,20 @@ def force_close(self):
         self.must_close = True
 
     def should_close(self):
-        if self.must_close or self.req.should_close():
+        if self.must_close:
+            # example: worker shutting down
             return True
-        if self.response_length is not None or self.chunked:
-            return False
-        if self.req.method == 'HEAD':
-            return False
-        if self.status_code < 200 or self.status_code in (204, 304):
-            return False
-        return True
+        if self.req.should_close():
+            # example: connection close or upgrade header
+            return True
+        if self.upgrade:
+            # close after the new protocol terminates
+            return True
+        if self.response_length is None and not self.chunked:
+            # close if there is a response body of unknown length
+            # the client cannot otherwise know when the response ends
+            return self.req.method != 'HEAD' and self.status_code not in (204, 304)
+        return False
 
     def start_response(self, status, headers, exc_info=None):
         if exc_info: