lazaret/crates/lazaret-core/tests/protocol.rs

use ed25519_dalek::SigningKey;
use lazaret_core::identity::{
    CapabilityId, CertificateDigest, ClientId, ConfigurationFingerprint, CoordinatorId, DomainId,
    GenerationId, RequestDigest, RequestId, SandboxId, SessionId, StreamDigest,
    TrustManifestDigest,
};
use lazaret_core::protocol::{
    MAX_RESULT_CHUNK_BYTES, MAX_UNACKNOWLEDGED_CHUNKS, ProtocolCodecError, ProtocolValidationError,
    RequestReplayGuard, ResultCollector, ResultSenderWindow, ResultWindowError, VerifiedSession,
    encode_signed, issue_protocol_credential, issue_session_capability, issue_trust_manifest,
    request_digest, stream_digest, verify_protocol_credential, verify_tls_credential,
    verify_trust_manifest,
};
use lazaret_core::records::{
    DeveloperDecision, DeveloperDecisionKind, EscalationOutcome, GenerationIdentity,
    MailboxCancellation, MailboxRequest, MailboxResult, MailboxResultAck, MailboxResultChunk,
    ProtocolCredential, ProtocolCredentialRole, ProtocolSubjectId, ProtocolVersion, ResultStream,
    SessionBinding, SessionRequestCapability, TlsCredential, TlsRole, TlsSubjectId, TrustManifest,
};

fn bytes(value: u8) -> [u8; 32] {
    [value; 32]
}

fn key(value: u8) -> SigningKey {
    SigningKey::from_bytes(&bytes(value))
}

fn generation(value: u8) -> GenerationIdentity {
    GenerationIdentity {
        agent_profile: format!("agent-{value}"),
        domain_id: DomainId::new(bytes(value)),
        generation_id: GenerationId::new(u64::from(value)),
        configuration_fingerprint: ConfigurationFingerprint::new(bytes(value + 1)),
        coordinator_id: CoordinatorId::new(bytes(value + 2)),
        sandbox_id: SandboxId::new(bytes(value + 3)),
    }
}

struct Fixture {
    generation: GenerationIdentity,
    tls_ca_digest: CertificateDigest,
    manifest: TrustManifest,
    manifest_digest: TrustManifestDigest,
    coordinator_credential: ProtocolCredential,
    sandbox_credential: ProtocolCredential,
    capability: SessionRequestCapability,
    binding: SessionBinding,
    coordinator_key: SigningKey,
    sandbox_key: SigningKey,
    anchor_key: SigningKey,
    session: VerifiedSession,
}

fn fixture() -> Fixture {
    let generation = generation(1);
    let tls_ca_digest = CertificateDigest::new(bytes(20));
    let issuer = key(21);
    let coordinator_key = key(22);
    let sandbox_key = key(23);
    let anchor_key = key(24);
    let (manifest, manifest_digest) =
        issue_trust_manifest(generation.clone(), tls_ca_digest, &issuer).unwrap();
    let coordinator_credential = issue_protocol_credential(
        generation.clone(),
        ProtocolCredentialRole::SessionCoordinator,
        ProtocolSubjectId::SessionCoordinator(generation.coordinator_id),
        &coordinator_key.verifying_key(),
        &issuer,
    )
    .unwrap();
    let sandbox_credential = issue_protocol_credential(
        generation.clone(),
        ProtocolCredentialRole::AgentSandbox,
        ProtocolSubjectId::AgentSandbox(generation.sandbox_id),
        &sandbox_key.verifying_key(),
        &issuer,
    )
    .unwrap();
    let binding = SessionBinding {
        generation_identity: generation.clone(),
        cli_id: ClientId::new(bytes(25)),
        session_id: SessionId::new(bytes(26)),
        session_anchor_public_key: anchor_key.verifying_key().to_bytes().to_vec(),
    };
    let capability = issue_session_capability(
        binding.clone(),
        coordinator_credential.credential_id,
        sandbox_credential.credential_id,
        &coordinator_key,
    )
    .unwrap();
    let session = VerifiedSession::new(
        manifest_digest,
        tls_ca_digest,
        &manifest,
        binding.clone(),
        &coordinator_credential,
        &sandbox_credential,
        capability.clone(),
    )
    .unwrap();
    Fixture {
        generation,
        tls_ca_digest,
        manifest,
        manifest_digest,
        coordinator_credential,
        sandbox_credential,
        capability,
        binding,
        coordinator_key,
        sandbox_key,
        anchor_key,
        session,
    }
}

