clusterflux-public/crates/disasmer-core/src/transport.rs

use serde::{Deserialize, Serialize};
use thiserror::Error;

use crate::{ArtifactId, Digest, NodeId, ProcessId, ProjectId, TenantId};

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum TransportKind {
    NativeQuic,
}

pub trait Transport {
    fn kind(&self) -> TransportKind;
    fn authenticated_direct_connections(&self) -> bool;
}

#[derive(Clone, Debug, Default)]
pub struct NativeQuicTransport;

impl Transport for NativeQuicTransport {
    fn kind(&self) -> TransportKind {
        TransportKind::NativeQuic
    }

    fn authenticated_direct_connections(&self) -> bool {
        true
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct DataPlaneScope {
    pub tenant: TenantId,
    pub project: ProjectId,
    pub process: ProcessId,
    pub object: DataPlaneObject,
    pub authorization_subject: String,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum DataPlaneObject {
    Artifact(ArtifactId),
    Blob(Digest),
    SourceSnapshot(Digest),
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct NodeEndpoint {
    pub node: NodeId,
    pub advertised_addr: String,
    pub public_key_fingerprint: Digest,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct RendezvousRequest {
    pub scope: DataPlaneScope,
    pub source: NodeEndpoint,
    pub destination: NodeEndpoint,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct DirectBulkTransferPlan {
    pub transport: TransportKind,
    pub scope: DataPlaneScope,
    pub source: NodeEndpoint,
    pub destination: NodeEndpoint,
    pub authorization_digest: Digest,
    pub coordinator_assisted_rendezvous: bool,
    pub coordinator_bulk_relay_allowed: bool,
}

#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum TransportError {
    #[error(
        "direct node-to-node connectivity is unavailable for scoped data-plane transfer: {reason}; coordinator bulk relay is disabled"
    )]
    DirectConnectivityUnavailable { reason: String },
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum BulkTransferDecision {
    DirectAuthenticated { scope: DataPlaneScope },
    FailClear { message: String },
}

#[derive(Clone, Debug, Error, PartialEq, Eq)]
#[error("bulk relay through coordinator is not allowed by default")]
pub struct BulkRelayDenied;

impl NativeQuicTransport {
    pub fn plan_authenticated_direct_bulk_transfer(
        &self,
        request: RendezvousRequest,
        direct_connectivity: bool,
        failure_reason: impl Into<String>,
    ) -> Result<DirectBulkTransferPlan, TransportError> {
        if !direct_connectivity {
            return Err(TransportError::DirectConnectivityUnavailable {
                reason: failure_reason.into(),
            });
        }

        let authorization_digest = data_plane_authorization_digest(&request);
        Ok(DirectBulkTransferPlan {
            transport: self.kind(),
            scope: request.scope,
            source: request.source,
            destination: request.destination,
            authorization_digest,
            coordinator_assisted_rendezvous: true,
            coordinator_bulk_relay_allowed: false,
        })
    }
}

pub fn direct_bulk_transfer_or_error(
    scope: DataPlaneScope,
    direct_connectivity: bool,
) -> BulkTransferDecision {
    if direct_connectivity {
        BulkTransferDecision::DirectAuthenticated { scope }
    } else {
        BulkTransferDecision::FailClear {
            message:
                "direct node-to-node connectivity is unavailable; coordinator bulk relay is disabled"
                    .to_owned(),
        }
    }
}

fn data_plane_authorization_digest(request: &RendezvousRequest) -> Digest {
    let object = match &request.scope.object {
        DataPlaneObject::Artifact(artifact) => format!("artifact:{artifact}"),
        DataPlaneObject::Blob(digest) => format!("blob:{}", digest.as_str()),
        DataPlaneObject::SourceSnapshot(digest) => format!("source:{}", digest.as_str()),
    };

    Digest::from_parts([
        b"dataplane-auth:v1".as_slice(),
        request.scope.tenant.as_str().as_bytes(),
        request.scope.project.as_str().as_bytes(),
        request.scope.process.as_str().as_bytes(),
        object.as_bytes(),
        request.scope.authorization_subject.as_bytes(),
        request.source.node.as_str().as_bytes(),
        request.source.public_key_fingerprint.as_str().as_bytes(),
        request.destination.node.as_str().as_bytes(),
        request
            .destination
            .public_key_fingerprint
            .as_str()
            .as_bytes(),
    ])
}

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

    fn endpoint(name: &str) -> NodeEndpoint {
        NodeEndpoint {
            node: NodeId::from(name),
            advertised_addr: format!("{name}.mesh.invalid:4433"),
            public_key_fingerprint: Digest::sha256(format!("{name}-public-key")),
        }
    }

    fn scope(project: &str) -> DataPlaneScope {
        DataPlaneScope {
            tenant: TenantId::from("tenant"),
            project: ProjectId::from(project),
            process: ProcessId::from("process"),
            object: DataPlaneObject::Artifact(ArtifactId::from("artifact")),
            authorization_subject: "node-a-to-node-b".to_owned(),
        }
    }

    #[test]
    fn failed_direct_transfer_does_not_silently_relay() {
        let decision = direct_bulk_transfer_or_error(scope("project"), false);

        assert!(matches!(decision, BulkTransferDecision::FailClear { .. }));
    }

    #[test]
    fn native_quic_rendezvous_plan_is_scoped_and_disallows_coordinator_bulk_relay() {
        let transport = NativeQuicTransport;
        let request = RendezvousRequest {
            scope: scope("project"),
            source: endpoint("node-a"),
            destination: endpoint("node-b"),
        };

        let plan = transport
            .plan_authenticated_direct_bulk_transfer(request.clone(), true, "")
            .unwrap();
        let changed_scope_plan = transport
            .plan_authenticated_direct_bulk_transfer(
                RendezvousRequest {
                    scope: scope("other-project"),
                    ..request
                },
                true,
                "",
            )
            .unwrap();

        assert_eq!(plan.transport, TransportKind::NativeQuic);
        assert_eq!(plan.scope.tenant, TenantId::from("tenant"));
        assert_eq!(plan.scope.project, ProjectId::from("project"));
        assert_eq!(plan.scope.process, ProcessId::from("process"));
        assert_eq!(
            plan.scope.object,
            DataPlaneObject::Artifact(ArtifactId::from("artifact"))
        );
        assert_eq!(plan.source.node, NodeId::from("node-a"));
        assert_eq!(plan.destination.node, NodeId::from("node-b"));
        assert!(plan.coordinator_assisted_rendezvous);
        assert!(!plan.coordinator_bulk_relay_allowed);
        assert_ne!(
            plan.authorization_digest,
            changed_scope_plan.authorization_digest
        );
    }

    #[test]
    fn failed_direct_rendezvous_reports_clear_error_instead_of_relaying() {
        let error = NativeQuicTransport
            .plan_authenticated_direct_bulk_transfer(
                RendezvousRequest {
                    scope: scope("project"),
                    source: endpoint("node-a"),
                    destination: endpoint("node-b"),
                },
                false,
                "nat traversal failed",
            )
            .unwrap_err();

        assert!(error.to_string().contains("nat traversal failed"));
        assert!(error
            .to_string()
            .contains("coordinator bulk relay is disabled"));
    }
}