bpo-35370: Add _PyEval_SetTrace() function (GH-18975)

* sys.settrace(), sys.setprofile() and _lsprof.Profiler.enable() now
  properly report PySys_Audit() error if "sys.setprofile" or
  "sys.settrace" audit event is denied.
* Add _PyEval_SetProfile() and _PyEval_SetTrace() function: similar
  to PyEval_SetProfile() and PyEval_SetTrace() but take a tstate
  parameter and return -1 on error.
* Add _PyObject_FastCallTstate() function.

2 files changed, 102 insertions, 44 deletions

diff --git a/Python/ceval.c b/Python/ceval.c
index ccd1c06a39..fc4f718de2 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -4586,51 +4586,85 @@ maybe_call_line_trace(Py_tracefunc func, PyObject *obj,
     return result;
 }
 
-void
-PyEval_SetProfile(Py_tracefunc func, PyObject *arg)
+int
+_PyEval_SetProfile(PyThreadState *tstate, Py_tracefunc func, PyObject *arg)
 {
+    assert(tstate != NULL);
+    /* The caller must hold the GIL */
+    assert(PyGILState_Check());
+
+    /* Call PySys_Audit() in the context of the current thread state,
+       even if tstate is not the current thread state. */
     if (PySys_Audit("sys.setprofile", NULL) < 0) {
-        return;
+        return -1;
     }
 
-    PyThreadState *tstate = _PyThreadState_GET();
-    PyObject *temp = tstate->c_profileobj;
-    Py_XINCREF(arg);
+    PyObject *profileobj = tstate->c_profileobj;
+
     tstate->c_profilefunc = NULL;
     tstate->c_profileobj = NULL;
-    /* Must make sure that tracing is not ignored if 'temp' is freed */
+    /* Must make sure that tracing is not ignored if 'profileobj' is freed */
     tstate->use_tracing = tstate->c_tracefunc != NULL;
-    Py_XDECREF(temp);
-    tstate->c_profilefunc = func;
+    Py_XDECREF(profileobj);
+
+    Py_XINCREF(arg);
     tstate->c_profileobj = arg;
+    tstate->c_profilefunc = func;
+
     /* Flag that tracing or profiling is turned on */
     tstate->use_tracing = (func != NULL) || (tstate->c_tracefunc != NULL);
+    return 0;
 }
 
 void
-PyEval_SetTrace(Py_tracefunc func, PyObject *arg)
+PyEval_SetProfile(Py_tracefunc func, PyObject *arg)
+{
+    PyThreadState *tstate = _PyThreadState_GET();
+    (void)_PyEval_SetProfile(tstate, func, arg);
+}
+
+int
+_PyEval_SetTrace(PyThreadState *tstate, Py_tracefunc func, PyObject *arg)
 {
+    assert(tstate != NULL);
+    /* The caller must hold the GIL */
+    assert(PyGILState_Check());
+
+    /* Call PySys_Audit() in the context of the current thread state,
+       even if tstate is not the current thread state. */
     if (PySys_Audit("sys.settrace", NULL) < 0) {
-        return;
+        return -1;
     }
 
-    _PyRuntimeState *runtime = &_PyRuntime;
-    PyThreadState *tstate = _PyRuntimeState_GetThreadState(runtime);
-    PyObject *temp = tstate->c_traceobj;
-    runtime->ceval.tracing_possible += (func != NULL) - (tstate->c_tracefunc != NULL);
-    Py_XINCREF(arg);
+    struct _ceval_runtime_state *ceval = &tstate->interp->runtime->ceval;
+    PyObject *traceobj = tstate->c_traceobj;
+    ceval->tracing_possible += (func != NULL) - (tstate->c_tracefunc != NULL);
+
     tstate->c_tracefunc = NULL;
     tstate->c_traceobj = NULL;
-    /* Must make sure that profiling is not ignored if 'temp' is freed */
-    tstate->use_tracing = tstate->c_profilefunc != NULL;
-    Py_XDECREF(temp);
-    tstate->c_tracefunc = func;
+    /* Must make sure that profiling is not ignored if 'traceobj' is freed */
+    tstate->use_tracing = (tstate->c_profilefunc != NULL);
+    Py_XDECREF(traceobj);
+
+    Py_XINCREF(arg);
     tstate->c_traceobj = arg;
+    tstate->c_tracefunc = func;
+
     /* Flag that tracing or profiling is turned on */
     tstate->use_tracing = ((func != NULL)
                            || (tstate->c_profilefunc != NULL));
+
+    return 0;
+}
+
+void
+PyEval_SetTrace(Py_tracefunc func, PyObject *arg)
+{
+    PyThreadState *tstate = _PyThreadState_GET();
+    (void)_PyEval_SetTrace(tstate, func, arg);
 }
 
