improved scgi.forkpool.SocketFileWrapper a lot

1 files changed, 44 insertions, 11 deletions

diff --git a/wsgitools/scgi/forkpool.py b/wsgitools/scgi/forkpool.py
index 8171172..3106d05 100644
--- a/wsgitools/scgi/forkpool.py
+++ b/wsgitools/scgi/forkpool.py
@@ -5,26 +5,57 @@ import sys
 
 class SocketFileWrapper:
     """Wraps a socket to a wsgi-compliant file-like object."""
-    def __init__(self, sock):
+    def __init__(self, sock, toread):
         """@param sock: is a socket.socket()"""
         self.sock = sock
         self.buff = ""
+        self.toread = toread
+
+    def _recv(self, size=None):
+        """internal method for receiving and counting incoming data"""
+        if size is None:
+            data = self.sock.recv()
+            self.readbytes += len(data)
+            return data
+        data = self.sock.recv(size)
+        self.toread -= len(data)
+        return data
+
+    def close(self):
+        """Does not close the socket, because it might still be needed. It
+        reads all data that should have been read as given by CONTENT_LENGTH."""
+        try:
+            while self.toread > 0:
+                if self.toread > 4096:
+                    self._recv(4096)
+                else:
+                    self._recv(self.toread)
+        except socket.error:
+            pass
+
     def read(self, size=None):
         """see pep333"""
         if size is None:
             try:
-                return self.buff + self.sock.recv()
-            except: raise # Yes, this is necessary.
-            else:
-                self.buff = ""
+                data = self._recv()
+            except socket.error:
+                if self.buff:
+                    ret, self.buff = self.buff, ""
+                    return ret
+                raise
+            return self.buff + data
         if size <= len(self.buff):
             ret, self.buff = self.buff[:size], self.buff[size:]
             return ret
         try:
-            return self.buff + self.sock.recv(size)
-        except: raise # Yes, this is necessary.
-        else:
-            self.buff = ""
+            data = self._recv(size - len(self.buff))
+        except socket.error:
+            if self.buff:
+                ret, self.buff = self.buff, ""
+                return self.buff
+            raise
+        return self.buff + data
+
     def readline(self):
         """see pep333"""
         while True:
@@ -33,7 +64,7 @@ class SocketFileWrapper:
                 ret, self.buff = self.buff[:split], self.buff[split:]
                 return ret
             except ValueError:
-                data = self.sock.recv(4096)
+                data = self._recv(4096)
                 if not data:
                     ret, self.buff = self.buff, ""
                     return ret
@@ -219,9 +250,10 @@ class SCGIServer:
                       '\r\n'.join(map("%s: %s".__mod__, headers))))
             return con.send
 
+        sfw = SocketFileWrapper(con, long(environ["CONTENT_LENGTH"]))
         environ.update({
             "wsgi.version": (1, 0),
-            "wsgi.input": SocketFileWrapper(con),
+            "wsgi.input": sfw,
             "wsgi.errors": self.error,
             "wsgi.url_scheme": "http",
             "wsgi.multithread": False,
@@ -242,4 +274,5 @@ class SCGIServer:
             con.send(data)
         if hasattr(result, "close"):
             result.close()
+        sfw.close()
         con.close()