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 mut 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 res = &mut stop_rx => {
452 if res.is_err() {
453 let _ = cancel.cancelled().await;
454 }
455 }
456 _ = cancel.cancelled() => {}
457 }
458
459 cancel.cancel();
460 receiver_tasks.shutdown().await;
461 state.sockets.write().await.clear();
462 }
463 });
464
465 Ok(MulticastHandle {
466 state,
467 cancel,
468 task: Arc::new(Mutex::new(Some(task))),
469 })
470}
471
472async fn receive_loop(
474 socket: BoundMulticastSocket,
475 state: Arc<MulticastEngineState>,
476 event_tx: mpsc::Sender<MulticastDiscoveryEvent>,
477 cancel: CancellationToken,
478) {
479 let mut buffer = vec![0u8; RECEIVE_BUFFER_SIZE];
480 let mut consecutive_errors = 0;
481
482 loop {
483 let received = tokio::select! {
484 _ = cancel.cancelled() => return,
485 res = socket.socket.recv_from(&mut buffer) => res,
486 };
487
488 let (size, source) = match received {
489 Ok(res) => {
490 consecutive_errors = 0;
491 res
492 }
493 Err(err) => {
494 consecutive_errors += 1;
495 tracing::warn!(
496 "Multicast recv error on interface {}: {err:#}",
497 socket.description
498 );
499 if consecutive_errors >= MAX_CONSECUTIVE_RECEIVE_ERRORS {
500 tracing::error!(
501 "Exceeded max consecutive receive errors on {}. Terminating socket loop.",
502 socket.description
503 );
504 return;
505 }
506 continue;
507 }
508 };
509
510 let message = match serde_json::from_slice::<MulticastMessageV2>(&buffer[..size]) {
511 Ok(msg) => msg,
512 Err(err) => {
513 tracing::warn!("Failed to deserialize multicast message from {source}: {err:#}");
514 continue;
515 }
516 };
517
518 if !state.node_info.fingerprint.is_empty()
520 && message.fingerprint == state.node_info.fingerprint
521 {
522 tracing::trace!(
523 "Suppressed self-multicast announcement with fingerprint {}",
524 message.fingerprint
525 );
526 continue;
527 }
528
529 let event = MulticastDiscoveryEvent::PeerDiscovered {
530 remote_ip: source.ip(),
531 scope_id: match &source {
532 SocketAddr::V6(v6) if v6.scope_id() != 0 => Some(v6.scope_id()),
533 _ => None,
534 },
535 message,
536 };
537
538 if event_tx.send(event).await.is_err() {
539 tracing::debug!("Discovery event receiver dropped. Shutting down receive loop.");
540 return;
541 }
542 }
543}
544
545#[cfg(test)]
546mod tests {
547 use super::*;
548 use std::sync::atomic::{AtomicU16, Ordering};
549
550 static NEXT_TEST_PORT: AtomicU16 = AtomicU16::new(54100);
551
552 const TEST_GROUP_V4: Ipv4Addr = Ipv4Addr::new(224, 0, 0, 169);
553 const TEST_GROUP_V6: Ipv6Addr = Ipv6Addr::new(0xff12, 0, 0, 0, 0, 0, 0xfd3a, 0xe422);
554
555 #[test]
556 fn test_node_info_to_message() {
557 let info = MulticastNodeInfo {
558 alias: "Test Node".to_string(),
559 version: "2.0".to_string(),
560 device_model: Some("Linux".to_string()),
561 device_type: Some(DeviceType::Desktop),
562 fingerprint: "FP123456".to_string(),
563 port: 53317,
564 protocol: ProtocolType::Https,
565 download: true,
566 };
567
568 let msg = info.to_message(true);
569 assert_eq!(msg.alias, "Test Node");
570 assert_eq!(msg.fingerprint, "FP123456");
571 assert!(msg.announce);
572 assert!(msg.download);
573
574 let json = serde_json::to_string(&msg).unwrap();
575 assert!(json.contains("\"announce\":true"));
576 assert!(json.contains("\"alias\":\"Test Node\""));
577 }
578
579 #[test]
580 fn test_multicast_config_default() {
581 let config = MulticastConfig::default();
582 assert_eq!(config.group_v4, DEFAULT_MULTICAST_IPV4);
583 assert_eq!(config.group_v6, Some(DEFAULT_MULTICAST_IPV6));
584 assert_eq!(config.port, DEFAULT_MULTICAST_PORT);
585 }
586
587 #[tokio::test]
588 async fn test_multi_node_discovery_exchange() {
589 let port = NEXT_TEST_PORT.fetch_add(1, Ordering::Relaxed);
590
591 let (node1_tx, _node1_rx) = mpsc::channel(32);
592 let (_node1_stop_tx, node1_stop_rx) = oneshot::channel();
593
594 let cfg1 = MulticastConfig {
595 group_v4: TEST_GROUP_V4,
596 group_v6: Some(TEST_GROUP_V6),
597 port,
598 interface_filter: InterfaceFilter::default(),
599 node_info: MulticastNodeInfo {
600 alias: "Node One".to_string(),
601 version: "2.0".to_string(),
602 device_model: Some("Linux".to_string()),
603 device_type: Some(DeviceType::Desktop),
604 fingerprint: "FP_NODE_ONE".to_string(),
605 port,
606 protocol: ProtocolType::Https,
607 download: true,
608 },
609 allow_loopback_fallback: false,
610 };
611
612 let handle1 = match start_multicast_engine(cfg1, node1_tx, node1_stop_rx).await {
613 Ok(h) => h,
614 Err(e) => {
615 eprintln!("Skipping test due to environment without multicast: {e}");
616 return;
617 }
618 };
619
620 let (node2_tx, mut node2_rx) = mpsc::channel(32);
621 let (_node2_stop_tx, node2_stop_rx) = oneshot::channel();
622
623 let cfg2 = MulticastConfig {
624 group_v4: TEST_GROUP_V4,
625 group_v6: Some(TEST_GROUP_V6),
626 port,
627 interface_filter: InterfaceFilter::default(),
628 node_info: MulticastNodeInfo {
629 alias: "Node Two".to_string(),
630 version: "2.0".to_string(),
631 device_model: Some("Android".to_string()),
632 device_type: Some(DeviceType::Mobile),
633 fingerprint: "FP_NODE_TWO".to_string(),
634 port,
635 protocol: ProtocolType::Https,
636 download: false,
637 },
638 allow_loopback_fallback: false,
639 };
640
641 let handle2 = match start_multicast_engine(cfg2, node2_tx, node2_stop_rx).await {
642 Ok(h) => h,
643 Err(e) => {
644 handle1.stop().await;
645 eprintln!("Skipping test due to environment without multicast: {e}");
646 return;
647 }
648 };
649
650 handle1.send_single_announcement().await;
652
653 let received_by_node2 = tokio::time::timeout(Duration::from_secs(3), async {
655 while let Some(evt) = node2_rx.recv().await {
656 match evt {
657 MulticastDiscoveryEvent::PeerDiscovered { message, .. }
658 if message.fingerprint == "FP_NODE_ONE" =>
659 {
660 return Some(message);
661 }
662 _ => {}
663 }
664 }
665 None
666 })
667 .await;
668
669
670
671 if let Ok(Some(msg)) = received_by_node2 {
672 assert_eq!(msg.alias, "Node One");
673 assert_eq!(msg.fingerprint, "FP_NODE_ONE");
674 assert_eq!(msg.device_type, Some(DeviceType::Desktop));
675 } else {
676 eprintln!("Warning: Multicast packet not received (likely filtered by network interface)");
677 }
678
679 handle1.stop().await;
681 handle2.stop().await;
682 }
683
684 #[tokio::test]
685 async fn test_loopback_suppression() {
686 let port = NEXT_TEST_PORT.fetch_add(1, Ordering::Relaxed);
687
688 let (node_tx, mut node_rx) = mpsc::channel(32);
689 let (_stop_tx, stop_rx) = oneshot::channel();
690
691 let cfg = MulticastConfig {
692 group_v4: TEST_GROUP_V4,
693 group_v6: Some(TEST_GROUP_V6),
694 port,
695 interface_filter: InterfaceFilter::default(),
696 node_info: MulticastNodeInfo {
697 alias: "Loopback Check".to_string(),
698 version: "2.0".to_string(),
699 device_model: None,
700 device_type: Some(DeviceType::Server),
701 fingerprint: "FP_SELF_ONLY".to_string(),
702 port,
703 protocol: ProtocolType::Https,
704 download: false,
705 },
706 allow_loopback_fallback: false,
707 };
708
709 let handle = match start_multicast_engine(cfg, node_tx, stop_rx).await {
710 Ok(h) => h,
711 Err(e) => {
712 eprintln!("Skipping test: {e}");
713 return;
714 }
715 };
716
717 handle.send_single_announcement().await;
719
720 let own_received = tokio::time::timeout(Duration::from_millis(500), async {
722 while let Some(evt) = node_rx.recv().await {
723 if matches!(
724 &evt,
725 MulticastDiscoveryEvent::PeerDiscovered { message, .. }
726 if message.fingerprint == "FP_SELF_ONLY"
727 ) {
728 return true;
729 }
730 }
731 false
732 })
733 .await;
734
735 assert!(
736 !own_received.unwrap_or(false),
737 "Own announcement must be suppressed by loopback filter"
738 );
739
740 handle.stop().await;
741 }
742
743
744 #[tokio::test]
745 async fn test_port_released_after_stop() {
746 let port = NEXT_TEST_PORT.fetch_add(1, Ordering::Relaxed);
747
748 let (tx1, _rx1) = mpsc::channel(32);
749 let (stop_tx1, stop_rx1) = oneshot::channel();
750
751 let cfg1 = MulticastConfig {
752 group_v4: TEST_GROUP_V4,
753 group_v6: Some(TEST_GROUP_V6),
754 port,
755 interface_filter: InterfaceFilter::default(),
756 node_info: MulticastNodeInfo {
757 alias: "Instance 1".to_string(),
758 version: "2.0".to_string(),
759 device_model: None,
760 device_type: None,
761 fingerprint: "FP_1".to_string(),
762 port,
763 protocol: ProtocolType::Https,
764 download: false,
765 },
766 allow_loopback_fallback: false,
767 };
768
769 let handle1 = match start_multicast_engine(cfg1, tx1, stop_rx1).await {
770 Ok(h) => h,
771 Err(e) => {
772 eprintln!("Skipping test: {e}");
773 return;
774 }
775 };
776
777 let _ = stop_tx1.send(());
778 handle1.wait_stopped().await;
779
780 let (tx2, _rx2) = mpsc::channel(32);
782 let (_stop_tx2, stop_rx2) = oneshot::channel();
783
784 let cfg2 = MulticastConfig {
785 group_v4: TEST_GROUP_V4,
786 group_v6: Some(TEST_GROUP_V6),
787 port,
788 interface_filter: InterfaceFilter::default(),
789 node_info: MulticastNodeInfo {
790 alias: "Instance 2".to_string(),
791 version: "2.0".to_string(),
792 device_model: None,
793 device_type: None,
794 fingerprint: "FP_2".to_string(),
795 port,
796 protocol: ProtocolType::Https,
797 download: false,
798 },
799 allow_loopback_fallback: false,
800 };
801
802 let handle2 = start_multicast_engine(cfg2, tx2, stop_rx2).await;
803 assert!(handle2.is_ok(), "Port should be reusable after shutdown");
804 if let Ok(h) = handle2 {
805 h.stop().await;
806 }
807 }
808}
809