add functional tests for write callback and behavior when no content length header is supplied

5 files changed, 205 insertions, 21 deletions

diff --git a/waitress/task.py b/waitress/task.py
index 0221458..4e58a76 100644
--- a/waitress/task.py
+++ b/waitress/task.py
@@ -247,32 +247,26 @@ class HTTPTask(object):
             for chunk in app_iter:
                 if first_chunk_len is None:
                     first_chunk_len = len(chunk)
+                    # Set a Content-Length header if one is not supplied.
+                    # start_response may not have been called until first
+                    # iteration as per PEP, so we must reinterrogate
+                    # self.content_length here
+                    cl = self.content_length
+                    has_content_length = cl != -1
+                    if not has_content_length and app_iter_len == 1:
+                        cl = self.content_length = first_chunk_len
                 # transmit headers only after first iteration of the iterable
                 # that returns a non-empty bytestring (PEP 3333)
                 if not chunk:
                     continue
-                # Set a Content-Length header if one is not supplied.
-                # start_response may not have been called until first iteration
-                # as per PEP
-                cl = self.content_length
-                has_content_length = cl != -1
-                if not has_content_length and app_iter_len == 1:
-                    cl = self.content_length = first_chunk_len
-                towrite = chunk
-                if has_content_length:
-                    towrite = chunk[:cl-self.content_bytes_written]
-                self.write(towrite)
-                if towrite != chunk:
-                    self.channel.server.log_info(
-                        'app_iter content exceeded the number of bytes '
-                        'specified by Content-Length header (%s)' % cl)
-                    break
+                self.write(chunk)
 
             cl = self.content_length
             if cl != -1:
                 if self.content_bytes_written != cl:
                     # close the connection so the client isn't sitting around
