motorcortex-rust 0.5.0

use crate::connection::{PipeEvent, PipeEventHandler};
use crate::error::{MotorcortexError, Result};
use crate::ConnectionOptions;
use nng_c_sys::{
    nng_close, nng_dialer, nng_pipe, nng_pipe_ev, nng_socket, nng_tls_config,
    nng_tls_config_server_name, nng_url_free, nng_url_parse,
};
use std::ffi::{CStr, CString};
use std::ptr;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

/// Wrapped context the NNG pipe-event callback operates on.
///
/// Both the initial connect wait (via `dialed`) and the driver's
/// post-connect state transitions (via `handler`) are served from
/// here — the callback just translates the C event code into a
/// [`PipeEvent`] and fans it out.
struct PipeContext {
    /// Flipped to `true` on the first `ADD_POST`. The initial
    /// `connect()` call polls this to unblock its wait loop once the
    /// dial finishes.
    dialed: AtomicBool,
    /// Handler installed by the driver; invoked on every
    /// `ADD_POST` / `REM_POST` event after the initial connect has
    /// returned.
    handler: PipeEventHandler,
}

extern "C" fn my_pipe_callback(
    _: nng_pipe,
    event: nng_pipe_ev::Type,
    context: *mut core::ffi::c_void,
) {
    if context.is_null() {
        return;
    }
    // SAFETY: the context is the leaked `Arc<PipeContext>` installed
    // by `connect`. We reconstruct the Arc, peek, then `forget` it so
    // the refcount isn't bumped down. The matching `Arc::from_raw` in
    // the ConnectionManager's Drop balances the original `into_raw`.
    let ctx = unsafe { Arc::from_raw(context as *const PipeContext) };
    match event {
        nng_pipe_ev::NNG_PIPE_EV_ADD_POST => {
            ctx.dialed.store(true, Ordering::SeqCst);
            (ctx.handler)(PipeEvent::Added);
        }
        nng_pipe_ev::NNG_PIPE_EV_REM_POST => {
            (ctx.handler)(PipeEvent::Removed);
        }
        _ => {}
    }
    // Keep the Arc's ref count intact — `into_raw` already bumped
    // the count; `from_raw` dropped it; we forget to restore.
    std::mem::forget(ctx);
}

pub struct ConnectionManager {
    /// NNG socket instance for communications
    pub sock: Option<nng_socket>,
    /// Optional TLS configuration for secure communication
    pub tls_cfg: Option<*mut nng_tls_config>,
    /// Optional NNG dialer to manage connection lifecycle
    pub dialer: Option<nng_dialer>,
    /// Raw pointer to the leaked `Arc<PipeContext>` registered as
    /// the pipe-notify context. Owned by this struct; freed on Drop
    /// after `disconnect()` has closed the socket.
    pipe_context_ptr: Option<*mut core::ffi::c_void>,
}

// SAFETY: ConnectionManager's raw pointers (nng_tls_config) are heap allocations
// managed by NNG that are safe to access from any thread. Sole ownership is
// maintained — no concurrent access is possible through move semantics.
unsafe impl Send for ConnectionManager {}

impl ConnectionManager {
    pub fn new() -> Self {
        Self {
            sock: None,
            tls_cfg: None,
            dialer: None,
            pipe_context_ptr: None,
        }
    }