+
 void
 _PyEval_SetCoroutineOriginTrackingDepth(PyThreadState *tstate, int new_depth)
 {
diff --git a/Python/sysmodule.c b/Python/sysmodule.c
index 2a8d12c03c..c78a627380 100644
--- a/Python/sysmodule.c
+++ b/Python/sysmodule.c
@@ -876,7 +876,7 @@ trace_init(void)
 
 
 static PyObject *
-call_trampoline(PyObject* callback,
+call_trampoline(PyThreadState *tstate, PyObject* callback,
                 PyFrameObject *frame, int what, PyObject *arg)
 {
     if (PyFrame_FastToLocalsWithError(frame) < 0) {
@@ -889,7 +889,7 @@ call_trampoline(PyObject* callback,
     stack[2] = (arg != NULL) ? arg : Py_None;
 
     /* call the Python-level function */
-    PyObject *result = _PyObject_FastCall(callback, stack, 3);
+    PyObject *result = _PyObject_FastCallTstate(tstate, callback, stack, 3);
 
     PyFrame_LocalsToFast(frame, 1);
     if (result == NULL) {
@@ -903,15 +903,17 @@ static int
 profile_trampoline(PyObject *self, PyFrameObject *frame,
                    int what, PyObject *arg)
 {
-    PyObject *result;
-
-    if (arg == NULL)
+    if (arg == NULL) {
         arg = Py_None;
-    result = call_trampoline(self, frame, what, arg);
+    }
+
+    PyThreadState *tstate = _PyThreadState_GET();
+    PyObject *result = call_trampoline(tstate, self, frame, what, arg);
     if (result == NULL) {
-        PyEval_SetProfile(NULL, NULL);
+        _PyEval_SetProfile(tstate, NULL, NULL);
         return -1;
     }
+
     Py_DECREF(result);
     return 0;
 }
@@ -921,20 +923,24 @@ trace_trampoline(PyObject *self, PyFrameObject *frame,
                  int what, PyObject *arg)
 {
     PyObject *callback;
-    PyObject *result;
-
-    if (what == PyTrace_CALL)
+    if (what == PyTrace_CALL) {
         callback = self;
-    else
+    }
+    else {
         callback = frame->f_trace;
-    if (callback == NULL)
+    }
+    if (callback == NULL) {
         return 0;
-    result = call_trampoline(callback, frame, what, arg);
+    }
+
+    PyThreadState *tstate = _PyThreadState_GET();
+    PyObject *result = call_trampoline(tstate, callback, frame, what, arg);
     if (result == NULL) {
-        PyEval_SetTrace(NULL, NULL);
+        _PyEval_SetTrace(tstate, NULL, NULL);
         Py_CLEAR(frame->f_trace);
         return -1;
     }
+
     if (result != Py_None) {
         Py_XSETREF(frame->f_trace, result);
     }
@@ -947,12 +953,21 @@ trace_trampoline(PyObject *self, PyFrameObject *frame,
 static PyObject *
 sys_settrace(PyObject *self, PyObject *args)
 {
-    if (trace_init() == -1)
+    if (trace_init() == -1) {
         return NULL;
-    if (args == Py_None)
-        PyEval_SetTrace(NULL, NULL);
-    else
-        PyEval_SetTrace(trace_trampoline, args);
+    }
+
+    PyThreadState *tstate = _PyThreadState_GET();
+    if (args == Py_None) {
+        if (_PyEval_SetTrace(tstate, NULL, NULL) < 0) {
+            return NULL;
+        }
+    }
+    else {
+        if (_PyEval_SetTrace(tstate, trace_trampoline, args) < 0) {
+            return NULL;
+        }
+    }
     Py_RETURN_NONE;
 }
 
@@ -987,12 +1002,21 @@ sys_gettrace_impl(PyObject *module)
 static PyObject *
 sys_setprofile(PyObject *self, PyObject *args)
 {
-    if (trace_init() == -1)
+    if (trace_init() == -1) {
         return NULL;
-    if (args == Py_None)
-        PyEval_SetProfile(NULL, NULL);
-    else
-        PyEval_SetProfile(profile_trampoline, args);
+    }
+
+    PyThreadState *tstate = _PyThreadState_GET();
+    if (args == Py_None) {
+        if (_PyEval_SetProfile(tstate, NULL, NULL) < 0) {
+            return NULL;
+        }
+    }
+    else {
+        if (_PyEval_SetProfile(tstate, profile_trampoline, args) < 0) {
+            return NULL;
+        }
+    }
     Py_RETURN_NONE;
 }