HeaderReader class to replace read_headers

cheroot/server.py

@@ -131,53 +131,71 @@ def bton(b, encoding='ISO-8859-1'):
     logging.statistics = {}
 
 
-def read_headers(rfile, hdict=None):
-    """Read headers from the given stream into the given header dict.
+class HeaderReader(object):
 
-    If hdict is None, a new header dict is created. Returns the populated
-    header dict.
+    def __call__(self, rfile, hdict=None):
+        """
+        Read headers from the given stream into the given header dict.
 
-    Headers which are repeated are folded together using a comma if their
-    specification so dictates.
+        If hdict is None, a new header dict is created. Returns the populated
+        header dict.
 
-    This function raises ValueError when the read bytes violate the HTTP spec.
-    You should probably return "400 Bad Request" if this happens.
-    """
-    if hdict is None:
-        hdict = {}
+        Headers which are repeated are folded together using a comma if their
+        specification so dictates.
 
-    while True:
-        line = rfile.readline()
-        if not line:
-            # No more data--illegal end of headers
-            raise ValueError('Illegal end of headers.')
-
-        if line == CRLF:
-            # Normal end of headers
-            break
-        if not line.endswith(CRLF):
-            raise ValueError('HTTP requires CRLF terminators')
-
-        if line[0] in (SPACE, TAB):
-            # It's a continuation line.
-            v = line.strip()
-        else:
-            try:
-                k, v = line.split(COLON, 1)
-            except ValueError:
-                raise ValueError('Illegal header line.')
-            # TODO: what about TE and WWW-Authenticate?
-            k = k.strip().title()
-            v = v.strip()
-            hname = k
+        This function raises ValueError when the read bytes violate the HTTP spec.
+        You should probably return "400 Bad Request" if this happens.
+        """
+        if hdict is None:
+            hdict = {}
+
+        while True:
+            line = rfile.readline()
+            if not line:
+                # No more data--illegal end of headers
+                raise ValueError('Illegal end of headers.')
+
+            if line == CRLF:
+                # Normal end of headers
+                break
+            if not line.endswith(CRLF):
+                raise ValueError('HTTP requires CRLF terminators')
+
+            if line[0] in (SPACE, TAB):
+                # It's a continuation line.
+                v = line.strip()
+            else:
+                try:
+                    k, v = line.split(COLON, 1)
+                except ValueError:
+                    raise ValueError('Illegal header line.')
+                v = v.strip()
+                k = self._transform_key(k)
+                hname = k
+
+            if not self._allow_header(k):
+                continue
 
-        if k in comma_separated_headers:
-            existing = hdict.get(hname)
-            if existing:
-                v = b', '.join((existing, v))
-        hdict[hname] = v
+            if k in comma_separated_headers:
+                existing = hdict.get(hname)
+                if existing:
+                    v = b', '.join((existing, v))
+            hdict[hname] = v
 
-    return hdict
+        return hdict
+
+    def _allow_header(self, key_name):
+        return key_name
+
+    def _transform_key(self, key_name):
+        # TODO: what about TE and WWW-Authenticate?
+        return key_name.strip().title()
+
+
+class DropUnderscoreHeaderReader(HeaderReader):
+    def _allow_header(self, key_name):
+        orig = super(DropUnderscoreHeaderReader, self)._allow_header(key_name)
+        return orig and '_' not in key_name
 
 
 class SizeCheckWrapper(object):
@@ -479,6 +497,11 @@ class HTTPRequest(object):
 
     This value is set automatically inside send_headers."""
 
+    header_reader = HeaderReader()
+    """
+    A HeaderReader instance or compatible reader.
+    """
+
     def __init__(self, server, conn):
         self.server = server
         self.conn = conn
@@ -637,7 +660,7 @@ def read_request_headers(self):
         """Read self.rfile into self.inheaders. Return success."""
         # then all the http headers
         try:
-            read_headers(self.rfile, self.inheaders)
+            self.header_reader(self.rfile, self.inheaders)
         except ValueError as ex:
             self.simple_response('400 Bad Request', ex.args[0])
             return False