    pub fn connect(
        &mut self,
        url: &str,
        connection_options: ConnectionOptions,
        socket_opener: unsafe extern "C" fn(*mut nng_c_sys::nng_socket) -> core::ffi::c_int,
        on_pipe_event: PipeEventHandler,
    ) -> Result<()> {
        // Open the socket
        let sock = unsafe {
            let mut sock: nng_c_sys::nng_socket = std::mem::zeroed();
            let rv = socket_opener(&mut sock);
            if rv != 0 {
                return Err(MotorcortexError::Connection(format!(
                    "Failed to open socket: {} ({})",
                    self.nng_error_to_string(nng_c_sys::nng_strerror(rv)),
                    rv
                )));
            }
            sock
        };
        self.sock = Some(sock);

        unsafe {
            // Set connection timeout for the dialer
            let rv_dialer_timeout = nng_c_sys::nng_setopt_ms(
                self.sock.unwrap(),
                nng_c_sys::NNG_OPT_RECVTIMEO.as_ptr() as *const core::ffi::c_char,
                connection_options.conn_timeout_ms as i32,
            );
            if rv_dialer_timeout != 0 {
                return Err(MotorcortexError::Connection(format!(
                    "Failed to set recv timeout: {} ({})",
                    self.nng_error_to_string(nng_c_sys::nng_strerror(rv_dialer_timeout)),
                    rv_dialer_timeout,
                )));
            }
        };

        if !connection_options.certificate.is_empty() {
            // Reject embedded NUL bytes instead of panicking — the
            // cert path comes from user code and an interior NUL
            // would make nng_tls_config_ca_file() read past the
            // intended C string anyway.
            let ca_file = CString::new(connection_options.certificate).map_err(|_| {
                MotorcortexError::Connection(
                    "certificate path contains an embedded NUL byte".into(),
                )
            })?;
            self.tls_cfg = unsafe {
                // Create tls config
                let mut tls_cfg: *mut nng_c_sys::nng_tls_config = ptr::null_mut();
                let mut rv = nng_c_sys::nng_tls_config_alloc(
                    &mut tls_cfg,
                    nng_c_sys::nng_tls_mode::NNG_TLS_MODE_CLIENT,
                );
                if rv != 0 {
                    return Err(MotorcortexError::Connection(format!(
                        "Failed to allocate certificate: {} ({})",
                        self.nng_error_to_string(nng_c_sys::nng_strerror(rv)),
                        rv
                    )));
                }

                rv = nng_c_sys::nng_tls_config_ca_file(tls_cfg, ca_file.as_ptr());
                if rv != 0 {
                    return Err(MotorcortexError::Connection(format!(
                        "Failed to load certificate from path {}, error: {} ({})",
                        ca_file.into_string().unwrap(),
                        self.nng_error_to_string(nng_c_sys::nng_strerror(rv)),
                        rv
                    )));
                }

                let mut nurl: *mut nng_c_sys::nng_url = ptr::null_mut();
                // Same rationale as the cert path above — interior
                // NUL would misbehave in the C API. Propagate as a
                // Connection error.
                let c_url = CString::new(url).map_err(|_| {
                    MotorcortexError::Connection(
                        "URL contains an embedded NUL byte".into(),
                    )
                })?;
                rv = nng_url_parse(&mut nurl, c_url.as_ptr());
                if rv == 0 && !nurl.is_null() {
                    nng_tls_config_server_name(tls_cfg, (*nurl).u_hostname);
                    nng_url_free(nurl);
                } else {
                    // URL parse failures surface via the matching
                    // dialer creation below; there's nothing useful
                    // to print here, and stdout in a library is
                    // impolite.
                }

                rv = nng_c_sys::nng_socket_set_ptr(
                    self.sock.unwrap(),
                    nng_c_sys::NNG_OPT_TLS_CONFIG.as_ptr() as *const core::ffi::c_char,
                    tls_cfg as *mut _,
                );

                if rv != 0 {
                    return Err(MotorcortexError::Connection(format!(
                        "Failed to apply certificate to the socket: {} ({})",
                        self.nng_error_to_string(nng_c_sys::nng_strerror(rv)),
                        rv
                    )));
                }
                Some(tls_cfg)
            };
        }

        // Free any previous pipe context from an earlier connect on
        // this same ConnectionManager before installing a fresh one.
        self.free_pipe_context();
        let pipe_ctx = Arc::new(PipeContext {
            dialed: AtomicBool::new(false),
            handler: on_pipe_event,
        });
        let connection_established = Arc::clone(&pipe_ctx);
        let pipe_context_ptr = Arc::into_raw(pipe_ctx) as *mut core::ffi::c_void;
        self.pipe_context_ptr = Some(pipe_context_ptr);
        self.configure_pipe_notifications(self.sock.unwrap(), pipe_context_ptr)?;

        self.dialer = unsafe {
            let mut dialer: nng_c_sys::nng_dialer = std::mem::zeroed();
            let c_url = CString::new(url).map_err(|_| {
                MotorcortexError::Connection(
                    "URL contains an embedded NUL byte".into(),
                )
            })?;
            let rv = nng_c_sys::nng_dialer_create(&mut dialer, self.sock.unwrap(), c_url.as_ptr());
            if rv != 0 {
                return Err(MotorcortexError::Connection(format!(
                    "Failed to create dialer: {} ({})",
                    self.nng_error_to_string(nng_c_sys::nng_strerror(rv)),
                    rv
                )));
            }

            // Configure NNG's built-in redial window. When
            // `reconnect == false` we pin both min and max to 0 so
            // NNG treats a failed dial as terminal and doesn't
            // retry in the background.
            let (min_ms, max_ms): (i32, i32) = if connection_options.reconnect {
                (
                    connection_options.reconnect_min.as_millis() as i32,
                    connection_options.reconnect_max.as_millis() as i32,
                )
            } else {
                (0, 0)
            };
            let _ = nng_c_sys::nng_dialer_setopt_ms(
                dialer,
                nng_c_sys::NNG_OPT_RECONNMINT.as_ptr() as *const core::ffi::c_char,
                min_ms,
            );
            let _ = nng_c_sys::nng_dialer_setopt_ms(
                dialer,
                nng_c_sys::NNG_OPT_RECONNMAXT.as_ptr() as *const core::ffi::c_char,
                max_ms,
            );

            let rv = nng_c_sys::nng_dialer_start(dialer, nng_c_sys::NNG_FLAG_NONBLOCK);
            if rv != 0 {
                return Err(MotorcortexError::Connection(format!(
                    "Failed to start dialer: {} ({})",
                    self.nng_error_to_string(nng_c_sys::nng_strerror(rv)),
                    rv
                )));
            }

            let timeout =
                Instant::now() + Duration::from_millis(connection_options.conn_timeout_ms as u64);
            while !connection_established.dialed.load(Ordering::SeqCst) {
                if Instant::now() > timeout {
                    return Err(MotorcortexError::Connection(
                        "Connection timeout".to_string(),
                    ));
                }
                std::thread::sleep(Duration::from_millis(50));
            }
            Some(dialer)
        };

        Ok(())
    }

