1use lodestone_wire::{
2 BinaryValue, Digest, ExecuteStateReadRequest, Identity, MutationOutcome, MutationResolver,
3 PrepareStateWriteRequest, ResolveStateWriteRequest, Resource, StateCommandResult, StateOperation,
4 StateOptionalValue, StateSetCondition, StateWriteResolution, SubmitStateWriteRequest, U64,
5};
6
7use crate::{CallOptions, Client, ClientError, SubmissionDisposition};
8
9const MAX_STATE_KEY_BYTES: usize = 254;
10const MAX_STATE_BATCH_KEYS: usize = 128;
11
12#[derive(Clone, Debug, Eq, PartialEq)]
14pub struct StateIndex {
15 name: String,
16 incarnation: Identity,
17}
18
19impl StateIndex {
20 pub fn new(name: impl Into<String>, incarnation: Identity) -> Result<Self, ClientError> {
21 let name = name.into();
22 crate::transport::validate_index_name(&name)?;
23 if !valid_identity(&incarnation.value) {
24 return Err(ClientError::InvalidStateIndex);
25 }
26 Ok(Self { name, incarnation })
27 }
28
29 #[must_use]
30 pub fn name(&self) -> &str {
31 &self.name
32 }
33
34 #[must_use]
35 pub fn incarnation(&self) -> &Identity {
36 &self.incarnation
37 }
38
39 fn resource(&self) -> Resource {
40 Resource {
41 kind: "index".to_owned(),
42 name: self.name.clone(),
43 incarnation: self.incarnation.clone(),
44 }
45 }
46}
47
48#[derive(Clone, Debug, Eq, PartialEq)]
50pub struct StateActionId(Digest);
51
52impl StateActionId {
53 pub fn new(value: Digest) -> Result<Self, ClientError> {
54 if !valid_digest(&value.value) {
55 return Err(ClientError::InvalidStateActionId);
56 }
57 Ok(Self(value))
58 }
59
60 pub fn generate() -> Result<Self, ClientError> {
61 let mut bytes = [0_u8; 32];
62 getrandom::fill(&mut bytes).map_err(|_| ClientError::RandomUnavailable)?;
63 if bytes == [0; 32] {
64 return Err(ClientError::RandomUnavailable);
65 }
66 Self::new(Digest {
67 value: lower_hex(&bytes),
68 })
69 }
70
71 #[must_use]
72 pub fn as_digest(&self) -> &Digest {
73 &self.0
74 }
75}
76
77#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd)]
80pub struct AbsoluteExpiry(u64);
81
82impl AbsoluteExpiry {
83 #[must_use]
84 pub const fn from_unix_millis(value: u64) -> Self {
85 Self(value)
86 }
87
88 #[must_use]
89 pub const fn unix_millis(self) -> u64 {
90 self.0
91 }
92
93 fn wire(self) -> U64 {
94 U64 {
95 value: self.0.to_string(),
96 }
97 }
98}
99
100#[derive(Clone, Debug, PartialEq)]
102pub enum StateWriteOutcome<T> {
103 Applied(T),
104 Resolved {
105 value: T,
106 sequence: U64,
107 },
108 NotSubmitted,
109 KnownRetired {
110 retired_through: U64,
111 },
112 StateReset {
113 submitted_reset_epoch: U64,
114 current_reset_epoch: U64,
115 },
116 Indeterminate {
117 observed_high_water: Option<U64>,
118 current_high_water: Option<U64>,
119 retired_through: Option<U64>,
120 },
121}
122
123impl Client {
124 pub async fn get(
126 &self,
127 index: &StateIndex,
128 key: &[u8],
129 options: CallOptions,
130 ) -> Result<Option<Vec<u8>>, ClientError> {
131 validate_key(key)?;
132 let result = self
133 .execute_state_read(
134 ExecuteStateReadRequest {
135 index: index.resource(),
136 operation: StateOperation::Get { key: binary(key) },
137 },
138 options,
139 )
140 .await?;
141 match result.result {
142 StateCommandResult::Get { value } => optional_value(value),
143 _ => Err(ClientError::UnexpectedStateResult),
144 }
145 }
146
147 pub async fn mget(
149 &self,
150 index: &StateIndex,
151 keys: &[Vec<u8>],
152 options: CallOptions,
153 ) -> Result<Vec<Option<Vec<u8>>>, ClientError> {
154 validate_keys(keys)?;
155 let result = self
156 .execute_state_read(
157 ExecuteStateReadRequest {
158 index: index.resource(),
159 operation: StateOperation::MultiGet {
160 keys: keys.iter().map(|key| binary(key)).collect(),
161 },
162 },
163 options,
164 )
165 .await?;
166 match result.result {
167 StateCommandResult::MultiGet { values } if values.len() == keys.len() => {
168 values.into_iter().map(optional_value).collect()
169 }
170 StateCommandResult::MultiGet { .. } => Err(ClientError::UnexpectedStateResult),
171 _ => Err(ClientError::UnexpectedStateResult),
172 }
173 }
174
175 pub async fn set(
177 &self,
178 index: &StateIndex,
179 action_id: &StateActionId,
180 key: &[u8],
181 value: &[u8],
182 expires_at: Option<AbsoluteExpiry>,
183 options: CallOptions,
184 ) -> Result<StateWriteOutcome<bool>, ClientError> {
185 validate_key(key)?;
186 let operation = StateOperation::Set {
187 key: binary(key),
188 value: binary(value),
189 condition: StateSetCondition::Always,
190 expires_at_millis: expires_at.map(AbsoluteExpiry::wire),
191 };
192 self.execute_state_write(index, action_id, operation, options, |result| match result {
193 StateCommandResult::Set { applied } => Ok(applied),
194 _ => Err(ClientError::UnexpectedStateResult),
195 })
196 .await
197 }
198
199 pub async fn delete(
201 &self,
202 index: &StateIndex,
203 action_id: &StateActionId,
204 keys: &[Vec<u8>],
205 options: CallOptions,
206 ) -> Result<StateWriteOutcome<u32>, ClientError> {
207 validate_keys(keys)?;
208 let operation = StateOperation::Delete {
209 keys: keys.iter().map(|key| binary(key)).collect(),
210 };
211 self.execute_state_write(index, action_id, operation, options, |result| match result {
212 StateCommandResult::Delete { removed } => Ok(removed),
213 _ => Err(ClientError::UnexpectedStateResult),
214 })
215 .await
216 }
217
218 async fn execute_state_write<T, F>(
219 &self,
220 index: &StateIndex,
221 action_id: &StateActionId,
222 operation: StateOperation,
223 options: CallOptions,
224 decode: F,
225 ) -> Result<StateWriteOutcome<T>, ClientError>
226 where
227 F: Fn(StateCommandResult) -> Result<T, ClientError>,
228 {
229 let prepared = self
230 .prepare_state_write(
231 PrepareStateWriteRequest {
232 index: index.resource(),
233 action_id: action_id.as_digest().clone(),
234 operation: operation.clone(),
235 },
236 options.clone(),
237 )
238 .await?;
239 let submitted = self
240 .submit_state_write(
241 SubmitStateWriteRequest {
242 index: index.resource(),
243 action_id: action_id.as_digest().clone(),
244 operation: operation.clone(),
245 locator: prepared.locator.clone(),
246 semantic_digest: prepared.semantic_digest.clone(),
247 },
248 options.clone(),
249 )
250 .await;
251 match submitted {
252 Ok(result) => match (result.result, result.resolution, result.mutation_outcome) {
253 (Some(value), None, MutationOutcome::Terminal) => Ok(StateWriteOutcome::Applied(decode(value)?)),
254 (
255 None,
256 Some(resolution),
257 MutationOutcome::ResolutionRequired {
258 resolver: MutationResolver::ResolveStateWrite,
259 },
260 ) => decode_resolution(resolution, decode),
261 _ => Err(ClientError::UnexpectedStateResult),
262 },
263 Err(error)
264 if matches!(
265 error.submission_disposition(),
266 SubmissionDisposition::Submitted | SubmissionDisposition::Indeterminate
267 ) =>
268 {
269 let resolved = self
270 .resolve_state_write(
271 ResolveStateWriteRequest {
272 index: index.resource(),
273 action_id: action_id.as_digest().clone(),
274 operation,
275 locator: prepared.locator,
276 semantic_digest: prepared.semantic_digest,
277 },
278 options.for_resolution(),
279 )
280 .await;
281 let resolved = match resolved {
282 Ok(resolved) => resolved,
283 Err(error) if error.is_deadline_exceeded() => {
284 return Ok(StateWriteOutcome::Indeterminate {
285 observed_high_water: None,
286 current_high_water: None,
287 retired_through: None,
288 });
289 }
290 Err(error) => return Err(error),
291 };
292 decode_resolution(resolved.resolution, decode)
293 }
294 Err(error) => Err(error),
295 }
296 }
297}
298
299fn decode_resolution<T, F>(resolution: StateWriteResolution, decode: F) -> Result<StateWriteOutcome<T>, ClientError>
300where
301 F: Fn(StateCommandResult) -> Result<T, ClientError>,
302{
303 match resolution {
304 StateWriteResolution::NotSubmitted => Ok(StateWriteOutcome::NotSubmitted),
305 StateWriteResolution::Resolved { result, sequence } => Ok(StateWriteOutcome::Resolved {
306 value: decode(result)?,
307 sequence,
308 }),
309 StateWriteResolution::KnownRetired { retired_through } => {
310 Ok(StateWriteOutcome::KnownRetired { retired_through })
311 }
312 StateWriteResolution::StateReset {
313 submitted_reset_epoch,
314 current_reset_epoch,
315 } => Ok(StateWriteOutcome::StateReset {
316 submitted_reset_epoch,
317 current_reset_epoch,
318 }),
319 StateWriteResolution::Indeterminate {
320 observed_high_water,
321 current_high_water,
322 retired_through,
323 } => Ok(StateWriteOutcome::Indeterminate {
324 observed_high_water,
325 current_high_water,
326 retired_through,
327 }),
328 }
329}
330
331fn validate_key(key: &[u8]) -> Result<(), ClientError> {
332 if key.is_empty() || key.len() > MAX_STATE_KEY_BYTES {
333 return Err(ClientError::InvalidStateKey);
334 }
335 Ok(())
336}
337
338fn validate_keys(keys: &[Vec<u8>]) -> Result<(), ClientError> {
339 if keys.is_empty() || keys.len() > MAX_STATE_BATCH_KEYS {
340 return Err(ClientError::InvalidStateBatch);
341 }
342 keys.iter().try_for_each(|key| validate_key(key))
343}
344
345fn optional_value(value: StateOptionalValue) -> Result<Option<Vec<u8>>, ClientError> {
346 match value {
347 StateOptionalValue::Missing => Ok(None),
348 StateOptionalValue::Value { value } => decode_binary(&value).map(Some),
349 }
350}
351
352fn binary(value: &[u8]) -> BinaryValue {
353 BinaryValue {
354 bytes_base64: encode_base64(value),
355 }
356}
357
358fn encode_base64(bytes: &[u8]) -> String {
359 const ALPHABET: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
360 let mut output = String::with_capacity(bytes.len().div_ceil(3).saturating_mul(4));
361 for chunk in bytes.chunks(3) {
362 let first = chunk[0];
363 let second = chunk.get(1).copied().unwrap_or(0);
364 let third = chunk.get(2).copied().unwrap_or(0);
365 output.push(char::from(ALPHABET[usize::from(first >> 2)]));
366 output.push(char::from(ALPHABET[usize::from(((first & 0x03) << 4) | (second >> 4))]));
367 output.push(if chunk.len() > 1 {
368 char::from(ALPHABET[usize::from(((second & 0x0f) << 2) | (third >> 6))])
369 } else {
370 '='
371 });
372 output.push(if chunk.len() > 2 {
373 char::from(ALPHABET[usize::from(third & 0x3f)])
374 } else {
375 '='
376 });
377 }
378 output
379}
380
381fn decode_binary(value: &BinaryValue) -> Result<Vec<u8>, ClientError> {
382 let input = value.bytes_base64.as_bytes();
383 if input.is_empty() {
384 return Ok(Vec::new());
385 }
386 if !input.len().is_multiple_of(4) {
387 return Err(ClientError::UnexpectedStateResult);
388 }
389 let mut output = Vec::with_capacity(input.len() / 4 * 3);
390 for (ordinal, chunk) in input.as_chunks::<4>().0.iter().enumerate() {
391 let last = ordinal + 1 == input.len() / 4;
392 let a = base64_digit(chunk[0])?;
393 let b = base64_digit(chunk[1])?;
394 let c = if chunk[2] == b'=' { 0 } else { base64_digit(chunk[2])? };
395 let d = if chunk[3] == b'=' { 0 } else { base64_digit(chunk[3])? };
396 let padding_before_last_chunk = !last && (chunk[2] == b'=' || chunk[3] == b'=');
397 let invalid_padding_order = chunk[2] == b'=' && chunk[3] != b'=';
398 let nonzero_double_padding_bits = chunk[2] == b'=' && b & 0x0f != 0;
399 let nonzero_single_padding_bits = chunk[3] == b'=' && chunk[2] != b'=' && c & 0x03 != 0;
400 if padding_before_last_chunk
401 || invalid_padding_order
402 || nonzero_double_padding_bits
403 || nonzero_single_padding_bits
404 {
405 return Err(ClientError::UnexpectedStateResult);
406 }
407 output.push((a << 2) | (b >> 4));
408 if chunk[2] != b'=' {
409 output.push((b << 4) | (c >> 2));
410 }
411 if chunk[3] != b'=' {
412 output.push((c << 6) | d);
413 }
414 }
415 if encode_base64(&output) != value.bytes_base64 {
416 return Err(ClientError::UnexpectedStateResult);
417 }
418 Ok(output)
419}
420
421fn base64_digit(value: u8) -> Result<u8, ClientError> {
422 match value {
423 b'A'..=b'Z' => Ok(value - b'A'),
424 b'a'..=b'z' => Ok(value - b'a' + 26),
425 b'0'..=b'9' => Ok(value - b'0' + 52),
426 b'+' => Ok(62),
427 b'/' => Ok(63),
428 _ => Err(ClientError::UnexpectedStateResult),
429 }
430}
431
432fn valid_identity(value: &str) -> bool {
433 value.len() == 32 && value != "00000000000000000000000000000000" && value.bytes().all(is_lower_hex)
434}
435
436fn valid_digest(value: &str) -> bool {
437 value.len() == 64
438 && value != "0000000000000000000000000000000000000000000000000000000000000000"
439 && value.bytes().all(is_lower_hex)
440}
441
442const fn is_lower_hex(value: u8) -> bool {
443 value.is_ascii_digit() || matches!(value, b'a'..=b'f')
444}
445
446fn lower_hex(bytes: &[u8]) -> String {
447 const HEX: &[u8; 16] = b"0123456789abcdef";
448 let mut output = String::with_capacity(bytes.len().saturating_mul(2));
449 for byte in bytes {
450 output.push(char::from(HEX[usize::from(byte >> 4)]));
451 output.push(char::from(HEX[usize::from(byte & 0x0f)]));
452 }
453 output
454}
455
456#[cfg(test)]
457mod tests {
458 use super::*;
459
460 #[test]
461 fn state_value_base64_is_canonical_and_round_trips() -> Result<(), ClientError> {
462 for value in [b"".as_slice(), b"a", b"ab", b"abc", &[0, 1, 254, 255]] {
463 let encoded = binary(value);
464 assert_eq!(decode_binary(&encoded)?, value);
465 }
466 assert!(
467 decode_binary(&BinaryValue {
468 bytes_base64: "YR==".to_owned(),
469 })
470 .is_err()
471 );
472 Ok(())
473 }
474
475 #[test]
476 fn state_semantic_inputs_refuse_ambiguous_identities_and_keys() {
477 assert!(StateIndex::new("state", Identity { value: "0".repeat(32) }).is_err());
478 assert!(StateActionId::new(Digest { value: "0".repeat(64) }).is_err());
479 assert!(validate_key(&[]).is_err());
480 assert!(validate_key(&vec![0; 255]).is_err());
481 assert_eq!(AbsoluteExpiry::from_unix_millis(42).unix_millis(), 42);
482 }
483}