From 8037e4a1b83b2b44875bd2eaf12cc0864aa55c10 Mon Sep 17 00:00:00 2001 From: Khyber Sen Date: Mon, 6 Jan 2025 11:37:58 -0800 Subject: [PATCH] runtime/type-registry: add new `type-registry` library for the type confusion APIs This is written in Rust like `memory-map` is, so I made the `CMakeLists.txt` very similar. Since the library may be called concurrently, it uses a hash map which is non-trivial in C, and we already have Rust libraries like `memory-map`, Rust seemed like a good choice here. --- runtime/CMakeLists.txt | 1 + runtime/tracer/CMakeLists.txt | 1 + runtime/type-registry/CMakeLists.txt | 21 ++++ runtime/type-registry/Cargo.lock | 7 ++ runtime/type-registry/Cargo.toml | 20 +++ runtime/type-registry/src/lib.rs | 175 +++++++++++++++++++++++++++ 6 files changed, 225 insertions(+) create mode 100644 runtime/type-registry/CMakeLists.txt create mode 100644 runtime/type-registry/Cargo.lock create mode 100644 runtime/type-registry/Cargo.toml create mode 100644 runtime/type-registry/src/lib.rs diff --git a/runtime/CMakeLists.txt b/runtime/CMakeLists.txt index da60ba148..ac91382b1 100644 --- a/runtime/CMakeLists.txt +++ b/runtime/CMakeLists.txt @@ -4,3 +4,4 @@ project(runtime) add_subdirectory(libia2) add_subdirectory(partition-alloc) add_subdirectory(tracer) +add_subdirectory(type-registry) diff --git a/runtime/tracer/CMakeLists.txt b/runtime/tracer/CMakeLists.txt index 3f2fb14ff..60e2ddf26 100644 --- a/runtime/tracer/CMakeLists.txt +++ b/runtime/tracer/CMakeLists.txt @@ -8,6 +8,7 @@ else() set(CARGO_TARGET_FLAG "") set(CARGO_ARCH_SUFFIX "") endif() + add_custom_command( OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/${CARGO_ARCH_SUFFIX}/release/libmemory_map.so COMMAND CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR} cargo build ${CARGO_TARGET_FLAG} --manifest-path ${CMAKE_CURRENT_SOURCE_DIR}/memory-map/Cargo.toml --release diff --git a/runtime/type-registry/CMakeLists.txt b/runtime/type-registry/CMakeLists.txt new file mode 100644 index 000000000..fe33ed475 --- /dev/null +++ b/runtime/type-registry/CMakeLists.txt @@ -0,0 +1,21 @@ +cmake_minimum_required(VERSION 3.12) +project(type-registry) + +if (LIBIA2_AARCH64) +set(CARGO_TARGET_FLAG "--target=aarch64-unknown-linux-gnu" "--config" "target.aarch64-unknown-linux-gnu.linker=\\\"aarch64-linux-gnu-gcc\\\"") +set(CARGO_ARCH_SUFFIX "aarch64-unknown-linux-gnu") +else() +set(CARGO_TARGET_FLAG "") +set(CARGO_ARCH_SUFFIX "") +endif() + +add_custom_command( + OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/${CARGO_ARCH_SUFFIX}/release/libtype_registry.so + COMMAND CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR} cargo build ${CARGO_TARGET_FLAG} --manifest-path ${CMAKE_CURRENT_SOURCE_DIR}/Cargo.toml --release + DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/Cargo.toml ${CMAKE_CURRENT_SOURCE_DIR}/src/lib.rs +) + +add_library(type-registry STATIC IMPORTED) +set_property(TARGET type-registry PROPERTY IMPORTED_LOCATION ${CMAKE_CURRENT_BINARY_DIR}/${CARGO_ARCH_SUFFIX}/release/libtype_registry.so) +add_custom_target(type-registry-tgt DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/${CARGO_ARCH_SUFFIX}/release/libtype_registry.so) +add_dependencies(type-registry type-registry-tgt) diff --git a/runtime/type-registry/Cargo.lock b/runtime/type-registry/Cargo.lock new file mode 100644 index 000000000..911018852 --- /dev/null +++ b/runtime/type-registry/Cargo.lock @@ -0,0 +1,7 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 3 + +[[package]] +name = "type-registry" +version = "0.1.0" diff --git a/runtime/type-registry/Cargo.toml b/runtime/type-registry/Cargo.toml new file mode 100644 index 000000000..d21323133 --- /dev/null +++ b/runtime/type-registry/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "type-registry" +version = "0.1.0" +edition = "2021" +rust-version = "1.84" + +[lib] +crate-type = ["cdylib"] + +[dependencies] + +# We always need `panic = "abort"`, +# as we don't want panics to cross the C/Rust boundary, +# and we want such panics for these security checks to always be fatal. + +[profile.dev] +panic = "abort" + +[profile.release] +panic = "abort" diff --git a/runtime/type-registry/src/lib.rs b/runtime/type-registry/src/lib.rs new file mode 100644 index 000000000..000437f1c --- /dev/null +++ b/runtime/type-registry/src/lib.rs @@ -0,0 +1,175 @@ +use std::collections::HashMap; +use std::fmt; +use std::fmt::Debug; +use std::fmt::Display; +use std::fmt::Formatter; +use std::ptr; +use std::sync::LazyLock; +use std::sync::RwLock; + +pub type Ptr = *const (); + +/// A pointer's address without provenance, used purely as an address for comparisons. +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +pub struct PtrAddr(usize); + +impl Display for PtrAddr { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "{:?}", ptr::without_provenance::<()>(self.0)) + } +} + +impl Debug for PtrAddr { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "{}", self) + } +} + +/// A unique ID for a type. +/// +/// This can be anything, as long as it's unique for a type in a program. +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +#[repr(transparent)] +pub struct TypeId(usize); + +impl Display for TypeId { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "Type{}", self.0) + } +} + +impl Debug for TypeId { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + write!(f, "{}", self) + } +} + +static GLOBAL_TYPE_REGISTRY: LazyLock = LazyLock::new(TypeRegistry::default); + +/// See [`TypeRegistry::construct`]. +#[no_mangle] +pub extern "C-unwind" fn ia2_type_registry_construct(ptr: Ptr, type_id: TypeId) { + GLOBAL_TYPE_REGISTRY.construct(ptr, type_id); +} + +/// See [`TypeRegistry::destruct`]. +#[no_mangle] +pub extern "C-unwind" fn ia2_type_registry_destruct(ptr: Ptr, expected_type_id: TypeId) { + GLOBAL_TYPE_REGISTRY.destruct(ptr, expected_type_id); +} + +/// See [`TypeRegistry::check`]. +#[no_mangle] +pub extern "C-unwind" fn ia2_type_registry_check(ptr: Ptr, expected_type_id: TypeId) { + GLOBAL_TYPE_REGISTRY.check(ptr, expected_type_id); +} + +#[derive(Default)] +pub struct TypeRegistry { + map: RwLock>, +} + +impl TypeRegistry { + /// Called after an object is constructed to register it and its type. + /// + /// Panics if `ptr` is already constructed. + #[track_caller] + pub fn construct(&self, ptr: Ptr, type_id: TypeId) { + let ptr = PtrAddr(ptr.addr()); + let guard = &mut *self.map.write().unwrap(); + let prev_type_id = guard.insert(ptr, type_id); + assert_eq!(prev_type_id, None); + } + + /// Called after an object is destructured to unregister it and its type. + /// + /// Panics if `ptr` has a different type. + #[track_caller] + pub fn destruct(&self, ptr: Ptr, expected_type_id: TypeId) { + let ptr = PtrAddr(ptr.addr()); + let guard = &mut *self.map.write().unwrap(); + let type_id = guard.remove(&ptr); + assert_eq!(type_id, Some(expected_type_id)); + } + + /// Called to check a pointer has the expected type. + /// + /// Panics if `ptr` is not registered or has a different type. + #[track_caller] + pub fn check(&self, ptr: Ptr, expected_type_id: TypeId) { + let ptr = PtrAddr(ptr.addr()); + let guard = &*self.map.read().unwrap(); + let type_id = guard.get(&ptr).copied(); + assert_eq!(type_id, Some(expected_type_id)); + } +} + +impl Drop for TypeRegistry { + fn drop(&mut self) { + match self.map.get_mut() { + Ok(map) => { + if !map.is_empty() { + eprintln!("warning: leak: {map:?}"); + } + } + Err(e) => { + eprintln!("poison: {e}") + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn ptr(addr: usize) -> Ptr { + core::ptr::without_provenance(addr) + } + + #[test] + fn normal() { + let registry = TypeRegistry::default(); + registry.construct(ptr(0xA), TypeId(1)); + registry.check(ptr(0xA), TypeId(1)); + registry.construct(ptr(0xB), TypeId(2)); + registry.check(ptr(0xA), TypeId(1)); + registry.check(ptr(0xB), TypeId(2)); + registry.destruct(ptr(0xA), TypeId(1)); + registry.check(ptr(0xB), TypeId(2)); + registry.destruct(ptr(0xB), TypeId(2)); + } + + #[test] + #[should_panic] + fn destruct_non_existent_ptr() { + let registry = TypeRegistry::default(); + registry.destruct(ptr(0xA), TypeId(1)); + } + + #[test] + #[should_panic] + fn check_non_existent_ptr() { + let registry = TypeRegistry::default(); + registry.check(ptr(0xA), TypeId(1)); + } + + #[test] + #[should_panic] + fn check_freed_ptr() { + let registry = TypeRegistry::default(); + registry.construct(ptr(0xA), TypeId(1)); + registry.check(ptr(0xA), TypeId(1)); + registry.destruct(ptr(0xA), TypeId(1)); + registry.check(ptr(0xA), TypeId(1)); + } + + #[test] + #[should_panic] + fn construct_existing_ptr() { + let registry = TypeRegistry::default(); + registry.construct(ptr(0xA), TypeId(1)); + registry.check(ptr(0xA), TypeId(1)); + registry.construct(ptr(0xA), TypeId(1)); + } +}