From 07e89e7a7b603537a96e8c3771d7963bbce812da Mon Sep 17 00:00:00 2001
From: James Roberts <jamesroberts.dev@gmail.com>
Date: Wed, 9 Feb 2022 23:54:43 +0200
Subject: [PATCH] Fixes response iterator bug

---
 fastwsgi.py        |  2 ++
 fastwsgi/request.c | 28 ++++++++++++----------------
 setup.py           |  2 +-
 3 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/fastwsgi.py b/fastwsgi.py
index 183ad8a..05e1b87 100644
--- a/fastwsgi.py
+++ b/fastwsgi.py
@@ -12,6 +12,7 @@ PORT = 5000
 BACKLOG = 1024
 LOGGING = 0
 
+
 def run_multi_process_server(app):
     workers = []
     for _ in range(NUM_WORKERS):
@@ -83,5 +84,6 @@ def run_from_cli(host, port, wsgi_app_import_string, logging):
 def run(wsgi_app, host=HOST, port=PORT, backlog=1024):
     print_server_details(host, port)
     print(f"Server listening at http://{host}:{port}")
+    print(f"Running on PID:", os.getpid())
     _fastwsgi.run_server(wsgi_app, host, port, backlog, LOGGING)
     # run_multi_process_server(wsgi_app)
diff --git a/fastwsgi/request.c b/fastwsgi/request.c
index 36f5772..c7487ac 100644
--- a/fastwsgi/request.c
+++ b/fastwsgi/request.c
@@ -168,26 +168,22 @@ void build_response(PyObject* wsgi_response, StartResponse* response, llhttp_t*
     logger("building response");
     Request* request = (Request*)parser->data;
 
-    PyObject* iter = NULL;
+    PyObject* iterator = NULL;
     PyObject* result = NULL;
+    PyObject* item = NULL;
 
     if (PyBytes_Check(wsgi_response))
         result = wsgi_response;
-    else if (PyIter_Check(wsgi_response))
-        iter = wsgi_response;
-    else
-        iter = PyObject_GetIter(wsgi_response);
+    else {
+        iterator = PyObject_GetIter(wsgi_response);
+        result = PyBytes_FromString("");
+        while (item = PyIter_Next(iterator)) {
+            PyBytes_ConcatAndDel(&result, item);
+        }
+    }
 
     int response_has_no_content = 0;
 
-    if (result == NULL) {
-        PyObject* next = PyIter_Next(iter);
-        if (next == NULL)
-            response_has_no_content = 1;
-        else
-            result = next;
-    }
-
     PyObject* status = PyUnicode_AsUTF8String(response->status);
     char* status_code = PyBytes_AS_STRING(status);
     if (strncmp(status_code, "204", 3) == 0 || strncmp(status_code, "304", 3) == 0) {
@@ -259,15 +255,15 @@ void build_response(PyObject* wsgi_response, StartResponse* response, llhttp_t*
     request->response_buffer.base = buf;
     request->response_buffer.len = strlen(buf);
 
-    if (iter != NULL && PyObject_HasAttrString(iter, "close")) {
-        PyObject* close = PyObject_GetAttrString(iter, "close");
+    if (iterator != NULL && PyObject_HasAttrString(iterator, "close")) {
+        PyObject* close = PyObject_GetAttrString(iterator, "close");
         if (close != NULL) {
             PyObject* close_result = PyObject_CallObject(close, NULL);
             Py_XDECREF(close_result);
         }
         Py_XDECREF(close);
     }
-    Py_XDECREF(iter);
+    Py_XDECREF(iterator);
     Py_XDECREF(result);
     result = NULL;
 }
diff --git a/setup.py b/setup.py
index 895b69d..86f31ca 100644
--- a/setup.py
+++ b/setup.py
@@ -10,7 +10,7 @@ module = Extension(
     "_fastwsgi",
     sources=SOURCES,
     include_dirs=["llhttp/include", "libuv/include"],
-    extra_compile_args=["-O3", "-fno-strict-aliasing", "-fcommon"],
+    extra_compile_args=["-O3", "-fno-strict-aliasing", "-fcommon", "-g", "-Wall"],
 )
 
 with open("README.md", "r", encoding="utf-8") as read_me: