1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
//! Deterministic borrow guards for account data.
//!
//! `Ref` and `RefMut` provide RAII borrow tracking on the `borrow_state`
//! field of `RuntimeAccount`. When dropped, they restore the borrow
//! state, preventing use-after-free and double-mutable-borrow bugs.
//!
//! These replace `core::cell::RefCell` without requiring alloc.
use crate::NOT_BORROWED;
/// Shared (immutable) borrow guard for account data.
///
/// On drop, decrements the borrow count in `RuntimeAccount.borrow_state`.
pub struct Ref<'a, T: ?Sized> {
value: &'a T,
state: *mut u8,
}
impl<'a, T: ?Sized> Ref<'a, T> {
/// Create a new shared borrow guard.
///
/// The caller must have already incremented `*state` to reflect
/// the new shared borrow.
#[inline(always)]
pub(crate) fn new(value: &'a T, state: *mut u8) -> Self {
Self { value, state }
}
/// Create a shared borrow guard from raw parts.
///
/// # Safety
///
/// The caller must ensure:
/// - The borrow state at `state` was already incremented
/// - `value` is valid for lifetime `'a`
/// - `state` points to a valid `RuntimeAccount.borrow_state`
#[inline(always)]
pub unsafe fn from_raw_parts(value: &'a T, state: *mut u8) -> Self {
Self { value, state }
}
/// Decompose into raw parts without running the destructor.
///
/// The caller takes responsibility for eventually releasing the
/// borrow (decrementing `*state`).
#[inline(always)]
pub fn into_raw_parts(self) -> (&'a T, *mut u8) {
let value = self.value;
let state = self.state;
core::mem::forget(self);
(value, state)
}
}
impl<T: ?Sized> core::ops::Deref for Ref<'_, T> {
type Target = T;
#[inline(always)]
fn deref(&self) -> &T {
self.value
}
}
impl<T: ?Sized> Drop for Ref<'_, T> {
fn drop(&mut self) {
// SAFETY: state points to RuntimeAccount.borrow_state in the
// BPF input buffer. We decrement the shared borrow count,
// restoring NOT_BORROWED when the last shared borrow is released.
unsafe {
let current = *self.state;
if current == 1 {
*self.state = NOT_BORROWED;
} else {
*self.state = current - 1;
}
}
}
}
/// Exclusive (mutable) borrow guard for account data.
///
/// On drop, restores `RuntimeAccount.borrow_state` to `NOT_BORROWED`.
pub struct RefMut<'a, T: ?Sized> {
value: &'a mut T,
state: *mut u8,
}
impl<'a, T: ?Sized> RefMut<'a, T> {
/// Create a new exclusive borrow guard.
///
/// The caller must have already set `*state = 0` to indicate
/// exclusive borrow.
#[inline(always)]
pub(crate) fn new(value: &'a mut T, state: *mut u8) -> Self {
Self { value, state }
}
/// Create an exclusive borrow guard from raw parts.
///
/// # Safety
///
/// The caller must ensure:
/// - The borrow state at `state` was set to 0 (exclusive)
/// - `value` is valid and unique for lifetime `'a`
/// - `state` points to a valid `RuntimeAccount.borrow_state`
#[inline(always)]
pub unsafe fn from_raw_parts(value: &'a mut T, state: *mut u8) -> Self {
Self { value, state }
}
/// Decompose into raw parts without running the destructor.
///
/// The caller takes responsibility for eventually releasing the
/// borrow (restoring `*state` to `NOT_BORROWED`).
#[inline(always)]
pub fn into_raw_parts(self) -> (&'a mut T, *mut u8) {
let manual = core::mem::ManuallyDrop::new(self);
let value = unsafe { core::ptr::read(&manual.value) };
let state = manual.state;
(value, state)
}
}
impl<T: ?Sized> core::ops::Deref for RefMut<'_, T> {
type Target = T;
#[inline(always)]
fn deref(&self) -> &T {
self.value
}
}
impl<T: ?Sized> core::ops::DerefMut for RefMut<'_, T> {
#[inline(always)]
fn deref_mut(&mut self) -> &mut T {
self.value
}
}
impl<T: ?Sized> Drop for RefMut<'_, T> {
fn drop(&mut self) {
// SAFETY: state points to RuntimeAccount.borrow_state.
// Restore to NOT_BORROWED when the exclusive borrow is released.
unsafe {
*self.state = NOT_BORROWED;
}
}
}