fn request(fixture: &Fixture, id: u8) -> MailboxRequest {
    MailboxRequest::new(
        fixture.binding.clone(),
        RequestId::new(bytes(id)),
        fixture.capability.clone(),
        fixture.sandbox_credential.credential_id,
        vec![b"/usr/bin/id".to_vec(), b"-u".to_vec()],
        "show the numeric user identity".into(),
    )
    .unwrap()
}

#[test]
fn exact_session_tuple_is_authenticated_and_request_ids_do_not_replay() {
    let fixture = fixture();
    let request = request(&fixture, 30);
    let encoded = encode_signed(&request, &fixture.sandbox_key).unwrap();
    let mut replay = RequestReplayGuard::default();
    let (accepted, digest) = replay.admit(&fixture.session, &encoded).unwrap();
    assert_eq!(accepted, request);
    assert_eq!(
        digest,
        request_digest(&postcard::to_stdvec(&request).unwrap())
    );
    assert!(matches!(
        replay.admit(&fixture.session, &encoded),
        Err(ProtocolValidationError::RequestReplay)
    ));

    let mut tampered = encoded.clone();
    let last = tampered.last_mut().unwrap();
    *last ^= 1;
    assert!(matches!(
        fixture.session.verify_request(&tampered),
        Err(ProtocolValidationError::Codec(
            ProtocolCodecError::InvalidSignature
        ))
    ));

    let wrong_key = encode_signed(&request, &key(31)).unwrap();
    assert!(matches!(
        fixture.session.verify_request(&wrong_key),
        Err(ProtocolValidationError::Codec(
            ProtocolCodecError::InvalidSignature
        ))
    ));

    let mut stale_binding = fixture.binding.clone();
    stale_binding.generation_identity.generation_id = GenerationId::new(2);
    let stale = MailboxRequest::new(
        stale_binding,
        RequestId::new(bytes(32)),
        fixture.capability.clone(),
        fixture.sandbox_credential.credential_id,
        vec![b"/usr/bin/true".to_vec()],
        String::new(),
    )
    .unwrap();
    let stale = encode_signed(&stale, &fixture.sandbox_key).unwrap();
    assert!(matches!(
        fixture.session.verify_request(&stale),
        Err(ProtocolValidationError::SessionBindingMismatch)
    ));
}

#[test]
fn trust_credentials_refuse_key_role_generation_and_digest_confusion() {
    let fixture = fixture();
    let issuer = verify_trust_manifest(
        &fixture.manifest,
        fixture.manifest_digest,
        &fixture.generation,
        fixture.tls_ca_digest,
    )
    .unwrap();
    assert!(matches!(
        verify_trust_manifest(
            &fixture.manifest,
            TrustManifestDigest::new(bytes(99)),
            &fixture.generation,
            fixture.tls_ca_digest,
        ),
        Err(ProtocolValidationError::ManifestDigestMismatch)
    ));
    assert!(matches!(
        verify_protocol_credential(
            &fixture.sandbox_credential,
            &issuer,
            &fixture.generation,
            ProtocolCredentialRole::SessionCoordinator,
            ProtocolSubjectId::SessionCoordinator(fixture.generation.coordinator_id),
        ),
        Err(ProtocolValidationError::RoleMismatch)
    ));
    let stale_generation = generation(40);
    assert!(matches!(
        verify_protocol_credential(
            &fixture.sandbox_credential,
            &issuer,
            &stale_generation,
            ProtocolCredentialRole::AgentSandbox,
            ProtocolSubjectId::AgentSandbox(fixture.generation.sandbox_id),
        ),
        Err(ProtocolValidationError::GenerationMismatch)
    ));
    let mut tampered_credential = fixture.sandbox_credential.clone();
    tampered_credential.public_key[0] ^= 1;
    assert!(matches!(
        verify_protocol_credential(
            &tampered_credential,
            &issuer,
            &fixture.generation,
            ProtocolCredentialRole::AgentSandbox,
            ProtocolSubjectId::AgentSandbox(fixture.generation.sandbox_id),
        ),
        Err(ProtocolValidationError::CredentialDigestMismatch)
    ));
    let mut tampered_capability = fixture.capability.clone();
    tampered_capability.capability_id = CapabilityId::new(bytes(98));
    assert!(matches!(
        VerifiedSession::new(
            fixture.manifest_digest,
            fixture.tls_ca_digest,
            &fixture.manifest,
            fixture.binding.clone(),
            &fixture.coordinator_credential,
            &fixture.sandbox_credential,
            tampered_capability,
        ),
        Err(ProtocolValidationError::CapabilityDigestMismatch)
    ));
}

