mdns-sd-discovery 0.2.0

use std::collections::HashSet;
use std::net::IpAddr;
use std::num::NonZeroU32;
use std::str::FromStr;

use futures_util::stream::StreamExt;
use log::{trace, warn};
use tokio::sync::mpsc::unbounded_channel;
use tokio::task::JoinHandle;
use zbus::message::Type as MessageType;
use zbus::{Connection, MatchRule, MessageStream};

use super::dbus::*;
use crate::browse::{
    BrowseEvent, BrowseEventReceiver, BrowseEventSender, DiscoveredService, RemovedService,
    ServiceBrowseError, parse_txt_entry,
};

/// Resolves a single service instance to a connectable endpoint using a
/// dedicated, short-lived D-Bus connection. Returns
/// [`ServiceBrowseError::ResolveFailed`] if the instance no longer responds,
/// which callers use as a liveness probe.
pub(crate) async fn resolve_once(
    name: &str,
    service_type: &str,
    domain: &str,
    interface_index: Option<NonZeroU32>,
) -> Result<DiscoveredService, ServiceBrowseError> {
    let interface = interface_to_avahi(interface_index)?;

    let conn = Connection::system().await.map_err(|err| {
        ServiceBrowseError::DnsSdUnavailable(format!("failed to connect to system D-Bus: {err}"))
    })?;
    let server = AvahiProxy::new(&conn).await.map_err(|err| {
        ServiceBrowseError::DnsSdUnavailable(format!("failed to connect to Avahi via D-Bus: {err}"))
    })?;

    match server
        .resolve_service(
            interface,
            AVAHI_PROTO_UNSPEC,
            name,
            service_type,
            domain,
            AVAHI_PROTO_UNSPEC,
            0,
        )
        .await
    {
        Ok(resolved) => Ok(resolved_to_service(resolved)),
        Err(err) => Err(ServiceBrowseError::ResolveFailed(
            name.to_string(),
            err.to_string(),
        )),
    }
}

/// Builds a [`DiscoveredService`] from the tuple returned by
/// `Server.ResolveService`.
fn resolved_to_service(resolved: ResolvedService) -> DiscoveredService {
    let (iface, _proto, name, service_type, domain, host, _aproto, address, port, txt, _flags) =
        resolved;
    let addresses: Vec<IpAddr> = IpAddr::from_str(&address).ok().into_iter().collect();
    let txt_records = txt.iter().map(|entry| parse_txt_entry(entry)).collect();
    DiscoveredService {
        name,
        service_type,
        domain,
        host_name: host,
        port,
        addresses,
        txt_records,
        interface_index: avahi_interface_to_index(iface),
    }
}

const SERVICE_BROWSER_INTERFACE: &str = "org.freedesktop.Avahi.ServiceBrowser";
const SERVICE_TYPE_BROWSER_INTERFACE: &str = "org.freedesktop.Avahi.ServiceTypeBrowser";

/// Guard returned alongside the event receiver. Dropping it aborts the root
/// browse task, which in turn drops (and thereby aborts) any child browse and
/// resolver tasks it owns.
pub(crate) struct BrowseGuard {
    handle: JoinHandle<()>,
}

impl Drop for BrowseGuard {
    fn drop(&mut self) {
        self.handle.abort();
    }
}

/// Aborts the wrapped task when dropped. Used to tie child task lifetimes to the
/// parent task that owns them.
struct AbortOnDrop(JoinHandle<()>);

impl Drop for AbortOnDrop {
    fn drop(&mut self) {
        self.0.abort();
    }
}

pub(crate) async fn browse_start(
    service_type: &Option<String>,
    domain: &Option<String>,
    interface_index: Option<NonZeroU32>,
) -> Result<(BrowseEventReceiver, BrowseGuard), ServiceBrowseError> {
    // Validate Avahi is reachable up front so the error surfaces from `browse()`
    // rather than asynchronously through the event stream.
    let conn = Connection::system().await.map_err(|err| {
        ServiceBrowseError::DnsSdUnavailable(format!("failed to connect to system D-Bus: {err}"))
    })?;
    AvahiProxy::new(&conn).await.map_err(|err| {
        ServiceBrowseError::DnsSdUnavailable(format!("failed to connect to Avahi via D-Bus: {err}"))
    })?;
    drop(conn);

    let interface = interface_to_avahi(interface_index)?;
    let domain = domain.clone().unwrap_or_default();
    let (tx, rx) = unbounded_channel();

    let handle = match service_type {
        Some(service_type) => {
            let service_type = service_type.clone();
            tokio::spawn(browse_one_type(interface, service_type, domain, tx))
        }
        None => tokio::spawn(browse_all_types(interface, domain, tx)),
    };

    Ok((rx, BrowseGuard { handle }))
}

