From 209936880df4d63730ea77cf4b6f84e5f2f56d7c Mon Sep 17 00:00:00 2001
From: Damien George <damien.p.george@gmail.com>
Date: Tue, 20 Feb 2018 18:00:44 +1100
Subject: [PATCH] py/builtinimport: Add compile-time option to disable external
 imports.

The new option is MICROPY_ENABLE_EXTERNAL_IMPORT and is enabled by default
so that the default behaviour is the same as before.  With it disabled
import is only supported for built-in modules, not for external files nor
frozen modules.  This allows to support targets that have no filesystem of
any kind and that only have access to pre-supplied built-in modules
implemented natively.
---
 py/builtinimport.c | 39 +++++++++++++++++++++++++++++++++++++++
 py/mpconfig.h      |  7 +++++++
 py/runtime.c       |  9 +++++++++
 3 files changed, 55 insertions(+)

diff --git a/py/builtinimport.c b/py/builtinimport.c
index 19bc9fc954..b8ed096caf 100644
--- a/py/builtinimport.c
+++ b/py/builtinimport.c
@@ -44,6 +44,8 @@
 #define DEBUG_printf(...) (void)0
 #endif
 
+#if MICROPY_ENABLE_EXTERNAL_IMPORT
+
 #define PATH_SEP_CHAR '/'
 
 bool mp_obj_is_package(mp_obj_t module) {
@@ -473,4 +475,41 @@ mp_obj_t mp_builtin___import__(size_t n_args, const mp_obj_t *args) {
     // Otherwise, we need to return top-level package
     return top_module_obj;
 }
+
+#else // MICROPY_ENABLE_EXTERNAL_IMPORT
+
+mp_obj_t mp_builtin___import__(size_t n_args, const mp_obj_t *args) {
+    // Check that it's not a relative import
+    if (n_args >= 5 && MP_OBJ_SMALL_INT_VALUE(args[4]) != 0) {
+        mp_raise_NotImplementedError("relative import");
+    }
+
+    // Check if module already exists, and return it if it does
+    qstr module_name_qstr = mp_obj_str_get_qstr(args[0]);
+    mp_obj_t module_obj = mp_module_get(module_name_qstr);
+    if (module_obj != MP_OBJ_NULL) {
+        return module_obj;
+    }
+
+    #if MICROPY_MODULE_WEAK_LINKS
+    // Check if there is a weak link to this module
+    mp_map_elem_t *el = mp_map_lookup((mp_map_t*)&mp_builtin_module_weak_links_map, MP_OBJ_NEW_QSTR(module_name_qstr), MP_MAP_LOOKUP);
+    if (el != NULL) {
+        // Found weak-linked module
+        mp_module_call_init(module_name_qstr, el->value);
+        return el->value;
+    }
+    #endif
+
+    // Couldn't find the module, so fail
+    if (MICROPY_ERROR_REPORTING == MICROPY_ERROR_REPORTING_TERSE) {
+        mp_raise_msg(&mp_type_ImportError, "module not found");
+    } else {
+        nlr_raise(mp_obj_new_exception_msg_varg(&mp_type_ImportError,
+            "no module named '%q'", module_name_qstr));
+    }
+}
+
+#endif // MICROPY_ENABLE_EXTERNAL_IMPORT
+
 MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(mp_builtin___import___obj, 1, 5, mp_builtin___import__);
diff --git a/py/mpconfig.h b/py/mpconfig.h
index b8a96f0b0a..e9071df2f7 100644
--- a/py/mpconfig.h
+++ b/py/mpconfig.h
@@ -405,6 +405,13 @@
 /*****************************************************************************/
 /* Python internal features                                                  */
 
+// Whether to enable import of external modules
+// When disabled, only importing of built-in modules is supported
+// When enabled, a port must implement mp_import_stat (among other things)
+#ifndef MICROPY_ENABLE_EXTERNAL_IMPORT
+#define MICROPY_ENABLE_EXTERNAL_IMPORT (1)
+#endif
+
 // Whether to use the POSIX reader for importing files
 #ifndef MICROPY_READER_POSIX
 #define MICROPY_READER_POSIX (0)
diff --git a/py/runtime.c b/py/runtime.c
index 9dff9847a0..65d0df639e 100644
--- a/py/runtime.c
+++ b/py/runtime.c
@@ -1337,6 +1337,8 @@ import_error:
         return dest[0];
     }
 
+    #if MICROPY_ENABLE_EXTERNAL_IMPORT
+
     // See if it's a package, then can try FS import
     if (!mp_obj_is_package(module)) {
         goto import_error;
@@ -1363,6 +1365,13 @@ import_error:
 
     // TODO lookup __import__ and call that instead of going straight to builtin implementation
     return mp_builtin___import__(5, args);
+
+    #else
+
+    // Package import not supported with external imports disabled
+    goto import_error;
+
+    #endif
 }
 
 void mp_import_all(mp_obj_t module) {