impl registry_keys

2 files changed, 135 insertions, 73 deletions

diff --git a/core/src/xmake/winos/registry_keys.c b/core/src/xmake/winos/registry_keys.c
index 42ad24f98..234cdf035 100644
--- a/core/src/xmake/winos/registry_keys.c
+++ b/core/src/xmake/winos/registry_keys.c
@@ -31,124 +31,114 @@
 #include "prefix.h"
 
 /* //////////////////////////////////////////////////////////////////////////////////////
- * private implementation
+ * types
  */
 
-/* //////////////////////////////////////////////////////////////////////////////////////
- * implementation
- */
-
-/* get registry keys
- *
- * local count, errors = winos.registry_keys("HKEY_LOCAL_MACHINE",
- *                                             "SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion\\AeDebug",
- *                                             function (key_path)
- *                                                 return true -- continue or break
- *                                             end)
- */
-tb_int_t xm_winos_registry_keys(lua_State* lua)
+// the enum info type
+typedef struct __xm_winos_registry_enum_info_t
 {
-    // check
-    tb_assert_and_check_return_val(lua, 0);
+    lua_State*          lua;
+    HKEY                key;
+    tb_int_t            ok;
+    tb_char_t const*    error;
+    tb_int_t            count;
+    tb_wchar_t          key_name[1024];
+    tb_char_t           key_path_a[TB_PATH_MAXN];
 
-    // get the arguments
-    tb_char_t const* rootkey = luaL_checkstring(lua, 1);
-    tb_char_t const* rootdir = luaL_checkstring(lua, 2);
-    tb_bool_t is_function    = lua_isfunction(lua, 3);
-    tb_check_return_val(rootkey && rootdir && is_function, 0);
+}xm_winos_registry_enum_info_t;
 
