bpo-38644: Add _PySys_Audit() which takes tstate (GH-19180)

Add _PySys_Audit() function to the internal C API: similar to
PySys_Audit(), but requires a mandatory tstate parameter.

Cleanup sys_audit_tstate() code: remove code path for NULL tstate,
since the function exits at entry if tstate is NULL. Remove also code
path for NULL tstate->interp: should_audit() now ensures that it is
not NULL (even if tstate->interp cannot be NULL in practice).

PySys_AddAuditHook() now checks if tstate is not NULL to decide if
tstate can be used or not, and tstate is set to NULL if the runtime
is not initialized yet.

Use _PySys_Audit() in sysmodule.c.

Include/cpython/sysmodule.h

@@ -13,7 +13,10 @@ PyAPI_FUNC(size_t) _PySys_GetSizeOf(PyObject *);
 
 typedef int(*Py_AuditHookFunction)(const char *, PyObject *, void *);
 
-PyAPI_FUNC(int) PySys_Audit(const char*, const char *, ...);
+PyAPI_FUNC(int) PySys_Audit(
+    const char *event,
+    const char *argFormat,
+    ...);
 PyAPI_FUNC(int) PySys_AddAuditHook(Py_AuditHookFunction, void*);
 
 #ifdef __cplusplus

Include/internal/pycore_sysmodule.h

@@ -0,0 +1,24 @@
+#ifndef Py_INTERNAL_SYSMODULE_H
+#define Py_INTERNAL_SYSMODULE_H
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#ifndef Py_BUILD_CORE
+#  error "this header requires Py_BUILD_CORE define"
+#endif
+
+PyAPI_FUNC(int) _PySys_Audit(
+    PyThreadState *tstate,
+    const char *event,
+    const char *argFormat,
+    ...);
+
+/* We want minimal exposure of this function, so use extern rather than
+   PyAPI_FUNC() to not export the symbol. */
+extern void _PySys_ClearAuditHooks(PyThreadState *tstate);
+
+#ifdef __cplusplus
+}
+#endif
+#endif /* !Py_INTERNAL_SYSMODULE_H */

Makefile.pre.in

@@ -1101,6 +1101,7 @@ PYTHON_HEADERS= \
 		$(srcdir)/Include/internal/pycore_pylifecycle.h \
 		$(srcdir)/Include/internal/pycore_pymem.h \
 		$(srcdir)/Include/internal/pycore_pystate.h \
+		$(srcdir)/Include/internal/pycore_sysmodule.h \
 		$(srcdir)/Include/internal/pycore_traceback.h \
 		$(srcdir)/Include/internal/pycore_tupleobject.h \
 		$(srcdir)/Include/internal/pycore_warnings.h \

PCbuild/pythoncore.vcxproj

@@ -183,6 +183,7 @@
     <ClInclude Include="..\Include\internal\pycore_pylifecycle.h" />
     <ClInclude Include="..\Include\internal\pycore_pymem.h" />
     <ClInclude Include="..\Include\internal\pycore_pystate.h" />
+    <ClInclude Include="..\Include\internal\pycore_sysmodule.h" />
     <ClInclude Include="..\Include\internal\pycore_traceback.h" />
     <ClInclude Include="..\Include\internal\pycore_tupleobject.h" />
     <ClInclude Include="..\Include\internal\pycore_warnings.h" />

PCbuild/pythoncore.vcxproj.filters

@@ -252,6 +252,9 @@
     <ClInclude Include="..\Include\internal\pycore_pystate.h">
       <Filter>Include</Filter>
     </ClInclude>
+    <ClInclude Include="..\Include\internal\pycore_sysmodule.h">
+      <Filter>Include</Filter>
+    </ClInclude>
     <ClInclude Include="..\Include\internal\pycore_traceback.h">
       <Filter>Include</Filter>
     </ClInclude>

Python/pylifecycle.c

@@ -16,6 +16,7 @@
 #include "pycore_pylifecycle.h"
 #include "pycore_pymem.h"
 #include "pycore_pystate.h"
+#include "pycore_sysmodule.h"
 #include "pycore_traceback.h"
 #include "grammar.h"
 #include "node.h"
@@ -1409,11 +1410,7 @@ Py_FinalizeEx(void)
     _PyGC_CollectIfEnabled();
 
     /* Clear all loghooks */