/// Subscribes to all signals of `interface` on a freshly created, dedicated
/// connection *before* any browser object is created, so the initial burst of
/// cached-entry signals is not missed. Returns the connection, the message
/// stream, and an Avahi server proxy on the connection.
async fn connect_and_subscribe(
    interface: &str,
) -> Result<(Connection, MessageStream, AvahiProxy<'static>), ServiceBrowseError> {
    let conn = Connection::system().await.map_err(|err| {
        ServiceBrowseError::DnsSdUnavailable(format!("failed to connect to system D-Bus: {err}"))
    })?;

    let rule = MatchRule::builder()
        .msg_type(MessageType::Signal)
        .sender("org.freedesktop.Avahi")
        .and_then(|b| b.interface(interface))
        .map_err(|err| ServiceBrowseError::BrowseFailed(err.to_string()))?
        .build();

    let messages = MessageStream::for_match_rule(rule, &conn, None)
        .await
        .map_err(|err| ServiceBrowseError::BrowseFailed(err.to_string()))?;

    let server = AvahiProxy::new(&conn)
        .await
        .map_err(|err| ServiceBrowseError::DnsSdUnavailable(err.to_string()))?;

    Ok((conn, messages, server))
}

/// Browses the DNS-SD service-type meta-query and starts a per-type instance
/// browse for each newly discovered type.
async fn browse_all_types(interface: i32, domain: String, tx: BrowseEventSender) {
    // `_conn` is unused directly but kept alive so its message stream stays open.
    let (_conn, mut messages, server) =
        match connect_and_subscribe(SERVICE_TYPE_BROWSER_INTERFACE).await {
            Ok(parts) => parts,
            Err(err) => {
                let _ = tx.send(Err(err));
                return;
            }
        };

    let type_browser = match server
        .service_type_browser_new(interface, AVAHI_PROTO_UNSPEC, &domain, 0)
        .await
    {
        Ok(browser) => browser,
        Err(err) => {
            let _ = tx.send(Err(ServiceBrowseError::BrowseFailed(format!(
                "ServiceTypeBrowserNew failed: {err}"
            ))));
            return;
        }
    };

    // Dedup discovered types (a type may be announced on multiple interfaces) and
    // keep each per-type instance browse alive for our lifetime.
    let mut seen: HashSet<(String, String)> = HashSet::new();
    let mut child_browsers: Vec<AbortOnDrop> = Vec::new();

    while let Some(msg) = messages.next().await {
        let msg = match msg {
            Ok(msg) => msg,
            Err(err) => {
                warn!("service type browser message error: {err}");
                continue;
            }
        };
        let member = msg.header().member().map(|m| m.as_str().to_owned());
        match member.as_deref() {
            Some("ItemNew") => {
                let (_iface, _proto, service_type, item_domain, _flags): (
                    i32,
                    i32,
                    String,
                    String,
                    u32,
                ) = match msg.body().deserialize() {
                    Ok(args) => args,
                    Err(err) => {
                        warn!("malformed service type ItemNew: {err}");
                        continue;
                    }
                };
                if seen.insert((service_type.clone(), item_domain.clone())) {
                    trace!("discovered service type {service_type:?} in domain {item_domain:?}");
                    let handle = tokio::spawn(browse_one_type(
                        interface,
                        service_type,
                        item_domain,
                        tx.clone(),
                    ));
                    child_browsers.push(AbortOnDrop(handle));
                }
            }
            Some("Failure") => {
                let err: String = msg.body().deserialize().unwrap_or_default();
                let _ = tx.send(Err(ServiceBrowseError::BrowseFailed(format!(
                    "service type browser failure: {err}"
                ))));
            }
            _ => {} // ItemRemove (of a type), AllForNow, CacheExhausted
        }
    }

    let _ = type_browser.free().await;
}

