Fix bug: When using dynamic linking, auth will be loaded repeatedly.

c/ckb_auth.h

@@ -68,17 +68,17 @@ typedef struct {
 // TODO: when ready, move it into ckb-c-stdlib
 typedef struct CkbAuthType {
     uint8_t algorithm_id;
-    uint8_t content[20];
+    uint8_t content[AUTH160_SIZE];
 } CkbAuthType;
 
 enum EntryCategoryType {
     EntryCategoryExec = 0,
-    EntryCategoryDynamicLinking = 1,
+    EntryCategoryDynamicLibrary = 1,
     EntryCategorySpawn = 2,
 };
 
 typedef struct CkbEntryType {
-    uint8_t code_hash[32];
+    uint8_t code_hash[BLAKE2B_BLOCK_SIZE];
     uint8_t hash_type;
     uint8_t entry_category;
 } CkbEntryType;
@@ -119,34 +119,108 @@ typedef int (*ckb_auth_validate_t)(uint8_t auth_algorithm_id,
                                    uint32_t message_size, uint8_t *pubkey_hash,
                                    uint32_t pubkey_hash_size);
 
-static uint8_t g_code_buff[300 * 1024] __attribute__((aligned(RISCV_PGSIZE)));
+#ifndef CKB_AUTH_DISABLE_DYNAMIC_LIB
+
+#ifndef CKB_AUTH_DL_BUFF_SIZE
+#define CKB_AUTH_DL_BUFF_SIZE 1024 * 200
+#endif  // CKB_AUTH_DL_MAX_COUNT
+
+#ifndef CKB_AUTH_DL_MAX_COUNT
+#define CKB_AUTH_DL_MAX_COUNT 8
+#endif  // CKB_AUTH_DL_MAX_COUNT
+
+static uint8_t g_dl_code_buffer[CKB_AUTH_DL_BUFF_SIZE]
+    __attribute__((aligned(RISCV_PGSIZE))) = {0};
+
+typedef struct {
+    uint8_t *code_ptr;
+    size_t code_ptr_size;
+
+    uint8_t code_hash[BLAKE2B_BLOCK_SIZE];
+    uint8_t hash_type;
+
+    void *handle;
+    ckb_auth_validate_t func;
+} CkbDLCache;
+static CkbDLCache g_dl_cache[CKB_AUTH_DL_MAX_COUNT];
+static size_t g_dl_cache_count = 0;
+
+int get_dl_func_by_code_hash(const uint8_t *code_hash, uint8_t hash_type,
+                             ckb_auth_validate_t *out_func) {
+    // Find from cache
+    for (size_t i = 0; i < g_dl_cache_count; i++) {
+        CkbDLCache *cache = &g_dl_cache[i];
+        if (memcmp(cache->code_hash, code_hash, BLAKE2B_BLOCK_SIZE) == 0 &&
+            hash_type == cache->hash_type) {
+            *out_func = cache->func;
+            return 0;
+        }
+    }
+
+    // get by cache
+    size_t buf_offset = 0;
+    if (g_dl_cache_count) {
+        CkbDLCache *cache = &g_dl_cache[g_dl_cache_count - 1];
+        buf_offset =
+            (size_t)(cache->code_ptr + cache->code_ptr_size - g_dl_code_buffer);
+        if (buf_offset % RISCV_PGSIZE != 0) {
+            buf_offset += RISCV_PGSIZE - buf_offset % RISCV_PGSIZE;
+        }
+    }
+
+    // load function by cell
+    CkbDLCache *cache = &g_dl_cache[g_dl_cache_count];
+    memset(cache, 0, sizeof(CkbDLCache));
+    cache->code_ptr = g_dl_code_buffer + buf_offset;
+
+    int err = ckb_dlopen2(code_hash, hash_type, cache->code_ptr,
+                          sizeof(g_dl_code_buffer) - buf_offset, &cache->handle,
+                          &cache->code_ptr_size);
+    if (err != 0) {
+        return err;
+    }
+    cache->func =
+        (ckb_auth_validate_t)ckb_dlsym(cache->handle, "ckb_auth_validate");
+    if (cache->func == 0) {
+        return CKB_INVALID_DATA;
+    }
+
+    *out_func = cache->func;
+    memcpy(cache->code_hash, code_hash, BLAKE2B_BLOCK_SIZE);
+    cache->hash_type = hash_type;
+
+    g_dl_cache_count += 1;
+    return 0;
+}
+
+#endif  // CKB_AUTH_DISABLE_DYNAMIC_LIB
 
 int ckb_auth(CkbEntryType *entry, CkbAuthType *id, const uint8_t *signature,
              uint32_t signature_size, const uint8_t *message32) {
     int err = 0;
-    if (entry->entry_category == EntryCategoryDynamicLinking) {
-        void *handle = NULL;
-        size_t consumed_size = 0;
-        err = ckb_dlopen2(entry->code_hash, entry->hash_type, g_code_buff,
-                          sizeof(g_code_buff), &handle, &consumed_size);
-        if (err != 0) return err;
-
-        ckb_auth_validate_t func =
-            (ckb_auth_validate_t)ckb_dlsym(handle, "ckb_auth_validate");
-        if (func == 0) {
-            return CKB_INVALID_DATA;
+    if (entry->entry_category == EntryCategoryDynamicLibrary) {
+#ifdef CKB_AUTH_DISABLE_DYNAMIC_LIB
+        // Disable DynamicLibrary via macro can save memory
+        return ERROR_INVALID_ARG;
+#else   // CKB_AUTH_DISABLE_DYNAMIC_LIB
+        ckb_auth_validate_t func = NULL;
+        err =
+            get_dl_func_by_code_hash(entry->code_hash, entry->hash_type, &func);
+        if (err) {
+            return err;
         }
-        return func(id->algorithm_id, signature, signature_size, message32, 32,
-                    id->content, 20);
+        return func(id->algorithm_id, signature, signature_size, message32,
+                    BLAKE2B_BLOCK_SIZE, id->content, AUTH160_SIZE);
+#endif  // CKB_AUTH_DISABLE_DYNAMIC_LIB
     } else if (entry->entry_category == EntryCategoryExec ||
                entry->entry_category == EntryCategorySpawn) {
         char algorithm_id_str[2 + 1];
         if (signature_size > 1024 * 8) {
             return CKB_INVALID_DATA;
         }
         char signature_str[signature_size * 2 + 1];
-        char message_str[32 * 2 + 1];
-        char pubkey_hash_str[20 * 2 + 1];
+        char message_str[BLAKE2B_BLOCK_SIZE * 2 + 1];
+        char pubkey_hash_str[AUTH160_SIZE * 2 + 1];
 
         uint32_t bin2hex_output_len = 0;
         if (ckb_bin2hex(&id->algorithm_id, 1, algorithm_id_str,
@@ -159,12 +233,12 @@ int ckb_auth(CkbEntryType *entry, CkbAuthType *id, const uint8_t *signature,
                           sizeof(signature_str), &bin2hex_output_len, true)) {
             return CKB_INVALID_DATA;
         }
-        if (ckb_bin2hex(message32, 32, message_str, sizeof(message_str),
+        if (ckb_bin2hex(message32, BLAKE2B_BLOCK_SIZE, message_str, sizeof(message_str),
                           &bin2hex_output_len, true)) {
             return CKB_INVALID_DATA;
         }
 
-        if (ckb_bin2hex(id->content, 20, pubkey_hash_str,
+        if (ckb_bin2hex(id->content, AUTH160_SIZE, pubkey_hash_str,
                           sizeof(pubkey_hash_str), &bin2hex_output_len, true)) {
             return CKB_INVALID_DATA;
         }

ckb-auth-rs/Cargo.toml

@@ -6,11 +6,18 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [features]
-default = []
+default = ["dynamic-library-memory-200"]
 ckb2023 = ["ckb-std/ckb2023"]
 
+enable-dynamic-library = ["lazy_static"]
+dynamic-library-memory-200 = ["enable-dynamic-library"]
+# enable these features when memory is not enough
+dynamic-library-memory-400 = ["enable-dynamic-library"]
+dynamic-library-memory-600 = ["enable-dynamic-library"]
+
 [dependencies]
 ckb-std = "0.14.3"
+lazy_static = { version = "1.4.0", optional = true, features = ["spin_no_std"] }
 
 [target.'cfg(target_arch = "riscv64")'.dependencies]
 hex = { version = "0.4.3", default-features = false, features = ["alloc"]}

ckb-auth-rs/src/ckb_auth.rs

@@ -1,19 +1,18 @@
 extern crate alloc;
 
 use crate::{CkbAuthError, CkbAuthType, CkbEntryType, EntryCategoryType};
-use alloc::collections::BTreeMap;
 use alloc::ffi::CString;
 use alloc::format;
-use alloc::vec::Vec;
 use ckb_std::high_level::exec_cell;
+use hex::encode;
+
+#[cfg(feature = "ckb2023")]
+use alloc::vec::Vec;
 #[cfg(feature = "ckb2023")]
 use ckb_std::high_level::spawn_cell;
-use ckb_std::{
-    ckb_types::core::ScriptHashType,
-    dynamic_loading_c_impl::{CKBDLContext, Library, Symbol},
-};
-use core::mem::size_of_val;
-use hex::encode;
+
+#[cfg(feature = "enable-dynamic-library")]
+use super::ckb_auth_dl::ckb_auth_dl;
 
 pub fn ckb_auth(
     entry: &CkbEntryType,
@@ -23,7 +22,8 @@ pub fn ckb_auth(
 ) -> Result<(), CkbAuthError> {
     match entry.entry_category {
         EntryCategoryType::Exec => ckb_auth_exec(entry, id, signature, message),
-        EntryCategoryType::DynamicLinking => ckb_auth_dl(entry, id, signature, message),
+        #[cfg(feature = "enable-dynamic-library")]
+        EntryCategoryType::DynamicLibrary => ckb_auth_dl(entry, id, signature, message),
         #[cfg(feature = "ckb2023")]
         EntryCategoryType::Spawn => ckb_auth_spawn(entry, id, signature, message),
     }
@@ -73,117 +73,3 @@ fn ckb_auth_exec(
     exec_cell(&entry.code_hash, entry.hash_type, &args)?;
     Ok(())
 }
-
-type DLContext = CKBDLContext<[u8; 512 * 1024]>;
-type CkbAuthValidate = unsafe extern "C" fn(
-    auth_algorithm_id: u8,
-    signature: *const u8,
-    signature_size: u32,
-    message: *const u8,
-    message_size: u32,
-    pubkey_hash: *mut u8,
-    pubkey_hash_size: u32,
-) -> i32;
-
-const EXPORTED_FUNC_NAME: &str = "ckb_auth_validate";
-
-struct CKBDLLoader {
-    pub context: DLContext,
-    pub context_used: usize,
-    pub loaded_lib: BTreeMap<[u8; 33], Library>,
-}
-
-static mut G_CKB_DL_LOADER: Option<CKBDLLoader> = None;
-impl CKBDLLoader {
-    pub fn get() -> &'static mut Self {
-        unsafe {
-            match G_CKB_DL_LOADER.as_mut() {
-                Some(v) => v,
-                None => {
-                    G_CKB_DL_LOADER = Some(Self::new());
-                    G_CKB_DL_LOADER.as_mut().unwrap()
-                }
-            }
-        }
-    }
-
-    fn new() -> Self {
-        Self {
-            context: unsafe { DLContext::new() },
-            context_used: 0,
-            loaded_lib: BTreeMap::new(),
-        }
-    }
-
-    fn get_lib(
-        &mut self,
-        code_hash: &[u8; 32],
-        hash_type: ScriptHashType,
-    ) -> Result<&Library, CkbAuthError> {
-        let mut lib_key = [0u8; 33];
-        lib_key[..32].copy_from_slice(code_hash);
-        lib_key[32] = hash_type as u8;
-
-        let has_lib = match self.loaded_lib.get(&lib_key) {
-            Some(_) => true,
-            None => false,
-        };
-
-        if !has_lib {
-            let size = size_of_val(&self.context);
-            let lib = self
-                .context
-                .load_with_offset(code_hash, hash_type, self.context_used, size)
-                .map_err(|_| CkbAuthError::LoadDLError)?;
-            self.context_used += lib.consumed_size();
-            self.loaded_lib.insert(lib_key.clone(), lib);
-        };
-        Ok(self.loaded_lib.get(&lib_key).unwrap())
-    }
-
-    pub fn get_validate_func<T>(
-        &mut self,
-        code_hash: &[u8; 32],
-        hash_type: ScriptHashType,
-        func_name: &str,
-    ) -> Result<Symbol<T>, CkbAuthError> {
-        let lib = self.get_lib(code_hash, hash_type)?;
-
-        let func: Option<Symbol<T>> = unsafe { lib.get(func_name.as_bytes()) };
-        if func.is_none() {
-            return Err(CkbAuthError::LoadDLFuncError);
-        }
-        Ok(func.unwrap())
-    }
-}
-
-fn ckb_auth_dl(
-    entry: &CkbEntryType,
-    id: &CkbAuthType,
-    signature: &[u8],
-    message: &[u8; 32],
-) -> Result<(), CkbAuthError> {
-    let func: Symbol<CkbAuthValidate> = CKBDLLoader::get().get_validate_func(
-        &entry.code_hash,
-        entry.hash_type,
-        EXPORTED_FUNC_NAME,
-    )?;
-
-    let mut pub_key = id.pubkey_hash.clone();
-    let rc_code = unsafe {
-        func(
-            id.algorithm_id.clone().into(),
-            signature.as_ptr(),
-            signature.len() as u32,
-            message.as_ptr(),
-            message.len() as u32,
-            pub_key.as_mut_ptr(),
-            pub_key.len() as u32,
-        )
-    };
-
-    match rc_code {
-        0 => Ok(()),
-        _ => Err(CkbAuthError::RunDLError),
-    }
-}