[pypy-commit] pypy merge-2.7.2: The underlying stream of a BufferedIO can call arbitrary code,

[amauryfa]

Author: Amaury Forgeot d'Arc <amauryfa at gmail.com>
Branch: merge-2.7.2
Changeset: r51696:624143b36c8d
Date: 2012-01-23 22:13 +0100
http://bitbucket.org/pypy/pypy/changeset/624143b36c8d/

Log:	The underlying stream of a BufferedIO can call arbitrary code, or
	yield to some signal handler; protect against reentrant code.

diff --git a/pypy/module/_io/interp_bufferedio.py b/pypy/module/_io/interp_bufferedio.py
--- a/pypy/module/_io/interp_bufferedio.py
+++ b/pypy/module/_io/interp_bufferedio.py
@@ -11,6 +11,7 @@
     W_IOBase, DEFAULT_BUFFER_SIZE, convert_size,
     check_readable_w, check_writable_w, check_seekable_w)
 from pypy.module._io.interp_io import W_BlockingIOError
+from pypy.module.thread import ll_thread
 import errno
 
 STATE_ZERO, STATE_OK, STATE_DETACHED = range(3)
@@ -28,6 +29,29 @@
         return False
 
 
+class TryLock(object):
+    "A Lock that raises RuntimeError when acquired twice by the same thread"
+    def __init__(self, space):
+        ## XXX cannot free a Lock?
+        ## if self.lock:
+        ##     self.lock.free()
+        self.lock = space.allocate_lock()
+        self.owner = 0
+        self.operr = OperationError(space.w_RuntimeError,
+                                    space.wrap("reentrant call"))
+
+    def __enter__(self):
+        if not self.lock.acquire(False):
+            if self.owner == ll_thread.get_ident():
+                raise self.operr
+            self.lock.acquire(True)
+        self.owner = ll_thread.get_ident()
+    
+    def __exit__(self,*args):
+        self.owner = 0
+        self.lock.release()
+
+
 class BlockingIOError(Exception):
     pass
 
@@ -130,10 +154,7 @@
 
         self.buffer = ['\0'] * self.buffer_size
 
-        ## XXX cannot free a Lock?
-        ## if self.lock:
-        ##     self.lock.free()
-        self.lock = space.allocate_lock()
+        self.lock = TryLock(space)
 
         try:
             self._raw_tell(space)
diff --git a/pypy/module/_io/test/test_bufferedio.py b/pypy/module/_io/test/test_bufferedio.py
--- a/pypy/module/_io/test/test_bufferedio.py
+++ b/pypy/module/_io/test/test_bufferedio.py
@@ -1,6 +1,9 @@
 from pypy.conftest import gettestobjspace, option
 from pypy.interpreter.gateway import interp2app
 from pypy.tool.udir import udir
+from pypy.module._io import interp_bufferedio
+from pypy.interpreter.error import OperationError
+import py.test
 
 class AppTestBufferedReader:
     spaceconfig = dict(usemodules=['_io'])
@@ -217,7 +220,7 @@
 
 class AppTestBufferedWriter:
     def setup_class(cls):
-        cls.space = gettestobjspace(usemodules=['_io'])
+        cls.space = gettestobjspace(usemodules=['_io', 'thread'])
         tmpfile = udir.join('tmpfile')
         cls.w_tmpfile = cls.space.wrap(str(tmpfile))
         if option.runappdirect:
@@ -425,6 +428,22 @@
         bufio.flush()
         assert rawio.count == 3
 
+    def test_reentrant_write(self):
+        import thread  # Reentrant-safe is only enabled with threads
+        import _io, errno
+        class MockRawIO(_io._RawIOBase):
+            def writable(self):
+                return True
+            def write(self, data):
+                bufio.write("something else")
+                return len(data)
+
+        rawio = MockRawIO()
+        bufio = _io.BufferedWriter(rawio)
+        bufio.write("test")
+        exc = raises(RuntimeError, bufio.flush)
+        assert "reentrant" in str(exc.value)  # And not e.g. recursion limit.
+
 class AppTestBufferedRWPair:
     def test_pair(self):
         import _io
@@ -494,3 +513,15 @@
                 expected[j] = 2
                 expected[i] = 1
                 assert raw.getvalue() == str(expected)
+        
+
+class TestNonReentrantLock:
+    def test_trylock(self):
+        space = gettestobjspace(usemodules=['thread'])
+        lock = interp_bufferedio.TryLock(space)
+        with lock:
+            pass
+        with lock:
+            exc = py.test.raises(OperationError, "with lock: pass")
+        assert exc.value.match(space, space.w_RuntimeError)
+