[3.7] bpo-17288: Prevent jumps from 'return' and 'exception' trace events. (GH-5928)

(cherry picked from commit e32bbaf376a09c149fa7c7f2919d7c9ce4e2a055)

Co-authored-by: xdegaye <xdegaye@gmail.com>

3 files changed, 94 insertions, 15 deletions

diff --git a/Lib/test/test_sys_settrace.py b/Lib/test/test_sys_settrace.py
index a05fbbd..46593cf 100644
--- a/Lib/test/test_sys_settrace.py
+++ b/Lib/test/test_sys_settrace.py
@@ -510,20 +510,35 @@ class RaisingTraceFuncTestCase(unittest.TestCase):
 class JumpTracer:
     """Defines a trace function that jumps from one place to another."""
 
-    def __init__(self, function, jumpFrom, jumpTo):
-        self.function = function
+    def __init__(self, function, jumpFrom, jumpTo, event='line',
+                 decorated=False):
+        self.code = function.__code__
         self.jumpFrom = jumpFrom
         self.jumpTo = jumpTo
+        self.event = event
+        self.firstLine = None if decorated else self.code.co_firstlineno
         self.done = False
 
     def trace(self, frame, event, arg):
-        if not self.done and frame.f_code == self.function.__code__:
-            firstLine = frame.f_code.co_firstlineno
-            if event == 'line' and frame.f_lineno == firstLine + self.jumpFrom:
+        if self.done:
+            return
+        # frame.f_code.co_firstlineno is the first line of the decorator when
+        # 'function' is decorated and the decorator may be written using
+        # multiple physical lines when it is too long. Use the first line
+        # trace event in 'function' to find the first line of 'function'.
+        if (self.firstLine is None and frame.f_code == self.code and
+                event == 'line'):
+            self.firstLine = frame.f_lineno - 1
+        if (event == self.event and self.firstLine and
+                frame.f_lineno == self.firstLine + self.jumpFrom):
+            f = frame
+            while f is not None and f.f_code != self.code:
+                f = f.f_back
+            if f is not None:
                 # Cope with non-integer self.jumpTo (because of
                 # no_jump_to_non_integers below).
                 try:
-                    frame.f_lineno = firstLine + self.jumpTo
+                    frame.f_lineno = self.firstLine + self.jumpTo
                 except TypeError:
                     frame.f_lineno = self.jumpTo
                 self.done = True
@@ -563,8 +578,9 @@ class JumpTestCase(unittest.TestCase):
                        "Expected: " + repr(expected) + "\n" +
                        "Received: " + repr(received))
 
-    def run_test(self, func, jumpFrom, jumpTo, expected, error=None):
-        tracer = JumpTracer(func, jumpFrom, jumpTo)
+    def run_test(self, func, jumpFrom, jumpTo, expected, error=None,
+                 event='line', decorated=False):
+        tracer = JumpTracer(func, jumpFrom, jumpTo, event, decorated)
         sys.settrace(tracer.trace)
         output = []
         if error is None:
@@ -575,15 +591,15 @@ class JumpTestCase(unittest.TestCase):
         sys.settrace(None)
         self.compare_jump_output(expected, output)
 
-    def jump_test(jumpFrom, jumpTo, expected, error=None):
+    def jump_test(jumpFrom, jumpTo, expected, error=None, event='line'):
         """Decorator that creates a test that makes a jump
         from one place to another in the following code.
         """
         def decorator(func):
             @wraps(func)
             def test(self):
-                # +1 to compensate a decorator line
-                self.run_test(func, jumpFrom+1, jumpTo+1, expected, error)
+                self.run_test(func, jumpFrom, jumpTo, expected,
+                              error=error, event=event, decorated=True)
             return test
         return decorator
 
@@ -1058,6 +1074,36 @@ output.append(4)
         sys.settrace(None)
         self.compare_jump_output([2, 3, 2, 3, 4], namespace["output"])
 