    /// Pin the dialer's reconnect window to zero so NNG treats the
    /// next failed dial as terminal and stops redialling in the
    /// background. Used by the driver's max-reconnect-attempts
    /// guard.
    ///
    /// Safe to call at any time — if no dialer is registered (e.g.
    /// never connected), this is a no-op.
    pub fn disable_dialer_reconnect(&self) {
        if let Some(dialer) = self.dialer {
            // SAFETY: `dialer` is a handle (an `nng_dialer` struct
            // containing an id); setting these two integer options
            // to 0 takes no caller-owned pointers. Safe to call on a
            // live or closed dialer per NNG docs — the closed case
            // returns an error we ignore.
            unsafe {
                let _ = nng_c_sys::nng_dialer_setopt_ms(
                    dialer,
                    nng_c_sys::NNG_OPT_RECONNMINT.as_ptr() as *const core::ffi::c_char,
                    0,
                );
                let _ = nng_c_sys::nng_dialer_setopt_ms(
                    dialer,
                    nng_c_sys::NNG_OPT_RECONNMAXT.as_ptr() as *const core::ffi::c_char,
                    0,
                );
            }
        }
    }

    pub fn disconnect(&mut self) -> Result<()> {
        // SAFETY: each call takes a handle this struct owns
        // exclusively (the `.take()` enforces single-close).
        // `nng_close` on a live socket joins its internal threads
        // and guarantees no callbacks fire after it returns, which
        // is what `free_pipe_context` relies on.
        unsafe {
            // Close the socket if it exists
            if let Some(sock) = self.sock.take() {
                let rv = nng_close(sock);
                if rv != 0 {
                    return Err(MotorcortexError::Connection(format!(
                        "Failed to close socket: {} ({})",
                        self.nng_error_to_string(nng_c_sys::nng_strerror(rv)),
                        rv
                    )));
                }
            }

            // Free the TLS configuration if it exists
            if let Some(tls_cfg) = self.tls_cfg.take() {
                nng_c_sys::nng_tls_config_free(tls_cfg);
            }
        }
        Ok(())
    }

    /// Release the leaked `Arc<PipeContext>` installed during
    /// `connect()`. Must be called after the socket is closed so NNG
    /// can't fire any more callbacks into the dropped context.
    fn free_pipe_context(&mut self) {
        if let Some(ptr) = self.pipe_context_ptr.take() {
            unsafe {
                // Reconstruct the Arc to drop it — balances the
                // `Arc::into_raw` in `connect()`.
                let _ = Arc::from_raw(ptr as *const PipeContext);
            }
        }
    }