/// Browses instances of a single service type, resolving each as it appears.
async fn browse_one_type(
    interface: i32,
    service_type: String,
    domain: String,
    tx: BrowseEventSender,
) {
    let (conn, mut messages, server) = match connect_and_subscribe(SERVICE_BROWSER_INTERFACE).await
    {
        Ok(parts) => parts,
        Err(err) => {
            let _ = tx.send(Err(err));
            return;
        }
    };

    let browser = match server
        .service_browser_new(interface, AVAHI_PROTO_UNSPEC, &service_type, &domain, 0)
        .await
    {
        Ok(browser) => browser,
        Err(err) => {
            let _ = tx.send(Err(ServiceBrowseError::BrowseFailed(format!(
                "ServiceBrowserNew failed for {service_type}: {err}"
            ))));
            return;
        }
    };

    // In-flight resolver tasks; aborted when this task ends.
    let mut resolvers: Vec<AbortOnDrop> = Vec::new();

    while let Some(msg) = messages.next().await {
        let msg = match msg {
            Ok(msg) => msg,
            Err(err) => {
                warn!("service browser message error: {err}");
                continue;
            }
        };
        let member = msg.header().member().map(|m| m.as_str().to_owned());
        match member.as_deref() {
            Some("ItemNew") => {
                let (iface, protocol, name, item_type, item_domain, _flags): (
                    i32,
                    i32,
                    String,
                    String,
                    String,
                    u32,
                ) = match msg.body().deserialize() {
                    Ok(args) => args,
                    Err(err) => {
                        warn!("malformed ItemNew: {err}");
                        continue;
                    }
                };
                resolvers.retain(|r| !r.0.is_finished());
                let handle = tokio::spawn(resolve_and_emit(
                    conn.clone(),
                    iface,
                    protocol,
                    name,
                    item_type,
                    item_domain,
                    tx.clone(),
                ));
                resolvers.push(AbortOnDrop(handle));
            }
            Some("ItemRemove") => {
                let (iface, _protocol, name, item_type, item_domain, _flags): (
                    i32,
                    i32,
                    String,
                    String,
                    String,
                    u32,
                ) = match msg.body().deserialize() {
                    Ok(args) => args,
                    Err(err) => {
                        warn!("malformed ItemRemove: {err}");
                        continue;
                    }
                };
                let removed = RemovedService {
                    name,
                    service_type: item_type,
                    domain: item_domain,
                    interface_index: avahi_interface_to_index(iface),
                };
                if tx.send(Ok(BrowseEvent::Removed(removed))).is_err() {
                    break;
                }
            }
            Some("Failure") => {
                let err: String = msg.body().deserialize().unwrap_or_default();
                let _ = tx.send(Err(ServiceBrowseError::BrowseFailed(format!(
                    "service browser failure for {service_type}: {err}"
                ))));
            }
            _ => {} // AllForNow, CacheExhausted
        }
    }

    let _ = browser.free().await;
}

/// Resolves a single discovered service instance via the synchronous
/// `Server.ResolveService` method and emits a `Found` event (or a
/// `ResolveFailed` error). Using the method (rather than the signal-based
/// `ServiceResolverNew`) avoids the subscribe-after-create signal race.
async fn resolve_and_emit(
    conn: Connection,
    interface: i32,
    protocol: i32,
    name: String,
    service_type: String,
    domain: String,
    tx: BrowseEventSender,
) {
    let server = match AvahiProxy::new(&conn).await {
        Ok(server) => server,
        Err(err) => {
            let _ = tx.send(Err(ServiceBrowseError::ResolveFailed(
                name,
                err.to_string(),
            )));
            return;
        }
    };

    match server
        .resolve_service(
            interface,
            protocol,
            &name,
            &service_type,
            &domain,
            AVAHI_PROTO_UNSPEC,
            0,
        )
        .await
    {
        Ok(resolved) => {
            let _ = tx.send(Ok(BrowseEvent::Found(resolved_to_service(resolved))));
        }
        Err(err) => {
            let _ = tx.send(Err(ServiceBrowseError::ResolveFailed(
                name,
                err.to_string(),
            )));
        }
    }
}

