Skip to content

Commit

Permalink
Fix bug: When using dynamic linking, auth will be loaded repeatedly. (#…
Browse files Browse the repository at this point in the history 
…27)

* ckb-auth-rs: Add features to control dynamic linking and CKBDLContext in static memory
* C language: supports dynamic loading of multiple ckb-auth
* Fix bug: When using dynamic linking, auth will be loaded repeatedly.
* Update tests/README.md
  • Loading branch information
@joii2020
joii2020 authored Dec 19, 2023
1 parent cb41284 commit f381f9b
Show file tree
Hide file tree
Showing 24 changed files with 2,106 additions and 202 deletions.
116 changes: 95 additions & 21 deletions c/ckb_auth.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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,
Expand All @@ -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;
}
Expand Down
9 changes: 8 additions & 1 deletion ckb-auth-rs/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"]}
Expand Down
132 changes: 9 additions & 123 deletions ckb-auth-rs/src/ckb_auth.rs
View file
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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),
}
Expand Down Expand Up @@ -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),
}
}
Loading

0 comments on commit f381f9b

Please sign in to comment.