lodestone-proto 0.0.16

//! Bounded orchestrator-to-placement intent wire.
//!
//! Orchestrators describe desired operations; they never carry a controller
//! token or a generic metastore command. Placement binds an accepted intent
//! to its current replicated claim and exact target preconditions.

use std::error::Error;
use std::fmt::{self, Display, Formatter};

use lodestone_types::{
    ClusterRevision, CollectionIndexRef, NodeLifecycleGeneration, NodeRole, NodeToken, RoutingHash, RoutingRangeEnd,
    RoutingRevision,
};
use serde::{Deserialize, Serialize};

pub const ORCHESTRATOR_INTENT_FORMAT_VERSION: u16 = 2;
pub const ORCHESTRATOR_INTENT_RESULT_FORMAT_VERSION: u16 = 1;
pub const MAX_ORCHESTRATOR_INTENT_BYTES: usize = 64 * 1024;
pub const MAX_ORCHESTRATOR_PROVIDER_ID_BYTES: usize = 256;
pub const MAX_ORCHESTRATOR_ENDPOINT_BYTES: usize = 2_048;
pub const MAX_ORCHESTRATOR_TOPOLOGY_BYTES: usize = 256;
pub const MAX_ORCHESTRATOR_CAPACITY_BYTES: usize = 256;
pub const MAX_ORCHESTRATOR_RESULT_NODES: usize = 16;

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct NodeIntentProposal {
    pub requested_role: NodeRole,
    pub enrollment_provider_id: Vec<u8>,
    pub provider_subject_digest: [u8; 32],
    pub public_key_digest: [u8; 32],
    pub endpoint: Vec<u8>,
    pub topology: Vec<u8>,
    pub capacity: Vec<u8>,
}

impl NodeIntentProposal {
    fn validate(&self) -> Result<(), OrchestratorIntentError> {
        if self.enrollment_provider_id.is_empty()
            || self.enrollment_provider_id.len() > MAX_ORCHESTRATOR_PROVIDER_ID_BYTES
            || self.provider_subject_digest == [0; 32]
            || self.public_key_digest == [0; 32]
            || self.endpoint.is_empty()
            || self.endpoint.len() > MAX_ORCHESTRATOR_ENDPOINT_BYTES
            || self.topology.len() > MAX_ORCHESTRATOR_TOPOLOGY_BYTES
            || self.capacity.len() > MAX_ORCHESTRATOR_CAPACITY_BYTES
        {
            return Err(OrchestratorIntentError::InvalidProposal);
        }
        Ok(())
    }
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
pub enum OrchestratorIntentKind {
    RequestNodeJoin {
        proposal: NodeIntentProposal,
    },
    RequestNodeDrain {
        target: NodeToken,
        expected_lifecycle_generation: NodeLifecycleGeneration,
    },
    RequestNodeReplacement {
        target: NodeToken,
        expected_lifecycle_generation: NodeLifecycleGeneration,
        replacement: NodeIntentProposal,
    },
    RequestVoterReplacement {
        remove: NodeToken,
        expected_lifecycle_generation: NodeLifecycleGeneration,
        replacement: NodeIntentProposal,
    },
    RequestSplit {
        collection_index_ref: CollectionIndexRef,
        expected_routing_revision: RoutingRevision,
        expected_range_authority_digest: [u8; 32],
        source_start: RoutingHash,
        source_end: RoutingRangeEnd,
        split_at: RoutingHash,
    },
    RequestMerge {
        collection_index_ref: CollectionIndexRef,
        expected_routing_revision: RoutingRevision,
        expected_range_authority_digest: [u8; 32],
        left_start: RoutingHash,
        boundary: RoutingHash,
        right_end: RoutingRangeEnd,
    },
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct OrchestratorIntent {
    pub format_version: u16,
    pub operation_id: [u8; 32],
    pub expected_cluster_revision: ClusterRevision,
    pub kind: OrchestratorIntentKind,
    pub semantic_digest: [u8; 32],
}

impl OrchestratorIntent {
    pub fn validate(&self) -> Result<(), OrchestratorIntentError> {
        if self.format_version != ORCHESTRATOR_INTENT_FORMAT_VERSION {
            return Err(OrchestratorIntentError::UnsupportedVersion(self.format_version));
        }
        if self.operation_id == [0; 32] || self.semantic_digest == [0; 32] {
            return Err(OrchestratorIntentError::InvalidOperationIdentity);
        }
        match &self.kind {
            OrchestratorIntentKind::RequestNodeJoin { proposal } => {
                proposal.validate()?;
            }
            OrchestratorIntentKind::RequestNodeDrain { .. } => {}
            OrchestratorIntentKind::RequestNodeReplacement { replacement, .. } => {
                replacement.validate()?;
            }
            OrchestratorIntentKind::RequestVoterReplacement { replacement, .. } => {
                replacement.validate()?;
                if replacement.requested_role != NodeRole::MetastoreVoter {
                    return Err(OrchestratorIntentError::InvalidProposal);
                }
            }
            OrchestratorIntentKind::RequestSplit {
                expected_range_authority_digest,
                source_start,
                source_end,
                split_at,
                ..
            } => {
                if *expected_range_authority_digest == [0; 32]
                    || split_at.get() <= source_start.get()
                    || u64::from(split_at.get()) >= source_end.get()
                {
                    return Err(OrchestratorIntentError::InvalidRange);
                }
            }
            OrchestratorIntentKind::RequestMerge {
                expected_range_authority_digest,
                left_start,
                boundary,
                right_end,
                ..
            } => {
                if *expected_range_authority_digest == [0; 32]
                    || boundary.get() <= left_start.get()
                    || u64::from(boundary.get()) >= right_end.get()
                {
                    return Err(OrchestratorIntentError::InvalidRange);
                }
            }
        }
        if self.semantic_digest != self.derived_semantic_digest()? {
            return Err(OrchestratorIntentError::SemanticDigestMismatch);
        }
        Ok(())
    }

