multicast_discovery_socket/
protocol.rs

1use std::borrow::Cow;
2use log::debug;
3use sha2_const_stable::Sha256;
4use crate::AdvertisementData;
5
6/// Message kind for `DiscoveryMessage` type
7#[derive(Copy, Clone)]
8pub enum DiscoveryMessageKind {
9    Discovery,
10    Announce,
11    ExtendAnnouncements,
12}
13
14impl DiscoveryMessageKind {
15    const fn str_name(&self) -> &'static [u8] {
16        match self {
17            DiscoveryMessageKind::Discovery => b"discovery",
18            DiscoveryMessageKind::Announce => b"announce",
19            DiscoveryMessageKind::ExtendAnnouncements => b"extend-announcements",
20        }
21    }
22    const fn pattern(&self) -> [u8; 32] {
23        Sha256::new().update(self.str_name()).finalize()
24    }
25
26    pub fn try_from_pattern(pattern: &[u8]) -> Option<Self> {
27        if pattern == Self::Discovery.pattern() {
28            Some(Self::Discovery)
29        } else if pattern == Self::Announce.pattern() {
30            Some(Self::Announce)
31        } else if pattern == Self::ExtendAnnouncements.pattern() {
32            Some(Self::ExtendAnnouncements)
33        } else {
34            None
35        }
36    }
37}
38
39#[derive(Clone)]
40pub enum DiscoveryMessage<'a, D: AdvertisementData> {
41    /// Ping packet used to trigger other endpoints to send Announce packet back
42    Discovery,
43    /// Tell other endpoints that we are running and available for making connection
44    Announce {
45        service_port: u16,
46        discover_id: u32,
47        disconnected: bool,
48        adv_data: Cow<'a, D>,
49    },
50    /// Request for endpoints on Primary and Backup ports to extend their announcements scope to Backup ports as well
51    ExtendAnnouncements,
52}
53
54impl<D: AdvertisementData> DiscoveryMessage<'_, D> {
55    fn msg_type(&self) -> DiscoveryMessageKind {
56        match self {
57            DiscoveryMessage::Discovery => DiscoveryMessageKind::Discovery,
58            DiscoveryMessage::Announce { .. } => DiscoveryMessageKind::Announce,
59            DiscoveryMessage::ExtendAnnouncements => DiscoveryMessageKind::ExtendAnnouncements,
60        }
61    }
62    pub(crate) fn try_parse(msg: &[u8]) -> Option<Self> {
63        if msg.len() < 32 {
64            debug!("Packet is too small, ignoring...");
65            return None;
66        }
67        let msg_pattern = &msg[..32];
68        let Some(msg_type) = DiscoveryMessageKind::try_from_pattern(msg_pattern) else {
69            debug!("Unknown discovery message pattern, ignoring...");
70            return None;
71        };
72
73        let msg_body = &msg[32..];
74        match msg_type {
75            DiscoveryMessageKind::Announce if msg_body.len() >= 2+4+1 => {
76                let service_port = u16::from_be_bytes(msg_body[0..2].try_into().unwrap());
77                let discover_id = u32::from_be_bytes(msg_body[2..6].try_into().unwrap());
78                let disconnected = msg_body[6] != 0;
79                let adv_data_body = &msg_body[7..];
80                let adv_data = D::try_decode(adv_data_body)?;
81
82                Some(DiscoveryMessage::Announce {
83                    adv_data: Cow::Owned(adv_data),
84                    service_port,
85                    disconnected,
86                    discover_id
87                })
88            }
89            DiscoveryMessageKind::Discovery if msg_body.is_empty() => {
90                Some(DiscoveryMessage::Discovery)
91            }
92            DiscoveryMessageKind::ExtendAnnouncements if msg_body.is_empty() => {
93                Some(DiscoveryMessage::ExtendAnnouncements)
94            }
95            _ => {
96                debug!("Discovery message has invalid length, ignoring...");
97                None
98            }
99        }
100    }
101    pub(crate) fn gen_message(&self) -> Vec<u8> {
102        let pattern = self.msg_type().pattern();
103        match self {
104            DiscoveryMessage::Discovery => pattern.to_vec(),
105            DiscoveryMessage::Announce { service_port, discover_id, disconnected, adv_data } => {
106                let mut hello_msg = pattern.to_vec();
107                hello_msg.extend_from_slice(service_port.to_be_bytes().as_ref());
108                hello_msg.extend_from_slice(discover_id.to_be_bytes().as_ref());
109                hello_msg.push(*disconnected as u8);
110                hello_msg.extend_from_slice(&adv_data.encode_to_bytes());
111                hello_msg
112            }
113            DiscoveryMessage::ExtendAnnouncements => pattern.to_vec(),
114        }
115    }
116}