Skip to main content

lodestone_client/
state.rs

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/// Exact public identity of one state-index incarnation.
13#[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/// Stable caller-owned identity for one prepared state mutation.
49#[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/// An absolute Unix epoch instant. This is deliberately not a duration:
78/// retries reuse the same instant and cannot extend a value's lifetime.
79#[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/// Stable outcome of the prepare/submit/resolve state-write workflow.
101#[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    /// Reads one exact key. Missing values return `None`.
125    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    /// Reads a colocated key set. Core refuses keys that do not share a slot.
148    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    /// Sets one value through the stable prepare/submit/resolve workflow.
176    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    /// Deletes one or more keys and returns the authoritative number removed.
200    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}