    pub fn derive_semantic_digest(&mut self) -> Result<(), OrchestratorIntentError> {
        self.semantic_digest = self.derived_semantic_digest()?;
        Ok(())
    }

    fn derived_semantic_digest(&self) -> Result<[u8; 32], OrchestratorIntentError> {
        let bytes = serde_json::to_vec(&(
            self.format_version,
            self.operation_id,
            self.expected_cluster_revision,
            &self.kind,
        ))
        .map_err(|error| OrchestratorIntentError::Codec(error.to_string()))?;
        let mut hasher = blake3::Hasher::new();
        hasher.update(b"lodestone/p6/orchestrator-intent/v1");
        hasher.update(&bytes);
        Ok(*hasher.finalize().as_bytes())
    }
}

#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum OrchestratorIntentResultKind {
    Completed,
    Accepted,
    Unavailable,
    Conflict,
    Refused,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct OrchestratorIntentResult {
    pub format_version: u16,
    pub operation_id: [u8; 32],
    pub semantic_digest: [u8; 32],
    pub kind: OrchestratorIntentResultKind,
    pub committed_revision: Option<ClusterRevision>,
    pub result_digest: [u8; 32],
    pub affected_nodes: Vec<NodeToken>,
    pub retryable: bool,
}

impl OrchestratorIntentResult {
    pub fn validate(&self) -> Result<(), OrchestratorIntentError> {
        if self.format_version != ORCHESTRATOR_INTENT_RESULT_FORMAT_VERSION {
            return Err(OrchestratorIntentError::UnsupportedResultVersion(self.format_version));
        }
        if self.operation_id == [0; 32]
            || self.semantic_digest == [0; 32]
            || self.result_digest == [0; 32]
            || self.affected_nodes.len() > MAX_ORCHESTRATOR_RESULT_NODES
            || (self.retryable && self.kind != OrchestratorIntentResultKind::Unavailable)
        {
            return Err(OrchestratorIntentError::InvalidResult);
        }
        Ok(())
    }
}

pub fn encode_orchestrator_intent(intent: &OrchestratorIntent) -> Result<Vec<u8>, OrchestratorIntentError> {
    intent.validate()?;
    encode_bounded(intent)
}

pub fn decode_orchestrator_intent(bytes: &[u8]) -> Result<OrchestratorIntent, OrchestratorIntentError> {
    let intent: OrchestratorIntent = decode_bounded(bytes)?;
    intent.validate()?;
    Ok(intent)
}

pub fn encode_orchestrator_intent_result(
    result: &OrchestratorIntentResult,
) -> Result<Vec<u8>, OrchestratorIntentError> {
    result.validate()?;
    encode_bounded(result)
}

pub fn decode_orchestrator_intent_result(bytes: &[u8]) -> Result<OrchestratorIntentResult, OrchestratorIntentError> {
    let result: OrchestratorIntentResult = decode_bounded(bytes)?;
    result.validate()?;
    Ok(result)
}

fn encode_bounded<T: Serialize>(value: &T) -> Result<Vec<u8>, OrchestratorIntentError> {
    let bytes = serde_json::to_vec(value).map_err(|error| OrchestratorIntentError::Codec(error.to_string()))?;
    if bytes.len() > MAX_ORCHESTRATOR_INTENT_BYTES {
        return Err(OrchestratorIntentError::MessageTooLarge(bytes.len()));
    }
    Ok(bytes)
}

fn decode_bounded<T>(bytes: &[u8]) -> Result<T, OrchestratorIntentError>
where
    T: for<'de> Deserialize<'de> + Serialize,
{
    if bytes.len() > MAX_ORCHESTRATOR_INTENT_BYTES {
        return Err(OrchestratorIntentError::MessageTooLarge(bytes.len()));
    }
    let value = serde_json::from_slice(bytes).map_err(|error| OrchestratorIntentError::Codec(error.to_string()))?;
    let canonical = serde_json::to_vec(&value).map_err(|error| OrchestratorIntentError::Codec(error.to_string()))?;
    if canonical != bytes {
        return Err(OrchestratorIntentError::NonCanonical);
    }
    Ok(value)
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum OrchestratorIntentError {
    Codec(String),
    InvalidOperationIdentity,
    InvalidProposal,
    InvalidRange,
    InvalidResult,
    MessageTooLarge(usize),
    NonCanonical,
    SemanticDigestMismatch,
    UnsupportedResultVersion(u16),
    UnsupportedVersion(u16),
}

impl Display for OrchestratorIntentError {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
        write!(formatter, "{self:?}")
    }
}

impl Error for OrchestratorIntentError {}

#[cfg(test)]
mod tests {
    use super::*;
    use lodestone_types::{CollectionId, IndexId, NodeId, NodeIncarnation, RecoveryEra};

