lodestone-client 0.0.4

use lodestone_wire::{
    BinaryValue, Digest, ExecuteStateReadRequest, Identity, PrepareStateWriteRequest, ResolveStateWriteRequest,
    Resource, StateCommandResult, StateOperation, StateOptionalValue, StateSetCondition, StateWriteResolution,
    SubmitStateWriteRequest, U64,
};

use crate::{CallOptions, Client, ClientError, SubmissionDisposition};

const MAX_STATE_KEY_BYTES: usize = 254;
const MAX_STATE_BATCH_KEYS: usize = 128;

/// Exact public identity of one state-index incarnation.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct StateIndex {
    name: String,
    incarnation: Identity,
}

impl StateIndex {
    pub fn new(name: impl Into<String>, incarnation: Identity) -> Result<Self, ClientError> {
        let name = name.into();
        crate::transport::validate_index_name(&name)?;
        if !valid_identity(&incarnation.value) {
            return Err(ClientError::InvalidStateIndex);
        }
        Ok(Self { name, incarnation })
    }

    #[must_use]
    pub fn name(&self) -> &str {
        &self.name
    }

    #[must_use]
    pub fn incarnation(&self) -> &Identity {
        &self.incarnation
    }

    fn resource(&self) -> Resource {
        Resource {
            kind: "index".to_owned(),
            name: self.name.clone(),
            incarnation: self.incarnation.clone(),
        }
    }
}

/// Stable caller-owned identity for one prepared state mutation.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct StateActionId(Digest);

impl StateActionId {
    pub fn new(value: Digest) -> Result<Self, ClientError> {
        if !valid_digest(&value.value) {
            return Err(ClientError::InvalidStateActionId);
        }
        Ok(Self(value))
    }

    pub fn generate() -> Result<Self, ClientError> {
        let mut bytes = [0_u8; 32];
        getrandom::fill(&mut bytes).map_err(|_| ClientError::RandomUnavailable)?;
        if bytes == [0; 32] {
            return Err(ClientError::RandomUnavailable);
        }
        Self::new(Digest {
            value: lower_hex(&bytes),
        })
    }

    #[must_use]
    pub fn as_digest(&self) -> &Digest {
        &self.0
    }
}

/// An absolute Unix epoch instant. This is deliberately not a duration:
/// retries reuse the same instant and cannot extend a value's lifetime.
#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd)]
pub struct AbsoluteExpiry(u64);

impl AbsoluteExpiry {
    #[must_use]
    pub const fn from_unix_millis(value: u64) -> Self {
        Self(value)
    }

    #[must_use]
    pub const fn unix_millis(self) -> u64 {
        self.0
    }

    fn wire(self) -> U64 {
        U64 {
            value: self.0.to_string(),
        }
    }
}

/// Stable outcome of the prepare/submit/resolve state-write workflow.
#[derive(Clone, Debug, PartialEq)]
pub enum StateWriteOutcome<T> {
    Applied(T),
    Resolved {
        value: T,
        sequence: U64,
    },
    NotSubmitted,
    KnownRetired {
        retired_through: U64,
    },
    StateReset {
        submitted_reset_epoch: U64,
        current_reset_epoch: U64,
    },
    Indeterminate {
        observed_high_water: Option<U64>,
        current_high_water: Option<U64>,
        retired_through: Option<U64>,
    },
}

impl Client {
    /// Reads one exact key. Missing values return `None`.
    pub async fn get(
        &self,
        index: &StateIndex,
        key: &[u8],
        options: CallOptions,
    ) -> Result<Option<Vec<u8>>, ClientError> {
        validate_key(key)?;
        let result = self
            .execute_state_read(
                ExecuteStateReadRequest {
                    index: index.resource(),
                    operation: StateOperation::Get { key: binary(key) },
                },
                options,
            )
            .await?;
        match result.result {
            StateCommandResult::Get { value } => optional_value(value),
            _ => Err(ClientError::UnexpectedStateResult),
        }
    }

    /// Reads a colocated key set. Core refuses keys that do not share a slot.
    pub async fn mget(
        &self,
        index: &StateIndex,
        keys: &[Vec<u8>],
        options: CallOptions,
    ) -> Result<Vec<Option<Vec<u8>>>, ClientError> {
        validate_keys(keys)?;
        let result = self
            .execute_state_read(
                ExecuteStateReadRequest {
                    index: index.resource(),
                    operation: StateOperation::MultiGet {
                        keys: keys.iter().map(|key| binary(key)).collect(),
                    },
                },
                options,
            )
            .await?;
        match result.result {
            StateCommandResult::MultiGet { values } if values.len() == keys.len() => {
                values.into_iter().map(optional_value).collect()
            }
            StateCommandResult::MultiGet { .. } => Err(ClientError::UnexpectedStateResult),
            _ => Err(ClientError::UnexpectedStateResult),
        }
    }

    /// Sets one value through the stable prepare/submit/resolve workflow.
    pub async fn set(
        &self,
        index: &StateIndex,
        action_id: &StateActionId,
        key: &[u8],
        value: &[u8],
        expires_at: Option<AbsoluteExpiry>,
        options: CallOptions,
    ) -> Result<StateWriteOutcome<bool>, ClientError> {
        validate_key(key)?;
        let operation = StateOperation::Set {
            key: binary(key),
            value: binary(value),
            condition: StateSetCondition::Always,
            expires_at_millis: expires_at.map(AbsoluteExpiry::wire),
        };
        self.execute_state_write(index, action_id, operation, options, |result| match result {
            StateCommandResult::Set { applied } => Ok(applied),
            _ => Err(ClientError::UnexpectedStateResult),
        })
        .await
    }

    /// Deletes one or more keys and returns the authoritative number removed.
    pub async fn delete(
        &self,
        index: &StateIndex,
        action_id: &StateActionId,
        keys: &[Vec<u8>],
        options: CallOptions,
    ) -> Result<StateWriteOutcome<u32>, ClientError> {
        validate_keys(keys)?;
        let operation = StateOperation::Delete {
            keys: keys.iter().map(|key| binary(key)).collect(),
        };
        self.execute_state_write(index, action_id, operation, options, |result| match result {
            StateCommandResult::Delete { removed } => Ok(removed),
            _ => Err(ClientError::UnexpectedStateResult),
        })
        .await
    }

    async fn execute_state_write<T, F>(
        &self,
        index: &StateIndex,
        action_id: &StateActionId,
        operation: StateOperation,
        options: CallOptions,
        decode: F,
    ) -> Result<StateWriteOutcome<T>, ClientError>
    where
        F: Fn(StateCommandResult) -> Result<T, ClientError>,
    {
        let prepared = self
            .prepare_state_write(
                PrepareStateWriteRequest {
                    index: index.resource(),
                    action_id: action_id.as_digest().clone(),
                    operation: operation.clone(),
                },
                options.clone(),
            )
            .await?;
        let submitted = self
            .submit_state_write(
                SubmitStateWriteRequest {
                    index: index.resource(),
                    action_id: action_id.as_digest().clone(),
                    operation: operation.clone(),
                    locator: prepared.locator.clone(),
                    semantic_digest: prepared.semantic_digest.clone(),
                },
                options.clone(),
            )
            .await;
        match submitted {
            Ok(result) => match (result.result, result.resolution) {
                (Some(value), None) if result.submitted => Ok(StateWriteOutcome::Applied(decode(value)?)),
                (None, Some(resolution)) if result.submitted => decode_resolution(resolution, decode),
                _ => Err(ClientError::UnexpectedStateResult),
            },
            Err(error)
                if matches!(
                    error.submission_disposition(),
                    SubmissionDisposition::Submitted | SubmissionDisposition::Indeterminate
                ) =>
            {
                let resolved = self
                    .resolve_state_write(
                        ResolveStateWriteRequest {
                            index: index.resource(),
                            action_id: action_id.as_digest().clone(),
                            operation,
                            locator: prepared.locator,
                            semantic_digest: prepared.semantic_digest,
                        },
                        options.for_resolution(),
                    )
                    .await;
                let resolved = match resolved {
                    Ok(resolved) => resolved,
                    Err(error) if error.is_deadline_exceeded() => {
                        return Ok(StateWriteOutcome::Indeterminate {
                            observed_high_water: None,
                            current_high_water: None,
                            retired_through: None,
                        });
                    }
                    Err(error) => return Err(error),
                };
                decode_resolution(resolved.resolution, decode)
            }
            Err(error) => Err(error),
        }
    }
}

fn decode_resolution<T, F>(resolution: StateWriteResolution, decode: F) -> Result<StateWriteOutcome<T>, ClientError>
where
    F: Fn(StateCommandResult) -> Result<T, ClientError>,
{
    match resolution {
        StateWriteResolution::NotSubmitted => Ok(StateWriteOutcome::NotSubmitted),
        StateWriteResolution::Resolved { result, sequence } => Ok(StateWriteOutcome::Resolved {
            value: decode(result)?,
            sequence,
        }),
        StateWriteResolution::KnownRetired { retired_through } => {
            Ok(StateWriteOutcome::KnownRetired { retired_through })
        }
        StateWriteResolution::StateReset {
            submitted_reset_epoch,
            current_reset_epoch,
        } => Ok(StateWriteOutcome::StateReset {
            submitted_reset_epoch,
            current_reset_epoch,
        }),
        StateWriteResolution::Indeterminate {
            observed_high_water,
            current_high_water,
            retired_through,
        } => Ok(StateWriteOutcome::Indeterminate {
            observed_high_water,
            current_high_water,
            retired_through,
        }),
    }
}

fn validate_key(key: &[u8]) -> Result<(), ClientError> {
    if key.is_empty() || key.len() > MAX_STATE_KEY_BYTES {
        return Err(ClientError::InvalidStateKey);
    }
    Ok(())
}

fn validate_keys(keys: &[Vec<u8>]) -> Result<(), ClientError> {
    if keys.is_empty() || keys.len() > MAX_STATE_BATCH_KEYS {
        return Err(ClientError::InvalidStateBatch);
    }
    keys.iter().try_for_each(|key| validate_key(key))
}

fn optional_value(value: StateOptionalValue) -> Result<Option<Vec<u8>>, ClientError> {
    match value {
        StateOptionalValue::Missing => Ok(None),
        StateOptionalValue::Value { value } => decode_binary(&value).map(Some),
    }
}

fn binary(value: &[u8]) -> BinaryValue {
    BinaryValue {
        bytes_base64: encode_base64(value),
    }
}

fn encode_base64(bytes: &[u8]) -> String {
    const ALPHABET: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    let mut output = String::with_capacity(bytes.len().div_ceil(3).saturating_mul(4));
    for chunk in bytes.chunks(3) {
        let first = chunk[0];
        let second = chunk.get(1).copied().unwrap_or(0);
        let third = chunk.get(2).copied().unwrap_or(0);
        output.push(char::from(ALPHABET[usize::from(first >> 2)]));
        output.push(char::from(ALPHABET[usize::from(((first & 0x03) << 4) | (second >> 4))]));
        output.push(if chunk.len() > 1 {
            char::from(ALPHABET[usize::from(((second & 0x0f) << 2) | (third >> 6))])
        } else {
            '='
        });
        output.push(if chunk.len() > 2 {
            char::from(ALPHABET[usize::from(third & 0x3f)])
        } else {
            '='
        });
    }
    output
}

