Fix: Stop TDI execution from editing caller's Python namespace

python/MDSplus/tests/data_case.py

@@ -26,6 +26,8 @@
 from numpy import ndarray, array, int32
 from math import log10
 import MDSplus as m
+import sys
+import __main__
 
 
 def _mimport(name, level=1):
@@ -42,7 +44,7 @@ class Tests(_common.Tests):
     TESTS = {
         'data', 'scalars', 'arrays', 'vms',
         'tdi', 'decompile', 'casts', 'tdipy',
-        'deserialize',
+        'deserialize', 'tdipy_namespace',
     }
 
     def _doThreeTest(self, tdiexpr, pyexpr, ans, **kwargs):
@@ -215,7 +217,7 @@ def doTest(suffix, cl, scl, ucl, **kwargs):
                         cl(0.08748866), # tand(a) (degrees)
                         cl(2.0), # anint(a/3)
                     ])
-            
+
             if real:
                 results.extend([
                     cl(1), # a mod b
@@ -224,7 +226,7 @@ def doTest(suffix, cl, scl, ucl, **kwargs):
                     False, # a < b
                     False, # a <= b
                 ])
-            
+
             m.Data.execute('_a=5%s,_b=2%s' % tuple([suffix]*2))
             a, b = cl(5), cl(2)
             with warnings.catch_warnings():
@@ -275,9 +277,9 @@ def results(cl, scl, ucl):
                     scl([5, 4, 2]), # abs1(-a)
                     scl([25, 16, 4]), # abssq(-a)
                 ]
-            
+
                 if almost:
