From ad21e9cfa15a0efd6c5560c38cddc3d8136adef8 Mon Sep 17 00:00:00 2001 From: weitengchen Date: Tue, 17 Jun 2025 10:30:43 -0700 Subject: [PATCH 1/3] support TLS --- litebox/src/platform/mod.rs | 15 ++ litebox_common_linux/Cargo.toml | 1 + litebox_common_linux/src/lib.rs | 81 +++++++++- litebox_platform_linux_userland/src/lib.rs | 169 +++++++++++++++++++-- 4 files changed, 248 insertions(+), 18 deletions(-) diff --git a/litebox/src/platform/mod.rs b/litebox/src/platform/mod.rs index 890a56096..60318e36c 100644 --- a/litebox/src/platform/mod.rs +++ b/litebox/src/platform/mod.rs @@ -539,3 +539,18 @@ pub trait SystemInfoProvider { /// execution context and transfer control to the syscall handler. fn get_syscall_entry_point(&self) -> usize; } + +/// A provider for thread-local storage. +pub trait ThreadLocalStorageProvider { + type ThreadLocalStorage; + + /// Get a thread-local storage value for the current thread. + fn get_thread_local(&self) -> *mut Self::ThreadLocalStorage; + + /// Set a thread-local storage value for the current thread. + fn set_thread_local(&self, value: Self::ThreadLocalStorage); + + /// Release the thread-local storage value for the current thread + /// and invoke the provided callback function on it. + fn release_thread_local(&self, f: fn(&mut Self::ThreadLocalStorage)); +} diff --git a/litebox_common_linux/Cargo.toml b/litebox_common_linux/Cargo.toml index 497ae9f5e..436706f3a 100644 --- a/litebox_common_linux/Cargo.toml +++ b/litebox_common_linux/Cargo.toml @@ -4,6 +4,7 @@ version = "0.1.0" edition = "2024" [dependencies] +bitfield = "0.19.1" bitflags = "2.9.0" cfg-if = "1.0.0" litebox = { path = "../litebox/", version = "0.1.0" } diff --git a/litebox_common_linux/src/lib.rs b/litebox_common_linux/src/lib.rs index a9247c015..6735fa38c 100644 --- a/litebox_common_linux/src/lib.rs +++ b/litebox_common_linux/src/lib.rs @@ -14,6 +14,8 @@ use syscalls::Sysno; pub mod errno; +extern crate alloc; + // TODO(jayb): Should errno::Errno be publicly re-exported? bitflags::bitflags! { @@ -681,13 +683,84 @@ pub unsafe fn wrfsbase(fs_base: usize) { } } +/// Reads the GS segment base address +/// +/// ## Safety +/// +/// If `CR4.FSGSBASE` is not set, this instruction will throw an `#UD`. +#[cfg(target_arch = "x86_64")] +pub unsafe fn rdgsbase() -> usize { + let ret: usize; + unsafe { + core::arch::asm!( + "rdgsbase {}", + out(reg) ret, + options(nostack, nomem) + ); + } + ret +} + +/// Writes the GS segment base address +/// +/// ## Safety +/// +/// If `CR4.FSGSBASE` is not set, this instruction will throw an `#UD`. +/// +/// The caller must ensure that this write operation has no unsafe side +/// effects, as the GS segment base address might be in use. +#[cfg(target_arch = "x86_64")] +pub unsafe fn wrgsbase(gs_base: usize) { + unsafe { + core::arch::asm!( + "wrgsbase {}", + in(reg) gs_base, + options(nostack, nomem) + ); + } +} + +/// Linux's `user_desc` struct used by the `set_thread_area` syscall. #[repr(C, packed)] #[derive(Debug, Clone)] pub struct UserDesc { - pub entry_number: i32, - pub base_addr: i32, - pub limit: i32, - pub flags: i32, + pub entry_number: u32, + pub base_addr: u32, + pub limit: u32, + pub flags: UserDescFlags, +} + +bitfield::bitfield! { + /// Flags for the `user_desc` struct. + #[derive(Clone, Copy)] + pub struct UserDescFlags(u32); + impl Debug; + /// 1 if the segment is 32-bit + pub seg_32bit, set_seg_32bit: 0; + /// Contents of the segment + pub contents, set_contents: 1, 2; + /// Read-exec only + pub read_exec_only, set_read_exec_only: 3; + /// Limit in pages + pub limit_in_pages, set_limit_in_pages: 4; + /// Segment not present + pub seg_not_present, set_seg_not_present: 5; + /// Usable by userland + pub useable, set_useable: 6; + /// 1 if the segment is 64-bit (x86_64 only) + pub lm, set_lm: 7; +} + +pub struct ThreadLocalStorage { + #[cfg(target_arch = "x86")] + pub self_ptr: *mut ThreadLocalStorage, + pub current_task: alloc::boxed::Box, + + pub __phantom: core::marker::PhantomData, +} + +pub struct Task { + pub tid: u32, } #[repr(C)] diff --git a/litebox_platform_linux_userland/src/lib.rs b/litebox_platform_linux_userland/src/lib.rs index 1d4578fdb..cd2257b29 100644 --- a/litebox_platform_linux_userland/src/lib.rs +++ b/litebox_platform_linux_userland/src/lib.rs @@ -574,6 +574,23 @@ fn set_fs_base_arch_prctl(fs_base: usize) -> Result, +) -> Result { + unsafe { syscalls::syscall1(syscalls::Sysno::set_thread_area, user_desc.as_usize()) }.map_err( + |err| match err { + syscalls::Errno::EFAULT => litebox_common_linux::errno::Errno::EFAULT, + syscalls::Errno::EINVAL => litebox_common_linux::errno::Errno::EINVAL, + syscalls::Errno::ENOSYS => litebox_common_linux::errno::Errno::ENOSYS, + syscalls::Errno::ESRCH => litebox_common_linux::errno::Errno::ESRCH, + _ => panic!("unexpected error {err}"), + }, + ) +} + pub struct PunchthroughToken { punchthrough: PunchthroughSyscall, } @@ -674,19 +691,7 @@ impl litebox::platform::PunchthroughToken for PunchthroughToken { } #[cfg(target_arch = "x86")] PunchthroughSyscall::SetThreadArea { user_desc } => { - use litebox::platform::RawConstPointer as _; - unsafe { - syscalls::syscall1(syscalls::Sysno::set_thread_area, user_desc.as_usize()) - } - .map_err(|err| { - litebox::platform::PunchthroughError::Failure(match err { - syscalls::Errno::EFAULT => litebox_common_linux::errno::Errno::EFAULT, - syscalls::Errno::EINVAL => litebox_common_linux::errno::Errno::EINVAL, - syscalls::Errno::ENOSYS => litebox_common_linux::errno::Errno::ENOSYS, - syscalls::Errno::ESRCH => litebox_common_linux::errno::Errno::ESRCH, - _ => panic!("unexpected error {err}"), - }) - }) + set_thread_area(user_desc).map_err(litebox::platform::PunchthroughError::Failure) } } } @@ -1207,12 +1212,120 @@ impl litebox::platform::SystemInfoProvider for LinuxUserland { } } +/// Similar to libc, we use fs/gs registers to store thread-local storage (TLS). +/// To avoid conflicts with libc's TLS, we choose to use gs on x86_64 and fs on x86 +/// as libc uses fs on x86_64 and gs on x86. +impl litebox::platform::ThreadLocalStorageProvider for LinuxUserland { + type ThreadLocalStorage = litebox_common_linux::ThreadLocalStorage; + + #[cfg(target_arch = "x86_64")] + fn get_thread_local(&self) -> *mut Self::ThreadLocalStorage { + let tls = unsafe { litebox_common_linux::rdgsbase() }; + if tls == 0 { + return core::ptr::null_mut(); + } + tls as *mut Self::ThreadLocalStorage + } + + #[cfg(target_arch = "x86")] + fn get_thread_local(&self) -> *mut Self::ThreadLocalStorage { + let mut fs_selector: u16; + unsafe { + core::arch::asm!( + "mov {0:x}, fs", + out(reg) fs_selector, + options(nostack, preserves_flags) + ); + } + if fs_selector == 0 { + return core::ptr::null_mut(); + } + + let mut addr: usize; + unsafe { + core::arch::asm!( + "mov {0}, fs:{offset}", + out(reg) addr, + offset = const core::mem::offset_of!(Self::ThreadLocalStorage, self_ptr), + options(nostack, preserves_flags) + ); + } + addr as *mut Self::ThreadLocalStorage + } + + #[cfg(target_arch = "x86_64")] + fn set_thread_local(&self, tls: Self::ThreadLocalStorage) { + let tls = Box::new(tls); + unsafe { litebox_common_linux::wrgsbase(Box::into_raw(tls) as usize) }; + } + + #[cfg(target_arch = "x86")] + fn set_thread_local(&self, tls: Self::ThreadLocalStorage) { + let mut tls = Box::new(tls); + tls.self_ptr = tls.as_mut(); + + let mut flags = litebox_common_linux::UserDescFlags(0); + flags.set_seg_32bit(true); + flags.set_useable(true); + let mut user_desc = litebox_common_linux::UserDesc { + entry_number: u32::MAX, + base_addr: Box::into_raw(tls) as u32, + limit: u32::try_from(core::mem::size_of::()).unwrap() - 1, + flags, + }; + let user_desc_ptr = litebox::platform::trivial_providers::TransparentMutPtr { + inner: &raw mut user_desc, + }; + set_thread_area(user_desc_ptr).expect("Failed to set thread area for TLS"); + + let new_fs_selector = ((user_desc.entry_number & 0xfff) << 3) | 0x3; // user mode + // set fs selector + unsafe { + core::arch::asm!( + "mov fs, {0:x}", + in(reg) new_fs_selector, + options(nostack, preserves_flags) + ); + } + } + + #[cfg(target_arch = "x86_64")] + fn release_thread_local(&self, f: fn(&mut Self::ThreadLocalStorage)) { + let tls = self.get_thread_local(); + if tls.is_null() { + return; + } + let mut tls = unsafe { Box::from_raw(tls) }; + f(&mut tls); + unsafe { + litebox_common_linux::wrgsbase(0); + } + } + + #[cfg(target_arch = "x86")] + fn release_thread_local(&self, f: fn(&mut Self::ThreadLocalStorage)) { + let tls = self.get_thread_local(); + if tls.is_null() { + return; + } + let mut tls = unsafe { Box::from_raw(tls) }; + f(&mut tls); + unsafe { + core::arch::asm!( + "mov fs, {0}", + in(reg) 0, + options(nostack, preserves_flags) + ); + } + } +} + #[cfg(test)] mod tests { use core::sync::atomic::AtomicU32; use std::thread::sleep; - use litebox::platform::RawMutex; + use litebox::platform::{RawMutex, ThreadLocalStorageProvider as _}; use crate::LinuxUserland; use litebox::platform::PageManagementProvider; @@ -1249,4 +1362,32 @@ mod tests { prev = page.end; } } + + #[test] + fn test_tls() { + let platform = LinuxUserland::new(None); + let tls = platform.get_thread_local(); + assert!(tls.is_null(), "TLS should be null in the main thread"); + platform.set_thread_local(litebox_common_linux::ThreadLocalStorage { + #[cfg(target_arch = "x86")] + self_ptr: core::ptr::null_mut(), + current_task: Box::new(litebox_common_linux::Task { tid: 0xffff }), + __phantom: core::marker::PhantomData, + }); + let tls = platform.get_thread_local(); + assert!(!tls.is_null(), "TLS should not be null after setting it"); + let tls = unsafe { &mut *tls }; + assert_eq!( + tls.current_task.tid, 0xffff, + "TLS should have the correct task ID" + ); + platform.release_thread_local(|tls| { + assert_eq!( + tls.current_task.tid, 0xffff, + "TLS should have the correct task ID" + ); + }); + let tls = platform.get_thread_local(); + assert!(tls.is_null(), "TLS should be null after releasing it"); + } } From 3730bd1a77cd79734d077dc355e77a038a4512fe Mon Sep 17 00:00:00 2001 From: weitengchen Date: Fri, 20 Jun 2025 18:08:03 -0700 Subject: [PATCH 2/3] fix interface --- litebox/src/platform/mod.rs | 27 +++-- litebox_common_linux/src/lib.rs | 18 ++++ litebox_platform_linux_userland/src/lib.rs | 111 +++++++++++++-------- 3 files changed, 107 insertions(+), 49 deletions(-) diff --git a/litebox/src/platform/mod.rs b/litebox/src/platform/mod.rs index 60318e36c..35198af63 100644 --- a/litebox/src/platform/mod.rs +++ b/litebox/src/platform/mod.rs @@ -544,13 +544,28 @@ pub trait SystemInfoProvider { pub trait ThreadLocalStorageProvider { type ThreadLocalStorage; - /// Get a thread-local storage value for the current thread. - fn get_thread_local(&self) -> *mut Self::ThreadLocalStorage; - /// Set a thread-local storage value for the current thread. - fn set_thread_local(&self, value: Self::ThreadLocalStorage); + /// + /// # Panics + /// + /// Panics if TLS is set already. + fn set_thread_local_storage(&self, value: Self::ThreadLocalStorage); + + /// Invokes the provided callback function with the thread-local storage value for the current thread. + /// + /// # Panics + /// + /// Panics if TLS is not set yet. + /// Panics if TLS is borrowed already (e.g., recursive call to [`Self::release_thread_local_storage`]). + fn with_thread_local_storage_mut(&self, f: F) -> R + where + F: FnOnce(&mut Self::ThreadLocalStorage) -> R; /// Release the thread-local storage value for the current thread - /// and invoke the provided callback function on it. - fn release_thread_local(&self, f: fn(&mut Self::ThreadLocalStorage)); + /// + /// # Panics + /// + /// Panics if TLS is not set yet. + /// Panics if TLS is being used by [`Self::with_thread_local_storage_mut`]. + fn release_thread_local_storage(&self) -> Self::ThreadLocalStorage; } diff --git a/litebox_common_linux/src/lib.rs b/litebox_common_linux/src/lib.rs index 6735fa38c..08f9de9a3 100644 --- a/litebox_common_linux/src/lib.rs +++ b/litebox_common_linux/src/lib.rs @@ -751,7 +751,11 @@ bitfield::bitfield! { pub lm, set_lm: 7; } +/// Struct for thread-local storage. pub struct ThreadLocalStorage { + /// Indicates whether the TLS is being borrowed. + pub borrowed: bool, + #[cfg(target_arch = "x86")] pub self_ptr: *mut ThreadLocalStorage, pub current_task: alloc::boxed::Box, @@ -759,7 +763,21 @@ pub struct ThreadLocalStorage { pub __phantom: core::marker::PhantomData, } +impl ThreadLocalStorage { + pub const fn new(task: alloc::boxed::Box) -> Self { + Self { + borrowed: false, + #[cfg(target_arch = "x86")] + self_ptr: core::ptr::null_mut(), + current_task: task, + __phantom: core::marker::PhantomData, + } + } +} + +/// A task associated with a thread in LiteBox. pub struct Task { + /// Thread identifier pub tid: u32, } diff --git a/litebox_platform_linux_userland/src/lib.rs b/litebox_platform_linux_userland/src/lib.rs index cd2257b29..4ca93348b 100644 --- a/litebox_platform_linux_userland/src/lib.rs +++ b/litebox_platform_linux_userland/src/lib.rs @@ -1212,23 +1212,18 @@ impl litebox::platform::SystemInfoProvider for LinuxUserland { } } -/// Similar to libc, we use fs/gs registers to store thread-local storage (TLS). -/// To avoid conflicts with libc's TLS, we choose to use gs on x86_64 and fs on x86 -/// as libc uses fs on x86_64 and gs on x86. -impl litebox::platform::ThreadLocalStorageProvider for LinuxUserland { - type ThreadLocalStorage = litebox_common_linux::ThreadLocalStorage; - +impl LinuxUserland { #[cfg(target_arch = "x86_64")] - fn get_thread_local(&self) -> *mut Self::ThreadLocalStorage { + fn get_thread_local_storage() -> *mut litebox_common_linux::ThreadLocalStorage { let tls = unsafe { litebox_common_linux::rdgsbase() }; if tls == 0 { return core::ptr::null_mut(); } - tls as *mut Self::ThreadLocalStorage + tls as *mut litebox_common_linux::ThreadLocalStorage } #[cfg(target_arch = "x86")] - fn get_thread_local(&self) -> *mut Self::ThreadLocalStorage { + fn get_thread_local_storage() -> *mut litebox_common_linux::ThreadLocalStorage { let mut fs_selector: u16; unsafe { core::arch::asm!( @@ -1246,21 +1241,40 @@ impl litebox::platform::ThreadLocalStorageProvider for LinuxUserland { core::arch::asm!( "mov {0}, fs:{offset}", out(reg) addr, - offset = const core::mem::offset_of!(Self::ThreadLocalStorage, self_ptr), + offset = const core::mem::offset_of!(litebox_common_linux::ThreadLocalStorage, self_ptr), options(nostack, preserves_flags) ); } - addr as *mut Self::ThreadLocalStorage + addr as *mut litebox_common_linux::ThreadLocalStorage } +} + +/// Similar to libc, we use fs/gs registers to store thread-local storage (TLS). +/// To avoid conflicts with libc's TLS, we choose to use gs on x86_64 and fs on x86 +/// as libc uses fs on x86_64 and gs on x86. +impl litebox::platform::ThreadLocalStorageProvider for LinuxUserland { + type ThreadLocalStorage = litebox_common_linux::ThreadLocalStorage; #[cfg(target_arch = "x86_64")] - fn set_thread_local(&self, tls: Self::ThreadLocalStorage) { + fn set_thread_local_storage(&self, tls: Self::ThreadLocalStorage) { + let old_gs_base = unsafe { litebox_common_linux::rdgsbase() }; + assert!(old_gs_base == 0, "TLS already set for this thread"); let tls = Box::new(tls); unsafe { litebox_common_linux::wrgsbase(Box::into_raw(tls) as usize) }; } #[cfg(target_arch = "x86")] - fn set_thread_local(&self, tls: Self::ThreadLocalStorage) { + fn set_thread_local_storage(&self, tls: Self::ThreadLocalStorage) { + let mut old_fs_selector: u16; + unsafe { + core::arch::asm!( + "mov {0:x}, fs", + out(reg) old_fs_selector, + options(nostack, preserves_flags) + ); + } + assert!(old_fs_selector == 0, "TLS already set for this thread"); + let mut tls = Box::new(tls); tls.self_ptr = tls.as_mut(); @@ -1290,26 +1304,22 @@ impl litebox::platform::ThreadLocalStorageProvider for LinuxUserland { } #[cfg(target_arch = "x86_64")] - fn release_thread_local(&self, f: fn(&mut Self::ThreadLocalStorage)) { - let tls = self.get_thread_local(); - if tls.is_null() { - return; - } - let mut tls = unsafe { Box::from_raw(tls) }; - f(&mut tls); + fn release_thread_local_storage(&self) -> Self::ThreadLocalStorage { + let tls = Self::get_thread_local_storage(); + assert!(!tls.is_null(), "TLS must be set before releasing it"); unsafe { litebox_common_linux::wrgsbase(0); } + + let tls = unsafe { Box::from_raw(tls) }; + assert!(!tls.borrowed, "TLS must not be borrowed when releasing it"); + *tls } #[cfg(target_arch = "x86")] - fn release_thread_local(&self, f: fn(&mut Self::ThreadLocalStorage)) { - let tls = self.get_thread_local(); - if tls.is_null() { - return; - } - let mut tls = unsafe { Box::from_raw(tls) }; - f(&mut tls); + fn release_thread_local_storage(&self) -> Self::ThreadLocalStorage { + let tls = Self::get_thread_local_storage(); + assert!(!tls.is_null(), "TLS must be set before releasing it"); unsafe { core::arch::asm!( "mov fs, {0}", @@ -1317,6 +1327,24 @@ impl litebox::platform::ThreadLocalStorageProvider for LinuxUserland { options(nostack, preserves_flags) ); } + + let tls = unsafe { Box::from_raw(tls) }; + assert!(!tls.borrowed, "TLS must not be borrowed when releasing it"); + *tls + } + + fn with_thread_local_storage_mut(&self, f: F) -> R + where + F: FnOnce(&mut Self::ThreadLocalStorage) -> R, + { + let tls = Self::get_thread_local_storage(); + assert!(!tls.is_null(), "TLS must be set before accessing it"); + let tls = unsafe { &mut *tls }; + assert!(!tls.borrowed, "TLS is already borrowed"); + tls.borrowed = true; // mark as borrowed + let ret = f(tls); + tls.borrowed = false; // mark as not borrowed anymore + ret } } @@ -1366,28 +1394,25 @@ mod tests { #[test] fn test_tls() { let platform = LinuxUserland::new(None); - let tls = platform.get_thread_local(); + let tls = LinuxUserland::get_thread_local_storage(); assert!(tls.is_null(), "TLS should be null in the main thread"); - platform.set_thread_local(litebox_common_linux::ThreadLocalStorage { - #[cfg(target_arch = "x86")] - self_ptr: core::ptr::null_mut(), - current_task: Box::new(litebox_common_linux::Task { tid: 0xffff }), - __phantom: core::marker::PhantomData, - }); - let tls = platform.get_thread_local(); - assert!(!tls.is_null(), "TLS should not be null after setting it"); - let tls = unsafe { &mut *tls }; - assert_eq!( - tls.current_task.tid, 0xffff, - "TLS should have the correct task ID" - ); - platform.release_thread_local(|tls| { + platform.set_thread_local_storage(litebox_common_linux::ThreadLocalStorage::new(Box::new( + litebox_common_linux::Task { tid: 0xffff }, + ))); + platform.with_thread_local_storage_mut(|tls| { assert_eq!( tls.current_task.tid, 0xffff, "TLS should have the correct task ID" ); + tls.current_task.tid = 0x1234; // Change the task ID }); - let tls = platform.get_thread_local(); + let tls = platform.release_thread_local_storage(); + assert_eq!( + tls.current_task.tid, 0x1234, + "TLS should have the correct task ID" + ); + + let tls = LinuxUserland::get_thread_local_storage(); assert!(tls.is_null(), "TLS should be null after releasing it"); } } From dbfd54bf67072b2b7cb429ad6353abcd425857a1 Mon Sep 17 00:00:00 2001 From: Weiteng Chen Date: Mon, 23 Jun 2025 10:05:43 -0700 Subject: [PATCH 3/3] Update litebox/src/platform/mod.rs Co-authored-by: Jay Bosamiya (Microsoft) --- litebox/src/platform/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litebox/src/platform/mod.rs b/litebox/src/platform/mod.rs index 35198af63..5d23690fc 100644 --- a/litebox/src/platform/mod.rs +++ b/litebox/src/platform/mod.rs @@ -556,7 +556,7 @@ pub trait ThreadLocalStorageProvider { /// # Panics /// /// Panics if TLS is not set yet. - /// Panics if TLS is borrowed already (e.g., recursive call to [`Self::release_thread_local_storage`]). + /// Panics if TLS is borrowed already (e.g., recursive call). fn with_thread_local_storage_mut(&self, f: F) -> R where F: FnOnce(&mut Self::ThreadLocalStorage) -> R;