alpine-protocol-sdk 0.1.13

use std::{
    fmt, io,
    net::{IpAddr, SocketAddr, UdpSocket},
    time::Duration,
};

use alpine::messages::{DiscoveryReply, DiscoveryRequest};
use rand::{rngs::OsRng, RngCore};
use serde_cbor;
use tracing::trace;

const DEFAULT_MULTICAST_IPV4: &str = "239.255.255.250:5555";
const DEFAULT_MULTICAST_IPV6: &str = "[ff12::1]:5555";
const DEFAULT_BROADCAST_IPV4: &str = "255.255.255.255:5555";

/// Options used to configure the blocking discovery helper.
pub struct DiscoveryClientOptions {
    pub remote_addr: SocketAddr,
    pub local_addr: SocketAddr,
    pub timeout: Duration,
    pub prefer_multicast: bool,
}

impl DiscoveryClientOptions {
    /// Creates options with the provided remote socket and a default timeout.
    pub fn new(remote_addr: SocketAddr, local_addr: SocketAddr, timeout: Duration) -> Self {
        Self {
            remote_addr,
            local_addr,
            timeout,
            prefer_multicast: true,
        }
    }

    pub fn disable_multicast(mut self) -> Self {
        self.prefer_multicast = false;
        self
    }
}

/// Errors that can happen while sending or receiving discovery payloads.
#[derive(Debug)]
pub enum DiscoveryError {
    Io(io::Error),
    Decode(serde_cbor::Error),
    Timeout,
    PermissionDenied,
    MulticastUnavailable,
    BroadcastBlocked,
}

impl fmt::Display for DiscoveryError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DiscoveryError::Io(err) => write!(f, "io error: {}", err),
            DiscoveryError::Decode(err) => write!(f, "cbors serialization error: {}", err),
            DiscoveryError::Timeout => write!(f, "discovery timed out"),
            DiscoveryError::PermissionDenied => {
                write!(f, "discovery channel permission denied")
            }
            DiscoveryError::MulticastUnavailable => {
                write!(f, "multicast discovery unavailable")
            }
            DiscoveryError::BroadcastBlocked => write!(f, "broadcast discovery blocked"),
        }
    }
}

impl std::error::Error for DiscoveryError {}

impl From<io::Error> for DiscoveryError {
    fn from(err: io::Error) -> Self {
        match err.kind() {
            io::ErrorKind::TimedOut | io::ErrorKind::WouldBlock => DiscoveryError::Timeout,
            _ => DiscoveryError::Io(err),
        }
    }
}

impl From<serde_cbor::Error> for DiscoveryError {
    fn from(err: serde_cbor::Error) -> Self {
        DiscoveryError::Decode(err)
    }
}

/// The outcome of a discovery request.
pub struct DiscoveryOutcome {
    pub reply: DiscoveryReply,
    pub peer: SocketAddr,
}

/// Stateless discovery helper that wraps the protocol request/response models.
pub struct DiscoveryClient {
    socket: UdpSocket,
    remote_addr: SocketAddr,
    prefer_multicast: bool,
}

impl DiscoveryClient {
    /// Creates a client that will send discovery packets to `remote_addr`.
    pub fn new(options: DiscoveryClientOptions) -> Result<Self, DiscoveryError> {
        let socket = UdpSocket::bind(options.local_addr)?;
        socket.set_broadcast(true)?;
        socket.set_read_timeout(Some(options.timeout))?;
        Ok(Self {
            socket,
            remote_addr: options.remote_addr,
            prefer_multicast: options.prefer_multicast,
        })
    }

    /// Sends a discovery payload with the requested capability names and waits for a reply.
    pub fn discover(&self, requested: &[String]) -> Result<DiscoveryOutcome, DiscoveryError> {
        let mut nonce = vec![0u8; 32];
        OsRng.fill_bytes(&mut nonce);
        let request = DiscoveryRequest::new(requested.to_vec(), nonce.clone());
        let payload = serde_cbor::to_vec(&request)?;

        let mut send_error: Option<DiscoveryError> = None;
        let mut sent = false;
        for target in self.discovery_targets() {
            match self.socket.send_to(&payload, target) {
                Ok(_) => {
                    trace!(target = %target, "discovery payload sent");
                    sent = true;
                    break;
                }
            Err(err) => {
                let kind = err.kind();
                let mapped = self.map_send_error(err, target);
                trace!(target = %target, error = %mapped, "discovery send failed");
                send_error = Some(mapped);
                if !self.should_continue_after_error(&kind) {
                    return Err(send_error.unwrap());
                }
            }
            }
        }

        if !sent {
            return Err(send_error.unwrap_or_else(|| DiscoveryError::PermissionDenied));
        }

        let mut buf = vec![0u8; 2048];
        let (len, peer) = match self.socket.recv_from(&mut buf) {
            Ok(res) => res,
            Err(err) => return Err(self.map_recv_error(err)),
        };
        let reply: DiscoveryReply = serde_cbor::from_slice(&buf[..len])?;
        Ok(DiscoveryOutcome { reply, peer })
    }

    fn discovery_targets(&self) -> Vec<SocketAddr> {
        let mut targets = Vec::new();

        if self.prefer_multicast {
            if let Ok(addr) = DEFAULT_MULTICAST_IPV4.parse() {
                targets.push(addr);
            }
            if let Ok(addr) = DEFAULT_MULTICAST_IPV6.parse() {
                targets.push(addr);
            }
        }

        if !targets.contains(&self.remote_addr) {
            targets.push(self.remote_addr);
        }

        if self.remote_addr.ip().is_ipv4() {
            if let Ok(addr) = DEFAULT_BROADCAST_IPV4.parse() {
                if !targets.contains(&addr) {
                    targets.push(addr);
                }
            }
        }

        targets
    }

    fn map_send_error(&self, err: io::Error, target: SocketAddr) -> DiscoveryError {
        match err.kind() {
            io::ErrorKind::PermissionDenied => {
                if target.ip().is_multicast() {
                    DiscoveryError::MulticastUnavailable
                } else if is_broadcast_addr(target.ip()) {
                    DiscoveryError::BroadcastBlocked
                } else {
                    DiscoveryError::PermissionDenied
                }
            }
            io::ErrorKind::ConnectionReset | io::ErrorKind::WouldBlock => {
                DiscoveryError::PermissionDenied
            }
            _ => DiscoveryError::Io(err),
        }
    }

    fn map_recv_error(&self, err: io::Error) -> DiscoveryError {
        match err.kind() {
            io::ErrorKind::TimedOut => DiscoveryError::Timeout,
            io::ErrorKind::PermissionDenied
            | io::ErrorKind::ConnectionReset
            | io::ErrorKind::WouldBlock => DiscoveryError::PermissionDenied,
            _ => DiscoveryError::Io(err),
        }
    }

    fn should_continue_after_error(&self, kind: &io::ErrorKind) -> bool {
        matches!(
            kind,
            io::ErrorKind::PermissionDenied
                | io::ErrorKind::WouldBlock
                | io::ErrorKind::ConnectionReset
        )
    }
}

fn is_broadcast_addr(ip: IpAddr) -> bool {
    matches!(ip, IpAddr::V4(addr) if addr.is_broadcast())
}