#[test]
fn tls_credentials_refuse_role_and_certificate_digest_confusion() {
    let fixture = fixture();
    let certificate = b"synthetic DER fixture".to_vec();
    let tls = TlsCredential::session_coordinator(
        fixture.generation.clone(),
        fixture.generation.coordinator_id,
        certificate.clone(),
        CertificateDigest::new(
            lazaret_core::protocol::stream_digest(&certificate)
                .as_bytes()
                .to_owned(),
        ),
    );
    verify_tls_credential(
        &tls,
        &fixture.generation,
        TlsRole::SessionCoordinator,
        TlsSubjectId::SessionCoordinator(fixture.generation.coordinator_id),
    )
    .unwrap();
    assert!(matches!(
        verify_tls_credential(
            &tls,
            &fixture.generation,
            TlsRole::ClientCli,
            TlsSubjectId::ClientCli(ClientId::new(bytes(1))),
        ),
        Err(ProtocolValidationError::RoleMismatch)
    ));
    let invalid_tls = TlsCredential::session_coordinator(
        fixture.generation.clone(),
        fixture.generation.coordinator_id,
        certificate,
        CertificateDigest::new(bytes(1)),
    );
    assert!(matches!(
        verify_tls_credential(
            &invalid_tls,
            &fixture.generation,
            TlsRole::SessionCoordinator,
            TlsSubjectId::SessionCoordinator(fixture.generation.coordinator_id),
        ),
        Err(ProtocolValidationError::CertificateDigestMismatch)
    ));
}

#[test]
fn session_anchor_and_sandbox_signatures_are_not_interchangeable() {
    let fixture = fixture();
    let request = request(&fixture, 50);
    let request_bytes = encode_signed(&request, &fixture.sandbox_key).unwrap();
    let (_, digest) = fixture.session.verify_request(&request_bytes).unwrap();
    let decision = DeveloperDecision {
        protocol_version: ProtocolVersion::V1,
        session_binding: fixture.binding.clone(),
        request_id: request.request_id(),
        request_sha256: digest,
        decision: DeveloperDecisionKind::Approve,
        approved_command_argv: request.command_argv().to_vec(),
    };
    let signed = encode_signed(&decision, &fixture.anchor_key).unwrap();
    assert_eq!(
        fixture
            .session
            .verify_developer_decision(&signed, request.request_id(), digest)
            .unwrap(),
        decision
    );
    let sandbox_signed = encode_signed(&decision, &fixture.sandbox_key).unwrap();
    assert!(matches!(
        fixture
            .session
            .verify_developer_decision(&sandbox_signed, request.request_id(), digest),
        Err(ProtocolValidationError::Codec(
            ProtocolCodecError::InvalidSignature
        ))
    ));
}

