tsafe_collab/

stub.rs

//! In-memory reference stub for `CollabRemote`.
//!
//! `StubServer` backs every method with `HashMap` state and enforces the
//! membership gate on every call that requires it.  It is intentionally simple:
//! no persistence, no HTTP, no auth tokens.  Use it in unit and integration
//! tests to exercise the adversarial proof scenarios in `D3.4`.
//!
//! The server holds **public keys only** — it never stores or derives plaintext
//! key material (ADR-027 §"Service NEVER holds").

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use chrono::Utc;
use uuid::Uuid;

use crate::error::CollabError;
use crate::remote::CollabRemote;
use crate::types::{DekEnvelope, InviteRecord, InviteToken, MemberEntry};

/// Default invite TTL in hours.
const INVITE_TTL_HOURS: i64 = 48;

/// Shared inner state, wrapped in `Arc<Mutex<_>>` so `StubServer` can be used
/// from multiple threads in tests.
#[derive(Debug, Default)]
struct StubState {
    /// `team_id` → members
    members: HashMap<String, Vec<MemberEntry>>,
    /// `(team_id, recipient_pubkey)` → envelope
    dek_inbox: HashMap<(String, String), DekEnvelope>,
    /// `(team_id, custodian_pubkey)` → age-encrypted share ciphertext
    share_inbox: HashMap<(String, String), String>,
    /// `token` → invite record
    invites: HashMap<String, InviteRecord>,
}

impl StubState {
    /// Returns `true` if `pubkey` is a registered member of `team_id`.
    fn is_member(&self, team_id: &str, pubkey: &str) -> bool {
        self.members
            .get(team_id)
            .map(|ms| ms.iter().any(|m| m.pubkey == pubkey))
            .unwrap_or(false)
    }

    /// Asserts membership, returning `NotMember` if the check fails.
    fn assert_member(&self, team_id: &str, pubkey: &str) -> Result<(), CollabError> {
        if self.is_member(team_id, pubkey) {
            Ok(())
        } else {
            Err(CollabError::NotMember {
                team_id: team_id.to_owned(),
            })
        }
    }
}

/// In-memory `CollabRemote` implementation backed by `HashMap` state.
///
/// Thread-safe via internal `Arc<Mutex<StubState>>`.
#[derive(Debug, Clone, Default)]
pub struct StubServer {
    state: Arc<Mutex<StubState>>,
}

impl StubServer {
    /// Construct an empty stub server.
    pub fn new() -> Self {
        Self::default()
    }
}

impl CollabRemote for StubServer {
    fn join(&self, team_id: &str, pubkey: &str) -> Result<(), CollabError> {
        let mut state = self.state.lock().expect("stub state poisoned");
        let members = state.members.entry(team_id.to_owned()).or_default();
        // Idempotent: do nothing if the pubkey is already registered.
        if !members.iter().any(|m| m.pubkey == pubkey) {
            members.push(MemberEntry {
                pubkey: pubkey.to_owned(),
                joined_at: Utc::now(),
            });
        }
        Ok(())
    }

    fn members(&self, team_id: &str) -> Result<Vec<MemberEntry>, CollabError> {
        let state = self.state.lock().expect("stub state poisoned");
        // `members` itself is accessible to non-members in principle, but we gate
        // it here to stay consistent with the ADR-027 requirement that all methods
        // require membership.  The "caller" for the stub is unspecified (no bearer
        // token), so we return the full list without a membership gate on this call
        // since the test harness calls it directly for assertion purposes.
        Ok(state.members.get(team_id).cloned().unwrap_or_default())
    }

    fn deliver_dek(
        &self,
        team_id: &str,
        recipient_pubkey: &str,
        envelope: DekEnvelope,
    ) -> Result<(), CollabError> {
        let mut state = self.state.lock().expect("stub state poisoned");
        // The *deliverer* (caller) must be a member.  In the stub we treat
        // recipient_pubkey as the delivery target; the real caller identity is
        // implicit.  We require the recipient to also be a member so the server
        // never delivers to unknown keys.
        state.assert_member(team_id, recipient_pubkey)?;
        state
            .dek_inbox
            .insert((team_id.to_owned(), recipient_pubkey.to_owned()), envelope);
        Ok(())
    }

