ghostlink_core/
mdns.rs

1//! mDNS-based peer discovery — an alternative path to the UDP broadcast
2//! fallback in `discovery.rs` for networks where broadcast traffic is
3//! filtered but multicast is not (some managed VLANs, cloud VPCs). Runs
4//! alongside UDP discovery rather than replacing it: callers typically merge
5//! results from both.
6//!
7//! A Ghostlink node advertises itself as an instance of `_ghostlink._tcp.local.`
8//! with its `NodeResources` encoded as TXT record key/value pairs. Peers
9//! browse for that service type and decode the TXT records back into
10//! `NodeResources`.
11
12use std::net::{IpAddr, SocketAddr};
13use std::sync::OnceLock;
14use std::time::{Duration, Instant};
15
16use mdns_sd::{ResolvedService, ServiceDaemon, ServiceEvent, ServiceInfo};
17
18use crate::protocol::NodeResources;
19
20/// mDNS service type Ghostlink nodes advertise themselves under.
21pub const MDNS_SERVICE_TYPE: &str = "_ghostlink._tcp.local.";
22
23static ADVERTISER: OnceLock<Option<ServiceDaemon>> = OnceLock::new();
24
25/// Registers this node under [`MDNS_SERVICE_TYPE`] so LAN peers can find it
26/// via mDNS browsing, idempotently (subsequent calls are no-ops). Best-effort:
27/// a failure (e.g. multicast unavailable in a sandboxed/headless environment)
28/// is logged and otherwise ignored — mirrors how UDP discovery failures never
29/// block server startup elsewhere in this crate.
30pub fn ensure_advertised(node: &NodeResources, port: u16) {
31    ADVERTISER.get_or_init(|| match advertise(node, port) {
32        Ok(daemon) => Some(daemon),
33        Err(err) => {
34            tracing::warn!("mdns: failed to advertise local node, continuing without it: {err}");
35            None
36        }
37    });
38}
39
40fn advertise(node: &NodeResources, port: u16) -> Result<ServiceDaemon, String> {
41    let daemon = ServiceDaemon::new().map_err(|e| format!("failed to start mDNS daemon: {e}"))?;
42
43    let properties = node_to_txt_properties(node);
44    let borrowed_properties: Vec<(&str, &str)> = properties
45        .iter()
46        .map(|(k, v)| (k.as_str(), v.as_str()))
47        .collect();
48
49    let instance_name = sanitize_instance_name(&node.id);
50    let host_name = format!("{instance_name}.local.");
51
52    // Empty ip + enable_addr_auto() lets the daemon discover and advertise
53    // this host's real interface addresses instead of us guessing one.
54    let service = ServiceInfo::new(
55        MDNS_SERVICE_TYPE,
56        &instance_name,
57        &host_name,
58        "",
59        port,
60        borrowed_properties.as_slice(),
61    )
62    .map_err(|e| format!("failed to build mDNS service info: {e}"))?
63    .enable_addr_auto();
64
65    daemon
66        .register(service)
67        .map_err(|e| format!("failed to register mDNS service: {e}"))?;
68
69    Ok(daemon)
70}
71
72/// Browses for other Ghostlink nodes over mDNS for up to `timeout`, returning
73/// each resolved peer's advertised resources and the address to reach it on.
74/// Blocks the calling thread for up to `timeout` — callers on an async
75/// runtime should run this via `spawn_blocking`, same as
76/// `discovery::broadcast_and_collect`.
77pub fn discover(timeout: Duration) -> Result<Vec<(NodeResources, SocketAddr)>, String> {
78    let daemon = ServiceDaemon::new().map_err(|e| format!("failed to start mDNS daemon: {e}"))?;
79    let receiver = daemon
80        .browse(MDNS_SERVICE_TYPE)
81        .map_err(|e| format!("failed to browse {MDNS_SERVICE_TYPE}: {e}"))?;
82
83    // Keyed by node id so a peer re-announcing during the browse window
84    // (mDNS resends resolved records periodically) doesn't produce
85    // duplicate entries.
86    let mut peers = std::collections::HashMap::new();
87    let deadline = Instant::now() + timeout;
88
89    while let Some(remaining) = deadline.checked_duration_since(Instant::now()) {
90        match receiver.recv_timeout(remaining) {
91            Ok(ServiceEvent::ServiceResolved(info)) => {
92                if let (Some(node), Some(addr)) =
93                    (node_from_service_info(&info), first_socket_addr(&info))
94                {
95                    peers.insert(node.id.clone(), (node, addr));
96                }
97            }
98            Ok(_) => {}
99            Err(_) => break, // timed out or the channel closed
100        }
101    }
102
103    let _ = daemon.stop_browse(MDNS_SERVICE_TYPE);
104    let _ = daemon.shutdown();
105
106    Ok(peers.into_values().collect())
107}
108
109fn first_socket_addr(info: &ResolvedService) -> Option<SocketAddr> {
110    info.get_addresses_v4()
111        .into_iter()
112        .next()
113        .map(|ip| SocketAddr::new(IpAddr::V4(ip), info.get_port()))
114}
115
116fn node_from_service_info(info: &ResolvedService) -> Option<NodeResources> {
117    node_from_txt_lookup(|key| info.get_property_val_str(key).map(str::to_owned))
118}
119
120fn node_to_txt_properties(node: &NodeResources) -> Vec<(String, String)> {
121    let mut props = vec![
122        ("id".to_string(), node.id.clone()),
123        ("vram_gb".to_string(), node.vram_gb.to_string()),
124        (
125            "system_memory_gb".to_string(),
126            node.system_memory_gb.to_string(),
127        ),
128        (
129            "compute_capability".to_string(),
130            node.compute_capability.clone(),
131        ),
132    ];
133    if let Some(gpu_name) = &node.gpu_name {
134        props.push(("gpu_name".to_string(), gpu_name.clone()));
135    }
136    if let Some(rpc_port) = node.rpc_port {
137        props.push(("rpc_port".to_string(), rpc_port.to_string()));
138    }
139    if let Some(rpc_build_id) = &node.rpc_build_id {
140        props.push(("rpc_build_id".to_string(), rpc_build_id.clone()));
141    }
142    props
143}
144
145fn node_from_txt_lookup(lookup: impl Fn(&str) -> Option<String>) -> Option<NodeResources> {
146    let id = lookup("id")?;
147    let vram_gb = lookup("vram_gb")?.parse().ok()?;
148    let system_memory_gb = lookup("system_memory_gb")?.parse().ok()?;
149    let compute_capability = lookup("compute_capability").unwrap_or_default();
150    let gpu_name = lookup("gpu_name");
151    let rpc_port = lookup("rpc_port").and_then(|v| v.parse().ok());
152    let rpc_build_id = lookup("rpc_build_id");
153    Some(NodeResources {
154        id,
155        vram_gb,
156        system_memory_gb,
157        compute_capability,
158        gpu_name,
159        rpc_port,
160        rpc_build_id,
161    })
162}
163
164/// mDNS instance names conventionally avoid raw dots (the DNS label
165/// separator) — collapse them so an id containing one doesn't produce a
166/// confusing on-wire label.
167fn sanitize_instance_name(id: &str) -> String {
168    id.replace('.', "-")
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174
175    fn lookup_from(props: &[(String, String)]) -> impl Fn(&str) -> Option<String> + '_ {
176        move |key: &str| props.iter().find(|(k, _)| k == key).map(|(_, v)| v.clone())
177    }
178
179    #[test]
180    fn txt_roundtrip_preserves_node_fields() {
181        let node = NodeResources::new("node-a", 24.0, 64.0, "8.9", Some("RTX 4090".to_string()));
182        let props = node_to_txt_properties(&node);
183        let decoded = node_from_txt_lookup(lookup_from(&props)).expect("decode");
184
185        assert_eq!(decoded.id, node.id);
186        assert_eq!(decoded.vram_gb, node.vram_gb);
187        assert_eq!(decoded.system_memory_gb, node.system_memory_gb);
188        assert_eq!(decoded.compute_capability, node.compute_capability);
189        assert_eq!(decoded.gpu_name, node.gpu_name);
190    }
191
192    #[test]
193    fn txt_roundtrip_without_gpu_name() {
194        let node = NodeResources::new("node-b", 8.0, 16.0, "7.5", None);
195        let props = node_to_txt_properties(&node);
196        let decoded = node_from_txt_lookup(lookup_from(&props)).expect("decode");
197
198        assert_eq!(decoded.gpu_name, None);
199    }
200
201    #[test]
202    fn txt_roundtrip_preserves_rpc_port() {
203        let node = NodeResources::new("node-c", 12.0, 32.0, "8.6", None).with_rpc_port(50052);
204        let props = node_to_txt_properties(&node);
205        let decoded = node_from_txt_lookup(lookup_from(&props)).expect("decode");
206
207        assert_eq!(decoded.rpc_port, Some(50052));
208    }
209
210    #[test]
211    fn txt_roundtrip_without_rpc_port() {
212        let node = NodeResources::new("node-d", 8.0, 16.0, "7.5", None);
213        let props = node_to_txt_properties(&node);
214        let decoded = node_from_txt_lookup(lookup_from(&props)).expect("decode");
215
216        assert_eq!(decoded.rpc_port, None);
217    }
218
219    #[test]
220    fn txt_roundtrip_preserves_rpc_build_id() {
221        let node =
222            NodeResources::new("node-e", 12.0, 32.0, "8.6", None).with_rpc_build_id("da296d6");
223        let props = node_to_txt_properties(&node);
224        let decoded = node_from_txt_lookup(lookup_from(&props)).expect("decode");
225
226        assert_eq!(decoded.rpc_build_id, Some("da296d6".to_string()));
227    }
228
229    #[test]
230    fn txt_roundtrip_without_rpc_build_id() {
231        let node = NodeResources::new("node-f", 8.0, 16.0, "7.5", None);
232        let props = node_to_txt_properties(&node);
233        let decoded = node_from_txt_lookup(lookup_from(&props)).expect("decode");
234
235        assert_eq!(decoded.rpc_build_id, None);
236    }
237
238    #[test]
239    fn decode_fails_gracefully_on_missing_required_field() {
240        // No "id" key present.
241        let props = vec![("vram_gb".to_string(), "8.0".to_string())];
242        assert!(node_from_txt_lookup(lookup_from(&props)).is_none());
243    }
244
245    #[test]
246    fn sanitize_instance_name_collapses_dots() {
247        assert_eq!(
248            sanitize_instance_name("workstation.local"),
249            "workstation-local"
250        );
251        assert_eq!(sanitize_instance_name("node-a"), "node-a");
252    }
253}
ghostlink_core/mdns.rs

ghostlink_core/
mdns.rs
