agentmux_srv/backend/lan_discovery/
lan_instances.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! A host-global record of this machine's LAN-enabled instances, so the one
5//! instance that holds the UDP probe port can list the others as `siblings`
6//! in its reply (`docs/specs/SPEC_LAN_FLEET_FEED_2026_10_03.md` §3, which
7//! links the mobile spec that owns the wire contract).
8//!
9//! One `<channel>.json` per instance under
10//! `registry::resolve_shared_lan_instances_dir()`, written only while LAN
11//! discovery runs. It carries the instance's `lan_key`, which that instance
12//! already broadcasts to the LAN over mDNS, and never the full `auth_key`.
13//! A channel with LAN off writes nothing, so it is never listed.
14
15use super::*;
16use std::path::{Path, PathBuf};
17
18/// How often a running instance rewrites its record.
19pub(super) const REWRITE_INTERVAL: std::time::Duration = std::time::Duration::from_secs(20);
20/// A record not rewritten for this long (three missed rewrites) is not listed.
21pub(super) const STALE_AFTER_MS: u64 = 60_000;
22/// Siblings listed in one reply, before the byte budget below applies.
23pub(super) const MAX_SIBLINGS: usize = 16;
24/// Serialized size a probe reply with siblings is kept under: a single
25/// unfragmented datagram on an Ethernet/Wi-Fi MTU of 1500, with headroom.
26pub(super) const MAX_REPLY_BYTES: usize = 1400;
27/// Directory entries looked at per read. Every file here is written by a
28/// local process, but a full disk of junk must not stall a probe reply.
29const MAX_FILES_SCANNED: usize = 64;
30/// A real record is about 200 bytes.
31const MAX_FILE_BYTES: u64 = 4096;
32const MAX_CHANNEL_LEN: usize = 64;
33const MAX_KEY_LEN: usize = 256;
34const MAX_VERSION_LEN: usize = 64;
35
36/// One instance's record, as stored on disk.
37#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
38pub(super) struct InstanceRecord {
39    pub channel: String,
40    pub port: u16,
41    pub lan_key: String,
42    pub version: String,
43    pub pid: u32,
44    pub updated_at_ms: u64,
45}
46
47/// The file a channel's record lives in: a readable stem, reduced to
48/// `[a-z0-9_-]` (plus other Unicode letters and digits) and capped so it can
49/// never name a path outside the directory, then 8 hex chars of a SHA-256 of
50/// the exact channel name. The slug alone is lossy (`Foo` and `foo`, or
51/// `foo bar` and `foo_bar`, share one), and two LAN-enabled channels sharing
52/// a file would overwrite each other's record (#4297). The real
53/// name stays in the file.
54pub(super) fn record_file_name(channel: &str) -> String {
55    use sha2::{Digest, Sha256};
56    let stem: String = agentmux_common::slug::file_stem(channel)
57        .chars()
58        .take(MAX_CHANNEL_LEN)
59        .collect();
60    let digest = Sha256::digest(channel.as_bytes());
61    let tag: String = digest[..4].iter().map(|b| format!("{b:02x}")).collect();
62    let stem = if stem.is_empty() { "_".to_string() } else { stem };
63    format!("{stem}-{tag}.json")
64}
65
66/// Write `record` atomically (temp file, then rename) with mode 0600 on Unix
67/// at open time, the same care `reactive::registry::write_entry_file` takes
68/// for the shared registry, which also holds keys. On Windows the file
69/// inherits the user-only ACL of the user's own profile directory.
70pub(super) fn write_record(dir: &Path, record: &InstanceRecord) -> std::io::Result<()> {
71    use std::io::Write;
72    std::fs::create_dir_all(dir)?;
73    let path = dir.join(record_file_name(&record.channel));
74    let tmp_path = path.with_extension("json.tmp");
75    let json = serde_json::to_vec(record).map_err(std::io::Error::other)?;
76    let mut opts = std::fs::OpenOptions::new();
77    opts.write(true).create(true).truncate(true);
78    #[cfg(unix)]
79    {
80        use std::os::unix::fs::OpenOptionsExt;
81        opts.mode(0o600);
82    }
83    let mut f = opts.open(&tmp_path)?;
84    if let Err(e) = f.write_all(&json) {
85        drop(f);
86        let _ = std::fs::remove_file(&tmp_path);
87        return Err(e);
88    }
89    drop(f);
90    std::fs::rename(&tmp_path, &path)
91}
92
93/// Remove `channel`'s record if it is this process's (`pid`). Another process
94/// on the same channel name owns its own record; this one leaves it alone.
95pub(super) fn remove_own_record(dir: &Path, channel: &str, pid: u32) {
96    let path = dir.join(record_file_name(channel));
97    if read_record(&path).is_some_and(|r| r.pid == pid) {
98        let _ = std::fs::remove_file(&path);
99    }
100}
101
102/// Read one record, refusing an oversized file or out-of-bounds fields.
103fn read_record(path: &Path) -> Option<InstanceRecord> {
104    if std::fs::metadata(path).ok()?.len() > MAX_FILE_BYTES {
105        return None;
106    }
107    let record: InstanceRecord = serde_json::from_slice(&std::fs::read(path).ok()?).ok()?;
108    let in_bounds = !record.channel.is_empty()
109        && record.channel.len() <= MAX_CHANNEL_LEN
110        && !record.lan_key.is_empty()
111        && record.lan_key.len() <= MAX_KEY_LEN
112        && record.version.len() <= MAX_VERSION_LEN
113        && record.port != 0;
114    in_bounds.then_some(record)
115}
116
117/// The other LAN-enabled instances on this host: records in `dir` that are
118/// fresh (rewritten within [`STALE_AFTER_MS`]), whose process is alive, and
119/// that are not this instance (neither `own_channel` nor `own_pid`). Sorted by
120/// channel, one per channel, at most [`MAX_SIBLINGS`].
121pub(super) fn read_siblings(
122    dir: &Path,
123    own_channel: &str,
124    own_pid: u32,
125    now_ms: u64,
126    pid_alive: impl Fn(u32) -> bool,
127) -> Vec<InstanceRecord> {
128    let Ok(entries) = std::fs::read_dir(dir) else {
129        return Vec::new();
130    };
131    let mut found: Vec<InstanceRecord> = entries
132        .flatten()
133        .map(|e| e.path())
134        .filter(|p| p.extension().and_then(|x| x.to_str()) == Some("json"))
135        .take(MAX_FILES_SCANNED)
136        .filter_map(|p| read_record(&p))
137        // A timestamp far in the future is as untrustworthy as an old one.
138        .filter(|r| now_ms.abs_diff(r.updated_at_ms) <= STALE_AFTER_MS)
139        .filter(|r| r.channel != own_channel && r.pid != own_pid)
140        .collect();
141    // Freshest first within a channel, so the dedup keeps the live writer if
142    // two files claim the same channel.
143    found.sort_by(|a, b| {
144        a.channel
145            .to_lowercase()
146            .cmp(&b.channel.to_lowercase())
147            .then_with(|| a.channel.cmp(&b.channel))
148            .then_with(|| b.updated_at_ms.cmp(&a.updated_at_ms))
149    });
150    found.dedup_by(|later, earlier| later.channel == earlier.channel);
151    // Last, because it is the one check that asks the OS.
152    found.retain(|r| pid_alive(r.pid));
153    found.truncate(MAX_SIBLINGS);
154    found
155}
156
157/// A sibling as it appears in the UDP reply. `auth_key` is that sibling's
158/// `lan_key`, under the field name every client already reads.
159pub(super) fn sibling_json(record: &InstanceRecord) -> serde_json::Value {
160    json!({
161        "channel": record.channel,
162        "port": record.port,
163        "auth_key": record.lan_key,
164        "version": record.version,
165    })
166}
167
168/// Set `response["siblings"]` to as many of `siblings` as keep the serialized
169/// reply within `budget` bytes. Always present, empty when there are none, so
170/// a client can tell "no siblings" from "a build that does not send them".
171pub(super) fn attach_siblings(
172    response: &mut serde_json::Value,
173    siblings: &[serde_json::Value],
174    budget: usize,
175) {
176    response["siblings"] = json!([]);
177    let mut size = serde_json::to_vec(response)
178        .map(|v| v.len())
179        .unwrap_or(usize::MAX);
180    let mut kept = Vec::new();
181    for sibling in siblings {
182        let extra = serde_json::to_vec(sibling)
183            .map(|v| v.len())
184            .unwrap_or(usize::MAX);
185        let separator = usize::from(!kept.is_empty());
186        if size.saturating_add(extra).saturating_add(separator) > budget {
187            break;
188        }
189        size += extra + separator;
190        kept.push(sibling.clone());
191    }
192    response["siblings"] = serde_json::Value::Array(kept);
193}
194
195/// Where this instance's record goes: `None` when the shared root can't be
196/// resolved, and always `None` under `cargo test`, whose `LanDiscovery::start`
197/// calls must not overwrite a real instance's record on the developer's
198/// machine. Tests that exercise the files pass a temp directory.
199pub(super) fn instances_dir() -> Option<PathBuf> {
200    if cfg!(test) {
201        return None;
202    }
203    crate::registry::resolve_shared_lan_instances_dir()
204}
205
206impl LanDiscovery {
207    fn own_record(&self) -> InstanceRecord {
208        InstanceRecord {
209            channel: crate::backend::reactive::registry::local_channel_id(),
210            port: self.port,
211            lan_key: self.auth_key.clone(),
212            version: self.version.clone(),
213            pid: std::process::id(),
214            updated_at_ms: agentmux_common::time::now_ms_u64(),
215        }
216    }
217
218    /// Write (or refresh) this instance's record, unless it has been withdrawn.
219    pub(super) fn publish_instance_record(&self) {
220        let live = self.instance_record_live.lock();
221        let (true, Some(dir)) = (*live, self.instances_dir.as_ref()) else {
222            return;
223        };
224        if let Err(e) = write_record(dir, &self.own_record()) {
225            tracing::debug!(error = %e, "LAN instance record not written");
226        }
227    }
228
229    /// Remove this instance's record and stop it being written again. The
230    /// flag is cleared under the same lock `publish_instance_record` holds, so
231    /// a rewrite racing a stop cannot bring the file back. Idempotent.
232    pub(super) fn withdraw_instance_record(&self) {
233        let mut live = self.instance_record_live.lock();
234        if !std::mem::replace(&mut *live, false) {
235            return;
236        }
237        if let Some(dir) = self.instances_dir.as_ref() {
238            let channel = crate::backend::reactive::registry::local_channel_id();
239            remove_own_record(dir, &channel, std::process::id());
240        }
241    }
242
243    /// Rewrites the record every [`REWRITE_INTERVAL`] until cancelled by
244    /// `shutdown()`. Holds its own `Arc`, like the other loops `start()` spawns.
245    pub(super) async fn instance_record_loop(self: Arc<Self>, mut cancel: oneshot::Receiver<()>) {
246        loop {
247            tokio::select! {
248                _ = tokio::time::sleep(REWRITE_INTERVAL) => self.publish_instance_record(),
249                _ = &mut cancel => return,
250            }
251        }
252    }
253
254    /// The `siblings` of a probe reply, as JSON, read fresh on each probe.
255    pub(super) fn sibling_entries(&self) -> Vec<serde_json::Value> {
256        let Some(dir) = self.instances_dir.as_ref() else {
257            return Vec::new();
258        };
259        read_siblings(
260            dir,
261            &crate::backend::reactive::registry::local_channel_id(),
262            std::process::id(),
263            agentmux_common::time::now_ms_u64(),
264            crate::backend::reactive::registry::pid_alive,
265        )
266        .iter()
267        .map(sibling_json)
268        .collect()
269    }
270}
271
272#[cfg(test)]
273mod tests {
274    use super::*;
275
276    const NOW: u64 = 1_000_000_000;
277
278    fn record(channel: &str, pid: u32, updated_at_ms: u64) -> InstanceRecord {
279        InstanceRecord {
280            channel: channel.to_string(),
281            port: 29700,
282            lan_key: format!("lan-key-{channel}"),
283            version: "0.59.7".to_string(),
284            pid,
285            updated_at_ms,
286        }
287    }
288
289    fn siblings(dir: &Path, alive: impl Fn(u32) -> bool) -> Vec<String> {
290        read_siblings(dir, "self-channel", 1, NOW, alive)
291            .into_iter()
292            .map(|r| r.channel)
293            .collect()
294    }
295
296    #[test]
297    fn a_written_record_reads_back() {
298        let dir = tempfile::tempdir().unwrap();
299        let r = record("stable", 42, NOW);
300        write_record(dir.path(), &r).unwrap();
301        assert_eq!(
302            read_siblings(dir.path(), "self-channel", 1, NOW, |_| true),
303            vec![r]
304        );
305        // No temp file left behind.
306        let names: Vec<_> = std::fs::read_dir(dir.path())
307            .unwrap()
308            .flatten()
309            .map(|e| e.file_name())
310            .collect();
311        assert_eq!(names, vec![std::ffi::OsString::from(record_file_name("stable"))]);
312    }
313
314    #[cfg(unix)]
315    #[test]
316    fn a_record_is_readable_by_its_owner_only() {
317        use std::os::unix::fs::PermissionsExt;
318        let dir = tempfile::tempdir().unwrap();
319        write_record(dir.path(), &record("stable", 42, NOW)).unwrap();
320        let mode = std::fs::metadata(dir.path().join(record_file_name("stable")))
321            .unwrap()
322            .permissions()
323            .mode();
324        assert_eq!(mode & 0o777, 0o600);
325    }
326
327    /// `<stem>-<8 hex>.json`: the readable stem, then the hash tag.
328    fn stem_of(name: &str) -> &str {
329        let base = name.strip_suffix(".json").expect("ends in .json");
330        let (stem, tag) = base.rsplit_once('-').expect("has a hash tag");
331        assert_eq!(tag.len(), 8);
332        assert!(tag.chars().all(|c| c.is_ascii_hexdigit()));
333        stem
334    }
335
336    #[test]
337    fn a_channel_cannot_name_a_path_outside_the_directory() {
338        assert_eq!(stem_of(&record_file_name("../../etc/passwd")), "______etc_passwd");
339        assert_eq!(stem_of(&record_file_name("C:\\Windows\\x")), "c__windows_x");
340        assert_eq!(stem_of(&record_file_name("local-main-b28b7a")), "local-main-b28b7a");
341        assert_eq!(stem_of(&record_file_name("")), "_");
342        let long = "x".repeat(500);
343        assert_eq!(stem_of(&record_file_name(&long)).len(), MAX_CHANNEL_LEN);
344        for name in ["../../etc/passwd", "C:\\Windows\\x", "a/b", ""] {
345            let f = record_file_name(name);
346            assert!(!f.contains('/') && !f.contains('\\') && !f.contains(".."), "{f}");
347        }
348
349        let dir = tempfile::tempdir().unwrap();
350        let nested = dir.path().join("instances");
351        write_record(&nested, &record("../escape", 42, NOW)).unwrap();
352        assert!(nested.join(record_file_name("../escape")).exists());
353        let outside: Vec<_> = std::fs::read_dir(dir.path())
354            .unwrap()
355            .flatten()
356            .map(|e| e.file_name())
357            .collect();
358        assert_eq!(outside, vec![std::ffi::OsString::from("instances")]);
359    }
360
361    #[test]
362    fn channels_with_the_same_slug_get_their_own_files() {
363        // The slug alone is lossy (#4297).
364        for (a, b) in [("Foo", "foo"), ("foo bar", "foo_bar"), ("dev/x", "dev_x")] {
365            assert_eq!(stem_of(&record_file_name(a)), stem_of(&record_file_name(b)));
366            assert_ne!(record_file_name(a), record_file_name(b), "{a} vs {b}");
367        }
368        assert_eq!(record_file_name("stable"), record_file_name("stable"), "stable per name");
369
370        let dir = tempfile::tempdir().unwrap();
371        write_record(dir.path(), &record("foo bar", 50, NOW)).unwrap();
372        write_record(dir.path(), &record("foo_bar", 51, NOW)).unwrap();
373        assert_eq!(siblings(dir.path(), |_| true), vec!["foo bar", "foo_bar"]);
374    }
375
376    #[test]
377    fn stale_dead_and_own_records_are_not_siblings() {
378        let dir = tempfile::tempdir().unwrap();
379        write_record(dir.path(), &record("fresh", 10, NOW - 5_000)).unwrap();
380        write_record(dir.path(), &record("stale", 11, NOW - STALE_AFTER_MS - 1)).unwrap();
381        write_record(dir.path(), &record("future", 12, NOW + STALE_AFTER_MS + 1)).unwrap();
382        write_record(dir.path(), &record("dead", 13, NOW)).unwrap();
383        write_record(dir.path(), &record("self-channel", 14, NOW)).unwrap();
384        write_record(dir.path(), &record("same-pid", 1, NOW)).unwrap();
385        assert_eq!(siblings(dir.path(), |pid| pid != 13), vec!["fresh"]);
386    }
387
388    #[test]
389    fn siblings_are_sorted_bounded_and_one_per_channel() {
390        let dir = tempfile::tempdir().unwrap();
391        for i in 0..(MAX_SIBLINGS + 4) {
392            write_record(
393                dir.path(),
394                &record(&format!("ch-{i:02}"), 100 + i as u32, NOW),
395            )
396            .unwrap();
397        }
398        let got = siblings(dir.path(), |_| true);
399        assert_eq!(got.len(), MAX_SIBLINGS);
400        let mut sorted = got.clone();
401        sorted.sort();
402        assert_eq!(got, sorted);
403
404        // Two files claiming one channel: the fresher is kept.
405        let dir = tempfile::tempdir().unwrap();
406        write_record(dir.path(), &record("dup", 20, NOW - 10_000)).unwrap();
407        let mut other = record("dup", 21, NOW);
408        other.port = 29704;
409        std::fs::write(
410            dir.path().join("other.json"),
411            serde_json::to_vec(&other).unwrap(),
412        )
413        .unwrap();
414        let got = read_siblings(dir.path(), "self-channel", 1, NOW, |_| true);
415        assert_eq!(got, vec![other]);
416    }
417
418    #[test]
419    fn malformed_oversized_and_out_of_bounds_files_are_skipped() {
420        let dir = tempfile::tempdir().unwrap();
421        std::fs::write(dir.path().join("junk.json"), b"not json").unwrap();
422        std::fs::write(
423            dir.path().join("big.json"),
424            vec![b' '; MAX_FILE_BYTES as usize + 1],
425        )
426        .unwrap();
427        std::fs::write(dir.path().join("ignored.txt"), b"{}").unwrap();
428        let mut no_key = record("nokey", 30, NOW);
429        no_key.lan_key.clear();
430        let mut zero_port = record("zeroport", 31, NOW);
431        zero_port.port = 0;
432        let long_channel = record(&"c".repeat(MAX_CHANNEL_LEN + 1), 32, NOW);
433        for (name, r) in [
434            ("nokey", no_key),
435            ("zeroport", zero_port),
436            ("long", long_channel),
437        ] {
438            std::fs::write(
439                dir.path().join(format!("{name}.json")),
440                serde_json::to_vec(&r).unwrap(),
441            )
442            .unwrap();
443        }
444        write_record(dir.path(), &record("good", 33, NOW)).unwrap();
445        assert_eq!(siblings(dir.path(), |_| true), vec!["good"]);
446    }
447
448    #[test]
449    fn only_the_owning_process_removes_a_record() {
450        let dir = tempfile::tempdir().unwrap();
451        write_record(dir.path(), &record("stable", 42, NOW)).unwrap();
452        remove_own_record(dir.path(), "stable", 43);
453        assert!(
454            dir.path().join(record_file_name("stable")).exists(),
455            "another pid's record stays"
456        );
457        remove_own_record(dir.path(), "stable", 42);
458        assert!(!dir.path().join(record_file_name("stable")).exists());
459    }
460
461    #[test]
462    fn a_sibling_carries_its_lan_key_and_nothing_else() {
463        let dir = tempfile::tempdir().unwrap();
464        // A record file with an extra field a forger (or a future build) added.
465        let mut raw = serde_json::to_value(record("stable", 42, NOW)).unwrap();
466        raw["auth_key"] = json!("FULL-AUTH-KEY-MUST-NOT-LEAK");
467        std::fs::write(
468            dir.path().join("stable.json"),
469            serde_json::to_vec(&raw).unwrap(),
470        )
471        .unwrap();
472        let got = read_siblings(dir.path(), "self-channel", 1, NOW, |_| true);
473        let entry = sibling_json(&got[0]);
474        assert_eq!(
475            entry,
476            json!({"channel": "stable", "port": 29700, "auth_key": "lan-key-stable", "version": "0.59.7"})
477        );
478        assert!(!entry.to_string().contains("FULL-AUTH-KEY"));
479    }
480
481    #[test]
482    fn siblings_stop_at_the_byte_budget() {
483        let mut response = probe_response_json("v0.59.7", "narko", "0.59.7", 29704, "k");
484        let many: Vec<_> = (0..MAX_SIBLINGS)
485            .map(|i| sibling_json(&record(&format!("local-main-channel-{i:02}"), 1, NOW)))
486            .collect();
487        attach_siblings(&mut response, &many, MAX_REPLY_BYTES);
488        let size = serde_json::to_vec(&response).unwrap().len();
489        let kept = response["siblings"].as_array().unwrap().len();
490        assert!(size <= MAX_REPLY_BYTES, "reply is {size} bytes");
491        assert!(kept > 0 && kept < MAX_SIBLINGS, "kept {kept}");
492
493        let mut response = probe_response_json("v0.59.7", "narko", "0.59.7", 29704, "k");
494        attach_siblings(&mut response, &[], MAX_REPLY_BYTES);
495        assert_eq!(response["siblings"], json!([]));
496    }
497
498    // -- On a `LanDiscovery` (idle mDNS daemon, none of `start()`'s tasks) --
499
500    fn discovery(dir: Option<PathBuf>) -> Arc<LanDiscovery> {
501        Arc::new(LanDiscovery {
502            daemon: mdns_sd::ServiceDaemon::new()
503                .expect("daemon construction (no register/browse)"),
504            instances: Arc::new(RwLock::new(HashMap::new())),
505            instance_id: "v0.59.7".to_string(),
506            event_bus: Arc::new(crate::backend::eventbus::EventBus::new()),
507            service_fullname: String::new(),
508            auth_key: "own-lan-key".to_string(),
509            hostname: "narko".to_string(),
510            version: "0.59.7".to_string(),
511            port: 29704,
512            udp_cancel: Mutex::new(None),
513            udp_peer_cancel: Mutex::new(None),
514            agent_names_cancel: Mutex::new(None),
515            instance_record_cancel: Mutex::new(None),
516            instances_dir: dir,
517            instance_record_live: Mutex::new(true),
518            announced_v4: Arc::new(Mutex::new(BTreeSet::new())),
519            monitored: true,
520            started_at: std::time::Instant::now(),
521        })
522    }
523
524    /// A pid that is alive for the whole test and is not ours: our parent
525    /// (the test harness or cargo).
526    fn another_live_pid() -> u32 {
527        let me = sysinfo::Pid::from_u32(std::process::id());
528        let mut sys = sysinfo::System::new();
529        sys.refresh_processes(sysinfo::ProcessesToUpdate::Some(&[me]), true);
530        let parent = sys
531            .process(me)
532            .and_then(|p| p.parent())
533            .expect("test process has a parent");
534        assert!(crate::backend::reactive::registry::pid_alive(
535            parent.as_u32()
536        ));
537        parent.as_u32()
538    }
539
540    /// The channel is read from the process environment, which other tests
541    /// change under this lock.
542    fn env_lock() -> std::sync::MutexGuard<'static, ()> {
543        crate::test_support::ISOLATED_AUTH_ENV_LOCK
544            .lock()
545            .unwrap_or_else(|e| e.into_inner())
546    }
547
548    #[test]
549    fn an_instance_publishes_and_withdraws_its_record() {
550        let _env = env_lock();
551        let dir = tempfile::tempdir().unwrap();
552        let d = discovery(Some(dir.path().to_path_buf()));
553        let path = dir.path().join(record_file_name(
554            &crate::backend::reactive::registry::local_channel_id(),
555        ));
556
557        d.publish_instance_record();
558        let r = read_record(&path).expect("written on publish");
559        assert_eq!(
560            (r.port, r.lan_key.as_str(), r.pid),
561            (29704, "own-lan-key", std::process::id())
562        );
563
564        d.withdraw_instance_record();
565        assert!(!path.exists(), "removed on withdraw");
566        d.publish_instance_record();
567        assert!(
568            !path.exists(),
569            "a rewrite after withdrawal does not bring it back"
570        );
571    }
572
573    #[test]
574    fn shutdown_withdraws_the_record() {
575        let _env = env_lock();
576        let dir = tempfile::tempdir().unwrap();
577        let d = discovery(Some(dir.path().to_path_buf()));
578        d.publish_instance_record();
579        assert_eq!(std::fs::read_dir(dir.path()).unwrap().count(), 1);
580        d.shutdown();
581        assert_eq!(std::fs::read_dir(dir.path()).unwrap().count(), 0);
582    }
583
584    #[test]
585    fn the_probe_reply_lists_live_siblings_but_the_desktop_reply_does_not() {
586        let _env = env_lock();
587        let dir = tempfile::tempdir().unwrap();
588        let now = agentmux_common::time::now_ms_u64();
589        write_record(
590            dir.path(),
591            &record("sibling-channel", another_live_pid(), now),
592        )
593        .unwrap();
594        write_record(dir.path(), &record("crashed-channel", u32::MAX, now)).unwrap();
595        let d = discovery(Some(dir.path().to_path_buf()));
596        d.publish_instance_record(); // this instance's own record is never listed
597
598        let reply = d.build_probe_response();
599        assert_eq!(
600            reply["siblings"],
601            json!([{
602                "channel": "sibling-channel",
603                "port": 29700,
604                "auth_key": "lan-key-sibling-channel",
605                "version": "0.59.7",
606            }])
607        );
608        assert_eq!(reply["auth_key"], "own-lan-key");
609        assert_eq!(reply["type"], UDP_RESPONSE_TYPE);
610        assert!(d.build_identity_response().get("siblings").is_none());
611    }
612
613    #[test]
614    fn the_reply_carries_install_id_and_channel_count_and_stays_one_datagram() {
615        let _env = env_lock();
616        let dir = tempfile::tempdir().unwrap();
617        let now = agentmux_common::time::now_ms_u64();
618        let pid = another_live_pid();
619        // As many siblings as a reply lists, each with the longest channel
620        // name a record may carry, so the byte budget is what limits them.
621        for i in 0..MAX_SIBLINGS {
622            let channel = format!("{i:02}{}", "c".repeat(MAX_CHANNEL_LEN - 2));
623            let mut r = record(&channel, pid, now);
624            r.lan_key = "k".repeat(64);
625            write_record(dir.path(), &r).unwrap();
626        }
627        let d = discovery(Some(dir.path().to_path_buf()));
628        let install_id = "mw3am46w5weex4a4fqrc3avnua";
629
630        let reply = d.probe_response_with(Some(install_id), Some(99), Some(29799));
631        let size = serde_json::to_vec(&reply).unwrap().len();
632        assert!(size <= MAX_REPLY_BYTES, "reply is {size} bytes");
633        assert_eq!(reply["install_id"], install_id);
634        assert_eq!(reply["channels_running"], 99);
635        assert_eq!(reply["viewer_port"], 29799);
636        let kept = reply["siblings"].as_array().unwrap().len();
637        assert!(kept > 0 && kept < MAX_SIBLINGS, "siblings are trimmed first: kept {kept}");
638
639        let identity = d.identity_response_with(Some(install_id), Some(3), Some(29702));
640        assert_eq!(identity["install_id"], install_id);
641        assert_eq!(identity["channels_running"], 3);
642        assert_eq!(identity["viewer_port"], 29702);
643
644        // Unknown: the fields are left out, not sent empty.
645        let bare = d.identity_response_with(None, None, None);
646        assert!(bare.get("install_id").is_none());
647        assert!(bare.get("channels_running").is_none());
648        assert!(bare.get("viewer_port").is_none());
649    }
650
651    fn probe_bytes() -> Vec<u8> {
652        serde_json::to_vec(&json!({"type": UDP_PROBE_TYPE, "v": UDP_PROTOCOL_VERSION})).unwrap()
653    }
654
655    #[tokio::test]
656    async fn the_responder_takes_over_the_port_once_its_holder_lets_go() {
657        let holder = std::net::UdpSocket::bind("0.0.0.0:0").unwrap();
658        let port = holder.local_addr().unwrap().port();
659        let d = discovery(None);
660        let (cancel_tx, cancel_rx) = oneshot::channel();
661        let task = tokio::spawn(d.clone().udp_responder_loop_on(
662            port,
663            std::time::Duration::from_millis(50),
664            cancel_rx,
665        ));
666        tokio::time::sleep(std::time::Duration::from_millis(150)).await;
667        assert!(!task.is_finished(), "a failed bind keeps retrying");
668        drop(holder);
669
670        let prober = UdpSocket::bind("127.0.0.1:0").await.unwrap();
671        let mut buf = [0u8; 2048];
672        let reply = tokio::time::timeout(std::time::Duration::from_secs(10), async {
673            loop {
674                let _ = prober.send_to(&probe_bytes(), ("127.0.0.1", port)).await;
675                let wait = std::time::Duration::from_millis(200);
676                if let Ok(Ok((len, _))) =
677                    tokio::time::timeout(wait, prober.recv_from(&mut buf)).await
678                {
679                    break serde_json::from_slice::<serde_json::Value>(&buf[..len]).unwrap();
680                }
681            }
682        })
683        .await
684        .expect("the responder answers once it holds the port");
685        assert_eq!(reply["type"], UDP_RESPONSE_TYPE);
686        assert_eq!(reply["siblings"], json!([]));
687
688        cancel_tx.send(()).unwrap();
689        tokio::time::timeout(std::time::Duration::from_secs(5), task)
690            .await
691            .unwrap()
692            .unwrap();
693    }
694
695    #[tokio::test]
696    async fn the_responder_stops_while_still_waiting_for_the_port() {
697        let holder = std::net::UdpSocket::bind("0.0.0.0:0").unwrap();
698        let port = holder.local_addr().unwrap().port();
699        let (cancel_tx, cancel_rx) = oneshot::channel();
700        let task = tokio::spawn(discovery(None).udp_responder_loop_on(
701            port,
702            std::time::Duration::from_secs(3600),
703            cancel_rx,
704        ));
705        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
706        cancel_tx.send(()).unwrap();
707        tokio::time::timeout(std::time::Duration::from_secs(5), task)
708            .await
709            .expect("cancel ends the retry wait")
710            .unwrap();
711        drop(holder);
712    }
713}
agentmux_srv/backend/lan_discovery/lan_instances.rs

agentmux_srv/backend/lan_discovery/
lan_instances.rs
