1use core::fmt;
2use core::sync::atomic::{AtomicUsize, Ordering};
3
4use cloud_sdk_sanitization::{SecretBuffer, sanitize_bytes};
5
6pub const MAX_CLIENT_WORKSPACE_LEASES: usize = usize::BITS as usize;
8
9#[derive(Clone, Copy, Debug, Eq, PartialEq)]
11pub enum WorkspacePoolError {
12 ZeroCapacity,
14 CapacityTooLarge,
16}
17
18impl_static_error!(WorkspacePoolError,
19 Self::ZeroCapacity => "client workspace pool capacity is zero",
20 Self::CapacityTooLarge => "client workspace pool capacity exceeds the atomic bound",
21);
22
23#[derive(Clone, Copy, Debug, Eq, PartialEq)]
25pub enum WorkspaceAcquireError {
26 Exhausted,
28}
29
30impl_static_error!(WorkspaceAcquireError,
31 Self::Exhausted => "client workspace pool is exhausted",
32);
33
34pub struct ClientWorkspace<'storage> {
48 target: SecretBuffer<'storage>,
49 request_body: SecretBuffer<'storage>,
50 response_body: SecretBuffer<'storage>,
51 response_headers: SecretBuffer<'storage>,
52}
53
54impl<'storage> ClientWorkspace<'storage> {
55 #[must_use]
57 pub fn new(
58 target: &'storage mut [u8],
59 request_body: &'storage mut [u8],
60 response_body: &'storage mut [u8],
61 response_headers: &'storage mut [u8],
62 ) -> Self {
63 sanitize_bytes(target);
64 sanitize_bytes(request_body);
65 sanitize_bytes(response_body);
66 sanitize_bytes(response_headers);
67 Self {
68 target: SecretBuffer::new(target),
69 request_body: SecretBuffer::new(request_body),
70 response_body: SecretBuffer::new(response_body),
71 response_headers: SecretBuffer::new(response_headers),
72 }
73 }
74
75 #[must_use]
77 pub fn capacities(&self) -> (usize, usize, usize, usize) {
78 (
79 self.target.as_slice().len(),
80 self.request_body.as_slice().len(),
81 self.response_body.as_slice().len(),
82 self.response_headers.as_slice().len(),
83 )
84 }
85
86 pub(crate) fn parts_mut(&mut self) -> ClientWorkspaceParts<'_> {
87 ClientWorkspaceParts {
88 target: self.target.as_mut_slice(),
89 request_body: self.request_body.as_mut_slice(),
90 response_body: self.response_body.as_mut_slice(),
91 response_headers: self.response_headers.as_mut_slice(),
92 }
93 }
94}
95
96impl fmt::Debug for ClientWorkspace<'_> {
97 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
98 formatter
99 .debug_struct("ClientWorkspace")
100 .field("capacities", &self.capacities())
101 .finish_non_exhaustive()
102 }
103}
104
105pub(crate) struct ClientWorkspaceParts<'storage> {
106 pub(crate) target: &'storage mut [u8],
107 pub(crate) request_body: &'storage mut [u8],
108 pub(crate) response_body: &'storage mut [u8],
109 pub(crate) response_headers: &'storage mut [u8],
110}
111
112impl ClientWorkspaceParts<'_> {
113 pub(crate) fn clear(&mut self) {
114 sanitize_bytes(self.target);
115 sanitize_bytes(self.request_body);
116 sanitize_bytes(self.response_body);
117 sanitize_bytes(self.response_headers);
118 }
119}
120
121pub struct ClientWorkspacePool<const N: usize> {
125 leased: AtomicUsize,
126}
127
128impl<const N: usize> ClientWorkspacePool<N> {
129 pub const fn new() -> Result<Self, WorkspacePoolError> {
131 if N == 0 {
132 return Err(WorkspacePoolError::ZeroCapacity);
133 }
134 if N > MAX_CLIENT_WORKSPACE_LEASES {
135 return Err(WorkspacePoolError::CapacityTooLarge);
136 }
137 Ok(Self {
138 leased: AtomicUsize::new(0),
139 })
140 }
141
142 pub fn try_acquire<'pool, 'storage>(
144 &'pool self,
145 workspace: ClientWorkspace<'storage>,
146 ) -> Result<ClientWorkspaceLease<'pool, 'storage, N>, WorkspaceAcquireError> {
147 let valid = valid_mask::<N>();
148 let mut observed = self.leased.load(Ordering::Acquire);
149 loop {
150 let available = !observed & valid;
151 if available == 0 {
152 return Err(WorkspaceAcquireError::Exhausted);
153 }
154 let index = available.trailing_zeros() as usize;
155 let bit = 1_usize << index;
156 match self.leased.compare_exchange_weak(
157 observed,
158 observed | bit,
159 Ordering::AcqRel,
160 Ordering::Acquire,
161 ) {
162 Ok(_) => {
163 return Ok(ClientWorkspaceLease {
164 workspace,
165 _slot: LeaseSlot { pool: self, bit },
166 });
167 }
168 Err(current) => observed = current,
169 }
170 }
171 }
172
173 #[must_use]
175 pub fn active_leases(&self) -> usize {
176 (self.leased.load(Ordering::Acquire) & valid_mask::<N>()).count_ones() as usize
177 }
178}
179
180impl<const N: usize> fmt::Debug for ClientWorkspacePool<N> {
181 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
182 formatter
183 .debug_struct("ClientWorkspacePool")
184 .field("capacity", &N)
185 .field("active_leases", &self.active_leases())
186 .finish()
187 }
188}
189
190pub struct ClientWorkspaceLease<'pool, 'storage, const N: usize> {
217 workspace: ClientWorkspace<'storage>,
218 _slot: LeaseSlot<'pool, N>,
219}
220
221impl<const N: usize> ClientWorkspaceLease<'_, '_, N> {
222 #[must_use]
224 pub fn capacities(&self) -> (usize, usize, usize, usize) {
225 self.workspace.capacities()
226 }
227
228 pub(crate) fn parts_mut(&mut self) -> ClientWorkspaceParts<'_> {
229 self.workspace.parts_mut()
230 }
231}
232
233impl<const N: usize> fmt::Debug for ClientWorkspaceLease<'_, '_, N> {
234 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
235 formatter
236 .debug_struct("ClientWorkspaceLease")
237 .field("capacities", &self.capacities())
238 .field("slot", &"[leased]")
239 .finish()
240 }
241}
242
243struct LeaseSlot<'pool, const N: usize> {
244 pool: &'pool ClientWorkspacePool<N>,
245 bit: usize,
246}
247
248impl<const N: usize> Drop for LeaseSlot<'_, N> {
249 fn drop(&mut self) {
250 self.pool.leased.fetch_and(!self.bit, Ordering::Release);
251 }
252}
253
254fn valid_mask<const N: usize>() -> usize {
255 let shift = match u32::try_from(N) {
256 Ok(value) => value,
257 Err(_) => return usize::MAX,
258 };
259 match 1_usize.checked_shl(shift) {
260 Some(upper_bit) => upper_bit.saturating_sub(1),
261 None => usize::MAX,
262 }
263}