wireguard-admin/src/app/first_delivery.rs

//! Session-local ownership of password-free first-delivery eligibility.
//!
//! The session keeps each generated private key beside the saved peer and
//! encrypted-record context that authorizes its single first-file publication.

use std::collections::BTreeMap;

use crate::delivery::ClientConfigPayload;
use crate::domain::PublicKey;
use crate::vault::{VaultBundle, WireGuardPrivateKey};
use crate::wireguard::ConfigDocument;

use super::{AppError, FirstIssueAttempt, FirstIssueDeliveryFailure, FirstIssueRetry};

#[derive(Clone, Eq, PartialEq)]
pub(super) struct IssueContext {
    peer_directives: String,
    encrypted_record: String,
}

pub(super) struct FirstDeliverySession {
    next_authorization: u64,
    issues: BTreeMap<PublicKey, InitialIssue>,
}

struct InitialIssue {
    private_key: WireGuardPrivateKey,
    context: IssueContext,
    stage: IssueStage,
}

#[derive(Clone, Copy, Eq, PartialEq)]
enum IssueStage {
    PendingDraft,
    AwaitingFirstFile,
    RetryAvailable,
    Authorizing(u64),
    RetryNeedsOverwriteConfirmation,
    RetryNeedsManualCleanup,
}

/// Read-only first-file availability with no access to the held private key,
/// issue context, or payload.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(super) enum SessionFirstDeliveryAvailability {
    InitialFile,
    Retry,
    RetryNeedsOverwriteConfirmation,
    RetryNeedsManualCleanup,
    Unavailable,
}

pub(super) struct FirstIssueAuthorization {
    identity: PublicKey,
    context: IssueContext,
    token: u64,
    payload: ClientConfigPayload,
}

impl FirstIssueAuthorization {
    pub(super) fn payload(&self) -> &ClientConfigPayload {
        &self.payload
    }
}

impl FirstDeliverySession {
    pub(super) fn new() -> Self {
        Self {
            next_authorization: 1,
            issues: BTreeMap::new(),
        }
    }

    pub(super) fn len(&self) -> usize {
        self.issues.len()
    }

    pub(super) fn read_model_availability(
        &self,
        identity: PublicKey,
        document: &ConfigDocument,
        vault: Option<&VaultBundle>,
    ) -> SessionFirstDeliveryAvailability {
        let Some(issue) = self.issues.get(&identity) else {
            return SessionFirstDeliveryAvailability::Unavailable;
        };
        let matching = vault
            .and_then(|vault| context_for(document, vault, identity).ok())
            .is_some_and(|actual| actual == issue.context);
        if !matching {
            return SessionFirstDeliveryAvailability::Unavailable;
        }
        match issue.stage {
            IssueStage::AwaitingFirstFile => SessionFirstDeliveryAvailability::InitialFile,
            IssueStage::RetryAvailable => SessionFirstDeliveryAvailability::Retry,
            IssueStage::RetryNeedsOverwriteConfirmation => {
                SessionFirstDeliveryAvailability::RetryNeedsOverwriteConfirmation
            }
            IssueStage::RetryNeedsManualCleanup => {
                SessionFirstDeliveryAvailability::RetryNeedsManualCleanup
            }
            IssueStage::PendingDraft | IssueStage::Authorizing(_) => {
                SessionFirstDeliveryAvailability::Unavailable
            }
        }
    }

    #[cfg(test)]
    pub(super) fn is_empty(&self) -> bool {
        self.issues.is_empty()
    }

    #[cfg(test)]
    pub(super) fn contains(&self, identity: &PublicKey) -> bool {
        self.issues.contains_key(identity)
    }

    pub(super) fn remember_pending(
        &mut self,
        identity: PublicKey,
        private_key: WireGuardPrivateKey,
        context: IssueContext,
    ) {
        self.issues.insert(
            identity,
            InitialIssue {
                private_key,
                context,
                stage: IssueStage::PendingDraft,
            },
        );
    }