    /// Utility function to convert NNG error messages to human-readable Rust strings.
    fn nng_error_to_string(&self, err_msg: *const core::ffi::c_char) -> String {
        // SAFETY: `err_msg` comes from `nng_strerror` which returns a
        // pointer to a static NUL-terminated C string. We null-check
        // it defensively and do not retain it past the `to_owned`.
        unsafe {
            if err_msg.is_null() {
                return "Unknown error".to_string();
            }
            let c_str = CStr::from_ptr(err_msg);
            c_str.to_string_lossy().into_owned()
        }
    }

    fn configure_pipe_notifications(
        &mut self,
        sock: nng_c_sys::nng_socket,
        state_update_context: *mut core::ffi::c_void,
    ) -> Result<()> {
        // SAFETY: `state_update_context` is an `Arc::into_raw`'d
        // `PipeContext` owned by this struct — it stays alive until
        // `free_pipe_context` runs in `Drop` (after the socket is
        // closed, so no in-flight callbacks can dereference freed
        // memory). `my_pipe_callback` is a `extern "C" fn` we own.
        unsafe {
            let mut er: i32;

            er = nng_c_sys::nng_pipe_notify(
                sock,
                nng_pipe_ev::NNG_PIPE_EV_ADD_PRE,
                Some(my_pipe_callback),
                state_update_context as *mut _,
            );
            if er != 0 {
                return Err(MotorcortexError::Connection(
                    "Failed to set pipe notification (ADD_PRE)".to_string(),
                ));
            }

            er = nng_c_sys::nng_pipe_notify(
                sock,
                nng_pipe_ev::NNG_PIPE_EV_ADD_POST,
                Some(my_pipe_callback),
                state_update_context as *mut _,
            );
            if er != 0 {
                return Err(MotorcortexError::Connection(
                    "Failed to set pipe notification (ADD_POST)".to_string(),
                ));
            }

            er = nng_c_sys::nng_pipe_notify(
                sock,
                nng_pipe_ev::NNG_PIPE_EV_REM_POST,
                Some(my_pipe_callback),
                state_update_context as *mut _,
            );
            if er != 0 {
                return Err(MotorcortexError::Connection(
                    "Failed to set pipe notification (REM_POST)".to_string(),
                ));
            }
        }
        Ok(())
    }
}

impl Drop for ConnectionManager {
    fn drop(&mut self) {
        // `disconnect` closes the socket (and the dialer with it),
        // which stops NNG from invoking any more pipe-notify
        // callbacks. Free the leaked context *after* that so no
        // callback ever sees freed memory.
        let _ = self.disconnect();
        self.free_pipe_context();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::connection::ConnectionOptions;
    use std::sync::Arc;

    fn noop_pipe_handler() -> PipeEventHandler {
        Arc::new(|_event| {})
    }

    #[test]
    fn connect_rejects_cert_path_with_nul_byte() {
        // An interior NUL in the cert path used to panic via
        // `CString::new(...).unwrap()`; it now surfaces as a clean
        // Connection error.
        let mut cm = ConnectionManager::new();
        let opts = ConnectionOptions::new("bad\0path.crt".into(), 100, 100);
        let err = cm
            .connect(
                "wss://127.0.0.1:0",
                opts,
                nng_c_sys::nng_req0_open,
                noop_pipe_handler(),
            )
            .expect_err("NUL byte must be rejected");
        assert!(
            matches!(err, MotorcortexError::Connection(ref msg) if msg.contains("NUL")),
            "expected Connection error mentioning NUL, got {err:?}",
        );
    }

    #[test]
    fn connect_rejects_url_with_nul_byte() {
        // Same story for the URL argument.
        let mut cm = ConnectionManager::new();
        let opts = ConnectionOptions::new(String::new(), 100, 100);
        let err = cm
            .connect(
                "wss://127.0.0.1\0:0",
                opts,
                nng_c_sys::nng_req0_open,
                noop_pipe_handler(),
            )
            .expect_err("NUL byte must be rejected");
        assert!(
            matches!(err, MotorcortexError::Connection(ref msg) if msg.contains("NUL")),
            "expected Connection error mentioning NUL, got {err:?}",
        );
    }
}