-    /* We want minimal exposure of this function, so define the extern
+    _PySys_ClearAuditHooks(tstate);
-     * here. The linker should discover the correct function without
-     * exporting a symbol. */
-    extern void _PySys_ClearAuditHooks(void);
-    _PySys_ClearAuditHooks();
 
     /* Destroy all modules */
     _PyImport_Cleanup(tstate);

Python/sysmodule.c

@@ -137,58 +137,67 @@ PySys_SetObject(const char *name, PyObject *v)
     return sys_set_object(tstate, name, v);
 }
 
+
 static int
-should_audit(PyThreadState *ts)
+should_audit(PyInterpreterState *is)
 {
-    if (!ts) {
+    /* tstate->interp cannot be NULL, but test it just in case
+       for extra safety */
+    assert(is != NULL);
+    if (!is) {
         return 0;
     }
-    PyInterpreterState *is = ts ? ts->interp : NULL;
+    return (is->runtime->audit_hook_head
-    return _PyRuntime.audit_hook_head
+            || is->audit_hooks
-        || (is && is->audit_hooks)
+            || PyDTrace_AUDIT_ENABLED());
-        || PyDTrace_AUDIT_ENABLED();
 }
 
-int
-PySys_Audit(const char *event, const char *argFormat, ...)
-{
-    PyObject *eventName = NULL;
-    PyObject *eventArgs = NULL;
-    PyObject *hooks = NULL;
-    PyObject *hook = NULL;
-    int res = -1;
-    PyThreadState *ts = _PyThreadState_GET();
 
+static int
+sys_audit_tstate(PyThreadState *ts, const char *event,
+                 const char *argFormat, va_list vargs)
+{
     /* N format is inappropriate, because you do not know
        whether the reference is consumed by the call.
        Assert rather than exception for perf reasons */
     assert(!argFormat || !strchr(argFormat, 'N'));
 
-    /* Early exit when no hooks are registered */
+    if (!ts) {
-    if (!should_audit(ts)) {
+        /* Audit hooks cannot be called with a NULL thread state */
         return 0;
     }
 
-    _Py_AuditHookEntry *e = _PyRuntime.audit_hook_head;
+    /* The current implementation cannot be called if tstate is not
+       the current Python thread state. */
+    assert(ts == _PyThreadState_GET());
+
+    /* Early exit when no hooks are registered */
+    PyInterpreterState *is = ts->interp;
+    if (!should_audit(is)) {
+        return 0;
+    }
+
+    PyObject *eventName = NULL;
+    PyObject *eventArgs = NULL;
+    PyObject *hooks = NULL;
+    PyObject *hook = NULL;
+    int res = -1;
+
     int dtrace = PyDTrace_AUDIT_ENABLED();
 
     PyObject *exc_type, *exc_value, *exc_tb;
-    if (ts) {
+    _PyErr_Fetch(ts, &exc_type, &exc_value, &exc_tb);
-        _PyErr_Fetch(ts, &exc_type, &exc_value, &exc_tb);
-    }
 
     /* Initialize event args now */
     if (argFormat && argFormat[0]) {
-        va_list args;
+        eventArgs = _Py_VaBuildValue_SizeT(argFormat, vargs);
-        va_start(args, argFormat);
-        eventArgs = _Py_VaBuildValue_SizeT(argFormat, args);
-        va_end(args);
         if (eventArgs && !PyTuple_Check(eventArgs)) {
             PyObject *argTuple = PyTuple_Pack(1, eventArgs);
             Py_DECREF(eventArgs);
             eventArgs = argTuple;
         }
-    } else {
+    }
+    else {
         eventArgs = PyTuple_New(0);
     }
     if (!eventArgs) {
@@ -196,6 +205,7 @@ PySys_Audit(const char *event, const char *argFormat, ...)
     }
 
     /* Call global hooks */
+    _Py_AuditHookEntry *e = is->runtime->audit_hook_head;
     for (; e; e = e->next) {
         if (e->hookCFunction(event, eventArgs, e->userData) < 0) {
             goto exit;
@@ -208,8 +218,7 @@ PySys_Audit(const char *event, const char *argFormat, ...)
     }
 
     /* Call interpreter hooks */
-    PyInterpreterState *is = ts ? ts->interp : NULL;
+    if (is->audit_hooks) {
-    if (is && is->audit_hooks) {
         eventName = PyUnicode_FromString(event);
         if (!eventName) {
             goto exit;
@@ -238,8 +247,8 @@ PySys_Audit(const char *event, const char *argFormat, ...)
                 ts->use_tracing = (ts->c_tracefunc || ts->c_profilefunc);
                 ts->tracing--;
             }
-            o = PyObject_CallFunctionObjArgs(hook, eventName,
+            PyObject* args[2] = {eventName, eventArgs};
-                                             eventArgs, NULL);
+            o = _PyObject_FastCallTstate(ts, hook, args, 2);
             if (canTrace) {
                 ts->tracing++;
                 ts->use_tracing = 0;
@@ -265,33 +274,66 @@ exit:
     Py_XDECREF(eventName);
     Py_XDECREF(eventArgs);
 
-    if (ts) {
+    if (!res) {
-        if (!res) {
+        _PyErr_Restore(ts, exc_type, exc_value, exc_tb);
-            _PyErr_Restore(ts, exc_type, exc_value, exc_tb);
+    }
-        } else {
+    else {
-            assert(_PyErr_Occurred(ts));
+        assert(_PyErr_Occurred(ts));
-            Py_XDECREF(exc_type);
+        Py_XDECREF(exc_type);
-            Py_XDECREF(exc_value);
+        Py_XDECREF(exc_value);
-            Py_XDECREF(exc_tb);
+        Py_XDECREF(exc_tb);
-        }
     }
 
     return res;
 }
 
+int
+_PySys_Audit(PyThreadState *tstate, const char *event,
+             const char *argFormat, ...)
+{
+    va_list vargs;
+#ifdef HAVE_STDARG_PROTOTYPES
+    va_start(vargs, argFormat);
+#else
+    va_start(vargs);
+#endif
+    int res = sys_audit_tstate(tstate, event, argFormat, vargs);
+    va_end(vargs);
+    return res;
+}
+
+int
+PySys_Audit(const char *event, const char *argFormat, ...)
+{
+    PyThreadState *tstate = _PyThreadState_GET();
+    va_list vargs;
+#ifdef HAVE_STDARG_PROTOTYPES
+    va_start(vargs, argFormat);
+#else
+    va_start(vargs);
+#endif
+    int res = sys_audit_tstate(tstate, event, argFormat, vargs);
+    va_end(vargs);
+    return res;
+}
+
 /* We expose this function primarily for our own cleanup during
  * finalization. In general, it should not need to be called,
- * and as such it is not defined in any header files.
+ * and as such the function is not exported.
- */
+ *
+ * Must be finalizing to clear hooks */
 void
-_PySys_ClearAuditHooks(void)
+_PySys_ClearAuditHooks(PyThreadState *ts)
 {
-    /* Must be finalizing to clear hooks */
+    assert(ts != NULL);
-    _PyRuntimeState *runtime = &_PyRuntime;
+    if (!ts) {
-    PyThreadState *ts = _PyRuntimeState_GetThreadState(runtime);
+        return;
+    }
+
+    _PyRuntimeState *runtime = ts->interp->runtime;
     PyThreadState *finalizing = _PyRuntimeState_GetFinalizing(runtime);
-    assert(!ts || finalizing == ts);
+    assert(finalizing == ts);
-    if (!ts || finalizing != ts) {
+    if (finalizing != ts) {
         return;
     }
 
@@ -302,11 +344,11 @@ _PySys_ClearAuditHooks(void)
 
     /* Hooks can abort later hooks for this event, but cannot
        abort the clear operation itself. */
-    PySys_Audit("cpython._PySys_ClearAuditHooks", NULL);
+    _PySys_Audit(ts, "cpython._PySys_ClearAuditHooks", NULL);
     _PyErr_Clear(ts);
 
-    _Py_AuditHookEntry *e = _PyRuntime.audit_hook_head, *n;
+    _Py_AuditHookEntry *e = runtime->audit_hook_head, *n;
-    _PyRuntime.audit_hook_head = NULL;
+    runtime->audit_hook_head = NULL;
     while (e) {
         n = e->next;
         PyMem_RawFree(e);
@@ -317,13 +359,21 @@ _PySys_ClearAuditHooks(void)
 int
 PySys_AddAuditHook(Py_AuditHookFunction hook, void *userData)
 {
+    /* tstate can be NULL, so access directly _PyRuntime:
+       PySys_AddAuditHook() can be called before Python is initialized. */
     _PyRuntimeState *runtime = &_PyRuntime;
-    PyThreadState *tstate = _PyRuntimeState_GetThreadState(runtime);
+    PyThreadState *tstate;
+    if (runtime->initialized) {
+        tstate = _PyRuntimeState_GetThreadState(runtime);
+    }
+    else {
+        tstate = NULL;
+    }
 
     /* Invoke existing audit hooks to allow them an opportunity to abort. */
     /* Cannot invoke hooks until we are initialized */
-    if (runtime->initialized) {
+    if (tstate != NULL) {
-        if (PySys_Audit("sys.addaudithook", NULL) < 0) {
+        if (_PySys_Audit(tstate, "sys.addaudithook", NULL) < 0) {
             if (_PyErr_ExceptionMatches(tstate, PyExc_RuntimeError)) {
                 /* We do not report errors derived from RuntimeError */
                 _PyErr_Clear(tstate);
@@ -333,10 +383,10 @@ PySys_AddAuditHook(Py_AuditHookFunction hook, void *userData)
         }
     }
 
-    _Py_AuditHookEntry *e = _PyRuntime.audit_hook_head;
+    _Py_AuditHookEntry *e = runtime->audit_hook_head;
     if (!e) {
         e = (_Py_AuditHookEntry*)PyMem_RawMalloc(sizeof(_Py_AuditHookEntry));
-        _PyRuntime.audit_hook_head = e;
+        runtime->audit_hook_head = e;
     } else {
         while (e->next) {
             e = e->next;
@@ -346,7 +396,7 @@ PySys_AddAuditHook(Py_AuditHookFunction hook, void *userData)
     }
 
     if (!e) {
-        if (runtime->initialized) {
+        if (tstate != NULL) {
             _PyErr_NoMemory(tstate);
         }
         return -1;
@@ -374,7 +424,7 @@ sys_addaudithook_impl(PyObject *module, PyObject *hook)
     PyThreadState *tstate = _PyThreadState_GET();
 
     /* Invoke existing audit hooks to allow them an opportunity to abort. */
-    if (PySys_Audit("sys.addaudithook", NULL) < 0) {
+    if (_PySys_Audit(tstate, "sys.addaudithook", NULL) < 0) {
         if (_PyErr_ExceptionMatches(tstate, PyExc_Exception)) {
             /* We do not report errors derived from Exception */
             _PyErr_Clear(tstate);
@@ -384,7 +434,6 @@ sys_addaudithook_impl(PyObject *module, PyObject *hook)
     }
 
     PyInterpreterState *is = tstate->interp;
-
     if (is->audit_hooks == NULL) {
         is->audit_hooks = PyList_New(0);
         if (is->audit_hooks == NULL) {
@@ -408,6 +457,7 @@ static PyObject *
 sys_audit(PyObject *self, PyObject *const *args, Py_ssize_t argc)
 {
     PyThreadState *tstate = _PyThreadState_GET();
+    assert(tstate != NULL);
 
     if (argc == 0) {
         _PyErr_SetString(tstate, PyExc_TypeError,
@@ -416,7 +466,7 @@ sys_audit(PyObject *self, PyObject *const *args, Py_ssize_t argc)
         return NULL;
     }
 
-    if (!should_audit(tstate)) {
+    if (!should_audit(tstate->interp)) {
         Py_RETURN_NONE;
     }
 
@@ -442,7 +492,7 @@ sys_audit(PyObject *self, PyObject *const *args, Py_ssize_t argc)
         return NULL;
     }
 
-    int res = PySys_Audit(event, "O", auditArgs);
+    int res = _PySys_Audit(tstate, event, "O", auditArgs);
     Py_DECREF(auditArgs);
 
     if (res < 0) {
@@ -1739,7 +1789,7 @@ sys__getframe_impl(PyObject *module, int depth)
     PyThreadState *tstate = _PyThreadState_GET();
     PyFrameObject *f = tstate->frame;
 
-    if (PySys_Audit("sys._getframe", "O", f) < 0) {
+    if (_PySys_Audit(tstate, "sys._getframe", "O", f) < 0) {
         return NULL;
     }