wireguard-admin/src/tui/state.rs

//! Secret-free navigation and session projections for the terminal interface.

use std::collections::BTreeSet;
use std::net::IpAddr;

use crate::app::{Application, ClientDisplay};
use crate::domain::{Client, PublicKey, RuntimeObservation};
use crate::runtime::RuntimeObservation as Statistics;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Screen {
    ClientList,
    ClientDetails,
    Help,
    Progress,
    Message,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Focus {
    ClientList,
    Details,
    Dialog,
    Help,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum DialogKind {
    Rename,
    Delete,
    Create,
    Commit,
    Restart,
    Exit,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum AppliedState {
    Unknown,
    Unavailable,
    Verified,
}

/// Public client values for both states. No configuration bytes or secrets are
/// retained in a view snapshot.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ClientStateRow {
    pub saved: Option<Client>,
    pub draft: Option<Client>,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct SessionView {
    pub generation: u64,
    pub clients: Vec<ClientDisplay>,
    pub client_states: std::collections::BTreeMap<PublicKey, ClientStateRow>,
    /// Number of peer sections in the installed server document, when its
    /// saved bytes can be projected. This includes nonmanaged peers.
    pub saved_peer_count: Option<usize>,
    pub document_issues: Vec<crate::wireguard::DocumentIssue>,
    pub applied: AppliedState,
    pub reserved_addresses: BTreeSet<IpAddr>,
    pub statistics: Statistics,
    pub has_draft: bool,
    pub vault_initialized: bool,
}

impl SessionView {
    pub fn from_application(app: &Application, generation: u64) -> Self {
        let states = app.states();
        let mut identities = BTreeSet::new();
        identities.extend(states.saved().clients().map(|client| client.identity));
        identities.extend(states.draft().clients().map(|client| client.identity));
        let client_states = identities
            .into_iter()
            .map(|identity| {
                (
                    identity,
                    ClientStateRow {
                        saved: states.saved().get(&identity).cloned(),
                        draft: states.draft().get(&identity).cloned(),
                    },
                )
            })
            .collect();
        let applied = match states.applied() {
            RuntimeObservation::Unknown => AppliedState::Unknown,
            RuntimeObservation::Unavailable => AppliedState::Unavailable,
            RuntimeObservation::Verified => AppliedState::Verified,
        };
        let projection = app.display_projection();
        let saved_peer_count =
            crate::wireguard::ConfigDocument::parse(&app.saved_bytes().server_config)
                .ok()
                .map(|document| document.peers().len());
        Self {
            generation,
            clients: projection.clients,
            client_states,
            saved_peer_count,
            document_issues: projection.document_issues,
            applied,
            reserved_addresses: states.reserved_addresses().clone(),
            statistics: Statistics::unknown(),
            has_draft: app.has_draft_changes(),
            vault_initialized: app.vault_initialized(),
        }
    }

    pub fn contains_identity(&self, identity: &PublicKey) -> bool {
        self.clients
            .iter()
            .any(|client| client.identity == Some(*identity))
    }

    /// Rows shown in the client table. Document issues are deliberately
    /// excluded because they are summaries, not client action targets.
    pub fn visible_records(&self) -> Vec<VisibleRecord<'_>> {
        let mut records = self
            .clients
            .iter()
            .map(VisibleRecord::Peer)
            .collect::<Vec<_>>();
        records.extend(self.client_states.iter().filter_map(|(identity, row)| {
            let client = row.saved.as_ref()?;
            if row.draft.is_some()
                || self
                    .clients
                    .iter()
                    .any(|peer| peer.identity == Some(*identity))
            {
                return None;
            }
            Some(VisibleRecord::DraftRemoved {
                identity: *identity,
                client,
            })
        }));
        records
    }
}

#[derive(Clone, Copy, Debug)]
pub enum VisibleRecord<'a> {
    Peer(&'a ClientDisplay),
    DraftRemoved {
        identity: PublicKey,
        client: &'a Client,
    },
}

impl VisibleRecord<'_> {
    pub fn identity(self) -> Option<PublicKey> {
        match self {
            Self::Peer(client) => client.identity,
            Self::DraftRemoved { identity, .. } => Some(identity),
        }
    }

    pub fn peer_index(self) -> Option<usize> {
        match self {
            Self::Peer(client) => Some(client.peer_index),
            Self::DraftRemoved { .. } => None,
        }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum OperationKind {
    CreateClient,
    RenameClient,
    DeleteClient,
    DiscardDraft,
    Commit,
    CreateBackup,
    ListBackups,
    ImportSelectedFile,
    InitialSetup,
    PreviewRestart,
    Restart,
    DeliverClient,
    Restore,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct PendingOperation {
    pub request_id: u64,
    pub kind: OperationKind,
    pub generation: u64,
}

/// Mutable UI navigation and latest accepted worker projection.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct UiState {
    pub screen: Screen,
    pub focus: Focus,
    pub dialog: Option<DialogKind>,
    pub selected_identity: Option<PublicKey>,
    /// Authoritative document index used only to inspect malformed peers that
    /// have no public key. Mutating actions must always use selected_identity.
    pub selected_peer_index: Option<usize>,
    pub scroll_offset: usize,
    pub generation: u64,
    pub view: Option<SessionView>,
    pub pending: Option<PendingOperation>,
    pub status_message: Option<String>,
}

impl Default for UiState {
    fn default() -> Self {
        Self {
            screen: Screen::ClientList,
            focus: Focus::ClientList,
            dialog: None,
            selected_identity: None,
            selected_peer_index: None,
            scroll_offset: 0,
            generation: 0,
            view: None,
            pending: None,
            status_message: None,
        }
    }
}

impl UiState {
    /// Replace the view only when it belongs to the current generation. The
    /// selected full public key remains stable through sorting and refreshes.
    pub fn apply_view(&mut self, view: SessionView) -> bool {
        if view.generation != self.generation {
            return false;
        }
        self.reconcile_selection(&view);
        self.view = Some(view);
        true
    }

    fn reconcile_selection(&mut self, view: &SessionView) {
        if let Some(identity) = self.selected_identity {
            let matching = view
                .clients
                .iter()
                .filter(|client| client.identity == Some(identity))
                .collect::<Vec<_>>();
            if matching.len() == 1 {
                self.selected_peer_index = Some(matching[0].peer_index);
                return;
            }
            if matching.len() > 1 {
                if self
                    .selected_peer_index
                    .is_some_and(|index| matching.iter().any(|client| client.peer_index == index))
                {
                    return;
                }
                // The old row disappeared and the key is now ambiguous. Do
                // not guess which peer the administrator meant.
                self.selected_identity = None;
                self.selected_peer_index = None;
                self.select_first_record(view);
                return;
            }
            let removed = view.client_states.get(&identity).is_some_and(|row| {
                row.saved.is_some()
                    && row.draft.is_none()
                    && !view
                        .clients
                        .iter()
                        .any(|client| client.identity == Some(identity))
            });
            if removed {
                self.selected_peer_index = None;
                return;
            }
            self.selected_identity = None;
            self.selected_peer_index = None;
        }

        if let Some(peer_index) = self.selected_peer_index {
            if let Some(client) = view
                .clients
                .iter()
                .find(|client| client.peer_index == peer_index)
            {
                self.selected_identity = client.identity;
                return;
            }
            self.selected_peer_index = None;
        }

        if self.selected_identity.is_some_and(|identity| {
            view.client_states
                .get(&identity)
                .is_some_and(|row| row.saved.is_some() && row.draft.is_none())
        }) {
            return;
        }

        self.select_first_record(view);
    }

    fn select_first_record(&mut self, view: &SessionView) {
        match view.visible_records().first().copied() {
            Some(VisibleRecord::Peer(client)) => {
                self.selected_peer_index = Some(client.peer_index);
                self.selected_identity = client.identity;
            }
            Some(VisibleRecord::DraftRemoved { identity, .. }) => {
                self.selected_peer_index = None;
                self.selected_identity = Some(identity);
            }
            None => {
                self.selected_peer_index = None;
                self.selected_identity = None;
            }
        }
    }

    /// Apply only statistics sampled from the still-current saved/draft view.
    /// This never changes the applied-state marker.
    pub fn apply_statistics(&mut self, generation: u64, statistics: Statistics) -> bool {
        if generation != self.generation {
            return false;
        }
        let Some(view) = self.view.as_mut() else {
            return false;
        };
        if view.generation != generation {
            return false;
        }
        view.statistics = statistics;
        true
    }

    pub fn begin_operation(&mut self, operation: PendingOperation) -> bool {
        if self.pending.is_some() || operation.generation < self.generation {
            return false;
        }
        self.generation = operation.generation;
        self.pending = Some(operation);
        self.screen = Screen::Progress;
        true
    }

    /// Close a form/dialog without cancelling work already submitted to the
    /// application worker. The pending progress state remains visible.
    pub fn cancel_dialog(&mut self) {
        self.dialog = None;
        if self.pending.is_some() {
            self.screen = Screen::Progress;
        } else {
            self.screen = if self.selected_identity.is_some() {
                Screen::ClientDetails
            } else {
                Screen::ClientList
            };
        }
    }

    pub fn finish_operation(&mut self, operation: PendingOperation) -> bool {
        if self.pending != Some(operation) {
            return false;
        }
        self.pending = None;
        self.screen = if self.dialog.is_some() {
            Screen::Message
        } else if self.selected_identity.is_some() {
            Screen::ClientDetails
        } else {
            Screen::ClientList
        };
        true
    }

    pub fn select(&mut self, identity: PublicKey) -> bool {
        let Some(view) = &self.view else { return false };
        let mut matching = view
            .clients
            .iter()
            .filter(|row| row.identity == Some(identity));
        let Some(client) = matching.next() else {
            return false;
        };
        if matching.next().is_some() {
            return false;
        }
        self.selected_identity = Some(identity);
        self.selected_peer_index = Some(client.peer_index);
        self.screen = Screen::ClientDetails;
        true
    }

    pub fn select_peer_index(&mut self, peer_index: usize) -> bool {
        if !self.highlight_peer_index(peer_index) {
            return false;
        }
        self.screen = Screen::ClientDetails;
        true
    }

    /// Move the client-list highlight without changing the active screen.
    pub fn highlight_peer_index(&mut self, peer_index: usize) -> bool {
        let Some(client) = self.view.as_ref().and_then(|view| {
            view.clients
                .iter()
                .find(|client| client.peer_index == peer_index)
        }) else {
            return false;
        };
        self.selected_peer_index = Some(peer_index);
        self.selected_identity = client.identity;
        true
    }

    pub fn selected_record(&self) -> Option<VisibleRecord<'_>> {
        let view = self.view.as_ref()?;
        let records = view.visible_records();
        self.selected_visible_index(&records)
            .and_then(|index| records.get(index).copied())
    }

    pub fn selected_visible_index(&self, records: &[VisibleRecord<'_>]) -> Option<usize> {
        self.selected_peer_index
            .and_then(|peer_index| {
                records
                    .iter()
                    .position(|record| record.peer_index() == Some(peer_index))
            })
            .or_else(|| {
                self.selected_identity.and_then(|identity| {
                    records
                        .iter()
                        .position(|record| record.identity() == Some(identity))
                })
            })
            .or_else(|| (!records.is_empty()).then_some(0))
    }

    pub fn select_visible_index(&mut self, selected: usize) -> bool {
        let Some(record) = self
            .view
            .as_ref()
            .and_then(|view| view.visible_records().get(selected).copied())
        else {
            return false;
        };
        match record {
            VisibleRecord::Peer(client) => {
                self.selected_peer_index = Some(client.peer_index);
                self.selected_identity = client.identity;
            }
            VisibleRecord::DraftRemoved { identity, .. } => {
                self.selected_peer_index = None;
                self.selected_identity = Some(identity);
            }
        }
        self.scroll_offset = 0;
        true
    }

    pub fn open_selected_record(&mut self) -> bool {
        let Some(view) = &self.view else { return false };
        let records = view.visible_records();
        let Some(index) = self.selected_visible_index(&records) else {
            return false;
        };
        self.select_visible_index(index);
        self.screen = Screen::ClientDetails;
        self.scroll_offset = 0;
        true
    }

    pub fn selected_client_mutable(&self) -> bool {
        let Some(identity) = self.selected_identity else {
            return false;
        };
        self.view.as_ref().is_some_and(|view| {
            let mut matching = view
                .clients
                .iter()
                .filter(|client| client.identity == Some(identity));
            let Some(client) = matching.next() else {
                return false;
            };
            matching.next().is_none()
                && self
                    .selected_peer_index
                    .is_none_or(|peer_index| peer_index == client.peer_index)
                && client.mutable_identity == Some(identity)
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::domain::PublicKey;
    use crate::wireguard::PeerIssue;

    fn row(identity: PublicKey, label: &str, mutable: bool) -> ClientDisplay {
        ClientDisplay {
            peer_index: 0,
            line: 1,
            identity: Some(identity),
            label: label.to_owned(),
            allowed_ips: vec!["10.0.0.2/32".to_owned()],
            has_subnet_routes: false,
            issues: if mutable {
                Vec::new()
            } else {
                vec![PeerIssue::MissingIpv4HostRoute]
            },
            mutable_identity: mutable.then_some(identity),
        }
    }

    fn view(generation: u64, clients: Vec<ClientDisplay>) -> SessionView {
        SessionView {
            generation,
            clients,
            client_states: Default::default(),
            saved_peer_count: None,
            document_issues: Vec::new(),
            applied: AppliedState::Unknown,
            reserved_addresses: BTreeSet::new(),
            statistics: Statistics::unknown(),
            has_draft: false,
            vault_initialized: true,
        }
    }

    #[test]
    fn refresh_preserves_selection_by_public_key_not_alias_or_row_position() {
        let a = PublicKey::from_bytes([1; 32]);
        let b = PublicKey::from_bytes([2; 32]);
        let mut state = UiState::default();
        assert!(state.apply_view(view(
            0,
            vec![row(a, "same alias", true), row(b, "same alias", true)]
        )));
        assert!(state.select(b));
        assert!(state.apply_view(view(
            0,
            vec![row(b, "renamed", true), row(a, "same alias", true)]
        )));
        assert_eq!(state.selected_identity, Some(b));
    }

    #[test]
    fn stale_view_is_rejected_and_unsupported_selected_peer_is_not_mutable() {
        let key = PublicKey::from_bytes([3; 32]);
        let mut state = UiState::default();
        assert!(state.apply_view(view(0, vec![row(key, "legacy", false)])));
        assert!(state.select(key));
        assert!(!state.selected_client_mutable());
        state.generation = 1;
        assert!(!state.apply_view(view(0, vec![row(key, "stale", true)])));
        assert_eq!(state.selected_identity, Some(key));
    }

    #[test]
    fn malformed_peer_can_be_inspected_by_document_location_without_mutation_identity() {
        let mut malformed = row(PublicKey::from_bytes([8; 32]), "malformed", false);
        malformed.identity = None;
        malformed.mutable_identity = None;
        malformed.peer_index = 3;
        let mut state = UiState::default();
        assert!(state.apply_view(view(0, vec![malformed])));
        assert!(state.select_peer_index(3));
        assert_eq!(state.selected_identity, None);
        assert_eq!(state.selected_peer_index, Some(3));
        assert!(!state.selected_client_mutable());
    }

    #[test]
    fn selection_uses_peer_position_for_duplicate_keys_and_disables_mutation() {
        let key = PublicKey::from_bytes([12; 32]);
        let mut first = row(key, "first duplicate", true);
        first.peer_index = 3;
        let mut second = row(key, "second duplicate", true);
        second.peer_index = 8;
        let mut state = UiState::default();
        assert!(state.apply_view(view(0, vec![first, second])));
        assert_eq!(state.selected_peer_index, Some(3));
        assert!(!state.selected_client_mutable());
        assert!(state.highlight_peer_index(8));
        assert_eq!(state.selected_record().unwrap().peer_index(), Some(8));
        assert_eq!(
            match state.selected_record().unwrap() {
                VisibleRecord::Peer(client) => client.label.as_str(),
                VisibleRecord::DraftRemoved { .. } => panic!("duplicate resolved to removal"),
            },
            "second duplicate"
        );
        assert!(!state.selected_client_mutable());

        let mut refreshed_first = row(key, "refreshed first", true);
        refreshed_first.peer_index = 3;
        let mut refreshed_second = row(key, "refreshed second", true);
        refreshed_second.peer_index = 8;
        assert!(state.apply_view(view(0, vec![refreshed_second, refreshed_first])));
        assert_eq!(state.selected_peer_index, Some(8));
        assert_eq!(
            match state.selected_record().unwrap() {
                VisibleRecord::Peer(client) => client.label.as_str(),
                VisibleRecord::DraftRemoved { .. } => panic!("duplicate resolved to removal"),
            },
            "refreshed second"
        );
        assert!(!state.selected_client_mutable());
    }

    #[test]
    fn removed_saved_client_is_visible_and_never_mutable() {
        let key = PublicKey::from_bytes([13; 32]);
        let removed = Client::new(
            key,
            "removed alias",
            crate::domain::ClientAddresses {
                ipv4: "10.0.0.9".parse().unwrap(),
                ipv6: None,
            },
        )
        .unwrap();
        let mut view = view(0, Vec::new());
        view.client_states.insert(
            key,
            ClientStateRow {
                saved: Some(removed),
                draft: None,
            },
        );
        view.document_issues
            .push(crate::wireguard::DocumentIssue::MalformedSectionHeader { line: 5 });
        let mut state = UiState::default();
        assert!(state.apply_view(view));
        assert_eq!(state.selected_identity, Some(key));
        assert_eq!(
            state.selected_visible_index(&state.view.as_ref().unwrap().visible_records()),
            Some(0)
        );
        assert!(matches!(
            state.selected_record(),
            Some(VisibleRecord::DraftRemoved { identity, .. }) if identity == key
        ));
        assert!(!state.selected_client_mutable());
        assert_eq!(state.view.as_ref().unwrap().visible_records().len(), 1);
    }

    #[test]
    fn cancelling_a_dialog_does_not_claim_submitted_work_was_cancelled() {
        let mut state = UiState {
            dialog: Some(DialogKind::Commit),
            ..UiState::default()
        };
        let operation = PendingOperation {
            request_id: 4,
            kind: OperationKind::Commit,
            generation: 1,
        };
        assert!(state.begin_operation(operation));
        state.cancel_dialog();
        assert_eq!(state.pending, Some(operation));
        assert_eq!(state.screen, Screen::Progress);
    }

    #[test]
    fn statistics_refresh_is_generation_bound_and_never_changes_applied_state() {
        let key = PublicKey::from_bytes([4; 32]);
        let mut state = UiState::default();
        assert!(state.apply_view(view(0, vec![row(key, "client", true)])));
        assert!(state.apply_statistics(0, Statistics::unknown()));
        assert_eq!(state.view.as_ref().unwrap().applied, AppliedState::Unknown);
        state.generation = 1;
        assert!(!state.apply_statistics(0, Statistics::unknown()));
    }
}