    pub(super) fn remove(&mut self, identity: &PublicKey) {
        self.issues.remove(identity);
    }

    pub(super) fn clear(&mut self) {
        self.issues.clear();
    }

    /// In-flight output can no longer be retried silently after a draft or
    /// runtime transition. Keep unrelated clients' rights intact.
    pub(super) fn invalidate_active_authorizations(&mut self) {
        for issue in self.issues.values_mut() {
            if matches!(issue.stage, IssueStage::Authorizing(_)) {
                issue.stage = IssueStage::RetryNeedsOverwriteConfirmation;
            }
        }
    }

    /// Keep only already committed first-file rights when a draft is dropped.
    pub(super) fn discard_draft(&mut self, document: &ConfigDocument, vault: Option<&VaultBundle>) {
        self.issues
            .retain(|_, issue| issue.stage == IssueStage::AwaitingFirstFile);
        self.retain_matching(document, vault);
    }

    pub(super) fn retain_matching(
        &mut self,
        document: &ConfigDocument,
        vault: Option<&VaultBundle>,
    ) {
        self.issues.retain(|identity, issue| {
            if issue.stage == IssueStage::PendingDraft {
                return true;
            }
            vault
                .and_then(|vault| context_for(document, vault, *identity).ok())
                .is_some_and(|actual| actual == issue.context)
        });
    }

    pub(super) fn promote_matching(
        &mut self,
        document: &ConfigDocument,
        vault: Option<&VaultBundle>,
    ) {
        self.issues.retain(|identity, issue| {
            let matches = vault
                .and_then(|vault| context_for(document, vault, *identity).ok())
                .is_some_and(|actual| actual == issue.context);
            if matches && issue.stage == IssueStage::PendingDraft {
                issue.stage = IssueStage::AwaitingFirstFile;
            }
            matches
        });
    }

    pub(super) fn begin_attempt(
        &mut self,
        identity: PublicKey,
        attempt: FirstIssueAttempt,
        document: &ConfigDocument,
        vault: Option<&VaultBundle>,
        render: impl FnOnce(&WireGuardPrivateKey) -> Result<ClientConfigPayload, AppError>,
    ) -> Result<FirstIssueAuthorization, AppError> {
        self.retain_matching(document, vault);
        let (is_retry, confirm_possible_overwrite, confirm_manual_cleanup) = match attempt {
            FirstIssueAttempt::Initial => (false, false, false),
            FirstIssueAttempt::Retry {
                confirm_possible_overwrite,
                confirm_manual_cleanup,
            } => (true, confirm_possible_overwrite, confirm_manual_cleanup),
        };
        let issue = self
            .issues
            .get(&identity)
            .ok_or(AppError::FirstIssueUnavailable)?;
        let allowed = match issue.stage {
            IssueStage::AwaitingFirstFile => !is_retry,
            IssueStage::RetryAvailable => is_retry,
            IssueStage::RetryNeedsOverwriteConfirmation => is_retry && confirm_possible_overwrite,
            IssueStage::RetryNeedsManualCleanup => is_retry && confirm_manual_cleanup,
            IssueStage::PendingDraft | IssueStage::Authorizing(_) => false,
        };
        if !allowed {
            return Err(AppError::FirstIssueUnavailable);
        }
        let current_context = context_for(
            document,
            vault.ok_or(AppError::FirstIssueUnavailable)?,
            identity,
        )?;
        if current_context != issue.context {
            self.issues.remove(&identity);
            return Err(AppError::FirstIssueUnavailable);
        }
        let payload = render(&issue.private_key)?;
        let context = issue.context.clone();
        let token = self.next_authorization;
        let next_token = token
            .checked_add(1)
            .ok_or(AppError::IssueAuthorizationExhausted)?;
        let issue = self
            .issues
            .get_mut(&identity)
            .ok_or(AppError::FirstIssueUnavailable)?;
        self.next_authorization = next_token;
        issue.stage = IssueStage::Authorizing(token);
        Ok(FirstIssueAuthorization {
            identity,
            context,
            token,
            payload,
        })
    }

