provide a real write callable, warn when cl exceeds or shorts specified, anticipate a case where start_response is not called until first iteration, set content-length header if len(app_iter) == 1 and none provided, raise an exception if start_response uncalled before a write is done

6 files changed, 108 insertions, 51 deletions

diff --git a/README.txt b/README.txt
index 6731fa3..e822647 100644
--- a/README.txt
+++ b/README.txt
@@ -136,9 +136,6 @@ known to work fairly well.
 Known Issues
 ------------
 
-- The server returns a ``write`` callable from ``start_response`` which
-  raises a ``NotImplementedError`` exception when called.
-
 - The server does not support the ``wsgi.file_wrapper`` protocol.
 
 Differences from ``zope.server``
@@ -169,5 +166,17 @@ Differences from ``zope.server``
 - Supports convenience ``waitress.serve`` function (e.g. ``from waitress
   import serve; serve(app)`` and convenience ``server.serve()`` function.
 
+- Returns a "real" write method from start_response.
+
+- Provides a getsockname method of the server FBO figuring out which port the
+  server is listening on when it's bound to port 0.
+
+- Warns when app_iter bytestream numbytes less than or greater than specified
+  Content-Length.
+
+- Set content-length header if len(app_iter) == 1 and none provided.
+
+- Raise an exception if start_response isnt called before any body write.
+
 .. _PasteDeploy: http://pythonpaste.org/deploy/
 
diff --git a/TODO.txt b/TODO.txt
index 3559b9d..035d03a 100644
--- a/TODO.txt
+++ b/TODO.txt
@@ -4,8 +4,6 @@
 
 - Documentation.
 
-- Figure out what to do when we have less data than Content-Length (close?)
-
 - Speed tweaking.
 
 - Coverage.
@@ -20,3 +18,4 @@
 
 - SERVER_IDENT -> ident
 
+- Logging.
diff --git a/waitress/task.py b/waitress/task.py
index 92e48df..354812e 100644
--- a/waitress/task.py
+++ b/waitress/task.py
@@ -140,8 +140,6 @@ class ThreadedTaskDispatcher(object):
 class HTTPTask(object):
     """An HTTP task accepts a request and writes to a channel.
 
-       Subclass this and override the execute() method.
-
        See ITask, IHeaderOutput.
     """
 
@@ -233,49 +231,69 @@ class HTTPTask(object):
                     self.content_length = int(v)
 
             # Return the write method used to write the response data.
-            return fake_write
+            return self.write
 
         # Call the application to handle the request and write a response
         app_iter = self.channel.server.application(env, start_response)
 
-        if not self.start_response_called:
-            raise RuntimeError('start_response was not called before app_iter '
-                               'returned')
-
-        # Set a Content-Length header if one is not supplied.
-        cl = self.content_length
-        if cl == -1:
+        app_iter_len = None
+        if hasattr(app_iter, '__len__'):
             app_iter_len = len(app_iter)
-            if app_iter_len == 1:
-                cl = self.content_length = len(app_iter[0])
-
-        has_content_length = cl != -1
-        bytes_written = 0
 
         try:
             # By iterating manually at this point, we execute task.write()
             # multiple times, allowing partial data to be sent.
-            for value in app_iter:
-                towrite = value
+            first_chunk_len = None
+            for chunk in app_iter:
+                if first_chunk_len is None:
+                    first_chunk_len = len(chunk)
+                # 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 = value[:cl-self.content_bytes_written]
-                bytes_written += len(towrite)
+                    towrite = chunk[:cl-self.content_bytes_written]
                 self.write(towrite)
-                if towrite != value:
-                    self.log_info(
-                        'warning: app_iter content exceeded the number '
-                        'of bytes specified by Content-Length header (%s)' % cl)
+                if towrite != chunk:
+                    self.channel.server.log_info(
+                        'app_iter content exceeded the number of bytes '
+                        'specified by Content-Length header (%s)' % cl)
                     break
 
-            if has_content_length:
-                if bytes_written != cl:
-                    self.log_info('warning: app_iter returned a number of '
-                                  'bytes (%s) too short for specified '
-                                  'Content-Length (%s)'  % (bytes_written, cl))
+            cl = self.content_length
+            if cl != -1:
+                if self.content_bytes_written != cl:
+                    self.close_on_finish = True
+                    self.channel.server.log_info(
+                        'app_iter returned a number of bytes (%s) too short '
+                        'for specified Content-Length (%s)' % (
+                            self.content_bytes_written,
+                            cl)
+                        )
         finally:
             if hasattr(app_iter, 'close'):
                 app_iter.close()
 
+    def write(self, data):
+        if not self.start_response_called:
+            raise RuntimeError('start_response was not called before body '
+                               'written')
+        channel = self.channel
+        if not self.wrote_header:
+            rh = self.build_response_header()
+            channel.write(rh)
+            self.wrote_header = True
+        if data:
+            self.content_bytes_written += len(data)
+            channel.write(data)
 
     def build_response_header(self):
         version = self.version
@@ -408,15 +426,6 @@ class HTTPTask(object):
         if not self.wrote_header:
             self.write(b'')
 
-    def write(self, data):
-        channel = self.channel
-        if not self.wrote_header:
-            rh = self.build_response_header()
-            channel.write(rh)
-            self.wrote_header = True
-        if data:
-            channel.write(data)
-
 def fake_write(body):
     raise NotImplementedError(
         "the waitress HTTP Server does not support the WSGI write() function.")
diff --git a/waitress/tests/fixtureapps/badcl.py b/waitress/tests/fixtureapps/badcl.py
new file mode 100644
index 0000000..b1009a4
--- /dev/null
+++ b/waitress/tests/fixtureapps/badcl.py
@@ -0,0 +1,17 @@
+def app(environ, start_response):
+    body = 'abcdefghi'
+    cl = len(body)
+    if environ['PATH_INFO'] == '/short':
+        cl = len(body) - 1
+    if environ['PATH_INFO'] == '/long':
+        cl = len(body) + 1
+    start_response(
+        '200 OK',
+        [('Content-Length', str(cl)), ('Content-Type', 'text/plain')]
+        )
+    return [body]
+
+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 94e7a82..65ffacd 100644
--- a/waitress/tests/test_functional.py
+++ b/waitress/tests/test_functional.py
@@ -311,6 +311,26 @@ class ExpectContinueTests(SubprocessTests, unittest.TestCase):
         self.assertEqual(length, len(response_body))
         self.assertEqual(response_body, tobytes(data))
 
+class BadContentLengthTests(SubprocessTests, unittest.TestCase):
+    def setUp(self):
+        echo = os.path.join(here, 'fixtureapps', 'badcl.py')
+        self.start_subprocess([self.exe, echo])
+
+    def tearDown(self):
+        self.stop_subprocess()
+
+    def test_short_cl(self):
+        self.conn.request('GET', '/short')
+        resp = self.getresponse()
+        self.assertEqual(resp.getheader('Content-Length'), '8')
+        resp.read()
+
+    def test_long_cl(self):
+        self.conn.request('GET', '/long')
+        resp = self.getresponse()
+        self.assertEqual(resp.getheader('Content-Length'), '10')
+        self.assertRaises(httplib.IncompleteRead, resp.read)
+
 def parse_headers(fp):
     """Parses only RFC2822 headers from a file pointer.
 
diff --git a/waitress/tests/test_task.py b/waitress/tests/test_task.py
index cf52510..c9a8861 100644
--- a/waitress/tests/test_task.py
+++ b/waitress/tests/test_task.py
@@ -105,6 +105,7 @@ class TestHTTPTask(unittest.TestCase):
         def execute():
             inst.executed = True
         inst.execute = execute
+        inst.start_response_called = True
         inst.service()
         self.assertTrue(inst.start_time)
         self.assertTrue(inst.channel.closed_when_done)
@@ -305,12 +306,12 @@ class TestHTTPTask(unittest.TestCase):
         self.assertTrue(lines[2].startswith(b'Date:'))
         self.assertEqual(lines[3], b'Server: hithere')
 
-    def test_getEnviron_already_cached(self):
+    def test_get_environment_already_cached(self):
         inst = self._makeOne()
         inst.environ = object()
         self.assertEqual(inst.get_environment(), inst.environ)
 
-    def test_getEnviron_path_startswith_more_than_one_slash(self):
+    def test_get_environment_path_startswith_more_than_one_slash(self):
         inst = self._makeOne()
         request_data = DummyParser()
         request_data.path = '///abc'
@@ -318,7 +319,7 @@ class TestHTTPTask(unittest.TestCase):
         environ = inst.get_environment()
         self.assertEqual(environ['PATH_INFO'], '/abc')
 
-    def test_getEnviron_path_empty(self):
+    def test_get_environment_path_empty(self):
         inst = self._makeOne()
         request_data = DummyParser()
         request_data.path = ''
@@ -326,14 +327,14 @@ class TestHTTPTask(unittest.TestCase):
         environ = inst.get_environment()
         self.assertEqual(environ['PATH_INFO'], '/')
 
-    def test_getEnviron_no_query(self):
+    def test_get_environment_no_query(self):
         inst = self._makeOne()
         request_data = DummyParser()
         inst.request_data = request_data
         environ = inst.get_environment()
         self.assertFalse('QUERY_STRING' in environ)
 
-    def test_getEnviron_with_query(self):
+    def test_get_environment_with_query(self):
         inst = self._makeOne()
         request_data = DummyParser()
         request_data.query = 'abc'
@@ -341,7 +342,7 @@ class TestHTTPTask(unittest.TestCase):
         environ = inst.get_environment()
         self.assertEqual(environ['QUERY_STRING'], 'abc')
 
-    def test_getEnviron_values(self):
+    def test_get_environment_values(self):
         import sys
         inst = self._makeOne()
         request_data = DummyParser()
@@ -379,6 +380,7 @@ class TestHTTPTask(unittest.TestCase):
     def test_finish_didnt_write_header(self):
         inst = self._makeOne()
         inst.wrote_header = False
+        inst.start_response_called = True
         inst.finish()
         self.assertTrue(inst.channel.written)
 
@@ -391,20 +393,21 @@ class TestHTTPTask(unittest.TestCase):
     def test_write_wrote_header(self):
         inst = self._makeOne()
         inst.wrote_header = True
+        inst.start_response_called = True
         inst.write(b'abc')
         self.assertEqual(inst.channel.written, b'abc')
 
     def test_write_header_not_written(self):
         inst = self._makeOne()
         inst.wrote_header = False
+        inst.start_response_called = True
         inst.write(b'abc')
         self.assertTrue(inst.channel.written)
         self.assertEqual(inst.wrote_header, True)
 
-    def test_execute_start_response_uncalled(self):
+    def test_write_start_response_uncalled(self):
         inst = self._makeOne()
-        inst.channel.server.application = lambda *arg: None
-        self.assertRaises(RuntimeError, inst.execute)
+        self.assertRaises(RuntimeError, inst.write, b'')
 
 
 class DummyTask(object):