Skip to main content

dns_sd_native/linux/
mod.rs

1use zbus::Connection;
2
3mod dbus;
4use dbus::*;
5use futures_util::stream::StreamExt;
6
7use log::{error, trace, warn};
8use std::num::NonZeroU32;
9
10use crate::{ServiceRegistrationError, TxtRecordValue};
11
12const AVAHI_IF_UNSPEC: i32 = -1;
13const AVAHI_PROTO_UNSPEC: i32 = -1;
14
15/// Reference to a registered service instance.
16///
17/// The service will be automatically unregistered when this value is dropped.
18pub struct ServiceRegistration {
19    entry_group: Option<EntryGroupProxy<'static>>,
20}
21
22impl ServiceRegistration {
23    pub(crate) async fn new(
24        service_type: &String,
25        port: u16,
26        name: &Option<String>,
27        host: &Option<String>,
28        domain: &Option<String>,
29        interface_index: Option<NonZeroU32>,
30        txt_record_values: &[(String, TxtRecordValue)],
31    ) -> Result<ServiceRegistration, ServiceRegistrationError> {
32        let conn = Connection::system().await.map_err(|err| {
33            ServiceRegistrationError::DnsSdUnavailable(format!(
34                "failed to connect to system D-Bus: {err}"
35            ))
36        })?;
37
38        let manager = AvahiProxy::new(&conn).await.map_err(|err| {
39            ServiceRegistrationError::DnsSdUnavailable(format!(
40                "failed to connect to Avahi via D-Bus: {err}"
41            ))
42        })?;
43        let entry_group = manager.entry_group_new().await.map_err(|err| {
44            ServiceRegistrationError::RegistrationError(format!(
45                "failed to create Avahi entry group: {err}"
46            ))
47        })?;
48
49        let protocol = AVAHI_PROTO_UNSPEC;
50        let flags = 0;
51
52        let interface = match interface_index {
53            Some(i) => {
54                let idx = i.get();
55                if idx > i32::MAX as u32 {
56                    return Err(ServiceRegistrationError::InvalidInterfaceIndex(idx));
57                }
58                idx as i32
59            }
60            None => AVAHI_IF_UNSPEC,
61        };
62        let domain = domain.as_deref().unwrap_or("");
63        let host = host.as_deref().unwrap_or("");
64        let name = if let Some(name) = name {
65            name.as_str()
66        } else {
67            &manager
68                .get_host_name()
69                .await
70                .map_err(|err| ServiceRegistrationError::HostnameUnavailable(err.to_string()))?
71        };
72        let txt: Vec<Vec<u8>> = txt_record_values
73            .iter()
74            .map(|(key, value)| {
75                let mut record = key.clone().into_bytes();
76                match value {
77                    TxtRecordValue::KeyOnly => {}
78                    TxtRecordValue::String(s) => {
79                        record.push(b'=');
80                        record.extend_from_slice(s.as_bytes());
81                    }
82                    TxtRecordValue::Binary(b) => {
83                        record.push(b'=');
84                        record.extend_from_slice(b);
85                    }
86                }
87                record
88            })
89            .collect();
90
91        let txt_refs: Vec<&[u8]> = txt.iter().map(|v| v.as_slice()).collect();
92
93        trace!(
94            "registering service with Avahi: interface={:?} protocol={:?} flags={:?} name={:?} type={:?} domain={:?} host={:?} port={:?} txt={:?}",
95            interface, protocol, flags, name, service_type, domain, host, port, txt
96        );
97        entry_group
98            .add_service(
99                interface,
100                protocol,
101                flags,
102                name,
103                service_type,
104                domain,
105                host,
106                port,
107                &txt_refs,
108            )
109            .await
110            .map_err(|err| {
111                ServiceRegistrationError::RegistrationError(format!(
112                    "Avahi add_service failed: {err}"
113                ))
114            })?;
115        // TODO: return ServiceRegistrationError::NameConflict if Avahi returns ErrorName == "org.freedesktop.Avahi.CollisionError"
116
117        entry_group.commit().await.map_err(|err| {
118            ServiceRegistrationError::RegistrationError(format!(
119                "Avahi entry group commit failed: {err}"
120            ))
121        })?;
122
123        match entry_group.get_state().await {
124            Ok(AVAHI_ENTRY_GROUP_ESTABLISHED) => {
125                trace!("service registration state: established");
126            }
127            Ok(AVAHI_ENTRY_GROUP_COLLISION) => {
128                return Err(ServiceRegistrationError::NameConflict);
129            }
130            Ok(AVAHI_ENTRY_GROUP_FAILURE) => {
131                return Err(ServiceRegistrationError::RegistrationFailed(
132                    "entry group entered failure state".into(),
133                ));
134            }
135            Err(err) => {
136                return Err(ServiceRegistrationError::RegistrationError(format!(
137                    "Avahi entry group get_state failed: {err}"
138                )));
139            }
140            Ok(state) => {
141                if state != AVAHI_ENTRY_GROUP_REGISTERING {
142                    warn!("service registration state: unknown state: {state}");
143                }
144                let mut state_stream = entry_group
145                    .receive_state_changed()
146                    .await
147                    .map_err(|err| ServiceRegistrationError::RegistrationError(err.to_string()))?;
148                while let Some(msg) = state_stream.next().await {
149                    let args = msg.args().map_err(|err| {
150                        ServiceRegistrationError::RegistrationError(err.to_string())
151                    })?;
152                    trace!(
153                        "state changed: state={:?} error={:?}",
154                        args.state, args.error
155                    );
156                    match args.state {
157                        AVAHI_ENTRY_GROUP_REGISTERING => continue,
158                        AVAHI_ENTRY_GROUP_ESTABLISHED => {
159                            trace!("service registration state: established");
160                        }
161                        AVAHI_ENTRY_GROUP_COLLISION => {
162                            return Err(ServiceRegistrationError::NameConflict);
163                        }
164                        AVAHI_ENTRY_GROUP_FAILURE => {
165                            return Err(ServiceRegistrationError::RegistrationFailed(format!(
166                                "entry group failure: {}",
167                                args.error
168                            )));
169                        }
170                        _ => {
171                            warn!("service registration state: unknown state: {args:?}");
172                        }
173                    }
174                    break;
175                }
176            }
177        }
178
179        Ok(Self {
180            entry_group: Some(entry_group),
181        })
182    }
183
184    /// Unregisters the service, notifying remote clients that the service is no longer available.
185    ///
186    /// Use this method instead of dropping the `ServiceRegistration` if you want to be notified of
187    /// any errors that occur during unregistration.
188    pub async fn unregister(mut self) -> Result<(), String> {
189        if let Some(entry_group) = self.entry_group.take() {
190            entry_group
191                .reset()
192                .await
193                .map_err(|err| format!("failed to unregister service: {err:?}"))
194        } else {
195            Ok(())
196        }
197    }
198}
199
200impl Drop for ServiceRegistration {
201    fn drop(&mut self) {
202        if let Some(entry_group) = self.entry_group.take() {
203            tokio::spawn(async move {
204                if let Err(err) = entry_group.reset().await {
205                    error!("failed to unregister service: {err}");
206                }
207            });
208        }
209    }
210}