1use lodestone_wire::{
2 BinaryValue, Digest, ExecuteStateReadRequest, Identity, PrepareStateWriteRequest, ResolveStateWriteRequest,
3 Resource, StateCommandResult, StateOperation, StateOptionalValue, StateSetCondition, StateWriteResolution,
4 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) {
253 (Some(value), None) if result.submitted => Ok(StateWriteOutcome::Applied(decode(value)?)),
254 (None, Some(resolution)) if result.submitted => decode_resolution(resolution, decode),
255 _ => Err(ClientError::UnexpectedStateResult),
256 },
257 Err(error)
258 if matches!(
259 error.submission_disposition(),
260 SubmissionDisposition::Submitted | SubmissionDisposition::Indeterminate
261 ) =>
262 {
263 let resolved = self
264 .resolve_state_write(
265 ResolveStateWriteRequest {
266 index: index.resource(),
267 action_id: action_id.as_digest().clone(),
268 operation,
269 locator: prepared.locator,
270 semantic_digest: prepared.semantic_digest,
271 },
272 options.for_resolution(),
273 )
274 .await;
275 let resolved = match resolved {
276 Ok(resolved) => resolved,
277 Err(error) if error.is_deadline_exceeded() => {
278 return Ok(StateWriteOutcome::Indeterminate {
279 observed_high_water: None,
280 current_high_water: None,
281 retired_through: None,
282 });
283 }
284 Err(error) => return Err(error),
285 };
286 decode_resolution(resolved.resolution, decode)
287 }
288 Err(error) => Err(error),
289 }
290 }
291}
292
293fn decode_resolution<T, F>(resolution: StateWriteResolution, decode: F) -> Result<StateWriteOutcome<T>, ClientError>
294where
295 F: Fn(StateCommandResult) -> Result<T, ClientError>,
296{
297 match resolution {
298 StateWriteResolution::NotSubmitted => Ok(StateWriteOutcome::NotSubmitted),
299 StateWriteResolution::Resolved { result, sequence } => Ok(StateWriteOutcome::Resolved {
300 value: decode(result)?,
301 sequence,
302 }),
303 StateWriteResolution::KnownRetired { retired_through } => {
304 Ok(StateWriteOutcome::KnownRetired { retired_through })
305 }
306 StateWriteResolution::StateReset {
307 submitted_reset_epoch,
308 current_reset_epoch,
309 } => Ok(StateWriteOutcome::StateReset {
310 submitted_reset_epoch,
311 current_reset_epoch,
312 }),
313 StateWriteResolution::Indeterminate {
314 observed_high_water,
315 current_high_water,
316 retired_through,
317 } => Ok(StateWriteOutcome::Indeterminate {
318 observed_high_water,
319 current_high_water,
320 retired_through,
321 }),
322 }
323}
324
325fn validate_key(key: &[u8]) -> Result<(), ClientError> {
326 if key.is_empty() || key.len() > MAX_STATE_KEY_BYTES {
327 return Err(ClientError::InvalidStateKey);
328 }
329 Ok(())
330}
331
332fn validate_keys(keys: &[Vec<u8>]) -> Result<(), ClientError> {
333 if keys.is_empty() || keys.len() > MAX_STATE_BATCH_KEYS {
334 return Err(ClientError::InvalidStateBatch);
335 }
336 keys.iter().try_for_each(|key| validate_key(key))
337}
338
339fn optional_value(value: StateOptionalValue) -> Result<Option<Vec<u8>>, ClientError> {
340 match value {
341 StateOptionalValue::Missing => Ok(None),
342 StateOptionalValue::Value { value } => decode_binary(&value).map(Some),
343 }
344}
345
346fn binary(value: &[u8]) -> BinaryValue {
347 BinaryValue {
348 bytes_base64: encode_base64(value),
349 }
350}
351
352fn encode_base64(bytes: &[u8]) -> String {
353 const ALPHABET: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
354 let mut output = String::with_capacity(bytes.len().div_ceil(3).saturating_mul(4));
355 for chunk in bytes.chunks(3) {
356 let first = chunk[0];
357 let second = chunk.get(1).copied().unwrap_or(0);
358 let third = chunk.get(2).copied().unwrap_or(0);
359 output.push(char::from(ALPHABET[usize::from(first >> 2)]));
360 output.push(char::from(ALPHABET[usize::from(((first & 0x03) << 4) | (second >> 4))]));
361 output.push(if chunk.len() > 1 {
362 char::from(ALPHABET[usize::from(((second & 0x0f) << 2) | (third >> 6))])
363 } else {
364 '='
365 });
366 output.push(if chunk.len() > 2 {
367 char::from(ALPHABET[usize::from(third & 0x3f)])
368 } else {
369 '='
370 });
371 }
372 output
373}
374
375fn decode_binary(value: &BinaryValue) -> Result<Vec<u8>, ClientError> {
376 let input = value.bytes_base64.as_bytes();
377 if input.is_empty() {
378 return Ok(Vec::new());
379 }
380 if !input.len().is_multiple_of(4) {
381 return Err(ClientError::UnexpectedStateResult);
382 }
383 let mut output = Vec::with_capacity(input.len() / 4 * 3);
384 for (ordinal, chunk) in input.chunks_exact(4).enumerate() {
385 let last = ordinal + 1 == input.len() / 4;
386 let a = base64_digit(chunk[0])?;
387 let b = base64_digit(chunk[1])?;
388 let c = if chunk[2] == b'=' { 0 } else { base64_digit(chunk[2])? };
389 let d = if chunk[3] == b'=' { 0 } else { base64_digit(chunk[3])? };
390 let padding_before_last_chunk = !last && (chunk[2] == b'=' || chunk[3] == b'=');
391 let invalid_padding_order = chunk[2] == b'=' && chunk[3] != b'=';
392 let nonzero_double_padding_bits = chunk[2] == b'=' && b & 0x0f != 0;
393 let nonzero_single_padding_bits = chunk[3] == b'=' && chunk[2] != b'=' && c & 0x03 != 0;
394 if padding_before_last_chunk
395 || invalid_padding_order
396 || nonzero_double_padding_bits
397 || nonzero_single_padding_bits
398 {
399 return Err(ClientError::UnexpectedStateResult);
400 }
401 output.push((a << 2) | (b >> 4));
402 if chunk[2] != b'=' {
403 output.push((b << 4) | (c >> 2));
404 }
405 if chunk[3] != b'=' {
406 output.push((c << 6) | d);
407 }
408 }
409 if encode_base64(&output) != value.bytes_base64 {
410 return Err(ClientError::UnexpectedStateResult);
411 }
412 Ok(output)
413}
414
415fn base64_digit(value: u8) -> Result<u8, ClientError> {
416 match value {
417 b'A'..=b'Z' => Ok(value - b'A'),
418 b'a'..=b'z' => Ok(value - b'a' + 26),
419 b'0'..=b'9' => Ok(value - b'0' + 52),
420 b'+' => Ok(62),
421 b'/' => Ok(63),
422 _ => Err(ClientError::UnexpectedStateResult),
423 }
424}
425
426fn valid_identity(value: &str) -> bool {
427 value.len() == 32 && value != "00000000000000000000000000000000" && value.bytes().all(is_lower_hex)
428}
429
430fn valid_digest(value: &str) -> bool {
431 value.len() == 64
432 && value != "0000000000000000000000000000000000000000000000000000000000000000"
433 && value.bytes().all(is_lower_hex)
434}
435
436const fn is_lower_hex(value: u8) -> bool {
437 value.is_ascii_digit() || matches!(value, b'a'..=b'f')
438}
439
440fn lower_hex(bytes: &[u8]) -> String {
441 const HEX: &[u8; 16] = b"0123456789abcdef";
442 let mut output = String::with_capacity(bytes.len().saturating_mul(2));
443 for byte in bytes {
444 output.push(char::from(HEX[usize::from(byte >> 4)]));
445 output.push(char::from(HEX[usize::from(byte & 0x0f)]));
446 }
447 output
448}
449
450#[cfg(test)]
451mod tests {
452 use super::*;
453
454 #[test]
455 fn state_value_base64_is_canonical_and_round_trips() -> Result<(), ClientError> {
456 for value in [b"".as_slice(), b"a", b"ab", b"abc", &[0, 1, 254, 255]] {
457 let encoded = binary(value);
458 assert_eq!(decode_binary(&encoded)?, value);
459 }
460 assert!(
461 decode_binary(&BinaryValue {
462 bytes_base64: "YR==".to_owned(),
463 })
464 .is_err()
465 );
466 Ok(())
467 }
468
469 #[test]
470 fn state_semantic_inputs_refuse_ambiguous_identities_and_keys() {
471 assert!(StateIndex::new("state", Identity { value: "0".repeat(32) }).is_err());
472 assert!(StateActionId::new(Digest { value: "0".repeat(64) }).is_err());
473 assert!(validate_key(&[]).is_err());
474 assert!(validate_key(&vec![0; 255]).is_err());
475 assert_eq!(AbsoluteExpiry::from_unix_millis(42).unix_millis(), 42);
476 }
477}