    fn proposal(role: NodeRole) -> NodeIntentProposal {
        NodeIntentProposal {
            requested_role: role,
            enrollment_provider_id: b"static".to_vec(),
            provider_subject_digest: [2; 32],
            public_key_digest: [3; 32],
            endpoint: b"127.0.0.1:1000".to_vec(),
            topology: Vec::new(),
            capacity: Vec::new(),
        }
    }

    fn token(id: u64) -> Result<NodeToken, Box<dyn Error>> {
        Ok(NodeToken::new(
            RecoveryEra::new(1)?,
            NodeId::new(id),
            NodeIncarnation::new(1)?,
        )?)
    }

    #[test]
    fn all_six_intents_round_trip_and_digest_is_exact() -> Result<(), Box<dyn Error>> {
        let scope = CollectionIndexRef::new(CollectionId::new(1)?, IndexId::new(2)?);
        let kinds = vec![
            OrchestratorIntentKind::RequestNodeJoin {
                proposal: proposal(NodeRole::Indexer),
            },
            OrchestratorIntentKind::RequestNodeDrain {
                target: token(1)?,
                expected_lifecycle_generation: NodeLifecycleGeneration::new(1),
            },
            OrchestratorIntentKind::RequestNodeReplacement {
                target: token(2)?,
                expected_lifecycle_generation: NodeLifecycleGeneration::new(1),
                replacement: proposal(NodeRole::Searcher),
            },
            OrchestratorIntentKind::RequestVoterReplacement {
                remove: token(3)?,
                expected_lifecycle_generation: NodeLifecycleGeneration::new(1),
                replacement: proposal(NodeRole::MetastoreVoter),
            },
            OrchestratorIntentKind::RequestSplit {
                collection_index_ref: scope,
                expected_routing_revision: RoutingRevision::new(1),
                expected_range_authority_digest: [7; 32],
                source_start: RoutingHash::new(0),
                source_end: RoutingRangeEnd::FULL,
                split_at: RoutingHash::new(10),
            },
            OrchestratorIntentKind::RequestMerge {
                collection_index_ref: scope,
                expected_routing_revision: RoutingRevision::new(1),
                expected_range_authority_digest: [7; 32],
                left_start: RoutingHash::new(0),
                boundary: RoutingHash::new(10),
                right_end: RoutingRangeEnd::FULL,
            },
        ];
        for (ordinal, kind) in kinds.into_iter().enumerate() {
            let mut intent = OrchestratorIntent {
                format_version: ORCHESTRATOR_INTENT_FORMAT_VERSION,
                operation_id: [u8::try_from(ordinal + 1)?; 32],
                expected_cluster_revision: ClusterRevision::new(1),
                kind,
                semantic_digest: [1; 32],
            };
            intent.derive_semantic_digest()?;
            let bytes = encode_orchestrator_intent(&intent)?;
            assert_eq!(decode_orchestrator_intent(&bytes)?, intent);
            let mut tampered = intent.clone();
            tampered.semantic_digest = [9; 32];
            assert!(matches!(
                decode_orchestrator_intent(&serde_json::to_vec(&tampered)?),
                Err(OrchestratorIntentError::SemanticDigestMismatch)
            ));
        }
        Ok(())
    }

