1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

use std::sync::Arc;
use tokio::{spawn, sync::RwLock};

use simple_dns::{header_buffer, Packet, PacketFlag, ResourceRecord};

use crate::{
    build_reply,
    resource_record_manager::ResourceRecordManager,
    socket_helper::{join_multicast, nonblocking, sender_socket},
    NetworkScope, SimpleMdnsError,
};

const FIVE_MINUTES: u32 = 60 * 5;

/// A simple mDNS responder aimed for service discovery.
/// In case you don't have a mDNS responder in your network, or for some reason don't want to use the ones available.
///
/// This responder will list for any mDNS query in the network via Multicast and will reply only to the resources that were added.
///
/// ```no_run
///     use simple_mdns::async_discovery::SimpleMdnsResponder;
///     use simple_dns::{Name, CLASS, ResourceRecord, rdata::{RData, A, SRV}};
///     use std::net::Ipv4Addr;
///
///     let mut responder = SimpleMdnsResponder::new(10);
///     let srv_name = Name::new_unchecked("_srvname._tcp.local");
///
///     responder.add_resource(ResourceRecord::new(
///         srv_name.clone(),
///         CLASS::IN,
///         10,
///         RData::A(A { address: Ipv4Addr::LOCALHOST.into() }),
///     ));
///
///     responder.add_resource(ResourceRecord::new(
///         srv_name.clone(),
///         CLASS::IN,
///         10,
///         RData::SRV(SRV {
///             port: 8080,
///             priority: 0,
///             weight: 0,
///             target: srv_name
///         })
///     ));
/// ```
///
/// This struct heavily relies on [`simple_dns`] crate and the same must be added as a dependency
pub struct SimpleMdnsResponder {
    resources: Arc<RwLock<ResourceRecordManager<'static>>>,
    rr_ttl: u32,
}

impl SimpleMdnsResponder {
    /// Creates a new SimpleMdnsResponder with specified ttl and IPV4 scope with UNSPECIFIED
    /// Interface
    pub fn new(rr_ttl: u32) -> Self {
        Self::new_with_scope(rr_ttl, NetworkScope::V4)
    }

    /// Creates a new SimpleMdnsResponder with specified ttl and network scope
    pub fn new_with_scope(rr_ttl: u32, scope: NetworkScope) -> Self {
        let responder = Self {
            resources: Arc::new(RwLock::new(ResourceRecordManager::new())),
            rr_ttl,
        };

        let resources = responder.resources.clone();
        spawn(async move {
            if let Err(err) = Self::responder_loop(resources, scope).await {
                log::error!("Dns Responder failed: {}", err);
            }
        });
        responder
    }

    /// Register a Resource Record
    pub async fn add_resource(&mut self, resource: ResourceRecord<'static>) {
        let mut resources = self.resources.write().await;
        resources.add_owned_resource(resource);
    }

    /// Remove a resource record
    pub async fn remove_resource_record(&mut self, resource: ResourceRecord<'static>) {
        let mut resources = self.resources.write().await;
        resources.remove_resource_record(&resource);
    }

    /// Remove all resource records
    pub async fn clear(&mut self) {
        let mut resources = self.resources.write().await;
        resources.clear();
    }

    async fn responder_loop(
        resources: Arc<RwLock<ResourceRecordManager<'_>>>,
        scope: NetworkScope,
    ) -> Result<(), SimpleMdnsError> {
        let mut recv_buffer = [0u8; 9000];
        let sender_socket = sender_socket(scope.is_v4()).and_then(nonblocking)?;

        let recv_socket = join_multicast(scope).and_then(nonblocking)?;

        loop {
            let (count, addr) = recv_socket.recv_from(&mut recv_buffer).await?;

            if header_buffer::has_flags(&recv_buffer[..count], PacketFlag::RESPONSE).unwrap_or(true)
            {
                continue;
            }

            match Packet::parse(&recv_buffer[..count]) {
                Ok(packet) => {
                    match build_reply(packet, &*resources.read().await) {
                        Some((reply_packet, unicast_response)) => {
                            let reply = match reply_packet.build_bytes_vec_compressed() {
                                Ok(reply) => reply,
                                Err(err) => {
                                    log::error!("Failed to build reply {err}");
                                    continue;
                                }
                            };

                            let reply_addr = if unicast_response {
                                addr
                            } else {
                                scope.socket_address()
                            };

                            sender_socket.send_to(&reply, reply_addr).await?;
                        }
                        None => {
                            log::trace!("No reply for query");
                            continue;
                        }
                    };
                }
                Err(err) => {
                    log::error!("Received Invalid packet {err}");
                }
            }
        }
    }

    /// Set the simple mdns responder's rr default ttl in seconds (defaults to 300).
    pub fn set_rr_ttl(&mut self, rr_default_ttl: u32) {
        self.rr_ttl = rr_default_ttl;
    }
}

impl Default for SimpleMdnsResponder {
    fn default() -> Self {
        Self::new(FIVE_MINUTES)
    }
}