oxdock_ssh_plugin/
types.rs1use std::fmt;
8use std::net::SocketAddr;
9use std::sync::{Arc, Mutex};
10
11use anyhow::{Result, bail};
12use oxdock_func_macro::oxdock_type;
13use tokio::sync::mpsc;
14
15use crate::state::{DownMsg, ServerState, UpMsg};
16
17#[oxdock_type(name = "SSH_SERVER")]
23#[derive(Debug, Clone)]
24pub struct SshServerTag {
25 state: Arc<ServerState>,
26}
27
28impl SshServerTag {
29 pub fn new(state: Arc<ServerState>) -> Self {
30 Self { state }
31 }
32
33 pub fn state(&self) -> &Arc<ServerState> {
34 &self.state
35 }
36}
37
38impl PartialEq for SshServerTag {
39 fn eq(&self, other: &Self) -> bool {
40 self.state.id() == other.state.id()
41 }
42}
43
44impl fmt::Display for SshServerTag {
45 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
46 write!(
47 formatter,
48 "SSH_SERVER({}, {})",
49 self.state.id(),
50 self.state.addr_text()
51 )
52 }
53}
54
55static SESSION_IDS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
57
58#[derive(Debug)]
62struct SshSessionInner {
63 id: u64,
64 command: Option<String>,
65 username: Option<String>,
66 peer_addr: Option<SocketAddr>,
67 pty_size: crate::state::SharedPtySize,
70 up_rx: Mutex<Option<mpsc::Receiver<UpMsg>>>,
71 down_tx: Mutex<Option<mpsc::Sender<DownMsg>>>,
72}
73
74#[oxdock_type(name = "SSH_SESSION")]
80#[derive(Debug, Clone)]
81pub struct SshSessionTag {
82 inner: Arc<SshSessionInner>,
83}
84
85impl SshSessionTag {
86 pub fn new(
87 command: Option<String>,
88 username: Option<String>,
89 peer_addr: Option<SocketAddr>,
90 pty_size: crate::state::SharedPtySize,
91 up_rx: mpsc::Receiver<UpMsg>,
92 down_tx: mpsc::Sender<DownMsg>,
93 ) -> Self {
94 Self {
95 inner: Arc::new(SshSessionInner {
96 id: SESSION_IDS.fetch_add(1, std::sync::atomic::Ordering::SeqCst),
97 command,
98 username,
99 peer_addr,
100 pty_size,
101 up_rx: Mutex::new(Some(up_rx)),
102 down_tx: Mutex::new(Some(down_tx)),
103 }),
104 }
105 }
106
107 pub fn command(&self) -> Option<String> {
108 self.inner.command.clone()
109 }
110
111 pub fn username(&self) -> Option<String> {
112 self.inner.username.clone()
113 }
114
115 pub fn peer_addr(&self) -> Option<SocketAddr> {
116 self.inner.peer_addr
117 }
118
119 pub fn pty_size(&self) -> crate::state::PtySize {
121 crate::state::snapshot_pty_size(&self.inner.pty_size).size
122 }
123
124 pub fn pty_size_handle(&self) -> crate::state::SharedPtySize {
126 Arc::clone(&self.inner.pty_size)
127 }
128
129 pub fn take_pump_ends(&self) -> Result<(mpsc::Receiver<UpMsg>, mpsc::Sender<DownMsg>)> {
133 let mut up = self.inner.up_rx.lock().unwrap_or_else(|p| p.into_inner());
134 let mut down = self.inner.down_tx.lock().unwrap_or_else(|p| p.into_inner());
135 if up.is_none() || down.is_none() {
136 bail!("SSH session already pumped");
137 }
138 Ok((up.take().unwrap(), down.take().unwrap()))
139 }
140}
141
142impl PartialEq for SshSessionTag {
143 fn eq(&self, other: &Self) -> bool {
144 self.inner.id == other.inner.id
145 }
146}
147
148impl fmt::Display for SshSessionTag {
149 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
150 let peer = self
151 .inner
152 .peer_addr
153 .map(|addr| addr.to_string())
154 .unwrap_or_default();
155 match &self.inner.username {
156 Some(user) => write!(
157 formatter,
158 "SSH_SESSION(sess-{} {}@{})",
159 self.inner.id, user, peer
160 ),
161 None => write!(formatter, "SSH_SESSION(sess-{} {})", self.inner.id, peer),
162 }
163 }
164}