    #[test]
    fn future_unknown_and_oversized_inputs_fail_closed() {
        assert!(matches!(
            decode_orchestrator_intent(&vec![b'x'; MAX_ORCHESTRATOR_INTENT_BYTES + 1]),
            Err(OrchestratorIntentError::MessageTooLarge(_))
        ));
        let unknown = br#"{"format_version":1,"operation_id":[1],"expected_cluster_revision":1,"kind":{"kind":"future"},"semantic_digest":[1]}"#;
        assert!(decode_orchestrator_intent(unknown).is_err());
    }

    #[test]
    fn range_intent_requires_an_exact_authority_digest() -> Result<(), Box<dyn Error>> {
        let scope = CollectionIndexRef::new(CollectionId::new(1)?, IndexId::new(2)?);
        let mut intent = OrchestratorIntent {
            format_version: ORCHESTRATOR_INTENT_FORMAT_VERSION,
            operation_id: [1; 32],
            expected_cluster_revision: ClusterRevision::new(1),
            kind: OrchestratorIntentKind::RequestSplit {
                collection_index_ref: scope,
                expected_routing_revision: RoutingRevision::new(1),
                expected_range_authority_digest: [0; 32],
                source_start: RoutingHash::new(0),
                source_end: RoutingRangeEnd::FULL,
                split_at: RoutingHash::new(10),
            },
            semantic_digest: [0; 32],
        };
        intent.derive_semantic_digest()?;
        assert!(matches!(intent.validate(), Err(OrchestratorIntentError::InvalidRange)));
        Ok(())
    }
}