/// Maps an optional interface index to Avahi's `i32` interface argument.
fn interface_to_avahi(interface_index: Option<NonZeroU32>) -> Result<i32, ServiceBrowseError> {
    match interface_index {
        Some(i) => {
            let idx = i.get();
            if idx > i32::MAX as u32 {
                return Err(ServiceBrowseError::InvalidInterfaceIndex(idx));
            }
            Ok(idx as i32)
        }
        None => Ok(AVAHI_IF_UNSPEC),
    }
}

/// Maps an Avahi `i32` interface value from a signal back to an interface index.
fn avahi_interface_to_index(interface: i32) -> Option<NonZeroU32> {
    if interface <= 0 {
        None
    } else {
        NonZeroU32::new(interface as u32)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn interface_to_avahi_none_is_unspecified() {
        assert_eq!(interface_to_avahi(None).unwrap(), AVAHI_IF_UNSPEC);
    }

    #[test]
    fn interface_to_avahi_passes_valid_index_through() {
        let idx = NonZeroU32::new(5).unwrap();
        assert_eq!(interface_to_avahi(Some(idx)).unwrap(), 5);
    }

    #[test]
    fn interface_to_avahi_accepts_i32_max() {
        let idx = NonZeroU32::new(i32::MAX as u32).unwrap();
        assert_eq!(interface_to_avahi(Some(idx)).unwrap(), i32::MAX);
    }

    #[test]
    fn interface_to_avahi_rejects_index_above_i32_max() {
        let idx = NonZeroU32::new(i32::MAX as u32 + 1).unwrap();
        match interface_to_avahi(Some(idx)) {
            Err(ServiceBrowseError::InvalidInterfaceIndex(i)) => {
                assert_eq!(i, i32::MAX as u32 + 1);
            }
            other => panic!("expected InvalidInterfaceIndex, got {other:?}"),
        }
    }

    #[test]
    fn avahi_interface_to_index_maps_positive_to_some() {
        assert_eq!(avahi_interface_to_index(2), NonZeroU32::new(2));
    }

    #[test]
    fn avahi_interface_to_index_zero_and_negative_are_none() {
        assert_eq!(avahi_interface_to_index(0), None);
        assert_eq!(avahi_interface_to_index(-1), None);
        assert_eq!(avahi_interface_to_index(AVAHI_IF_UNSPEC), None);
    }

    #[test]
    fn interface_round_trips_through_avahi_mapping() {
        let idx = NonZeroU32::new(9).unwrap();
        let avahi = interface_to_avahi(Some(idx)).unwrap();
        assert_eq!(avahi_interface_to_index(avahi), Some(idx));
    }

    #[test]
    fn resolved_to_service_maps_all_fields() {
        let resolved: ResolvedService = (
            3,
            AVAHI_PROTO_UNSPEC,
            "My Web Server".to_string(),
            "_http._tcp".to_string(),
            "local".to_string(),
            "macbook.local".to_string(),
            AVAHI_PROTO_UNSPEC,
            "192.168.1.10".to_string(),
            8080,
            vec![b"path=/index.html".to_vec(), b"flag".to_vec()],
            0,
        );
        let service = resolved_to_service(resolved);
        assert_eq!(service.name, "My Web Server");
        assert_eq!(service.service_type, "_http._tcp");
        assert_eq!(service.domain, "local");
        assert_eq!(service.host_name, "macbook.local");
        assert_eq!(service.port, 8080);
        assert_eq!(
            service.addresses,
            vec![IpAddr::from_str("192.168.1.10").unwrap()]
        );
        assert_eq!(service.txt("path"), Some(&b"/index.html"[..]));
        assert_eq!(service.txt("flag"), None);
        assert_eq!(service.interface_index, NonZeroU32::new(3));
    }

    #[test]
    fn resolved_to_service_unparsable_address_yields_no_addresses() {
        let resolved: ResolvedService = (
            0,
            AVAHI_PROTO_UNSPEC,
            "svc".to_string(),
            "_http._tcp".to_string(),
            "local".to_string(),
            "host.local".to_string(),
            AVAHI_PROTO_UNSPEC,
            String::new(),
            80,
            Vec::new(),
            0,
        );
        let service = resolved_to_service(resolved);
        assert!(service.addresses.is_empty());
        assert_eq!(service.interface_index, None);
    }
}