#[test]
fn cancellation_ack_chunks_and_terminal_result_bind_the_accepted_request() {
    let fixture = fixture();
    let request = request(&fixture, 60);
    let request_id = request.request_id();
    let (_, digest) = fixture
        .session
        .verify_request(&encode_signed(&request, &fixture.sandbox_key).unwrap())
        .unwrap();

    let cancellation = MailboxCancellation {
        protocol_version: ProtocolVersion::V1,
        session_binding: fixture.binding.clone(),
        request_id,
        session_capability: fixture.capability.clone(),
        sandbox_credential_id: fixture.sandbox_credential.credential_id,
    };
    let cancellation = encode_signed(&cancellation, &fixture.sandbox_key).unwrap();
    fixture
        .session
        .verify_cancellation(&cancellation, request_id)
        .unwrap();

    let ack = MailboxResultAck {
        protocol_version: ProtocolVersion::V1,
        session_binding: fixture.binding.clone(),
        request_id,
        request_sha256: digest,
        session_capability: fixture.capability.clone(),
        sandbox_credential_id: fixture.sandbox_credential.credential_id,
        highest_contiguous_sequence: 1,
    };
    let ack = encode_signed(&ack, &fixture.sandbox_key).unwrap();
    fixture
        .session
        .verify_result_ack(&ack, request_id, digest)
        .unwrap();

    let chunks = [
        (ResultStream::Stdout, b"hello".as_slice()),
        (ResultStream::Stderr, b"warning".as_slice()),
    ];
    let mut collector = ResultCollector::default();
    for (sequence, (stream, bytes)) in chunks.iter().enumerate() {
        let chunk = MailboxResultChunk {
            protocol_version: ProtocolVersion::V1,
            session_binding: fixture.binding.clone(),
            request_id,
            request_sha256: digest,
            coordinator_credential_id: fixture.coordinator_credential.credential_id,
            sequence: u64::try_from(sequence).unwrap(),
            stream: *stream,
            bytes: bytes.to_vec(),
        };
        let encoded = encode_signed(&chunk, &fixture.coordinator_key).unwrap();
        let verified = fixture
            .session
            .verify_result_chunk(&encoded, request_id, digest)
            .unwrap();
        collector.push(&verified).unwrap();
        if sequence == 0 {
            assert_eq!(
                collector.push(&verified),
                Err(ResultWindowError::UnexpectedSequence {
                    expected: 1,
                    actual: 0,
                })
            );
        }
    }
    let result = MailboxResult {
        protocol_version: ProtocolVersion::V1,
        session_binding: fixture.binding.clone(),
        request_id,
        request_sha256: digest,
        coordinator_credential_id: fixture.coordinator_credential.credential_id,
        outcome: EscalationOutcome::Completed,
        exit_code: Some(0),
        chunk_count: 2,
        stdout_sha256: stream_digest(b"hello"),
        stderr_sha256: stream_digest(b"warning"),
    };
    let encoded = encode_signed(&result, &fixture.coordinator_key).unwrap();
    let verified = fixture
        .session
        .verify_result(&encoded, request_id, digest)
        .unwrap();
    collector.finish(&verified).unwrap();

    let mut wrong_digest = result;
    wrong_digest.stdout_sha256 = StreamDigest::new(bytes(1));
    assert_eq!(
        collector.finish(&wrong_digest),
        Err(ResultWindowError::StdoutDigestMismatch)
    );
    assert!(matches!(
        fixture
            .session
            .verify_result(&encoded, request_id, RequestDigest::new(bytes(99))),
        Err(ProtocolValidationError::RequestMismatch)
    ));
}

#[test]
fn result_windows_are_bounded_and_contiguous() {
    let mut sender = ResultSenderWindow::default();
    for expected in 0..MAX_UNACKNOWLEDGED_CHUNKS {
        assert_eq!(sender.reserve(1).unwrap(), u64::try_from(expected).unwrap());
    }
    assert_eq!(sender.reserve(1), Err(ResultWindowError::WindowFull));
    sender.acknowledge(3).unwrap();
    assert_eq!(sender.outstanding(), 4);
    assert_eq!(sender.acknowledge(3), Err(ResultWindowError::AckRegression));
    assert_eq!(
        sender.acknowledge(99),
        Err(ResultWindowError::AckBeyondSent)
    );
    assert_eq!(
        sender.reserve(MAX_RESULT_CHUNK_BYTES + 1),
        Err(ResultWindowError::ChunkTooLarge {
            length: MAX_RESULT_CHUNK_BYTES + 1
        })
    );

    let fixture = fixture();
    let oversized = MailboxResultChunk {
        protocol_version: ProtocolVersion::V1,
        session_binding: fixture.binding,
        request_id: RequestId::new(bytes(70)),
        request_sha256: RequestDigest::new(bytes(71)),
        coordinator_credential_id: fixture.coordinator_credential.credential_id,
        sequence: 0,
        stream: ResultStream::Stdout,
        bytes: vec![0; MAX_RESULT_CHUNK_BYTES + 1],
    };
    let encoded = encode_signed(&oversized, &fixture.coordinator_key).unwrap();
    assert!(matches!(
        fixture.session.verify_result_chunk(
            &encoded,
            RequestId::new(bytes(70)),
            RequestDigest::new(bytes(71))
        ),
        Err(ProtocolValidationError::ChunkTooLarge { .. })
    ));
}

#[test]
fn complete_signed_record_limit_is_enforced() {
    let fixture = fixture();
    let oversized = MailboxRequest::new(
        fixture.binding,
        RequestId::new(bytes(80)),
        fixture.capability,
        fixture.sandbox_credential.credential_id,
        vec![b"/usr/bin/true".to_vec()],
        "x".repeat(65_536),
    )
    .unwrap();
    assert!(matches!(
        encode_signed(&oversized, &fixture.sandbox_key),
        Err(ProtocolCodecError::RecordTooLarge { .. })
    ));
}