--- a/cosocket.py	Fri May 01 16:46:31 2009 -0700
+++ b/cosocket.py	Fri May 01 17:07:18 2009 -0700
@@ -8,6 +8,7 @@
 import weakref
 
 DEFAULT_TIMEOUT = 90.0
+DEFAULT_MAX_DATA = 65536
 
 time_map = {}
 
@@ -26,8 +27,10 @@
     def __init__(self, coroutine, conn = None):
         asynchat.async_chat.__init__(self, conn)
         self.set_terminator(None)
+        self.__max_data = DEFAULT_MAX_DATA
         self.__coroutine = coroutine
         self.__data = []
+        self.__data_len = 0
         self.__coroutine_stack = []
         self.__timeout = 0
         self.__time_passed = 0
@@ -138,14 +141,21 @@
         time_map[self.__on_tick] = self.__on_tick
 
     def collect_incoming_data(self, data):
-        # TODO: Enforce some maximum data length.
         self.__data.append(data)
+        self.__data_len += len(data)
+        if self.__max_data and self.__data_len > self.__max_data:
+            self.handle_close()
+
+    def set_max_data(self, amount):
+        self.__max_data = amount
 
     def found_terminator(self):
-        self.set_terminator(None)
-        data = ''.join(self.__data)
-        self.__data = []
-        self.continue_from_yield(data)
+        if not (self.__max_data and self.__data_len > self.__max_data):
+            self.set_terminator(None)
+            data = ''.join(self.__data)
+            self.__data = []
+            self.__data_len = 0
+            self.continue_from_yield(data)
 
 class CoroutineSocketServer(asyncore.dispatcher):
     def __init__(self, addr, coroutineFactory):
@@ -175,17 +185,20 @@
 
 class until_received(object):
     def __init__(self, terminator = None, bytes = None,
-                 timeout = DEFAULT_TIMEOUT):
+                 timeout = DEFAULT_TIMEOUT, max_data = DEFAULT_MAX_DATA):
         self._timeout = timeout
         if terminator:
             self._terminator = terminator
+            self._max_data = max_data
         elif bytes:
             self._terminator = bytes
+            self._max_data = 0
         else:
             raise ValueError()
 
     def execute(self, dispatcher):
         dispatcher.set_timeout(self._timeout)
+        dispatcher.set_max_data(self._max_data)
         dispatcher.set_terminator(self._terminator)
 
 class until_sent(object):