1use crate::discovery::interface::{local_interfaces, InterfaceFilter};
13use lanflare_proto::discovery::{DeviceType, MulticastMessageV2, ProtocolType};
14use rand::Rng;
15use socket2::{Domain, Protocol, Socket, Type};
16use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
17use std::sync::Arc;
18use std::time::Duration;
19use tokio::net::UdpSocket;
20use tokio::sync::{mpsc, oneshot, Mutex, RwLock};
21use tokio_util::sync::CancellationToken;
22
23pub const DEFAULT_MULTICAST_IPV4: Ipv4Addr = Ipv4Addr::new(224, 0, 0, 167);
25
26pub const DEFAULT_MULTICAST_IPV6: Ipv6Addr =
28 Ipv6Addr::new(0xff12, 0, 0, 0, 0, 0, 0xfd3a, 0xe420);
29
30pub const DEFAULT_MULTICAST_PORT: u16 = 53317;
32
33pub const ANNOUNCE_BURST_DELAYS: [Duration; 3] = [
35 Duration::from_millis(100),
36 Duration::from_millis(500),
37 Duration::from_millis(2000),
38];
39
40pub const MAX_BURST_JITTER: Duration = Duration::from_millis(50);
42
43const RECEIVE_BUFFER_SIZE: usize = 65536;
45
46const MAX_CONSECUTIVE_RECEIVE_ERRORS: u32 = 10;
48
49#[derive(Clone, Debug, PartialEq, Eq)]
51pub struct MulticastNodeInfo {
52 pub alias: String,
54
55 pub version: String,
57
58 pub device_model: Option<String>,
60
61 pub device_type: Option<DeviceType>,
63
64 pub fingerprint: String,
67
68 pub port: u16,
70
71 pub protocol: ProtocolType,
73
74 pub download: bool,
76}
77
78impl MulticastNodeInfo {
79 pub fn to_message(&self, announce: bool) -> MulticastMessageV2 {
80 MulticastMessageV2 {
81 alias: self.alias.clone(),
82 version: self.version.clone(),
83 device_model: self.device_model.clone(),
84 device_type: self.device_type,
85 fingerprint: self.fingerprint.clone(),
86 port: self.port,
87 protocol: self.protocol,
88 download: self.download,
89 announce,
90 }
91 }
92}
93
94#[derive(Clone, Debug)]
96pub struct MulticastConfig {
97 pub group_v4: Ipv4Addr,
99
100 pub group_v6: Option<Ipv6Addr>,
102
103 pub port: u16,
105
106 pub interface_filter: InterfaceFilter,
108
109 pub node_info: MulticastNodeInfo,
111
112 pub allow_loopback_fallback: bool,
114}
115
116impl Default for MulticastConfig {
117 fn default() -> Self {
118 Self {
119 group_v4: DEFAULT_MULTICAST_IPV4,
120 group_v6: Some(DEFAULT_MULTICAST_IPV6),
121 port: DEFAULT_MULTICAST_PORT,
122 interface_filter: InterfaceFilter::default(),
123 node_info: MulticastNodeInfo {
124 alias: "LanFlare Node".to_string(),
125 version: "2.0".to_string(),
126 device_model: None,
127 device_type: Some(DeviceType::Desktop),
128 fingerprint: String::new(),
129 port: DEFAULT_MULTICAST_PORT,
130 protocol: ProtocolType::Https,
131 download: false,
132 },
133 allow_loopback_fallback: false,
134 }
135 }
136}
137
138#[derive(Clone, Debug, PartialEq, Eq)]
140pub enum MulticastDiscoveryEvent {
141 PeerDiscovered {
143 remote_ip: IpAddr,
145
146 scope_id: Option<u32>,
148
149 message: MulticastMessageV2,
151 },
152
153 AllSocketsFailed,
155}
156
157pub struct BoundMulticastSocket {
159 pub description: String,
160 pub target: SocketAddr,
161 pub socket: Arc<UdpSocket>,
162}
163
164pub fn bind_multicast_socket_v4(
166 group: Ipv4Addr,
167 port: u16,
168 interface: Ipv4Addr,
169) -> std::io::Result<UdpSocket> {
170 let socket = Socket::new(Domain::IPV4, Type::DGRAM, Some(Protocol::UDP))?;
171
172 socket.set_reuse_address(true)?;
174 #[cfg(all(unix, not(any(target_os = "solaris", target_os = "illumos"))))]
175 socket.set_reuse_port(true)?;
176
177 socket.bind(&SocketAddr::from(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, port)).into())?;
179
180 socket.join_multicast_v4(&group, &interface)?;
182
183 socket.set_multicast_if_v4(&interface)?;
185
186 socket.set_multicast_loop_v4(true)?;
188
189 socket.set_multicast_ttl_v4(1)?;
191
192 socket.set_nonblocking(true)?;
193
194 UdpSocket::from_std(socket.into())
195}
196
197pub fn bind_multicast_socket_v6(
199 group: Ipv6Addr,
200 port: u16,
201 interface_index: u32,
202) -> std::io::Result<UdpSocket> {
203 let socket = Socket::new(Domain::IPV6, Type::DGRAM, Some(Protocol::UDP))?;
204
205 socket.set_only_v6(true)?;
207
208 socket.set_reuse_address(true)?;
209 #[cfg(all(unix, not(any(target_os = "solaris", target_os = "illumos"))))]
210 socket.set_reuse_port(true)?;
211
212 socket.bind(&SocketAddr::from(SocketAddrV6::new(Ipv6Addr::UNSPECIFIED, port, 0, 0)).into())?;
213
214 socket.join_multicast_v6(&group, interface_index)?;
215
216 socket.set_multicast_if_v6(interface_index)?;
217
218 socket.set_multicast_loop_v6(true)?;
219
220 socket.set_multicast_hops_v6(1)?;
222
223 socket.set_nonblocking(true)?;
224
225 UdpSocket::from_std(socket.into())
226}
227
228pub fn bind_multicast_sockets(
230 config: &MulticastConfig,
231) -> std::io::Result<Vec<BoundMulticastSocket>> {
232 let mut interfaces = local_interfaces(&config.interface_filter)?;
233
234 if interfaces.v4.is_empty() && interfaces.v6.is_empty() && config.allow_loopback_fallback {
236 interfaces.v4.push(crate::discovery::interface::LocalInterfaceV4 {
237 name: "lo".to_string(),
238 address: Ipv4Addr::new(127, 0, 0, 1),
239 });
240 }
241
242 let mut sockets = Vec::new();
243
244 for interface in interfaces.v4 {
245 let description = format!("{} ({})", interface.name, interface.address);
246 match bind_multicast_socket_v4(config.group_v4, config.port, interface.address) {
247 Ok(socket) => {
248 tracing::info!(
249 "Bound UDP multicast socket v4 (interface: {}, group: {}, port: {})",
250 description,
251 config.group_v4,
252 config.port
253 );
254 sockets.push(BoundMulticastSocket {
255 description,
256 target: SocketAddr::from(SocketAddrV4::new(config.group_v4, config.port)),
257 socket: Arc::new(socket),
258 });
259 }
260 Err(err) => {
261 tracing::warn!(
262 "Could not bind UDP multicast socket v4 ({}, group: {}, port: {}): {err:#}",
263 description,
264 config.group_v4,
265 config.port
266 );
267 }
268 }
269 }
270
271 if let Some(group_v6) = config.group_v6 {
272 for interface in interfaces.v6 {
273 let description = format!("{} (IPv6, if_index: {})", interface.name, interface.index);
274 match bind_multicast_socket_v6(group_v6, config.port, interface.index) {
275 Ok(socket) => {
276 tracing::info!(
277 "Bound UDP multicast socket v6 (interface: {}, group: {}, port: {})",
278 description,
279 group_v6,
280 config.port
281 );
282 sockets.push(BoundMulticastSocket {
283 description,
284 target: SocketAddr::from(SocketAddrV6::new(
285 group_v6,
286 config.port,
287 0,
288 interface.index,
289 )),
290 socket: Arc::new(socket),
291 });
292 }
293 Err(err) => {
294 tracing::warn!(
295 "Could not bind UDP multicast socket v6 ({}, group: {}, port: {}): {err:#}",
296 description,
297 group_v6,
298 config.port
299 );
300 }
301 }
302 }
303 }
304
305 Ok(sockets)
306}
307
308#[derive(Clone)]
310struct SendTarget {
311 target: SocketAddr,
312 socket: Arc<UdpSocket>,
313}
314
315struct MulticastEngineState {
316 node_info: MulticastNodeInfo,
317 sockets: RwLock<Vec<SendTarget>>,
318}
319
320impl MulticastEngineState {
321 async fn broadcast_announcement(&self) {
323 let msg = self.node_info.to_message(true);
324 let payload = match serde_json::to_vec(&msg) {
325 Ok(p) => p,
326 Err(e) => {
327 tracing::error!("Failed to serialize multicast announcement: {e}");
328 return;
329 }
330 };
331
332 let sockets = self.sockets.read().await.clone();
333 for SendTarget { target, socket } in sockets {
334 if let Err(err) = socket.send_to(&payload, target).await {
335 tracing::warn!("Failed sending multicast packet to {target}: {err}");
336 }
337 }
338 }
339}
340
341#[derive(Clone)]
343pub struct MulticastHandle {
344 state: Arc<MulticastEngineState>,
345 cancel: CancellationToken,
346 task: Arc<Mutex<Option<tokio::task::JoinHandle<()>>>>,
347}
348
349impl MulticastHandle {
350 pub async fn announce_burst(&self) {
355 self.announce_custom_burst(&ANNOUNCE_BURST_DELAYS, MAX_BURST_JITTER)
356 .await;
357 }
358
359 pub async fn announce_custom_burst(&self, delays: &[Duration], max_jitter: Duration) {
361 for &base_delay in delays {
362 let jitter_millis = if !max_jitter.is_zero() {
363 rand::rng().random_range(0..max_jitter.as_millis().max(1) as u64)
364 } else {
365 0
366 };
367 let delay = base_delay + Duration::from_millis(jitter_millis);
368
369 tokio::select! {
370 _ = self.cancel.cancelled() => return,
371 _ = tokio::time::sleep(delay) => {}
372 }
373
374 tracing::debug!("Sending multicast announcement burst");
375 self.state.broadcast_announcement().await;
376 }
377 }
378
379 pub async fn send_single_announcement(&self) {
381 self.state.broadcast_announcement().await;
382 }
383
384 pub async fn stop(&self) {
386 self.cancel.cancel();
387 if let Some(task) = self.task.lock().await.take() {
388 let _ = task.await;
389 }
390 }
391
392 pub async fn wait_stopped(&self) {
394 if let Some(task) = self.task.lock().await.take() {
395 let _ = task.await;
396 }
397 }
398}
399
400pub async fn start_multicast_engine(
405 config: MulticastConfig,
406 event_tx: mpsc::Sender<MulticastDiscoveryEvent>,
407 stop_rx: oneshot::Receiver<()>,
408) -> anyhow::Result<MulticastHandle> {
409 let bound_sockets = bind_multicast_sockets(&config)?;
410 if bound_sockets.is_empty() {
411 anyhow::bail!(
412 "No network interface available for multicast on port {}",
413 config.port
414 );
415 }
416
417 let state = Arc::new(MulticastEngineState {
418 node_info: config.node_info,
419 sockets: RwLock::new(
420 bound_sockets
421 .iter()
422 .map(|s| SendTarget {
423 target: s.target,
424 socket: s.socket.clone(),
425 })
426 .collect(),
427 ),
428 });
429
430 let cancel = CancellationToken::new();
431
432 let task = tokio::spawn({
433 let state = state.clone();
434 let cancel = cancel.clone();
435 async move {
436 let mut receiver_tasks = tokio::task::JoinSet::new();
437 for s in bound_sockets {
438 receiver_tasks.spawn(receive_loop(
439 s,
440 state.clone(),
441 event_tx.clone(),
442 cancel.clone(),
443 ));
444 }
445
446 tokio::select! {
447 _ = async { while receiver_tasks.join_next().await.is_some() {} } => {
448 tracing::error!("All multicast sockets failed. Stopping multicast discovery.");
449 let _ = event_tx.try_send(MulticastDiscoveryEvent::AllSocketsFailed);
450 }
451 _ = stop_rx => {}
452 _ = cancel.cancelled() => {}
453 }
454
455 cancel.cancel();
456 receiver_tasks.shutdown().await;
457 state.sockets.write().await.clear();
458 }
459 });
460
461 Ok(MulticastHandle {
462 state,
463 cancel,
464 task: Arc::new(Mutex::new(Some(task))),
465 })
466}
467
468async fn receive_loop(
470 socket: BoundMulticastSocket,
471 state: Arc<MulticastEngineState>,
472 event_tx: mpsc::Sender<MulticastDiscoveryEvent>,
473 cancel: CancellationToken,
474) {
475 let mut buffer = vec![0u8; RECEIVE_BUFFER_SIZE];
476 let mut consecutive_errors = 0;
477
478 loop {
479 let received = tokio::select! {
480 _ = cancel.cancelled() => return,
481 res = socket.socket.recv_from(&mut buffer) => res,
482 };
483
484 let (size, source) = match received {
485 Ok(res) => {
486 consecutive_errors = 0;
487 res
488 }
489 Err(err) => {
490 consecutive_errors += 1;
491 tracing::warn!(
492 "Multicast recv error on interface {}: {err:#}",
493 socket.description
494 );
495 if consecutive_errors >= MAX_CONSECUTIVE_RECEIVE_ERRORS {
496 tracing::error!(
497 "Exceeded max consecutive receive errors on {}. Terminating socket loop.",
498 socket.description
499 );
500 return;
501 }
502 continue;
503 }
504 };
505
506 let message = match serde_json::from_slice::<MulticastMessageV2>(&buffer[..size]) {
507 Ok(msg) => msg,
508 Err(err) => {
509 tracing::warn!("Failed to deserialize multicast message from {source}: {err:#}");
510 continue;
511 }
512 };
513
514 if !state.node_info.fingerprint.is_empty()
516 && message.fingerprint == state.node_info.fingerprint
517 {
518 tracing::trace!(
519 "Suppressed self-multicast announcement with fingerprint {}",
520 message.fingerprint
521 );
522 continue;
523 }
524
525 let event = MulticastDiscoveryEvent::PeerDiscovered {
526 remote_ip: source.ip(),
527 scope_id: match &source {
528 SocketAddr::V6(v6) if v6.scope_id() != 0 => Some(v6.scope_id()),
529 _ => None,
530 },
531 message,
532 };
533
534 if event_tx.send(event).await.is_err() {
535 tracing::debug!("Discovery event receiver dropped. Shutting down receive loop.");
536 return;
537 }
538 }
539}
540
541#[cfg(test)]
542mod tests {
543 use super::*;
544 use std::sync::atomic::{AtomicU16, Ordering};
545
546 static NEXT_TEST_PORT: AtomicU16 = AtomicU16::new(54100);
547
548 const TEST_GROUP_V4: Ipv4Addr = Ipv4Addr::new(224, 0, 0, 169);
549 const TEST_GROUP_V6: Ipv6Addr = Ipv6Addr::new(0xff12, 0, 0, 0, 0, 0, 0xfd3a, 0xe422);
550
551 #[test]
552 fn test_node_info_to_message() {
553 let info = MulticastNodeInfo {
554 alias: "Test Node".to_string(),
555 version: "2.0".to_string(),
556 device_model: Some("Linux".to_string()),
557 device_type: Some(DeviceType::Desktop),
558 fingerprint: "FP123456".to_string(),
559 port: 53317,
560 protocol: ProtocolType::Https,
561 download: true,
562 };
563
564 let msg = info.to_message(true);
565 assert_eq!(msg.alias, "Test Node");
566 assert_eq!(msg.fingerprint, "FP123456");
567 assert!(msg.announce);
568 assert!(msg.download);
569
570 let json = serde_json::to_string(&msg).unwrap();
571 assert!(json.contains("\"announce\":true"));
572 assert!(json.contains("\"alias\":\"Test Node\""));
573 }
574
575 #[test]
576 fn test_multicast_config_default() {
577 let config = MulticastConfig::default();
578 assert_eq!(config.group_v4, DEFAULT_MULTICAST_IPV4);
579 assert_eq!(config.group_v6, Some(DEFAULT_MULTICAST_IPV6));
580 assert_eq!(config.port, DEFAULT_MULTICAST_PORT);
581 }
582
583 #[tokio::test]
584 async fn test_multi_node_discovery_exchange() {
585 let port = NEXT_TEST_PORT.fetch_add(1, Ordering::Relaxed);
586
587 let (node1_tx, _node1_rx) = mpsc::channel(32);
588 let (_node1_stop_tx, node1_stop_rx) = oneshot::channel();
589
590 let cfg1 = MulticastConfig {
591 group_v4: TEST_GROUP_V4,
592 group_v6: Some(TEST_GROUP_V6),
593 port,
594 interface_filter: InterfaceFilter::default(),
595 node_info: MulticastNodeInfo {
596 alias: "Node One".to_string(),
597 version: "2.0".to_string(),
598 device_model: Some("Linux".to_string()),
599 device_type: Some(DeviceType::Desktop),
600 fingerprint: "FP_NODE_ONE".to_string(),
601 port,
602 protocol: ProtocolType::Https,
603 download: true,
604 },
605 allow_loopback_fallback: false,
606 };
607
608 let handle1 = match start_multicast_engine(cfg1, node1_tx, node1_stop_rx).await {
609 Ok(h) => h,
610 Err(e) => {
611 eprintln!("Skipping test due to environment without multicast: {e}");
612 return;
613 }
614 };
615
616 let (node2_tx, mut node2_rx) = mpsc::channel(32);
617 let (_node2_stop_tx, node2_stop_rx) = oneshot::channel();
618
619 let cfg2 = MulticastConfig {
620 group_v4: TEST_GROUP_V4,
621 group_v6: Some(TEST_GROUP_V6),
622 port,
623 interface_filter: InterfaceFilter::default(),
624 node_info: MulticastNodeInfo {
625 alias: "Node Two".to_string(),
626 version: "2.0".to_string(),
627 device_model: Some("Android".to_string()),
628 device_type: Some(DeviceType::Mobile),
629 fingerprint: "FP_NODE_TWO".to_string(),
630 port,
631 protocol: ProtocolType::Https,
632 download: false,
633 },
634 allow_loopback_fallback: false,
635 };
636
637 let handle2 = match start_multicast_engine(cfg2, node2_tx, node2_stop_rx).await {
638 Ok(h) => h,
639 Err(e) => {
640 handle1.stop().await;
641 eprintln!("Skipping test due to environment without multicast: {e}");
642 return;
643 }
644 };
645
646 handle1.send_single_announcement().await;
648
649 let received_by_node2 = tokio::time::timeout(Duration::from_secs(3), async {
651 while let Some(evt) = node2_rx.recv().await {
652 match evt {
653 MulticastDiscoveryEvent::PeerDiscovered { message, .. }
654 if message.fingerprint == "FP_NODE_ONE" =>
655 {
656 return Some(message);
657 }
658 _ => {}
659 }
660 }
661 None
662 })
663 .await;
664
665
666
667 if let Ok(Some(msg)) = received_by_node2 {
668 assert_eq!(msg.alias, "Node One");
669 assert_eq!(msg.fingerprint, "FP_NODE_ONE");
670 assert_eq!(msg.device_type, Some(DeviceType::Desktop));
671 } else {
672 eprintln!("Warning: Multicast packet not received (likely filtered by network interface)");
673 }
674
675 handle1.stop().await;
677 handle2.stop().await;
678 }
679
680 #[tokio::test]
681 async fn test_loopback_suppression() {
682 let port = NEXT_TEST_PORT.fetch_add(1, Ordering::Relaxed);
683
684 let (node_tx, mut node_rx) = mpsc::channel(32);
685 let (_stop_tx, stop_rx) = oneshot::channel();
686
687 let cfg = MulticastConfig {
688 group_v4: TEST_GROUP_V4,
689 group_v6: Some(TEST_GROUP_V6),
690 port,
691 interface_filter: InterfaceFilter::default(),
692 node_info: MulticastNodeInfo {
693 alias: "Loopback Check".to_string(),
694 version: "2.0".to_string(),
695 device_model: None,
696 device_type: Some(DeviceType::Server),
697 fingerprint: "FP_SELF_ONLY".to_string(),
698 port,
699 protocol: ProtocolType::Https,
700 download: false,
701 },
702 allow_loopback_fallback: false,
703 };
704
705 let handle = match start_multicast_engine(cfg, node_tx, stop_rx).await {
706 Ok(h) => h,
707 Err(e) => {
708 eprintln!("Skipping test: {e}");
709 return;
710 }
711 };
712
713 handle.send_single_announcement().await;
715
716 let own_received = tokio::time::timeout(Duration::from_millis(500), async {
718 while let Some(evt) = node_rx.recv().await {
719 if matches!(
720 &evt,
721 MulticastDiscoveryEvent::PeerDiscovered { message, .. }
722 if message.fingerprint == "FP_SELF_ONLY"
723 ) {
724 return true;
725 }
726 }
727 false
728 })
729 .await;
730
731 assert!(
732 !own_received.unwrap_or(false),
733 "Own announcement must be suppressed by loopback filter"
734 );
735
736 handle.stop().await;
737 }
738
739
740 #[tokio::test]
741 async fn test_port_released_after_stop() {
742 let port = NEXT_TEST_PORT.fetch_add(1, Ordering::Relaxed);
743
744 let (tx1, _rx1) = mpsc::channel(32);
745 let (stop_tx1, stop_rx1) = oneshot::channel();
746
747 let cfg1 = MulticastConfig {
748 group_v4: TEST_GROUP_V4,
749 group_v6: Some(TEST_GROUP_V6),
750 port,
751 interface_filter: InterfaceFilter::default(),
752 node_info: MulticastNodeInfo {
753 alias: "Instance 1".to_string(),
754 version: "2.0".to_string(),
755 device_model: None,
756 device_type: None,
757 fingerprint: "FP_1".to_string(),
758 port,
759 protocol: ProtocolType::Https,
760 download: false,
761 },
762 allow_loopback_fallback: false,
763 };
764
765 let handle1 = match start_multicast_engine(cfg1, tx1, stop_rx1).await {
766 Ok(h) => h,
767 Err(e) => {
768 eprintln!("Skipping test: {e}");
769 return;
770 }
771 };
772
773 let _ = stop_tx1.send(());
774 handle1.wait_stopped().await;
775
776 let (tx2, _rx2) = mpsc::channel(32);
778 let (_stop_tx2, stop_rx2) = oneshot::channel();
779
780 let cfg2 = MulticastConfig {
781 group_v4: TEST_GROUP_V4,
782 group_v6: Some(TEST_GROUP_V6),
783 port,
784 interface_filter: InterfaceFilter::default(),
785 node_info: MulticastNodeInfo {
786 alias: "Instance 2".to_string(),
787 version: "2.0".to_string(),
788 device_model: None,
789 device_type: None,
790 fingerprint: "FP_2".to_string(),
791 port,
792 protocol: ProtocolType::Https,
793 download: false,
794 },
795 allow_loopback_fallback: false,
796 };
797
798 let handle2 = start_multicast_engine(cfg2, tx2, stop_rx2).await;
799 assert!(handle2.is_ok(), "Port should be reusable after shutdown");
800 if let Ok(h) = handle2 {
801 h.stop().await;
802 }
803 }
804}
805