agentmux_srv/backend/viewer/
mod.rs1pub mod cert;
22pub mod client;
23pub mod feed;
24pub mod pairing;
25pub mod panes;
26pub mod terminal;
27
28use std::collections::HashMap;
29use std::net::IpAddr;
30use std::path::PathBuf;
31use std::sync::atomic::{AtomicU16, Ordering};
32use std::sync::{Arc, OnceLock};
33
34use tokio_util::sync::CancellationToken;
35
36pub const MAX_FEEDS_PER_DEVICE: usize = 4;
38pub const MAX_FEEDS: usize = 16;
40pub const LAST_SEEN_EVERY_MS: i64 = 60_000;
42
43#[derive(Debug, Clone, Copy, PartialEq, Eq)]
45pub enum FeedRefused {
46 PerDevice,
47 Total,
48}
49
50#[derive(Default)]
51struct FeedRegistry {
52 by_device: HashMap<String, Vec<(u64, CancellationToken)>>,
53 next_id: u64,
54}
55
56impl FeedRegistry {
57 fn total(&self) -> usize {
58 self.by_device.values().map(Vec::len).sum()
59 }
60}
61
62pub struct FeedLease {
65 registry: Arc<parking_lot::Mutex<FeedRegistry>>,
66 device_id: String,
67 id: u64,
68 pub token: CancellationToken,
69}
70
71impl Drop for FeedLease {
72 fn drop(&mut self) {
73 let mut reg = self.registry.lock();
74 if let Some(list) = reg.by_device.get_mut(&self.device_id) {
75 list.retain(|(id, _)| *id != self.id);
76 if list.is_empty() {
77 reg.by_device.remove(&self.device_id);
78 }
79 }
80 }
81}
82
83static ADVERTISED_PORT: OnceLock<Arc<AtomicU16>> = OnceLock::new();
87
88pub fn advertised_viewer_port() -> Option<u16> {
90 ADVERTISED_PORT.get().map(|p| p.load(Ordering::Relaxed)).filter(|p| *p != 0)
91}
92
93pub struct ViewerService {
94 data_dir: Option<PathBuf>,
96 tls: parking_lot::Mutex<Option<Arc<cert::ViewerTls>>>,
97 pub pairing: pairing::Pairing,
98 pub hub: Arc<feed::FeedHub>,
99 feeds: Arc<parking_lot::Mutex<FeedRegistry>>,
100 max_per_device: usize,
101 max_total: usize,
102 advertised_port: Arc<AtomicU16>,
104 bound: parking_lot::Mutex<Vec<IpAddr>>,
106 reservation: parking_lot::Mutex<Option<(u16, socket2::Socket)>>,
111 last_seen_written: parking_lot::Mutex<HashMap<String, i64>>,
113}
114
115impl ViewerService {
116 pub fn new(data_dir: Option<PathBuf>, broker: &crate::backend::mps::Broker) -> Arc<Self> {
118 Self::with_limits(data_dir, broker, Arc::new(feed::FeedHub::default()), MAX_FEEDS_PER_DEVICE, MAX_FEEDS)
119 }
120
121 pub fn with_limits(
122 data_dir: Option<PathBuf>,
123 broker: &crate::backend::mps::Broker,
124 hub: Arc<feed::FeedHub>,
125 max_per_device: usize,
126 max_total: usize,
127 ) -> Arc<Self> {
128 hub.install(broker);
129 Arc::new(Self {
130 data_dir,
131 tls: parking_lot::Mutex::new(None),
132 pairing: pairing::Pairing::new(),
133 hub,
134 feeds: Arc::new(parking_lot::Mutex::new(FeedRegistry::default())),
135 max_per_device,
136 max_total,
137 advertised_port: Arc::new(AtomicU16::new(0)),
138 bound: parking_lot::Mutex::new(Vec::new()),
139 reservation: parking_lot::Mutex::new(None),
140 last_seen_written: parking_lot::Mutex::new(HashMap::new()),
141 })
142 }
143
144 pub fn install_global(&self) {
146 let _ = ADVERTISED_PORT.set(Arc::clone(&self.advertised_port));
147 }
148
149 pub fn advertised_port_handle(&self) -> Arc<AtomicU16> {
151 Arc::clone(&self.advertised_port)
152 }
153
154 pub fn advertised_port(&self) -> Option<u16> {
155 Some(self.advertised_port.load(Ordering::Relaxed)).filter(|p| *p != 0)
156 }
157
158 pub fn tls(&self) -> Result<Arc<cert::ViewerTls>, String> {
160 let mut slot = self.tls.lock();
161 if let Some(tls) = slot.as_ref() {
162 return Ok(Arc::clone(tls));
163 }
164 let tls = Arc::new(match &self.data_dir {
165 Some(dir) => cert::load_or_create(dir)?,
166 None => return Err("the viewer has no data directory for its certificate".to_string()),
167 });
168 *slot = Some(Arc::clone(&tls));
169 Ok(tls)
170 }
171
172 pub fn reserve_port(&self, candidates: impl IntoIterator<Item = u16>) -> Option<u16> {
175 let mut slot = self.reservation.lock();
176 if let Some((port, _)) = slot.as_ref() {
177 return Some(*port);
178 }
179 for port in candidates {
180 let Ok(socket) = socket2::Socket::new(socket2::Domain::IPV4, socket2::Type::STREAM, None) else {
181 continue;
182 };
183 let addr = std::net::SocketAddr::from(([127, 0, 0, 1], port));
184 if socket.bind(&addr.into()).is_ok() {
185 *slot = Some((port, socket));
186 return Some(port);
187 }
188 }
189 None
190 }
191
192 pub fn set_bound(&self, addrs: Vec<IpAddr>, port: u16) {
195 let up = !addrs.is_empty();
196 *self.bound.lock() = addrs;
197 self.advertised_port.store(if up { port } else { 0 }, Ordering::Relaxed);
198 }
199
200 pub fn pair_host(&self) -> Option<IpAddr> {
203 let bound = self.bound.lock().clone();
204 let primary = crate::backend::lan_listeners::primary_lan_ipv4();
205 primary
206 .filter(|p| bound.contains(p))
207 .or_else(|| bound.iter().copied().find(IpAddr::is_ipv4))
208 }
209
210 pub fn open_feed(&self, device_id: &str) -> Result<FeedLease, FeedRefused> {
212 let mut reg = self.feeds.lock();
213 if reg.total() >= self.max_total {
214 return Err(FeedRefused::Total);
215 }
216 if reg.by_device.get(device_id).map_or(0, Vec::len) >= self.max_per_device {
217 return Err(FeedRefused::PerDevice);
218 }
219 reg.next_id += 1;
220 let id = reg.next_id;
221 let token = CancellationToken::new();
222 reg.by_device.entry(device_id.to_string()).or_default().push((id, token.clone()));
223 Ok(FeedLease { registry: Arc::clone(&self.feeds), device_id: device_id.to_string(), id, token })
224 }
225
226 pub fn close_device_feeds(&self, device_id: &str) -> usize {
228 let reg = self.feeds.lock();
229 let Some(list) = reg.by_device.get(device_id) else { return 0 };
230 for (_, token) in list {
231 token.cancel();
232 }
233 list.len()
234 }
235
236 #[cfg(test)]
237 pub fn open_feeds(&self) -> usize {
238 self.feeds.lock().total()
239 }
240
241 pub fn last_seen_due(&self, device_id: &str, now_ms: i64) -> bool {
244 let mut written = self.last_seen_written.lock();
245 match written.get(device_id) {
246 Some(at) if now_ms - at < LAST_SEEN_EVERY_MS => false,
247 _ => {
248 written.insert(device_id.to_string(), now_ms);
249 true
250 }
251 }
252 }
253
254 pub fn forget_device(&self, device_id: &str) {
256 self.last_seen_written.lock().remove(device_id);
257 }
258}
259
260#[cfg(test)]
261mod tests {
262 use super::*;
263
264 fn service(per_device: usize, total: usize) -> Arc<ViewerService> {
265 let broker = crate::backend::mps::Broker::new();
266 ViewerService::with_limits(None, &broker, Arc::new(feed::FeedHub::default()), per_device, total)
267 }
268
269 #[test]
270 fn feeds_are_capped_per_device_and_in_all() {
271 let v = service(2, 3);
272 let a1 = v.open_feed("a").unwrap();
273 let _a2 = v.open_feed("a").unwrap();
274 assert_eq!(v.open_feed("a").err(), Some(FeedRefused::PerDevice));
275 let _b1 = v.open_feed("b").unwrap();
276 assert_eq!(v.open_feed("c").err(), Some(FeedRefused::Total));
277 drop(a1);
278 assert_eq!(v.open_feeds(), 2);
279 assert!(v.open_feed("c").is_ok(), "a closed feed frees its place");
280 }
281
282 #[test]
283 fn closing_a_devices_feeds_cancels_only_its_own() {
284 let v = service(4, 16);
285 let a1 = v.open_feed("a").unwrap();
286 let a2 = v.open_feed("a").unwrap();
287 let b1 = v.open_feed("b").unwrap();
288 assert_eq!(v.close_device_feeds("a"), 2);
289 assert!(a1.token.is_cancelled() && a2.token.is_cancelled());
290 assert!(!b1.token.is_cancelled());
291 assert_eq!(v.close_device_feeds("nobody"), 0);
292 }
293
294 #[test]
295 fn last_seen_is_written_at_most_once_a_minute() {
296 let v = service(4, 16);
297 assert!(v.last_seen_due("a", 1_000));
298 assert!(!v.last_seen_due("a", 1_000 + LAST_SEEN_EVERY_MS - 1));
299 assert!(v.last_seen_due("b", 1_000), "per device");
300 assert!(v.last_seen_due("a", 1_000 + LAST_SEEN_EVERY_MS));
301 }
302
303 #[test]
304 fn the_port_is_reserved_once_and_skips_a_taken_one() {
305 let taken = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
306 let taken_port = taken.local_addr().unwrap().port();
307 let free = std::net::TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap().port();
308 let v = service(4, 16);
309 assert_eq!(v.reserve_port([taken_port, free]), Some(free));
310 assert!(
311 std::net::TcpListener::bind(("127.0.0.1", free)).is_err(),
312 "the reservation holds the port against other instances"
313 );
314 assert_eq!(v.reserve_port([taken_port]), Some(free), "kept, not chosen again");
315 }
316
317 #[test]
318 fn the_port_is_advertised_only_while_bound() {
319 let v = service(4, 16);
320 assert_eq!(v.advertised_port(), None);
321 v.set_bound(vec![IpAddr::from([198, 51, 100, 7])], 29702);
322 assert_eq!(v.advertised_port(), Some(29702));
323 assert_eq!(v.pair_host(), Some(IpAddr::from([198, 51, 100, 7])));
324 v.set_bound(Vec::new(), 29702);
325 assert_eq!(v.advertised_port(), None);
326 assert_eq!(v.pair_host(), None);
327 }
328
329 #[test]
330 fn without_a_data_directory_there_is_no_certificate() {
331 assert!(service(4, 16).tls().is_err());
332 }
333}