From 09803ad149b1def32ebf12603af50b384eb56efd 2011-10-15 08:03:19
From: MinRK <benjaminrk@gmail.com>
Date: 2011-10-15 08:03:19
Subject: [PATCH] protect against bad return type of CustomTB

includes test
also expands a few docstrings
---

diff --git a/IPython/core/interactiveshell.py b/IPython/core/interactiveshell.py
index 08309e9..93f9e23 100644
--- a/IPython/core/interactiveshell.py
+++ b/IPython/core/interactiveshell.py
@@ -1484,13 +1484,42 @@ class InteractiveShell(SingletonConfigurable, Magic):
             print 'Exception value:',value
             print 'Traceback      :',tb
             #print 'Source code    :','\n'.join(self.buffer)
+        
+        def validate_stb(stb):
+            """validate structured traceback return type
+            
+            return type of CustomTB *should* be a list of strings, but allow
+            single strings or None, which are harmless.
+            
+            This function will *always* return a list of strings,
+            and will raise a TypeError if stb is inappropriate.
+            """
+            msg = "CustomTB must return list of strings, not %r" % stb
+            if stb is None:
+                return []
+            elif isinstance(stb, basestring):
+                return [stb]
+            elif not isinstance(stb, list):
+                raise TypeError(msg)
+            # it's a list
+            for line in stb:
+                # check every element
+                if not isinstance(line, basestring):
+                    raise TypeError(msg)
+            return stb
 
         if handler is None:
             wrapped = dummy_handler
         else:
             def wrapped(self,etype,value,tb,tb_offset=None):
+                """wrap CustomTB handler, to protect IPython from user code
+                
+                This makes it harder (but not impossible) for custom exception
+                handlers to crash IPython.
+                """
                 try:
-                    return handler(self,etype,value,tb,tb_offset=tb_offset)
+                    stb = handler(self,etype,value,tb,tb_offset=tb_offset)
+                    return validate_stb(stb)
                 except:
                     # clear custom handler immediately
                     self.set_custom_exc((), None)
@@ -1499,7 +1528,10 @@ class InteractiveShell(SingletonConfigurable, Magic):
                     stb = self.InteractiveTB.structured_traceback(*sys.exc_info())
                     print >> io.stdout, self.InteractiveTB.stb2text(stb)
                     print >> io.stdout, "The original exception:"
-                    self.showtraceback((etype,value,tb), tb_offset=tb_offset)
+                    stb = self.InteractiveTB.structured_traceback(
+                                            (etype,value,tb), tb_offset=tb_offset
+                    )
+                return stb
 
         self.CustomTB = types.MethodType(wrapped,self)
         self.custom_exceptions = exc_tuple
@@ -1570,11 +1602,7 @@ class InteractiveShell(SingletonConfigurable, Magic):
                 sys.last_value = value
                 sys.last_traceback = tb
                 if etype in self.custom_exceptions:
-                    # FIXME: Old custom traceback objects may just return a
-                    # string, in that case we just put it into a list
                     stb = self.CustomTB(etype, value, tb, tb_offset)
-                    if isinstance(ctb, basestring):
-                        stb = [stb]
                 else:
                     if exception_only:
                         stb = ['An exception has occurred, use %tb to see '
diff --git a/IPython/core/tests/test_interactiveshell.py b/IPython/core/tests/test_interactiveshell.py
index 9e7b848..a0b31b0 100644
--- a/IPython/core/tests/test_interactiveshell.py
+++ b/IPython/core/tests/test_interactiveshell.py
@@ -155,7 +155,7 @@ class InteractiveShellTestCase(unittest.TestCase):
         try:
             # capture stderr
             io.stderr = StringIO()
-            ip.set_custom_exc((IOError,),lambda etype,value,tb: None)
+            ip.set_custom_exc((IOError,), lambda etype,value,tb: 1/0)
             self.assertEquals(ip.custom_exceptions, (IOError,))
             ip.run_cell(u'raise IOError("foo")')
             self.assertEquals(ip.custom_exceptions, ())
@@ -163,4 +163,20 @@ class InteractiveShellTestCase(unittest.TestCase):
         finally:
             io.stderr = save_stderr
 
+    def test_bad_custom_tb_return(self):
+        """Check that InteractiveShell is protected from bad return types in custom exception handlers"""
+        ip = get_ipython()
+        from IPython.utils import io
+        save_stderr = io.stderr
+        try:
+            # capture stderr
+            io.stderr = StringIO()
+            ip.set_custom_exc((NameError,),lambda etype,value,tb, tb_offset=None: 1)
+            self.assertEquals(ip.custom_exceptions, (NameError,))
+            ip.run_cell(u'a=abracadabra')
+            self.assertEquals(ip.custom_exceptions, ())
+            self.assertTrue("Custom TB Handler failed" in io.stderr.getvalue())
+        finally:
+            io.stderr = save_stderr
+