-                    results.extend([    
+                    results.extend([
                         cl([148.413159,  54.5981500,  7.38905610]),  # exp(a)
                         cl([1.60943791,  1.38629436,  0.69314718]),  # log(a)
                         cl([-0.95892427, -0.7568025,  0.90929743]),  # sin(a)
@@ -298,7 +300,7 @@ def results(cl, scl, ucl):
                             cl([0.08748866,  0.06992681,  0.03492077]),  # tand(a) (degrees)
                             cl([2, 1, 1]),  # anint(a/3)
                         ])
-                
+
                 if real:
                     results.extend([
                         cl([1, 1, 2]), # a mod b
@@ -307,7 +309,7 @@ def results(cl, scl, ucl):
                         array([False, False, True]), # a < b
                         array([False, False, True]), # a <= b
                     ])
-                
+
                 return results
             """ test array """
             m.Data.execute('_a=[5%s,4%s,2%s],_b=[2%s,3%s,5%s]' %
@@ -526,5 +528,39 @@ def tdipy(self):
         if not self.inThread:
             self._doTdiTest("TEST()", m.Array([1, 2]))
 
+    def tdipy_namespace(self):
+        """
+        All Python TDI functions should run in their own namespaces.
+        Previously, they would overwrite variables in the caller's namespace.
+        """
+        # Py.py assigns _tb, MDSplus and Py at module level.
+        # These assignments should not escape to the caller
+        sentinel = object()
+        __main__._tb = sentinel
+        try:
+            before = set(vars(__main__))
+            file_before = __file__
+            self.assertTrue(file_before.endswith("data_case.py"), file_before)
+            self._doTdiTest("Py('a=1','a')", 1)
+            self.assertIs(__main__._tb, sentinel, "Sentinel value was overwritten")
+            self.assertEqual(set(vars(__main__)) - before, set())
+            self.assertEqual(__file__, file_before)
+        finally:
+            del __main__._tb
+        # Inside the TDI function, __file__ should refer to the source for that function
+        # however, it should also not escape to the caller's globals
+        internal_file = str(m.Data.execute(
+            "Py('import sys; a=sys.modules[\"tdi_functions.Py\"].__file__','a')"
+        ))
+        self.assertNotEqual(internal_file, file_before)
+        self.assertTrue(internal_file.endswith("Py.py"), internal_file)
+        self.assertEqual(__file__, file_before)
+        # All TDI functions should also be isolated from each other,
+        # so the _tb variable assigned in Py.py should not be visible in pyfun.py
+        self._doTdiTest("Py('a=123','a')", m.Int32(123))
+        self._doTdiTest("pyfun('str',*,123)", m.String("123"))
+        self.assertIn("_tb", vars(sys.modules["tdi_functions.Py"]))
+        self.assertNotIn("_tb", vars(sys.modules["tdi_functions.pyfun"]))
+
 
 Tests.main()

tdishr/TdiExtPython.c

@@ -90,6 +90,9 @@ static int (*PyCallable_Check)() = NULL;
 static void (*PyErr_Clear)() = NULL;
 static PyObject *(*PyImport_AddModule)() = NULL;
 static PyObject *(*PyModule_AddObject)() = NULL;
+static int (*PyObject_SetAttrString)() = NULL;
+static PyObject *(*PyModule_GetDict)() = NULL;
+static int (*PyDict_SetItemString)() = NULL;
 #ifdef MACOS_ARM64
 static PyObject *(*PyObject_CallFunctionObjArgs)(void *, ...) = NULL;
 #else
@@ -113,10 +116,10 @@ static PyObject *(*PyObject_CallFunction)() = NULL;
 static PyObject *(*PyTuple_New)() = NULL;
 static void (*PyTuple_SetItem)() = NULL;
 #ifdef USE_EXECFILE
-static int (*PyRun_SimpleStringFlags)() = NULL;
+static PyObject *(*PyRun_StringFlags)() = NULL;
 #else
 static FILE *(*_Py_fopen_obj)() = NULL;
-static int (*PyRun_SimpleFileExFlags)() = NULL;
+static PyObject *(*PyRun_FileExFlags)() = NULL;
 #endif
 static PyObject *pointerToObject = NULL;
 static PyObject *makeData = NULL;
@@ -139,7 +142,7 @@ inline static void initialize()
     envsym = "python2.7";
 #ifdef MACOS_ARM64
     const char *aspath = "/opt/local/lib/libpython2.7.dylib";   // (MW) TODO: for MacPorts version
-#else 
+#else
     const char *aspath = "/usr/lib/python2.7.so.1";
 #endif
     setenv("PyLib", envsym, B_FALSE);
@@ -231,6 +234,9 @@ inline static void initialize()
   loadrtn(PyErr_Clear, 1);
   loadrtn(PyImport_AddModule, 1);
   loadrtn(PyModule_AddObject, 1);
+  loadrtn(PyObject_SetAttrString, 1);
+  loadrtn(PyModule_GetDict, 1);
+  loadrtn(PyDict_SetItemString, 1);
   loadrtn(PyObject_CallFunctionObjArgs, 1);
   loadrtn(PyString_FromString, 0);
   if (!PyString_FromString)
@@ -248,10 +254,10 @@ inline static void initialize()
   loadrtn(PyTuple_New, 1);
   loadrtn(PyTuple_SetItem, 1);
 #ifdef USE_EXECFILE
-  loadrtn(PyRun_SimpleStringFlags, 1);
+  loadrtn(PyRun_StringFlags, 1);
 #else
   loadrtn(_Py_fopen_obj, 0);
-  loadrtn(PyRun_SimpleFileExFlags, 1);
+  loadrtn(PyRun_FileExFlags, 1);
 #endif
   loadrtn(PyGILState_Check, 0);
   loadrtn(PyGILState_Release, 1);
@@ -435,25 +441,28 @@ static inline int is_callable(const PyObject *const fun,
 }
 
 #ifdef USE_EXECFILE
-static inline PyObject *get_exec_file(const PyObject *const __main__)
+static inline PyObject *get_exec_file(const PyObject *const module,
+                                      PyObject *const globals)
 {
-  PyObject *execfile = PyObject_GetAttrString(__main__, "execfile");
+  PyObject *execfile = PyObject_GetAttrString(module, "_mds_execfile");
   if (!execfile)
   { // not defined yet, so we define it
     PyErr_Clear();
-    char def[] = "import __main__\ndef execfile(filename):\n with "
+    char def[] = "def _mds_execfile(filename,ns):\n with "
                  "open(filename,'r') as f:\n  "
-                 "exec(compile(f.read(),filename,'exec'),__main__.__dict__,__"
-                 "main__.__dict__)\0";
+                 "exec(compile(f.read(),filename,'exec'),ns,ns)\0";
     int flags = 0;
-    if (PyRun_SimpleStringFlags(def, &flags))
+    PyObject *ans = PyRun_StringFlags(def, Py_file_input, globals, globals,
+                                      &flags);
+    if (!ans)
     {
-      fprintf(stderr, "Error defining execfile\n");
+      fprintf(stderr, "Error defining _mds_execfile\n");
       if (PyErr_Occurred())
         PyErr_Print();
       return NULL;
     }
-    execfile = PyObject_GetAttrString(__main__, "execfile");
+    Py_DecRef(ans);
+    execfile = PyObject_GetAttrString(module, "_mds_execfile");
   }
   return execfile;
 }
@@ -475,41 +484,74 @@ static inline void add__file__fun(const PyObject *const tdi_functions,
   free(__file__fun);
 }
 
-static inline int load_python_fun(const char *const fullpath,
-                                  char **const funname)
+// Each TDI python file gets its own module, named after the function it defines,
+// so that its module level names cannot collide with the caller's globals or with
+// another TDI function's.
+static inline PyObject *get_fun_module(const char *const funname)
 {
-  // get __main__
-  PyObject *__main__ = PyImport_AddModule("__main__");
-  if (!__main__)
+  static const char prefix[] = "tdi_functions.";
+  char *modname = malloc(sizeof(prefix) + strlen(funname));
+  strcpy(modname, prefix);
+  strcat(modname, funname);
+  PyObject *module = PyImport_AddModule(modname);
+  if (!module)
   {
-    fprintf(stderr, "Error getting __main__ module'\n");
+    fprintf(stderr, "Error getting module '%s'\n", modname);
     if (PyErr_Occurred())
       PyErr_Print();
+  }
+  free(modname);
+  return module;
+}
+
+static inline int load_python_fun(const char *const fullpath,
+                                  char **const funname)
+{
+  // The function name is the basename without the trailing '.py', and is needed
+  // before the file runs in order to name its namespace
+  const char *c, *p = fullpath;
+  for (; (c = strchr(p, '/')); p = c + 1)
+    ;
+#ifdef _WIN32
+  for (; (c = strchr(p, '\\')); p = c + 1)
+    ;
+#endif
+  const size_t mlen = strlen(p) - 3;
+  *funname = memcpy(malloc(mlen + 1), p, mlen);
+  funname[0][mlen] = '\0';
+  PyObject *module = get_fun_module(*funname);
+  if (!module)
+  {
+    free(*funname);
+    *funname = NULL;
     return MDSplusERROR;
   }
+  PyObject *globals = PyModule_GetDict(module);
 #ifdef USE_EXECFILE
-  PyObject *execfile = get_exec_file(__main__);
+  PyObject *execfile = get_exec_file(module, globals);
   if (!execfile)
+  {
+    free(*funname);
+    *funname = NULL;
     return MDSplusERROR;
+  }
 #endif
-  // add __file__=<fullpath> to globals
-  if (PyModule_AddObject(
-          __main__, "__file__",
-          PyString_FromString(fullpath)))
-  { // no need to deref PyString
+  PyObject *__file__ = PyString_FromString(fullpath);
+  // add __file__=<fullpath> to the namespace, so the file sees its own path
+  if (PyObject_SetAttrString(module, "__file__", __file__))
+  {
     fprintf(stderr, "Failed adding __file__='%s'\n", fullpath);
     if (PyErr_Occurred())
       PyErr_Print();
   }
-  PyObject *__file__ = PyString_FromString(fullpath);
 #ifdef USE_EXECFILE
-  PyObject *ans = PyObject_CallFunctionObjArgs(execfile, __file__, NULL);
+  PyObject *ans = PyObject_CallFunctionObjArgs(execfile, __file__, globals, NULL);
   Py_DecRef(execfile);
   Py_DecRef(ans);
   if (!ans)
   {
 #else
-  int err;
+  PyObject *ans;
   INIT_AND_FCLOSE_ON_EXIT(fp);
   if (_Py_fopen_obj)
     fp = _Py_fopen_obj(__file__, "r");
@@ -519,31 +561,26 @@ static inline int load_python_fun(const char *const fullpath,
   {
     fprintf(stderr, "Error opening file '%s'\n", fullpath);
     Py_DecRef(__file__);
+    free(*funname);
+    *funname = NULL;
     return MDSplusERROR;
   }
   int flags = 0;
-  err = PyRun_SimpleFileExFlags(fp, fullpath, 1, &flags);
+  ans = PyRun_FileExFlags(fp, fullpath, Py_file_input, globals, globals, 1, &flags);
   FCLOSE_CANCEL(fp);
-  if (err)
+  Py_DecRef(ans);
+  if (!ans)
   {
 #endif
     fprintf(stderr, "Error compiling file '%s'\n", fullpath);
     if (PyErr_Occurred())
       PyErr_Print();
     Py_DecRef(__file__);
+    free(*funname);
+    *funname = NULL;
     return TdiUNKNOWN_VAR;
   }
-  const char *c, *p = fullpath;
-  for (; (c = strchr(p, '/')); p = c + 1)
-    ;
-#ifdef _WIN32
-  for (; (c = strchr(p, '\\')); p = c + 1)
-    ;
-#endif
-  const size_t mlen = strlen(p) - 3;
-  *funname = memcpy(malloc(mlen + 1), p, mlen);
-  funname[0][mlen] = '\0';
-  PyObject *pyFunction = PyObject_GetAttrString(__main__, *funname);
+  PyObject *pyFunction = PyObject_GetAttrString(module, *funname);
   if (!is_callable(pyFunction, *funname, fullpath))
   {
     free(*funname);
@@ -568,27 +605,13 @@ static inline int call_python_fun(const char *const filename, const int nargs,
                                   mdsdsc_xd_t *const out_ptr)
 {
   PyObject *tdi_functions = PyImport_AddModule("tdi_functions");
-  if (tdi_functions)
+  if (!tdi_functions)
   {
-    PyObject *__main__ = PyImport_AddModule("__main__");
-    char *__file__fun = malloc(strlen(filename) + 9);
-    strcpy(__file__fun, "__file__");
-    strcat(__file__fun, filename);
-    PyObject *__file__ = PyObject_GetAttrString(tdi_functions, __file__fun);
-    free(__file__fun);
-    if (__file__)
-      PyModule_AddObject(__main__, "__file__", __file__);
-    else
-    { // silently fail and set __file__ to None
-      PyModule_AddObject(__main__, "__file__", Py_None);
-      if (PyErr_Occurred())
-        PyErr_Clear();
-    }
-  }
-  else
     fprintf(stderr, "Failed getting module tdi_functions\n");
-  if (PyErr_Occurred())
-    PyErr_Print();
+    if (PyErr_Occurred())
+      PyErr_Print();
+    return MDSplusERROR;
+  }
   if ((strcasecmp("py", filename) == 0) && (MdsSandboxEnabled() == 1))
     return MDSplusSANDBOX;
   PyObject *pyFunction = PyObject_GetAttrString(tdi_functions, filename);
@@ -598,6 +621,24 @@ static inline int call_python_fun(const char *const filename, const int nargs,
       PyErr_Print();
     return MDSplusERROR;
   }
+  // Point __file__ at the function's source file.
+  char *__file__fun = malloc(strlen(filename) + 9);
+  strcpy(__file__fun, "__file__");
+  strcat(__file__fun, filename);
+  PyObject *__file__ = PyObject_GetAttrString(tdi_functions, __file__fun);
+  free(__file__fun);
+  if (__file__)
+  {
+    PyObject *fun_globals = PyObject_GetAttrString(pyFunction, "__globals__");
+    if (fun_globals)
+    {
+      PyDict_SetItemString(fun_globals, "__file__", __file__);
+      Py_DecRef(fun_globals);
+    }
+    Py_DecRef(__file__);
+  }
+  if (PyErr_Occurred())
+    PyErr_Clear();
   PyObject *pyArgs = args_to_tuple(nargs, args);
   PyObject *ans = PyObject_CallObject(pyFunction, pyArgs);
   Py_DecRef(pyFunction);