-                    # waiting for more data
+                    # waiting for more data when there are too few bytes
+                    # to service content-length
                     self.close_on_finish = True
                     self.channel.server.log_info(
                         'app_iter returned too few bytes (%s) '
@@ -294,8 +288,18 @@ class HTTPTask(object):
             channel.write(rh)
             self.wrote_header = True
         if data:
-            self.content_bytes_written += len(data)
-            channel.write(data)
+            towrite = data
+            cl = self.content_length
+            if cl != -1:
+                towrite = data[:cl-self.content_bytes_written]
+            if towrite != data:
+                # XXX warn instead of relog
+                self.channel.server.log_info(
+                    'written content exceeded the number of bytes '
+                    'specified by Content-Length header (%s)' % cl)
+            if towrite:
+                self.content_bytes_written += len(towrite)
+                channel.write(towrite)
 
     def build_response_header(self):
         version = self.version
@@ -323,7 +327,7 @@ class HTTPTask(object):
         if content_length_header is None and self.content_length != -1:
             content_length_header = str(self.content_length)
             self.response_headers.append(
-                ('Content-Length',content_length_header)
+                ('Content-Length', content_length_header)
                 )
 
         def close_on_finish():
diff --git a/waitress/tests/fixtureapps/badcl.py b/waitress/tests/fixtureapps/badcl.py
index 96948b6..e3cf2ee 100644
--- a/waitress/tests/fixtureapps/badcl.py
+++ b/waitress/tests/fixtureapps/badcl.py
@@ -1,5 +1,5 @@
 def app(environ, start_response):
-    body = 'abcdefghi'
+    body = b'abcdefghi'
     cl = len(body)
     if environ['PATH_INFO'] == '/short_body':
         cl = len(body) +1
diff --git a/waitress/tests/fixtureapps/nocl.py b/waitress/tests/fixtureapps/nocl.py
new file mode 100644
index 0000000..f19a9d2
--- /dev/null
+++ b/waitress/tests/fixtureapps/nocl.py
@@ -0,0 +1,14 @@
+def app(environ, start_response):
+    body = b'abcdefghi'
+    app_iter = [body]
+    if environ['PATH_INFO'] == '/generator':
+        def gen():
+            yield body
+        app_iter = gen()
+    start_response('200 OK', [])
+    return app_iter
+
+if __name__ == '__main__':
+    from waitress import serve
+    serve(app, port=61523, verbose=False)
+    
diff --git a/waitress/tests/fixtureapps/writecb.py b/waitress/tests/fixtureapps/writecb.py
new file mode 100644
index 0000000..7d4824c
--- /dev/null
+++ b/waitress/tests/fixtureapps/writecb.py
@@ -0,0 +1,19 @@
+def app(environ, start_response):
+    path_info = environ['PATH_INFO']
+    if path_info == '/no_content_length':
+        headers = []
+    else:
+        headers = [('Content-Length', '9')]
+    write = start_response('200 OK', headers)
+    if path_info == '/long_body':
+        write(b'abcdefghij')
+    elif path_info == '/short_body':
+        write(b'abcdefgh')
+    else:
+        write(b'abcdefghi')
+    return []
+
+if __name__ == '__main__':
+    from waitress import serve
+    serve(app, port=61523, verbose=False)
+    
diff --git a/waitress/tests/test_functional.py b/waitress/tests/test_functional.py
index 5c49c69..40fb98f 100644
--- a/waitress/tests/test_functional.py
+++ b/waitress/tests/test_functional.py
@@ -373,6 +373,153 @@ class BadContentLengthTests(SubprocessTests, unittest.TestCase):
         response_body = fp.read(content_length)
         self.assertEqual(int(status), 200)
 
+class NoContentLengthTests(SubprocessTests, unittest.TestCase):
+    def setUp(self):
+        echo = os.path.join(here, 'fixtureapps', 'nocl.py')
+        self.start_subprocess([self.exe, echo])
+
+    def tearDown(self):
+        self.stop_subprocess()
+
+    def test_generator(self):
+        self.conn.request("GET", "/generator",
+                          headers={"Connection": "Keep-Alive",
+                                   "Content-Length": "0"})
+        resp = self.getresponse()
+        self.assertEqual(resp.getheader('Content-Length'), None)
+        self.assertEqual(resp.getheader('Connection'), 'close')
+        self.assertEqual(resp.read(), b'abcdefghi')
+
+    def test_list(self):
+        self.conn.request("GET", "/list",
+                          headers={"Connection": "Keep-Alive",
+                                   "Content-Length": "0"})
+        resp = self.getresponse()
+        self.assertEqual(resp.getheader('Content-Length'), '9')
+        self.assertEqual(resp.getheader('Connection'), None)
+        self.assertEqual(resp.read(), b'abcdefghi')
+
+class WriteCallbackTests(SubprocessTests, unittest.TestCase):
+    def setUp(self):
+        echo = os.path.join(here, 'fixtureapps', 'writecb.py')
+        self.start_subprocess([self.exe, echo])
+
+    def tearDown(self):
+        self.stop_subprocess()
+
+    def test_short_body(self):
+        # check to see if server closes connection when body is too short
+        # for cl header
+        to_send = tobytes(
+            "GET /short_body HTTP/1.0\n"
+             "Connection: Keep-Alive\n"
+             "Content-Length: 0\n"
+             "\n"
+             )
+        self.sock.connect((self.host, self.port))
+        self.sock.send(to_send)
+        fp = self.sock.makefile('rb', 0)
+        line = fp.readline() # status line
+        version, status, reason = (x.strip() for x in line.split(None, 2))
+        headers = parse_headers(fp)
+        content_length = int(headers.get('content-length')) or None
+        self.assertEqual(content_length, 9)
+        response_body = fp.read(content_length)
+        self.assertEqual(int(status), 200)
+        self.assertNotEqual(content_length, len(response_body))
+        self.assertEqual(len(response_body), content_length-1)
+        self.assertEqual(response_body, tobytes('abcdefgh'))
+        # remote closed connection (despite keepalive header); not sure why
+        # first send succeeds
+        self.sock.send(to_send[:5])
+        self.assertRaises(socket.error, self.sock.send, to_send[5:])
+
+    def test_long_body(self):
+        # check server doesnt close connection when body is too long
+        # for cl header
+        to_send = tobytes(
+            "GET /long_body HTTP/1.0\n"
+             "Connection: Keep-Alive\n"
+             "Content-Length: 0\n"
+             "\n"
+             )
+        self.sock.connect((self.host, self.port))
+        self.sock.send(to_send)
+        fp = self.sock.makefile('rb', 0)
+        line = fp.readline() # status line
+        version, status, reason = (x.strip() for x in line.split(None, 2))
+        headers = parse_headers(fp)
+        content_length = int(headers.get('content-length')) or None
+        self.assertEqual(content_length, 9)
+        response_body = fp.read(content_length)
+        self.assertEqual(int(status), 200)
+        self.assertEqual(content_length, len(response_body))
+        self.assertEqual(response_body, tobytes('abcdefghi'))
+        # remote does not close connection (keepalive header)
+        self.sock.send(to_send)
+        fp = self.sock.makefile('rb', 0)
+        line = fp.readline() # status line
+        version, status, reason = (x.strip() for x in line.split(None, 2))
+        headers = parse_headers(fp)
+        content_length = int(headers.get('content-length')) or None
+        response_body = fp.read(content_length)
+        self.assertEqual(int(status), 200)
+
+    def test_equal_body(self):
+        # check server doesnt close connection when body is equal to
+        # cl header
+        to_send = tobytes(
+            "GET /equal_body HTTP/1.0\n"
+             "Connection: Keep-Alive\n"
+             "Content-Length: 0\n"
+             "\n"
+             )
+        self.sock.connect((self.host, self.port))
+        self.sock.send(to_send)
+        fp = self.sock.makefile('rb', 0)
+        line = fp.readline() # status line
+        version, status, reason = (x.strip() for x in line.split(None, 2))
+        headers = parse_headers(fp)
+        content_length = int(headers.get('content-length')) or None
+        self.assertEqual(content_length, 9)
+        response_body = fp.read(content_length)
+        self.assertEqual(int(status), 200)
+        self.assertEqual(content_length, len(response_body))
+        self.assertEqual(response_body, tobytes('abcdefghi'))
+        # remote does not close connection (keepalive header)
+        self.sock.send(to_send)
+        fp = self.sock.makefile('rb', 0)
+        line = fp.readline() # status line
+        version, status, reason = (x.strip() for x in line.split(None, 2))
+        headers = parse_headers(fp)
+        content_length = int(headers.get('content-length')) or None
+        response_body = fp.read(content_length)
+        self.assertEqual(int(status), 200)
+
+    def test_no_content_length(self):
+        # wtf happens when there's no content-length
+        to_send = tobytes(
+            "GET /no_content_length HTTP/1.0\n"
+             "Connection: Keep-Alive\n"
+             "Content-Length: 0\n"
+             "\n"
+             )
+        self.sock.connect((self.host, self.port))
+        self.sock.send(to_send)
+        fp = self.sock.makefile('rb', 0)
+        line = fp.readline() # status line
+        version, status, reason = (x.strip() for x in line.split(None, 2))
+        headers = parse_headers(fp)
+        content_length = headers.get('content-length')
+        self.assertEqual(content_length, None)
+        response_body = fp.read()
+        self.assertEqual(int(status), 200)
+        self.assertEqual(response_body, tobytes('abcdefghi'))
+        # remote closed connection (despite keepalive header); not sure why
+        # first send succeeds
+        self.sock.send(to_send[:5])
+        self.assertRaises(socket.error, self.sock.send, to_send[5:])
+
 def parse_headers(fp):
     """Parses only RFC2822 headers from a file pointer.