+    @jump_test(2, 3, [1], event='call', error=(ValueError, "can't jump from"
+               " the 'call' trace event of a new frame"))
+    def test_no_jump_from_call(output):
+        output.append(1)
+        def nested():
+            output.append(3)
+        nested()
+        output.append(5)
+
+    @jump_test(2, 1, [1], event='return', error=(ValueError,
+               "can only jump from a 'line' trace event"))
+    def test_no_jump_from_return_event(output):
+        output.append(1)
+        return
+
+    @jump_test(2, 1, [1], event='exception', error=(ValueError,
+               "can only jump from a 'line' trace event"))
+    def test_no_jump_from_exception_event(output):
+        output.append(1)
+        1 / 0
+
+    @jump_test(3, 2, [2], event='return', error=(ValueError,
+               "can't jump from a yield statement"))
+    def test_no_jump_from_yield(output):
+        def gen():
+            output.append(2)
+            yield 3
+        next(gen())
+        output.append(5)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/Misc/NEWS.d/next/Core and Builtins/2018-02-27-13-36-21.bpo-17288.Gdj24S.rst b/Misc/NEWS.d/next/Core and Builtins/2018-02-27-13-36-21.bpo-17288.Gdj24S.rst
new file mode 100644
index 0000000..ce9e84c
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2018-02-27-13-36-21.bpo-17288.Gdj24S.rst
@@ -0,0 +1 @@
+Prevent jumps from 'return' and 'exception' trace events.
diff --git a/Objects/frameobject.c b/Objects/frameobject.c
index 5c73e85..b7a16ad 100644
--- a/Objects/frameobject.c
+++ b/Objects/frameobject.c
@@ -56,6 +56,9 @@ frame_getlineno(PyFrameObject *f, void *closure)
  *  o 'try'/'for'/'while' blocks can't be jumped into because the blockstack
  *    needs to be set up before their code runs, and for 'for' loops the
  *    iterator needs to be on the stack.
+ *  o Jumps cannot be made from within a trace function invoked with a
+ *    'return' or 'exception' event since the eval loop has been exited at
+ *    that time.
  */
 static int
 frame_setlineno(PyFrameObject *f, PyObject* p_new_lineno)
@@ -91,13 +94,32 @@ frame_setlineno(PyFrameObject *f, PyObject* p_new_lineno)
         return -1;
     }
 
+    /* Upon the 'call' trace event of a new frame, f->f_lasti is -1 and
+     * f->f_trace is NULL, check first on the first condition.
+     * Forbidding jumps from the 'call' event of a new frame is a side effect
+     * of allowing to set f_lineno only from trace functions. */
+    if (f->f_lasti == -1) {
+        PyErr_Format(PyExc_ValueError,
+                     "can't jump from the 'call' trace event of a new frame");
+        return -1;
+    }
+
     /* You can only do this from within a trace function, not via
      * _getframe or similar hackery. */
-    if (!f->f_trace)
-    {
+    if (!f->f_trace) {
         PyErr_Format(PyExc_ValueError,
-                     "f_lineno can only be set by a"
-                     " line trace function");
+                     "f_lineno can only be set by a trace function");
+        return -1;
+    }
+
+    /* Forbid jumps upon a 'return' trace event (except after executing a
+     * YIELD_VALUE or YIELD_FROM opcode, f_stacktop is not NULL in that case)
+     * and upon an 'exception' trace event.
+     * Jumps from 'call' trace events have already been forbidden above for new
+     * frames, so this check does not change anything for 'call' events. */
+    if (f->f_stacktop == NULL) {
+        PyErr_SetString(PyExc_ValueError,
+                "can only jump from a 'line' trace event");
         return -1;
     }
 
@@ -156,6 +178,16 @@ frame_setlineno(PyFrameObject *f, PyObject* p_new_lineno)
 
     /* We're now ready to look at the bytecode. */
     PyBytes_AsStringAndSize(f->f_code->co_code, (char **)&code, &code_len);
+
+    /* The trace function is called with a 'return' trace event after the
+     * execution of a yield statement. */
+    assert(f->f_lasti != -1);
+    if (code[f->f_lasti] == YIELD_VALUE || code[f->f_lasti] == YIELD_FROM) {
+        PyErr_SetString(PyExc_ValueError,
+                "can't jump from a yield statement");
+        return -1;
+    }
+
     min_addr = Py_MIN(new_lasti, f->f_lasti);
     max_addr = Py_MAX(new_lasti, f->f_lasti);