    pub(super) fn validate_authorization(
        &mut self,
        authorization: &FirstIssueAuthorization,
        document: &ConfigDocument,
        vault: Option<&VaultBundle>,
    ) -> Result<(), AppError> {
        self.retain_matching(document, vault);
        let current = context_for(
            document,
            vault.ok_or(AppError::FirstIssueUnavailable)?,
            authorization.identity,
        )?;
        if current != authorization.context
            || !matches!(
                self.issues.get(&authorization.identity),
                Some(issue)
                    if issue.stage == IssueStage::Authorizing(authorization.token)
                        && issue.context == authorization.context
            )
        {
            return Err(AppError::FirstIssueUnavailable);
        }
        Ok(())
    }

    pub(super) fn require_overwrite_confirmation(&mut self, identity: PublicKey) {
        if let Some(issue) = self.issues.get_mut(&identity)
            && matches!(issue.stage, IssueStage::Authorizing(_))
        {
            issue.stage = IssueStage::RetryNeedsOverwriteConfirmation;
        }
    }

    pub(super) fn complete_publication(
        &mut self,
        authorization: FirstIssueAuthorization,
        document: &ConfigDocument,
        vault: Option<&VaultBundle>,
    ) -> Result<(), AppError> {
        self.validate_authorization(&authorization, document, vault)?;
        self.issues.remove(&authorization.identity);
        Ok(())
    }

    pub(super) fn fail_publication(
        &mut self,
        authorization: FirstIssueAuthorization,
        failure: FirstIssueDeliveryFailure,
        document: &ConfigDocument,
        vault: Option<&VaultBundle>,
    ) -> Result<FirstIssueRetry, AppError> {
        self.validate_authorization(&authorization, document, vault)?;
        let issue = self
            .issues
            .get_mut(&authorization.identity)
            .ok_or(AppError::FirstIssueUnavailable)?;
        let retry = match failure {
            FirstIssueDeliveryFailure::BeforePublication => {
                issue.stage = IssueStage::RetryAvailable;
                FirstIssueRetry::Available
            }
            FirstIssueDeliveryFailure::PublicationUncertain => {
                issue.stage = IssueStage::RetryNeedsOverwriteConfirmation;
                FirstIssueRetry::OverwriteConfirmationRequired
            }
            FirstIssueDeliveryFailure::CleanupRequired => {
                issue.stage = IssueStage::RetryNeedsManualCleanup;
                FirstIssueRetry::ManualCleanupRequired
            }
        };
        Ok(retry)
    }
}

pub(super) fn context_for(
    document: &ConfigDocument,
    vault: &VaultBundle,
    identity: PublicKey,
) -> Result<IssueContext, AppError> {
    let matching = document
        .peers()
        .iter()
        .filter(|peer| {
            peer.identity == Some(identity) || peer.public_key_candidates.contains(&identity)
        })
        .collect::<Vec<_>>();
    let [peer] = matching.as_slice() else {
        return Err(AppError::FirstIssueUnavailable);
    };
    if peer.identity != Some(identity) || !peer.issues.is_empty() {
        return Err(AppError::FirstIssueUnavailable);
    }
    // Alias/admin comments are descriptive metadata; only WireGuard directives
    // bind first-issue authority, so renames do not invalidate the creation.
    let peer_directives = document
        .source_slice(peer.section)
        .ok_or(AppError::FirstIssueUnavailable)?
        .lines()
        .map(str::trim)
        .filter(|line| !line.is_empty() && !line.starts_with(['#', ';']))
        .collect::<Vec<_>>()
        .join("\n");
    let encrypted_record = vault
        .record_context(identity)
        .ok_or(AppError::FirstIssueUnavailable)?
        .to_owned();
    Ok(IssueContext {
        peer_directives,
        encrypted_record,
    })
}