Skip to main content

lodestone_client/
state.rs

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/// 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, 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}