vrchat_osc 2.2.0

mod fetch;
mod mdns;
mod oscquery;

pub use oscquery::models;
pub use rosc;

pub use fetch::Error as FetchError;
pub use mdns::Error as MdnsError;
pub use oscquery::Error as OscQueryError;

use crate::fetch::fetch;
use crate::oscquery::OscQuery;

use futures::{stream, StreamExt};
use hickory_proto::rr::Name;
use oscquery::models::{HostInfo, OscNode, OscRootNode};
use rosc::OscPacket;
use std::str::FromStr;
use std::{
    collections::HashMap,
    net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4},
    sync::Arc,
};
use tokio::{
    net::UdpSocket,
    sync::{mpsc, RwLock},
    task::JoinHandle,
};
use wildmatch::WildMatch;

const OSC_PACKET_BUFFER_SIZE: usize = 65535; // Max UDP packet size

/// Defines the possible errors that can occur within the VRChatOSC library.
#[derive(thiserror::Error, Debug)]
pub enum Error {
    #[error("OSC error: {0}")]
    OscError(#[from] rosc::OscError),
    #[error("OSCQuery error: {0}")]
    OscQueryError(#[from] oscquery::Error),
    #[error("mDNS error: {0}")]
    MdnsError(#[from] mdns::Error),
    #[error("Hickory DNS protocol error: {0}")]
    HickoryError(#[from] hickory_proto::ProtoError),
    #[error("I/O error: {0}")]
    IoError(#[from] std::io::Error),
    #[error("Fetch error: {0}")]
    FetchError(#[from] fetch::Error),
    #[error("No valid network interface found: no non-loopback IPv4 address available")]
    NoValidInterface,
}

/// Holds handles related to a registered OSC service.
struct ServiceHandle {
    /// Join handle for the OSC listening task.
    osc: JoinHandle<()>,
    /// The OSCQuery server instance.
    osc_query: OscQuery,
}

/// Represents the type of OSC service discovered.
pub enum ServiceType {
    /// OSC service type.
    Osc(String, SocketAddr),
    /// OSCQuery service type.
    OscQuery(String, SocketAddr),
}

/// Main struct for managing VRChat OSC communication and service discovery.
pub struct VRChatOSC {
    /// Socket for sending OSC messages.
    send_socket: UdpSocket,
    /// The destination IP address for OSC messages (VRChat's IP address).
    osc_ip: Arc<RwLock<IpAddr>>,
    /// mDNS client instance for service discovery.
    mdns: mdns::Mdns,
    /// Stores registered service handles, mapping service name to its handle.
    service_handles: Arc<RwLock<HashMap<String, ServiceHandle>>>,
    /// Callback function to be executed when a new mDNS service is discovered.
    on_service_discovered_callback:
        Arc<RwLock<Option<Arc<dyn Fn(ServiceType) + Send + Sync + 'static>>>>,
}

/// Finds a non-loopback IPv4 interface address that is up.
///
/// This function enumerates network interfaces and selects one that is:
/// - IPv4
/// - Not a loopback address
/// - In an "up" state
///
/// # Returns
/// A non-loopback IPv4 address, or None if no suitable interface is found.
fn find_non_loopback_ipv4() -> Option<IpAddr> {
    if let Ok(interfaces) = if_addrs::get_if_addrs() {
        for iface in interfaces {
            if let std::net::IpAddr::V4(ipv4) = iface.addr.ip() {
                if !ipv4.is_loopback() {
                    return Some(IpAddr::V4(ipv4));
                }
            }
        }
    }
    None
}

/// Finds the local IP address that can reach the given destination IP.
///
/// This function connects a temporary UDP socket to the destination to let the
/// OS routing table determine the appropriate interface, then returns the local address.
///
/// # Arguments
/// * `dest_ip` - The destination IP address to reach.
///
/// # Returns
/// The local IP address on the interface that can reach the destination.
fn find_local_ip_for_destination(dest_ip: IpAddr) -> Result<IpAddr, Error> {
    // If the destination IP is a loopback address, return it directly.
    if dest_ip.is_loopback() {
        return Ok(dest_ip);
    }

    // Check if the destination IP is one of our local interface addresses.
    if let Ok(interfaces) = if_addrs::get_if_addrs() {
        for iface in interfaces {
            if iface.addr.ip() == dest_ip {
                return Ok(dest_ip);
            }
        }
    }

    // Create a UDP socket and connect to the destination (without sending data).
    // The OS will determine the correct interface based on routing table.
    let socket = match dest_ip {
        IpAddr::V4(_) => std::net::UdpSocket::bind("0.0.0.0:0")?,
        IpAddr::V6(_) => std::net::UdpSocket::bind("[::]:0")?,
    };

    socket.connect((dest_ip, 0))?;
    Ok(socket.local_addr()?.ip())
}

/// Sanitizes the service name to be compatible with mDNS and VRChat requirements.
///
/// Rules:
/// - Keep ASCII alphanumerics, hyphens, and all non-ASCII characters.
/// - Replace any other ASCII characters (including spaces and symbols) with '-'.
/// - Replace Unicode control characters with '-'.
/// - Replace uppercase ASCII characters with lowercase.
fn sanitize_service_name(name: &str) -> String {
    name.chars()
        .map(|c| {
            if (c.is_ascii() && !c.is_ascii_alphanumeric()) || c.is_control() {
                '-'
            } else if c.is_ascii_uppercase() {
                c.to_ascii_lowercase()
            } else {
                c
            }
        })
        .collect()
}

impl VRChatOSC {
    /// Creates a new `VRChatOSC` instance.
    ///
    /// # Arguments
    /// * `osc_ip` - Optional destination IP address for OSC messages.
    ///   If `None`, it will attempt to automatically find a suitable network interface.
    pub async fn new(osc_ip: Option<IpAddr>) -> Result<Arc<VRChatOSC>, Error> {
        let osc_ip = match osc_ip {
            Some(ip) => ip,
            None => find_non_loopback_ipv4().ok_or(Error::NoValidInterface)?,
        };
        let socket = UdpSocket::bind(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0)).await?;

        // Create an mpsc channel for notifying about discovered mDNS services.
        let (discover_notifier_tx, mut discover_notifier_rx) = mpsc::channel(8);

        // Derive the advertised IP (our local interface IP) from the OSC destination IP.
        // This finds the local IP address on the same network as the destination.
        let advertised_ip = find_local_ip_for_destination(osc_ip)?;

        // Initialize the mDNS client with our local interface IP for announcements.
        let mdns_client = mdns::Mdns::new(discover_notifier_tx, advertised_ip).await?;

        // Start following OSC services and OSCQuery JSON services on the local network.
        let _ = mdns_client
            .follow(Name::from_ascii("_osc._udp.local.")?)
            .await;
        let _ = mdns_client
            .follow(Name::from_ascii("_oscjson._tcp.local.")?)
            .await;

        // Prepare a shared storage for the service discovered callback.
        let on_service_discovered_callback = Arc::new(RwLock::new(
            None::<Arc<dyn Fn(ServiceType) + Send + Sync + 'static>>,
        ));
        let callback_arc_clone = on_service_discovered_callback.clone();

        // Spawn a new asynchronous task to listen for service discovery notifications.
        // This task will own the `discover_notifier_rx` (receiver end of the mpsc channel).
        tokio::spawn(async move {
            // Continuously try to receive messages from the discovery notification channel.
            loop {
                if let Some((service_name, socket_addr)) = discover_notifier_rx.recv().await {
                    let callback_guard = callback_arc_clone.read().await;
                    // If a callback is registered, invoke it with the service name and address.
                    if let Some(callback) = callback_guard.as_ref() {
                        if service_name.trim_to(3).to_utf8() == "_osc._udp.local." {
                            callback(ServiceType::Osc(service_name.to_utf8(), socket_addr));
                        } else if service_name.trim_to(3).to_utf8() == "_oscjson._tcp.local." {
                            callback(ServiceType::OscQuery(service_name.to_utf8(), socket_addr));
                        }
                    }
                }
            }
        });

        Ok(Arc::new(VRChatOSC {
            send_socket: socket,
            osc_ip: Arc::new(RwLock::new(osc_ip)),
            mdns: mdns_client,
            service_handles: Arc::new(RwLock::new(HashMap::new())),
            on_service_discovered_callback,
        }))
    }

    /// Sets the destination IP address for OSC messages.
    ///
    /// # Arguments
    /// * `ip` - The new destination IP address. If `None`, it will attempt to automatically find a suitable network interface.
    pub async fn set_osc_ip(&self, ip: Option<IpAddr>) -> Result<(), Error> {
        let ip = match ip {
            Some(ip) => ip,
            None => find_non_loopback_ipv4().ok_or(Error::NoValidInterface)?,
        };
        *self.osc_ip.write().await = ip;
        let advertised_ip = find_local_ip_for_destination(ip)?;
        self.mdns.set_advertised_ip(advertised_ip).await;
        Ok(())
    }

    /// Gets the currently configured OSC destination IP address.
    pub async fn get_osc_ip(&self) -> IpAddr {
        *self.osc_ip.read().await
    }

    /// Registers a callback to be invoked when a new OSC service is discovered on the network.
    ///
    /// # Arguments
    /// * `callback` - A function or closure called with the discovered service details.
    pub async fn on_connect<F>(&self, callback: F)
    where
        F: Fn(ServiceType) + Send + Sync + 'static,
    {
        let mut callback_guard = self.on_service_discovered_callback.write().await;
        *callback_guard = Some(Arc::new(callback));
    }

    /// Registers a local OSC service to be discoverable by other clients.
    ///
    /// # Arguments
    /// * `service_name` - The name of the service to register.
    /// * `parameters` - The OSC address space and parameters for this service.
    /// * `handler` - A function called when an OSC packet is received for this service.
    pub async fn register<F>(
        &self,
        service_name: &str,
        parameters: OscRootNode,
        handler: F,
    ) -> Result<(), Error>
    where
        F: Fn(OscPacket) + Send + 'static,
    {
        // Start OSC server (UDP listener)
        // Bind to an ephemeral port on all interfaces.
        let socket = UdpSocket::bind(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0)).await?;
        let osc_local_addr = socket.local_addr()?; // Get the actual address it bound to.

        // Spawn a task to handle incoming OSC packets.
        let osc_handle = tokio::spawn(async move {
            let mut buf = [0; OSC_PACKET_BUFFER_SIZE]; // Buffer for receiving OSC packets.
            loop {
                // Wait to receive data on the socket.
                match socket.recv_from(&mut buf).await {
                    Ok((len, addr)) => {
                        // Decode the received UDP data into an OSC packet.
                        if let Ok((_, packet)) = rosc::decoder::decode_udp(&buf[..len]) {
                            handler(packet); // Call the provided handler with the decoded packet.
                        } else {
                            log::debug!("Failed to decode OSC packet from {}", addr);
                        }
                    }
                    Err(e) => {
                        if e.kind() == std::io::ErrorKind::ConnectionReset
                            || e.kind() == std::io::ErrorKind::BrokenPipe
                        {
                            log::warn!("Socket connection error ({}). Task for {:?} might need to be restarted or interface is down.", e, socket.local_addr().ok());
                            break;
                        } else {
                            log::warn!(
                                "Failed to receive data on OSC socket {:?}: {}",
                                socket.local_addr().ok(),
                                e
                            );
                        }
                        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
                        continue;
                    }
                }
            }
        });

        // Start OSCQuery server (HTTP server)
        let host_info = HostInfo::new(
            service_name.to_string(),
            osc_local_addr.ip(),   // Use the IP of the OSC server.
            osc_local_addr.port(), // Use the port of the OSC server.
        );
        let mut osc_query = OscQuery::new(host_info, parameters);
        // Serve OSCQuery on an ephemeral port on all interfaces.
        let osc_query_local_addr = osc_query
            .serve(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0))
            .await?;

        // Create mDNS service announcements.
        let sanitized_service_name = sanitize_service_name(service_name);

        // Register the OSC and OSCQuery services with mDNS.
        self.mdns
            .register(
                Name::from_str(&format!("{}._osc._udp.local.", sanitized_service_name))?,
                osc_local_addr.port(),
            )
            .await?;
        self.mdns
            .register(
                Name::from_str(&format!("{}._oscjson._tcp.local.", sanitized_service_name))?,
                osc_query_local_addr.port(),
            )
            .await?;

        // Save service handles for later management (e.g., unregistering).
        let mut handles = self.service_handles.write().await;
        handles.insert(
            service_name.to_string(),
            ServiceHandle {
                osc: osc_handle,
                osc_query,
            },
        );
        Ok(())
    }

    /// Unregisters an OSC service and stops its servers.
    ///
    /// # Arguments
    /// * `service_name` - The name of the service to unregister.
    pub async fn unregister(&self, service_name: &str) -> Result<(), Error> {
        let sanitized_service_name = sanitize_service_name(service_name);
        // Remove the service from our tracking.
        let mut service_handles_map = self.service_handles.write().await;
        if let Some(mut service_handle_entry) = service_handles_map.remove(service_name) {
            // Unregister from mDNS.
            self.mdns
                .unregister(Name::from_str(&format!(
                    "{}._osc._udp.local.",
                    sanitized_service_name
                ))?)
                .await?;
            self.mdns
                .unregister(Name::from_str(&format!(
                    "{}._oscjson._tcp.local.",
                    sanitized_service_name
                ))?)
                .await?;

            // Stop the associated tasks/servers.
            service_handle_entry.osc.abort(); // Abort the OSC listening task.
            service_handle_entry.osc_query.shutdown(); // Gracefully shutdown the OSCQuery server.
        }
        Ok(())
    }

    /// Sends an OSC packet to services matching a name pattern.
    ///
    /// # Arguments
    /// * `packet` - The `OscPacket` to send.
    /// * `to` - A pattern (e.g., "VRChat-Client-*") to match against discovered service names.
    pub async fn send(&self, packet: OscPacket, to: &str) -> Result<(), Error> {
        // Find services matching the pattern. The matching logic is within `find_service`.
        // The closure provided to `find_service` determines if a service (by its Name) matches.
        let services = self
            .mdns
            .find_service(|name, _| {
                // `WildMatch` performs glob-style pattern matching.
                WildMatch::new(&format!("{}._osc._udp.local.", to)).matches(&name.to_utf8())
            })
            .await;

        if services.is_empty() {
            log::info!("No mDNS services found matching the expression: {}", to);
            return Ok(());
        }

        // Encode the OSC packet into bytes.
        let msg_buf = rosc::encoder::encode(&packet)?;
        // Send the packet to all found services.
        let send_futs = services
            .into_iter()
            .map(|(_, addr)| self.send_socket.send_to(&msg_buf, addr));
        let results = futures::future::join_all(send_futs).await;
        for res in results {
            res?;
        }

        Ok(())
    }

    /// Sends an OSC packet to a specific socket address.
    ///
    /// # Arguments
    /// * `packet` - The `OscPacket` to send.
    /// * `addr` - The `SocketAddr` to send the packet to.
    pub async fn send_to_addr(&self, packet: OscPacket, addr: SocketAddr) -> Result<(), Error> {
        let msg_buf = rosc::encoder::encode(&packet)?;
        self.send_socket.send_to(&msg_buf, addr).await?;
        Ok(())
    }

    /// Retrieves a specific parameter from services matching a name pattern.
    ///
    /// # Arguments
    /// * `method` - The OSC path of the parameter (e.g., "/avatar/parameters/MyParam").
    /// * `from` - A pattern (e.g., "VRChat-Client-*") to match against discovered service names.
    ///
    /// # Returns
    /// A list of matching service names and their corresponding parameter values.
    pub async fn get_parameter(
        &self,
        method: &str,
        from: &str,
    ) -> Result<Vec<(String, OscNode)>, Error> {
        // Find services matching the pattern. The matching logic is within `find_service`.
        // The closure provided to `find_service` determines if a service matches.
        let services = self
            .mdns
            .find_service(|name, _| {
                WildMatch::new(&format!("{}._oscjson._tcp.local.", from)).matches(&name.to_utf8())
            })
            .await;

        if services.is_empty() {
            log::info!(
                "No mDNS services found for get_parameter matching expression: {}",
                from
            );
            return Ok(Vec::new());
        }

        // Asynchronously fetch the parameter from all matching services.
        // `stream::iter` creates a stream from the services.
        // `map` transforms each service into a future that fetches the parameter.
        // `buffer_unordered(3)` allows up to 3 fetches to run concurrently.
        // `filter_map` discards any fetches that resulted in an error.
        // `collect` gathers all successful results into a Vec.
        let params = stream::iter(services)
            .map(|(name, addr)| async move {
                fetch::<_, OscNode>(addr, method)
                    .await
                    .map(|(param, _)| (name.to_utf8(), param))
            })
            .buffer_unordered(3)
            .filter_map(|res| async {
                if let Err(e) = &res {
                    log::warn!("Failed to fetch parameter: {:?}", e);
                }
                res.ok()
            })
            .collect::<Vec<_>>()
            .await;

        Ok(params)
    }

    /// Retrieves a specific parameter from a specific service address.
    ///
    /// # Arguments
    /// * `method` - The OSC path of the parameter (e.g., "/avatar/parameters/MyParam").
    /// * `addr` - The address of the service.
    ///
    /// # Returns
    /// The fetched parameter value.
    pub async fn get_parameter_from_addr(
        &self,
        method: &str,
        addr: SocketAddr,
    ) -> Result<OscNode, Error> {
        let (param, _url) = fetch::<_, OscNode>(addr, method).await?;
        Ok(param)
    }

    /// Shuts down all registered services and cleans up resources.
    /// This method should be called before the VRChatOSC instance is dropped
    /// to ensure graceful shutdown of asynchronous tasks and network services.
    pub async fn shutdown(&self) -> Result<(), Error> {
        let mut service_handles_map = self.service_handles.write().await;
        let service_names: Vec<String> = service_handles_map.keys().cloned().collect();

        for name in service_names {
            if let Some(mut handle) = service_handles_map.remove(&name) {
                let sanitized_service_name = sanitize_service_name(&name);
                // Attempt to unregister from mDNS. Errors are logged but not propagated to allow other services to shut down.
                if let Err(e) = self
                    .mdns
                    .unregister(Name::from_str(&format!(
                        "{}._osc._udp.local.",
                        sanitized_service_name
                    ))?)
                    .await
                {
                    log::error!("Failed to unregister OSC for {}: {}", name, e);
                }
                if let Err(e) = self
                    .mdns
                    .unregister(Name::from_str(&format!(
                        "{}._oscjson._tcp.local.",
                        sanitized_service_name
                    ))?)
                    .await
                {
                    log::error!("Failed to unregister OSCQuery for {}: {}", name, e);
                }

                handle.osc.abort();
                handle.osc_query.shutdown();
            }
        }
        Ok(())
    }

    /// Lists the names of all currently registered services.
    pub async fn list_services(&self) -> Vec<String> {
        let handles = self.service_handles.read().await;
        handles.keys().cloned().collect()
    }
}

impl Drop for VRChatOSC {
    fn drop(&mut self) {
        // Best-effort synchronous cleanup.
        // For robust cleanup, especially of async tasks and network resources,
        // the asynchronous `shutdown` method should be called explicitly.
        if let Ok(mut handles) = self.service_handles.try_write() {
            let service_names: Vec<String> = handles.keys().cloned().collect();
            for name in service_names {
                if let Some(mut service_handle) = handles.remove(&name) {
                    // mDNS unregistration cannot be reliably called here due to async and potential blocking.
                    service_handle.osc.abort();
                    // OscQuery::shutdown() is assumed to be synchronous or non-blocking here.
                    // If it's async, it cannot be .await-ed in drop.
                    service_handle.osc_query.shutdown();
                }
            }
        } else {
            // This might happen if the lock is poisoned or contended in a way not suitable for drop.
            // In a real application, this should be logged or handled appropriately.
            // Using log::error! or eprintln! here might be appropriate.
            // For now, we acknowledge that proper async shutdown is preferred.
            if !std::thread::panicking() {
                // Avoid double panic if already panicking
                log::warn!("VRChatOSC: Could not acquire lock on service_handles during drop. Explicitly call shutdown() for robust cleanup.");
            }
        }
    }
}