py/objmodule: Add a table of built-in modules with delegation.

This replaces the previous QSTR_null entry in the globals dict which could
leak out to Python (e.g. via iteration of mod.__dict__) and could lead to
crashes.

It results in smaller code size at the expense of turning a lookup into a
loop, but the list it is looping over likely only contains one or two
elements.

To allow a module to register its custom attr function it can use the new
`MP_REGISTER_MODULE_DELEGATION` macro.

This work was funded through GitHub Sponsors.

Signed-off-by: Jim Mussared <jim.mussared@gmail.com>

9 files changed, 67 insertions, 36 deletions

diff --git a/py/builtin.h b/py/builtin.h
index 81d0789802..57f275fb31 100644
--- a/py/builtin.h
+++ b/py/builtin.h
@@ -132,6 +132,8 @@ extern const mp_obj_module_t mp_module___main__;
 extern const mp_obj_module_t mp_module_builtins;
 extern const mp_obj_module_t mp_module_sys;
 
+void mp_module_sys_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest);
+
 // Modules needed by the parser when MICROPY_COMP_MODULE_CONST is enabled.
 extern const mp_obj_module_t mp_module_errno;
 extern const mp_obj_module_t mp_module_uctypes;
diff --git a/py/builtinhelp.c b/py/builtinhelp.c
index c60ed44e9c..9bd56cca09 100644
--- a/py/builtinhelp.c
+++ b/py/builtinhelp.c
@@ -151,12 +151,7 @@ STATIC void mp_help_print_obj(const mp_obj_t obj) {
     if (map != NULL) {
         for (uint i = 0; i < map->alloc; i++) {
             mp_obj_t key = map->table[i].key;
-            if (key != MP_OBJ_NULL
-                #if MICROPY_MODULE_ATTR_DELEGATION
-                // MP_MODULE_ATTR_DELEGATION_ENTRY entries have MP_QSTRnull as qstr key.
-                && key != MP_OBJ_NEW_QSTR(MP_QSTRnull)
-                #endif
-                ) {
+            if (key != MP_OBJ_NULL) {
                 mp_help_print_info_about_object(key, map->table[i].value);
             }
         }
diff --git a/py/makemoduledefs.py b/py/makemoduledefs.py
index 2f787905a9..5a1362a4c9 100644
--- a/py/makemoduledefs.py
+++ b/py/makemoduledefs.py
@@ -22,11 +22,16 @@ import io
 import argparse
 
 
-pattern = re.compile(
+register_pattern = re.compile(
     r"\s*(MP_REGISTER_MODULE|MP_REGISTER_EXTENSIBLE_MODULE)\(MP_QSTR_(.*?),\s*(.*?)\);",
     flags=re.DOTALL,
 )
 
+delegation_pattern = re.compile(
+    r"\s*(?:MP_REGISTER_MODULE_DELEGATION)\((.*?),\s*(.*?)\);",
+    flags=re.DOTALL,
+)
+
 
 def find_module_registrations(filename):
     """Find any MP_REGISTER_MODULE definitions in the provided file.
@@ -37,7 +42,8 @@ def find_module_registrations(filename):
     global pattern
 
     with io.open(filename, encoding="utf-8") as c_file_obj:
-        return set(re.findall(pattern, c_file_obj.read()))
+        c = c_file_obj.read()
+        return set(re.findall(register_pattern, c)), set(re.findall(delegation_pattern, c))
 
 
 def generate_module_table_header(modules):
@@ -50,7 +56,6 @@ def generate_module_table_header(modules):
     # Print header file for all external modules.
     mod_defs = set()
     extensible_mod_defs = set()
-    print("// Automatically generated by makemoduledefs.py.\n")
     for macro_name, module_name, obj_module in modules:
         mod_def = "MODULE_DEF_{}".format(module_name.upper())
         if macro_name == "MP_REGISTER_MODULE":
@@ -97,13 +102,27 @@ def generate_module_table_header(modules):
     print("// MICROPY_REGISTERED_EXTENSIBLE_MODULES")
 
 
+def generate_module_delegations(delegations):
+    print("\n#define MICROPY_MODULE_DELEGATIONS \\")
+    for obj_module, fun_name in delegations:
+        print(
+            "    {{ MP_ROM_PTR(&{obj_module}), {fun_name} }},".format(
+                obj_module=obj_module, fun_name=fun_name
+            )
+        )
+    print("// MICROPY_MODULE_DELEGATIONS")
+
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("file", nargs=1, help="file with MP_REGISTER_MODULE definitions")
     args = parser.parse_args()
 
-    modules = find_module_registrations(args.file[0])
+    print("// Automatically generated by makemoduledefs.py.\n")
+
+    modules, delegations = find_module_registrations(args.file[0])
     generate_module_table_header(sorted(modules))
+    generate_module_delegations(sorted(delegations))
 
 
 if __name__ == "__main__":
diff --git a/py/makeqstrdefs.py b/py/makeqstrdefs.py
index dd96513351..64249f76ce 100644
--- a/py/makeqstrdefs.py
+++ b/py/makeqstrdefs.py
@@ -93,7 +93,9 @@ def process_file(f):
     elif args.mode == _MODE_COMPRESS:
         re_match = re.compile(r'MP_COMPRESSED_ROM_TEXT\("([^"]*)"\)')
     elif args.mode == _MODE_MODULE:
-        re_match = re.compile(r"MP_REGISTER_(?:EXTENSIBLE_)?MODULE\(.*?,\s*.*?\);")
+        re_match = re.compile(
+            r"(?:MP_REGISTER_MODULE|MP_REGISTER_EXTENSIBLE_MODULE|MP_REGISTER_MODULE_DELEGATION)\(.*?,\s*.*?\);"
+        )
     elif args.mode == _MODE_ROOT_POINTER:
         re_match = re.compile(r"MP_REGISTER_ROOT_POINTER\(.*?\);")
     output = []
diff --git a/py/modsys.c b/py/modsys.c
index ddc732e005..a0ecb87b5f 100644
--- a/py/modsys.c
+++ b/py/modsys.c
@@ -207,7 +207,7 @@ STATIC const uint16_t sys_mutable_keys[] = {
     MP_QSTRnull,
 };
 
-STATIC void mp_module_sys_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
+void mp_module_sys_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
     MP_STATIC_ASSERT(MP_ARRAY_SIZE(sys_mutable_keys) == MP_SYS_MUTABLE_NUM + 1);
     MP_STATIC_ASSERT(MP_ARRAY_SIZE(MP_STATE_VM(sys_mutable)) == MP_SYS_MUTABLE_NUM);
     mp_module_generic_attr(attr, dest, sys_mutable_keys, MP_STATE_VM(sys_mutable));
@@ -280,11 +280,6 @@ STATIC const mp_rom_map_elem_t mp_module_sys_globals_table[] = {
     #if MICROPY_PY_SYS_ATEXIT
     { MP_ROM_QSTR(MP_QSTR_atexit), MP_ROM_PTR(&mp_sys_atexit_obj) },
     #endif
-
-    #if MICROPY_PY_SYS_ATTR_DELEGATION
-    // Delegation of attr lookup.
-    MP_MODULE_ATTR_DELEGATION_ENTRY(&mp_module_sys_attr),
-    #endif
 };
 
 STATIC MP_DEFINE_CONST_DICT(mp_module_sys_globals, mp_module_sys_globals_table);
@@ -317,6 +312,7 @@ MP_REGISTER_ROOT_POINTER(mp_obj_t sys_exitfunc);
 #if MICROPY_PY_SYS_ATTR_DELEGATION
 // Contains mutable sys attributes.
 MP_REGISTER_ROOT_POINTER(mp_obj_t sys_mutable[MP_SYS_MUTABLE_NUM]);
+MP_REGISTER_MODULE_DELEGATION(mp_module_sys, &mp_module_sys_attr);
 #endif
 
 #endif // MICROPY_PY_SYS
diff --git a/py/mpconfig.h b/py/mpconfig.h
index a004f548ad..504ad64b01 100644
--- a/py/mpconfig.h
+++ b/py/mpconfig.h
@@ -839,8 +839,8 @@ typedef double mp_float_t;
 #define MICROPY_STREAMS_POSIX_API (0)
 #endif
 
-// Whether modules can use MP_MODULE_ATTR_DELEGATION_ENTRY() to delegate failed
-// attribute lookups.
+// Whether modules can use MP_REGISTER_MODULE_DELEGATION() to delegate failed
+// attribute lookups to a custom handler function.
 #ifndef MICROPY_MODULE_ATTR_DELEGATION
 #define MICROPY_MODULE_ATTR_DELEGATION (MICROPY_CONFIG_ROM_LEVEL_AT_LEAST_EXTRA_FEATURES)
 #endif
diff --git a/py/obj.h b/py/obj.h
index 005383e9f3..de4c15b22c 100644
--- a/py/obj.h
+++ b/py/obj.h
@@ -437,6 +437,10 @@ typedef struct _mp_rom_obj_t { mp_const_obj_t o; } mp_rom_obj_t;
 // As above, but allow this module to be extended from the filesystem.
 #define MP_REGISTER_EXTENSIBLE_MODULE(module_name, obj_module)
 
+// Add a custom handler for a builtin module that will be called to delegate
+// failed attribute lookups.
+#define MP_REGISTER_MODULE_DELEGATION(obj_module, fun_name)
+
 // Declare a root pointer (to avoid garbage collection of a global static variable).
 // param variable_declaration: a valid C variable declaration
 #define MP_REGISTER_ROOT_POINTER(variable_declaration)
diff --git a/py/objmodule.c b/py/objmodule.c
index bf38366266..ad96a3be28 100644
--- a/py/objmodule.c
+++ b/py/objmodule.c
@@ -63,20 +63,7 @@ STATIC void module_print(const mp_print_t *print, mp_obj_t self_in, mp_print_kin
     mp_printf(print, "<module '%s'>", module_name);
 }
 
-STATIC void module_attr_try_delegation(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
-    #if MICROPY_MODULE_ATTR_DELEGATION
-    // Delegate lookup to a module's custom attr method (found in last lot of globals dict).
-    mp_obj_module_t *self = MP_OBJ_TO_PTR(self_in);
-    mp_map_t *map = &self->globals->map;
-    if (map->table[map->alloc - 1].key == MP_OBJ_NEW_QSTR(MP_QSTRnull)) {
-        ((mp_attr_fun_t)MP_OBJ_TO_PTR(map->table[map->alloc - 1].value))(self_in, attr, dest);
-    }
-    #else
-    (void)self_in;
-    (void)attr;
-    (void)dest;
-    #endif
-}
+STATIC void module_attr_try_delegation(mp_obj_t self_in, qstr attr, mp_obj_t *dest);
 
 STATIC void module_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
     mp_obj_module_t *self = MP_OBJ_TO_PTR(self_in);
@@ -177,6 +164,18 @@ STATIC const mp_rom_map_elem_t mp_builtin_extensible_module_table[] = {
 };
 MP_DEFINE_CONST_MAP(mp_builtin_extensible_module_map, mp_builtin_extensible_module_table);
 
+#if MICROPY_MODULE_ATTR_DELEGATION && defined(MICROPY_MODULE_DELEGATIONS)
+typedef struct _mp_module_delegation_entry_t {
+    mp_rom_obj_t mod;
+    mp_attr_fun_t fun;
+} mp_module_delegation_entry_t;
+
+STATIC const mp_module_delegation_entry_t mp_builtin_module_delegation_table[] = {
+    // delegation entries declared with MP_REGISTER_MODULE_DELEGATION()
+    MICROPY_MODULE_DELEGATIONS
+};
+#endif
+
 // Attempts to find (and initialise) a built-in, otherwise returns
 // MP_OBJ_NULL.
 mp_obj_t mp_module_get_builtin(qstr module_name, bool extensible) {
@@ -230,6 +229,23 @@ mp_obj_t mp_module_get_builtin(qstr module_name, bool extensible) {
     return elem->value;
 }
 
+STATIC void module_attr_try_delegation(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
+    #if MICROPY_MODULE_ATTR_DELEGATION
+    // Delegate lookup to a module's custom attr method.
+    size_t n = MP_ARRAY_SIZE(mp_builtin_module_delegation_table);
+    for (size_t i = 0; i < n; ++i) {
+        if (*(mp_obj_t *)(&mp_builtin_module_delegation_table[i].mod) == self_in) {
+            mp_builtin_module_delegation_table[i].fun(self_in, attr, dest);
+            break;
+        }
+    }
+    #else
+    (void)self_in;
+    (void)attr;
+    (void)dest;
+    #endif
+}
+
 void mp_module_generic_attr(qstr attr, mp_obj_t *dest, const uint16_t *keys, mp_obj_t *values) {
     for (size_t i = 0; keys[i] != MP_QSTRnull; ++i) {
         if (attr == keys[i]) {
diff --git a/py/objmodule.h b/py/objmodule.h
index 63ae3c3bdf..9cc9a2f102 100644
--- a/py/objmodule.h
+++ b/py/objmodule.h
@@ -28,9 +28,6 @@
 
 #include "py/obj.h"
 
-// Place at the very end of a module's globals_table.
-#define MP_MODULE_ATTR_DELEGATION_ENTRY(ptr) { MP_ROM_QSTR(MP_QSTRnull), MP_ROM_PTR(ptr) }
-
 extern const mp_map_t mp_builtin_module_map;
 extern const mp_map_t mp_builtin_extensible_module_map;