-    // query key-value
-    tb_bool_t   ok = tb_false;
-    tb_int_t    count = 0;
-    HKEY        key = tb_null;
+/* //////////////////////////////////////////////////////////////////////////////////////
+ * private implementation
+ */
+static tb_void_t xm_winos_registry_enum_keys(xm_winos_registry_enum_info_t* info, tb_wchar_t const* rootdir, tb_long_t recursion)
+{
+    // enum keys
     HKEY        keynew = tb_null;
-    tb_char_t*  value = tb_null;
+    lua_State*  lua = info->lua;
+    tb_wchar_t* key_path = tb_null;
+    tb_size_t   key_path_maxn = TB_PATH_MAXN;
     do
     {
-        // get registry rootkey
-        if (!tb_strcmp(rootkey, "HKEY_CLASSES_ROOT"))         key = HKEY_CLASSES_ROOT;
-        else if (!tb_strcmp(rootkey, "HKEY_CURRENT_CONFIG"))  key = HKEY_CURRENT_CONFIG;
-        else if (!tb_strcmp(rootkey, "HKEY_CURRENT_USER"))    key = HKEY_CURRENT_USER;
-        else if (!tb_strcmp(rootkey, "HKEY_LOCAL_MACHINE"))   key = HKEY_LOCAL_MACHINE;
-        else if (!tb_strcmp(rootkey, "HKEY_USERS"))           key = HKEY_USERS;
-        else
+        // open registry key
+        if (RegOpenKeyExW(info->key, rootdir, 0, KEY_READ, &keynew) != ERROR_SUCCESS && keynew)
         {
-            lua_pushnil(lua);
-            lua_pushfstring(lua, "invalid registry rootkey: %s", rootkey);
+            info->ok = -1;
+            info->error = "open registry key failed";
             break;
         }
 
-        // open registry key
-        if (RegOpenKeyExA(key, rootdir, 0, KEY_READ, &keynew) != ERROR_SUCCESS && keynew)
+        // query keys
+        DWORD key_name_num = 0;
+        DWORD key_name_maxn = 0;
+        if (RegQueryInfoKeyW(keynew, tb_null, tb_null, tb_null, &key_name_num, &key_name_maxn, tb_null, tb_null, tb_null, tb_null, tb_null, tb_null) != ERROR_SUCCESS)
         {
-            lua_pushnil(lua);
-            lua_pushfstring(lua, "open registry key failed: %s\\%s", rootkey, rootdir);
+            // cannot query this key, we ignore it
             break;
         }
+        key_name_maxn++; // add `\0`
 
-        // query keys
-        DWORD key_path_num = 0;
-        DWORD key_path_maxn = 0;
-        if (RegQueryInfoKeyW(keynew, tb_null, tb_null, tb_null, &key_path_num, &key_path_maxn, tb_null, tb_null, tb_null, tb_null, tb_null, tb_null) != ERROR_SUCCESS)
+        // ensure enough key path buffer
+        if (key_name_maxn > tb_arrayn(info->key_name))
         {
-            lua_pushnil(lua);
-            lua_pushfstring(lua, "query registry info failed: %s\\%s", rootkey, rootdir);
+            info->ok = -1;
+            info->error = "no enough key path buffer";
             break;
         }
-        key_path_maxn++; // add `\0`
 
-        // ensure enough key path buffer
-        tb_wchar_t key_path[TB_PATH_MAXN];
-        if (key_path_maxn > tb_arrayn(key_path))
+        // init key path
+        key_path = (tb_wchar_t*)tb_malloc(key_path_maxn * sizeof(tb_wchar_t));
+        if (!key_path)
         {
-            lua_pushnil(lua);
-            lua_pushfstring(lua, "no enough key path buffer: %s\\%s", rootkey, rootdir);
+            info->ok = -1;
+            info->error = "no enough key path buffer";
             break;
         }
 
         // get all keys
         DWORD i = 0;
-        tb_char_t key_path_a[TB_PATH_MAXN];
-        for (i = 0; i < key_path_num; i++)
+        for (i = 0; i < key_name_num && info->ok > 0; i++)
         {
-            // get key path
-            key_path[0] = L'\0';
-            DWORD key_path_size = tb_arrayn(key_path);
-            if (RegEnumKeyExW(keynew, i, key_path, &key_path_size, tb_null, tb_null, tb_null, tb_null) != ERROR_SUCCESS)
+            // get key name
+            info->key_name[0] = L'\0';
+            DWORD key_name_size = tb_arrayn(info->key_name);
+            if (RegEnumKeyExW(keynew, i, info->key_name, &key_name_size, tb_null, tb_null, tb_null, tb_null) != ERROR_SUCCESS)
             {
-                lua_pushnil(lua);
-                lua_pushfstring(lua, "get registry key path(%d) failed: %s\\%s", i, rootkey, rootdir);
+                info->ok = -1;
+                info->error = "get registry key failed";
                 break;
             }
 
+            // get key path
+            tb_swprintf(key_path, key_path_maxn, L"%s\\%s", rootdir, info->key_name);
+
             // get key path (mbs)
-            tb_size_t key_path_a_size = tb_wtoa(key_path_a, key_path, sizeof(key_path_a));
+            tb_size_t key_path_a_size = tb_wtoa(info->key_path_a, key_path, sizeof(info->key_path_a));
             if (key_path_a_size == -1)
             {
-                lua_pushnil(lua);
-                lua_pushfstring(lua, "convert registry key path(%d) failed: %s\\%s", i, rootkey, rootdir);
+                info->ok = -1;
+                info->error = "convert registry key path failed";
                 break;
             }
 
-            // do callback(key_path)
-            lua_pushvalue(lua, 3);
-            lua_pushlstring(lua, key_path_a, key_path_a_size);
+            // do callback(key_name)
+            lua_pushvalue(lua, 4);
+            lua_pushlstring(lua, info->key_path_a, key_path_a_size);
             lua_call(lua, 1, 1);
-            count++;
+            info->count++;
 
             // is continue?
             tb_bool_t is_continue = lua_toboolean(lua, -1);
             lua_pop(lua, 1);
             if (!is_continue)
             {
-                ok = tb_true;
+                info->ok = 0;
                 break;
             }
-        }
 
-        // ok
-        if (i == key_path_num)
-            ok = tb_true;
+            // enum all subkeys
+            if (recursion > 0 || recursion < 0)
+                xm_winos_registry_enum_keys(info, key_path, recursion > 0? recursion - 1 : recursion);
+        }
 
     } while (0);
 
