moshpits 0.2.0

A Rust implementation of in the same vein as Mosh, the mobile shell.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

// Copyright (c) 2025 moshpit developers
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or https://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

use std::{
    collections::{HashMap, VecDeque},
    fmt,
    sync::Arc,
};

use libmoshpit::{EncryptedFrame, TerminalMessage};
use tokio::sync::{Mutex, mpsc::Sender};
use tokio_util::sync::CancellationToken;
use uuid::Uuid;

/// Maximum bytes kept in the per-session scrollback ring buffer (64 KiB).
pub(crate) const SCROLLBACK_CAPACITY: usize = 65_536;

/// Replaceable output handle for a session.
#[derive(Debug)]
pub(crate) struct SessionOutputHandle {
    /// Per-connection UUID used to tag outbound [`EncryptedFrame::Bytes`] datagrams.
    pub kex_uuid: Uuid,
    /// Channel to the live [`crate::UdpSender`].  Set to `None` when the client has
    /// disconnected and the PTY is running headless.
    pub tx: Option<Sender<EncryptedFrame>>,
    /// Cancellation token for the current connection's UDP tasks.  Cancelled on resume
    /// to shut down the stale reader/sender pair.
    pub conn_token: Option<CancellationToken>,
    /// UDP port allocated for the current connection, returned to the pool when the PTY
    /// session ends.
    pub udp_port: Option<u16>,
}

/// Full state for one live PTY session.
pub(crate) struct SessionRecord {
    /// Forward keyboard / resize events from the connected client into this channel.
    pub term_tx: Sender<TerminalMessage>,
    /// Shared, replaceable output handle – updated on every reconnect.
    pub output_handle: Arc<Mutex<SessionOutputHandle>>,
    /// Ring buffer of raw PTY output bytes for scrollback replay on reconnect.
    pub scrollback: Arc<Mutex<VecDeque<u8>>>,
    /// Server-side vt100 emulator tracking current PTY screen state.
    /// Fed by the PTY reader thread; queried on reconnect and by the periodic
    /// screen-state sync task to produce [`libmoshpit::EncryptedFrame::ScreenState`] frames.
    pub server_emulator: Arc<Mutex<vt100::Parser>>,
}

impl fmt::Debug for SessionRecord {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("SessionRecord")
            .field("output_handle", &self.output_handle)
            .field("scrollback", &self.scrollback)
            .finish_non_exhaustive()
    }
}

/// Full session registry: maps stable session UUID → [`SessionRecord`].
pub(crate) type FullSessionRegistry = Arc<Mutex<HashMap<Uuid, SessionRecord>>>;

/// Create a new, empty [`FullSessionRegistry`].
pub(crate) fn new_full_registry() -> FullSessionRegistry {
    Arc::new(Mutex::new(HashMap::new()))
}

#[cfg(test)]
mod test {
    use std::{collections::VecDeque, sync::Arc};

    use libmoshpit::{EncryptedFrame, TerminalMessage};
    use tokio::sync::{Mutex, mpsc::channel};
    use uuid::Uuid;

    use super::{SCROLLBACK_CAPACITY, SessionOutputHandle, SessionRecord, new_full_registry};

    #[test]
    fn scrollback_capacity_is_64kib() {
        assert_eq!(SCROLLBACK_CAPACITY, 65_536);
    }

    #[tokio::test]
    async fn new_full_registry_is_empty() {
        let registry = new_full_registry();
        assert!(registry.lock().await.is_empty());
    }

    #[test]
    fn session_output_handle_debug() {
        let handle = SessionOutputHandle {
            kex_uuid: Uuid::nil(),
            tx: None::<tokio::sync::mpsc::Sender<EncryptedFrame>>,
            conn_token: None,
            udp_port: None,
        };
        let s = format!("{handle:?}");
        assert!(s.contains("SessionOutputHandle"));
        assert!(s.contains("kex_uuid"));
    }

    #[tokio::test]
    async fn session_record_debug() {
        let (term_tx, _term_rx) = channel::<TerminalMessage>(1);
        let (frame_tx, _frame_rx) = channel::<EncryptedFrame>(1);
        let output_handle = Arc::new(Mutex::new(SessionOutputHandle {
            kex_uuid: Uuid::nil(),
            tx: Some(frame_tx),
            conn_token: None,
            udp_port: None,
        }));
        let scrollback = Arc::new(Mutex::new(VecDeque::<u8>::new()));
        let server_emulator = Arc::new(Mutex::new(vt100::Parser::new(24, 80, 0)));
        let record = SessionRecord {
            term_tx,
            output_handle,
            scrollback,
            server_emulator,
        };
        let s = format!("{record:?}");
        assert!(s.contains("SessionRecord"));
        assert!(s.contains("output_handle"));
    }
}