axhal/
tls.rs

1//! Thread Local Storage (TLS) support.
2//!
3//! ## TLS layout for x86_64
4//!
5//! ```text
6//! aligned --> +-------------------------+- static_tls_offset
7//! allocation  |                         | \
8//!             | .tdata                  |  |
9//! | address   |                         |  |
10//! | grow up   + - - - - - - - - - - - - +   > Static TLS block
11//! v           |                         |  |  (length: static_tls_size)
12//!             | .tbss                   |  |
13//!             |                         |  |
14//!             +-------------------------+  |
15//!             | / PADDING / / / / / / / | /
16//!             +-------------------------+
17//!    tls_ptr -+-> self pointer (void *) | \
18//! (tp_offset) |                         |  |
19//!             | Custom TCB format       |   > Thread Control Block (TCB)
20//!             | (might be used          |  |  (length: TCB_SIZE)
21//!             |  by a libC)             |  |
22//!             |                         | /
23//!             +-------------------------+- (total length: tls_area_size)
24//! ```
25//!
26//! ## TLS layout for AArch64 and RISC-V
27//!
28//! ```text
29//!             +-------------------------+
30//!             |                         | \
31//!             | Custom TCB format       |  |
32//!             | (might be used          |   > Thread Control Block (TCB)
33//!             |  by a libC)             |  |  (length: TCB_SIZE)
34//!             |                         | /
35//!    tls_ptr -+-------------------------+
36//! (tp_offset) | GAP_ABOVE_TP            |
37//!             +-------------------------+- static_tls_offset
38//!             |                         | \
39//!             | .tdata                  |  |
40//!             |                         |  |
41//!             + - - - - - - - - - - - - +   > Static TLS block
42//!             |                         |  |  (length: static_tls_size)
43//!             | .tbss                   |  |
44//!             |                         | /
45//!             +-------------------------+- (total length: tls_area_size)
46//! ```
47//!
48//! Reference:
49//! 1. <https://github.com/unikraft/unikraft/blob/staging/arch/x86/x86_64/tls.c>
50//! 2. <https://github.com/unikraft/unikraft/blob/staging/arch/arm/arm64/tls.c>
51
52extern crate alloc;
53
54use memory_addr::align_up;
55
56use core::alloc::Layout;
57use core::ptr::NonNull;
58
59const TLS_ALIGN: usize = 0x10;
60
61cfg_if::cfg_if! {
62    if #[cfg(target_arch = "x86_64")] {
63        const TCB_SIZE: usize = 8; // to store TLS self pointer
64        const GAP_ABOVE_TP: usize = 0;
65    } else if #[cfg(target_arch = "aarch64")] {
66        const TCB_SIZE: usize = 0;
67        const GAP_ABOVE_TP: usize = 16;
68    } else if #[cfg(target_arch = "riscv64")] {
69        const TCB_SIZE: usize = 0;
70        const GAP_ABOVE_TP: usize = 0;
71    } else if #[cfg(target_arch = "loongarch64")] {
72        const TCB_SIZE: usize = 0;
73        const GAP_ABOVE_TP: usize = 0;
74    }
75}
76
77unsafe extern "C" {
78    fn _stdata();
79    fn _etdata();
80    fn _etbss();
81}
82
83/// The memory region for thread-local storage.
84pub struct TlsArea {
85    base: NonNull<u8>,
86    layout: Layout,
87}
88
89impl Drop for TlsArea {
90    fn drop(&mut self) {
91        unsafe {
92            alloc::alloc::dealloc(self.base.as_ptr(), self.layout);
93        }
94    }
95}
96
97impl TlsArea {
98    /// Returns the pointer to the TLS static area.
99    ///
100    /// One should set the hardware thread pointer register to this value.
101    pub fn tls_ptr(&self) -> *mut u8 {
102        unsafe { self.base.as_ptr().add(tp_offset()) }
103    }
104
105    /// Allocates the memory region for TLS, and initializes it.
106    pub fn alloc() -> Self {
107        let layout = Layout::from_size_align(tls_area_size(), TLS_ALIGN).unwrap();
108        let area_base = unsafe { alloc::alloc::alloc_zeroed(layout) };
109
110        let tls_load_base = _stdata as *mut u8;
111        let tls_load_size = _etbss as usize - _stdata as usize;
112        unsafe {
113            // copy data from .tbdata section
114            core::ptr::copy_nonoverlapping(
115                tls_load_base,
116                area_base.add(static_tls_offset()),
117                tls_load_size,
118            );
119            // initialize TCB
120            init_tcb(area_base);
121        }
122
123        Self {
124            base: NonNull::new(area_base).unwrap(),
125            layout,
126        }
127    }
128}
129
130fn static_tls_size() -> usize {
131    align_up(_etbss as usize - _stdata as usize, TLS_ALIGN)
132}
133
134fn static_tls_offset() -> usize {
135    if cfg!(target_arch = "x86_64") {
136        0
137    } else if cfg!(any(
138        target_arch = "aarch64",
139        target_arch = "riscv64",
140        target_arch = "loongarch64"
141    )) {
142        TCB_SIZE + GAP_ABOVE_TP
143    } else {
144        unreachable!()
145    }
146}
147
148fn tp_offset() -> usize {
149    if cfg!(target_arch = "x86_64") {
150        static_tls_size()
151    } else if cfg!(any(
152        target_arch = "aarch64",
153        target_arch = "riscv64",
154        target_arch = "loongarch64"
155    )) {
156        TCB_SIZE
157    } else {
158        unreachable!()
159    }
160}
161
162fn tls_area_size() -> usize {
163    if cfg!(target_arch = "x86_64") {
164        static_tls_size() + TCB_SIZE
165    } else if cfg!(any(
166        target_arch = "aarch64",
167        target_arch = "riscv64",
168        target_arch = "loongarch64"
169    )) {
170        TCB_SIZE + GAP_ABOVE_TP + static_tls_size()
171    } else {
172        unreachable!()
173    }
174}
175
176unsafe fn init_tcb(tls_area: *mut u8) {
177    if cfg!(target_arch = "x86_64") {
178        let tp_addr = tls_area.add(tp_offset()).cast::<usize>();
179        tp_addr.write(tp_addr as usize); // write self pointer
180    }
181}