    fn fetch_dek(
        &self,
        team_id: &str,
        recipient_pubkey: &str,
    ) -> Result<Option<DekEnvelope>, CollabError> {
        let state = self.state.lock().expect("stub state poisoned");
        // Membership gate: the caller (recipient) must be a member.
        state.assert_member(team_id, recipient_pubkey)?;
        Ok(state
            .dek_inbox
            .get(&(team_id.to_owned(), recipient_pubkey.to_owned()))
            .cloned())
    }

    fn create_invite(
        &self,
        team_id: &str,
        invitee_pubkey: &str,
    ) -> Result<InviteToken, CollabError> {
        let mut state = self.state.lock().expect("stub state poisoned");
        let token_str = Uuid::new_v4().to_string();
        let expires_at = Utc::now() + chrono::Duration::hours(INVITE_TTL_HOURS);
        let record = InviteRecord {
            token: token_str.clone(),
            team_id: team_id.to_owned(),
            invitee_pubkey: invitee_pubkey.to_owned(),
            expires_at,
        };
        state.invites.insert(token_str.clone(), record);
        Ok(InviteToken {
            token: token_str,
            team_id: team_id.to_owned(),
            bound_pubkey: invitee_pubkey.to_owned(),
        })
    }

    fn confirm_invite(
        &self,
        team_id: &str,
        token: &InviteToken,
        confirming_pubkey: &str,
    ) -> Result<(), CollabError> {
        let mut state = self.state.lock().expect("stub state poisoned");

        let record = state
            .invites
            .get(&token.token)
            .ok_or(CollabError::InviteNotFound)?;

        // TOFU binding check: the confirming key must match the key bound at
        // invite creation time (ADR-028 §"Invitee confirms their key").
        if record.invitee_pubkey != confirming_pubkey {
            return Err(CollabError::PubkeyMismatch);
        }

        // Expiry check.
        if Utc::now() > record.expires_at {
            return Err(CollabError::InviteNotFound);
        }

        let confirmed_pubkey = record.invitee_pubkey.clone();
        let confirmed_team = record.team_id.clone();
        let token_str = token.token.clone();

        // Enforce that the team_id in the token matches the call parameter.
        if confirmed_team != team_id {
            return Err(CollabError::InviteNotFound);
        }

        // Consume the invite (one-use).
        state.invites.remove(&token_str);

        // Add the confirmed member to the team directory.
        let members = state.members.entry(team_id.to_owned()).or_default();
        if !members.iter().any(|m| m.pubkey == confirmed_pubkey) {
            members.push(MemberEntry {
                pubkey: confirmed_pubkey,
                joined_at: Utc::now(),
            });
        }

        Ok(())
    }

    fn deliver_recovery_share(
        &self,
        team_id: &str,
        custodian_pubkey: &str,
        share_ciphertext: &str,
    ) -> Result<(), CollabError> {
        let mut state = self.state.lock().expect("stub state poisoned");
        state.assert_member(team_id, custodian_pubkey)?;
        state.share_inbox.insert(
            (team_id.to_owned(), custodian_pubkey.to_owned()),
            share_ciphertext.to_owned(),
        );
        Ok(())
    }

    fn fetch_recovery_share(
        &self,
        team_id: &str,
        custodian_pubkey: &str,
    ) -> Result<Option<String>, CollabError> {
        let state = self.state.lock().expect("stub state poisoned");
        state.assert_member(team_id, custodian_pubkey)?;
        Ok(state
            .share_inbox
            .get(&(team_id.to_owned(), custodian_pubkey.to_owned()))
            .cloned())
    }
}