@@ -157,6 +147,79 @@ tb_int_t xm_winos_registry_keys(lua_State* lua)
         RegCloseKey(keynew);
     keynew = tb_null;
 
+    // free key path
+    if (key_path) tb_free(key_path);
+    key_path = tb_null;
+}
+
+/* //////////////////////////////////////////////////////////////////////////////////////
+ * implementation
+ */
+
+/* get registry keys
+ *
+ * local count, errors = winos.registry_keys("HKEY_LOCAL_MACHINE",
+ *                                             "SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion\\AeDebug",
+ *                                             function (key_path)
+ *                                                 return true -- continue or break
+ *                                             end)
+ */
+tb_int_t xm_winos_registry_keys(lua_State* lua)
+{
+    // check
+    tb_assert_and_check_return_val(lua, 0);
+
+    // get the arguments
+    tb_char_t const* rootkey = luaL_checkstring(lua, 1);
+    tb_char_t const* rootdir = luaL_checkstring(lua, 2);
+    tb_long_t recursion      = lua_tointeger(lua, 3);
+    tb_bool_t is_function    = lua_isfunction(lua, 4);
+    tb_check_return_val(rootkey && rootdir && is_function, 0);
+
+    // enum keys
+    tb_bool_t   ok = tb_false;
+    tb_int_t    count = 0;
+    HKEY        key = tb_null;
+    do
+    {
+        // get registry rootkey
+        if (!tb_strcmp(rootkey, "HKEY_CLASSES_ROOT"))         key = HKEY_CLASSES_ROOT;
+        else if (!tb_strcmp(rootkey, "HKEY_CURRENT_CONFIG"))  key = HKEY_CURRENT_CONFIG;
+        else if (!tb_strcmp(rootkey, "HKEY_CURRENT_USER"))    key = HKEY_CURRENT_USER;
+        else if (!tb_strcmp(rootkey, "HKEY_LOCAL_MACHINE"))   key = HKEY_LOCAL_MACHINE;
+        else if (!tb_strcmp(rootkey, "HKEY_USERS"))           key = HKEY_USERS;
+        else
+        {
+            lua_pushnil(lua);
+            lua_pushfstring(lua, "invalid registry rootkey: %s", rootkey);
+            break;
+        }
+
+        // do enum
+        tb_wchar_t rootdir_w[TB_PATH_MAXN];
+        if (tb_atow(rootdir_w, rootdir, TB_PATH_MAXN) == -1)
+        {
+            lua_pushnil(lua);
+            lua_pushfstring(lua, "rootdir is too long: %s", rootdir);
+            break;
+        }
+        xm_winos_registry_enum_info_t info;
+        info.lua   = lua;
+        info.key   = key;
+        info.count = 0;
+        info.ok    = tb_true;
+        info.error = tb_null;
+        xm_winos_registry_enum_keys(&info, rootdir_w, recursion);
+        count = info.count;
+        ok    = info.ok >= 0;
+        if (!ok)
+        {
+            lua_pushnil(lua);
+            lua_pushfstring(lua, "%s: %s\\%s", info.error? info.error : "enum registry keys failed", rootkey, rootdir);
+        }
+
+    } while (0);
+
     // ok?
     if (ok)
     {
diff --git a/core/src/xmake/winos/registry_values.c b/core/src/xmake/winos/registry_values.c
index 7a9211007..17aefef6e 100644
--- a/core/src/xmake/winos/registry_values.c
+++ b/core/src/xmake/winos/registry_values.c
@@ -53,12 +53,11 @@ tb_int_t xm_winos_registry_values(lua_State* lua)
     tb_bool_t is_function    = lua_isfunction(lua, 3);
     tb_check_return_val(rootkey && rootdir && is_function, 0);
 
-    // query key-value
+    // enum values
     tb_bool_t   ok = tb_false;
     tb_int_t    count = 0;
     HKEY        key = tb_null;
     HKEY        keynew = tb_null;
-    tb_char_t*  value = tb_null;
     do
     {
         // get registry rootkey