fn decode_binary(value: &BinaryValue) -> Result<Vec<u8>, ClientError> {
    let input = value.bytes_base64.as_bytes();
    if input.is_empty() {
        return Ok(Vec::new());
    }
    if !input.len().is_multiple_of(4) {
        return Err(ClientError::UnexpectedStateResult);
    }
    let mut output = Vec::with_capacity(input.len() / 4 * 3);
    for (ordinal, chunk) in input.chunks_exact(4).enumerate() {
        let last = ordinal + 1 == input.len() / 4;
        let a = base64_digit(chunk[0])?;
        let b = base64_digit(chunk[1])?;
        let c = if chunk[2] == b'=' { 0 } else { base64_digit(chunk[2])? };
        let d = if chunk[3] == b'=' { 0 } else { base64_digit(chunk[3])? };
        let padding_before_last_chunk = !last && (chunk[2] == b'=' || chunk[3] == b'=');
        let invalid_padding_order = chunk[2] == b'=' && chunk[3] != b'=';
        let nonzero_double_padding_bits = chunk[2] == b'=' && b & 0x0f != 0;
        let nonzero_single_padding_bits = chunk[3] == b'=' && chunk[2] != b'=' && c & 0x03 != 0;
        if padding_before_last_chunk
            || invalid_padding_order
            || nonzero_double_padding_bits
            || nonzero_single_padding_bits
        {
            return Err(ClientError::UnexpectedStateResult);
        }
        output.push((a << 2) | (b >> 4));
        if chunk[2] != b'=' {
            output.push((b << 4) | (c >> 2));
        }
        if chunk[3] != b'=' {
            output.push((c << 6) | d);
        }
    }
    if encode_base64(&output) != value.bytes_base64 {
        return Err(ClientError::UnexpectedStateResult);
    }
    Ok(output)
}

fn base64_digit(value: u8) -> Result<u8, ClientError> {
    match value {
        b'A'..=b'Z' => Ok(value - b'A'),
        b'a'..=b'z' => Ok(value - b'a' + 26),
        b'0'..=b'9' => Ok(value - b'0' + 52),
        b'+' => Ok(62),
        b'/' => Ok(63),
        _ => Err(ClientError::UnexpectedStateResult),
    }
}

fn valid_identity(value: &str) -> bool {
    value.len() == 32 && value != "00000000000000000000000000000000" && value.bytes().all(is_lower_hex)
}

fn valid_digest(value: &str) -> bool {
    value.len() == 64
        && value != "0000000000000000000000000000000000000000000000000000000000000000"
        && value.bytes().all(is_lower_hex)
}

const fn is_lower_hex(value: u8) -> bool {
    value.is_ascii_digit() || matches!(value, b'a'..=b'f')
}

fn lower_hex(bytes: &[u8]) -> String {
    const HEX: &[u8; 16] = b"0123456789abcdef";
    let mut output = String::with_capacity(bytes.len().saturating_mul(2));
    for byte in bytes {
        output.push(char::from(HEX[usize::from(byte >> 4)]));
        output.push(char::from(HEX[usize::from(byte & 0x0f)]));
    }
    output
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn state_value_base64_is_canonical_and_round_trips() -> Result<(), ClientError> {
        for value in [b"".as_slice(), b"a", b"ab", b"abc", &[0, 1, 254, 255]] {
            let encoded = binary(value);
            assert_eq!(decode_binary(&encoded)?, value);
        }
        assert!(
            decode_binary(&BinaryValue {
                bytes_base64: "YR==".to_owned(),
            })
            .is_err()
        );
        Ok(())
    }

    #[test]
    fn state_semantic_inputs_refuse_ambiguous_identities_and_keys() {
        assert!(StateIndex::new("state", Identity { value: "0".repeat(32) }).is_err());
        assert!(StateActionId::new(Digest { value: "0".repeat(64) }).is_err());
        assert!(validate_key(&[]).is_err());
        assert!(validate_key(&vec![0; 255]).is_err());
        assert_eq!(AbsoluteExpiry::from_unix_millis(42).unix_millis(), 42);
    }
}