remoteingress/rust/crates/remoteingress-core/src/edge.rs

use std::collections::HashMap;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{mpsc, Mutex, RwLock};
use tokio::task::JoinHandle;
use tokio_rustls::TlsConnector;
use serde::{Deserialize, Serialize};

use remoteingress_protocol::*;

/// Edge configuration (hub-host + credentials only; ports come from hub).
#[derive(Debug, Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct EdgeConfig {
    pub hub_host: String,
    pub hub_port: u16,
    pub edge_id: String,
    pub secret: String,
}

/// Handshake config received from hub after authentication.
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct HandshakeConfig {
    listen_ports: Vec<u16>,
    #[serde(default = "default_stun_interval")]
    stun_interval_secs: u64,
}

fn default_stun_interval() -> u64 {
    300
}

/// Runtime config update received from hub via FRAME_CONFIG.
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct ConfigUpdate {
    listen_ports: Vec<u16>,
}

/// Events emitted by the edge.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
#[serde(tag = "type")]
pub enum EdgeEvent {
    TunnelConnected,
    TunnelDisconnected,
    #[serde(rename_all = "camelCase")]
    PublicIpDiscovered { ip: String },
    #[serde(rename_all = "camelCase")]
    PortsAssigned { listen_ports: Vec<u16> },
    #[serde(rename_all = "camelCase")]
    PortsUpdated { listen_ports: Vec<u16> },
}

/// Edge status response.
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct EdgeStatus {
    pub running: bool,
    pub connected: bool,
    pub public_ip: Option<String>,
    pub active_streams: usize,
    pub listen_ports: Vec<u16>,
}

/// The tunnel edge that listens for client connections and multiplexes them to the hub.
pub struct TunnelEdge {
    config: RwLock<EdgeConfig>,
    event_tx: mpsc::UnboundedSender<EdgeEvent>,
    event_rx: Mutex<Option<mpsc::UnboundedReceiver<EdgeEvent>>>,
    shutdown_tx: Mutex<Option<mpsc::Sender<()>>>,
    running: RwLock<bool>,
    connected: Arc<RwLock<bool>>,
    public_ip: Arc<RwLock<Option<String>>>,
    active_streams: Arc<AtomicU32>,
    next_stream_id: Arc<AtomicU32>,
    listen_ports: Arc<RwLock<Vec<u16>>>,
}

impl TunnelEdge {
    pub fn new(config: EdgeConfig) -> Self {
        let (event_tx, event_rx) = mpsc::unbounded_channel();
        Self {
            config: RwLock::new(config),
            event_tx,
            event_rx: Mutex::new(Some(event_rx)),
            shutdown_tx: Mutex::new(None),
            running: RwLock::new(false),
            connected: Arc::new(RwLock::new(false)),
            public_ip: Arc::new(RwLock::new(None)),
            active_streams: Arc::new(AtomicU32::new(0)),
            next_stream_id: Arc::new(AtomicU32::new(1)),
            listen_ports: Arc::new(RwLock::new(Vec::new())),
        }
    }

    /// Take the event receiver (can only be called once).
    pub async fn take_event_rx(&self) -> Option<mpsc::UnboundedReceiver<EdgeEvent>> {
        self.event_rx.lock().await.take()
    }

    /// Get the current edge status.
    pub async fn get_status(&self) -> EdgeStatus {
        EdgeStatus {
            running: *self.running.read().await,
            connected: *self.connected.read().await,
            public_ip: self.public_ip.read().await.clone(),
            active_streams: self.active_streams.load(Ordering::Relaxed) as usize,
            listen_ports: self.listen_ports.read().await.clone(),
        }
    }

    /// Start the edge: connect to hub, start listeners.
    pub async fn start(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        let config = self.config.read().await.clone();
        let (shutdown_tx, shutdown_rx) = mpsc::channel::<()>(1);
        *self.shutdown_tx.lock().await = Some(shutdown_tx);
        *self.running.write().await = true;

        let connected = self.connected.clone();
        let public_ip = self.public_ip.clone();
        let active_streams = self.active_streams.clone();
        let next_stream_id = self.next_stream_id.clone();
        let event_tx = self.event_tx.clone();
        let listen_ports = self.listen_ports.clone();

        tokio::spawn(async move {
            edge_main_loop(
                config,
                connected,
                public_ip,
                active_streams,
                next_stream_id,
                event_tx,
                listen_ports,
                shutdown_rx,
            )
            .await;
        });

        Ok(())
    }

    /// Stop the edge.
    pub async fn stop(&self) {
        if let Some(tx) = self.shutdown_tx.lock().await.take() {
            let _ = tx.send(()).await;
        }
        *self.running.write().await = false;
        *self.connected.write().await = false;
        self.listen_ports.write().await.clear();
    }
}

async fn edge_main_loop(
    config: EdgeConfig,
    connected: Arc<RwLock<bool>>,
    public_ip: Arc<RwLock<Option<String>>>,
    active_streams: Arc<AtomicU32>,
    next_stream_id: Arc<AtomicU32>,
    event_tx: mpsc::UnboundedSender<EdgeEvent>,
    listen_ports: Arc<RwLock<Vec<u16>>>,
    mut shutdown_rx: mpsc::Receiver<()>,
) {
    let mut backoff_ms: u64 = 1000;
    let max_backoff_ms: u64 = 30000;

    loop {
        // Try to connect to hub
        let result = connect_to_hub_and_run(
            &config,
            &connected,
            &public_ip,
            &active_streams,
            &next_stream_id,
            &event_tx,
            &listen_ports,
            &mut shutdown_rx,
        )
        .await;

        *connected.write().await = false;
        let _ = event_tx.send(EdgeEvent::TunnelDisconnected);
        active_streams.store(0, Ordering::Relaxed);
        listen_ports.write().await.clear();

        match result {
            EdgeLoopResult::Shutdown => break,
            EdgeLoopResult::Reconnect => {
                log::info!("Reconnecting in {}ms...", backoff_ms);
                tokio::select! {
                    _ = tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)) => {}
                    _ = shutdown_rx.recv() => break,
                }
                backoff_ms = (backoff_ms * 2).min(max_backoff_ms);
            }
        }
    }
}

enum EdgeLoopResult {
    Shutdown,
    Reconnect,
}

async fn connect_to_hub_and_run(
    config: &EdgeConfig,
    connected: &Arc<RwLock<bool>>,
    public_ip: &Arc<RwLock<Option<String>>>,
    active_streams: &Arc<AtomicU32>,
    next_stream_id: &Arc<AtomicU32>,
    event_tx: &mpsc::UnboundedSender<EdgeEvent>,
    listen_ports: &Arc<RwLock<Vec<u16>>>,
    shutdown_rx: &mut mpsc::Receiver<()>,
) -> EdgeLoopResult {
    // Build TLS connector that skips cert verification (auth is via secret)
    let tls_config = rustls::ClientConfig::builder()
        .dangerous()
        .with_custom_certificate_verifier(Arc::new(NoCertVerifier))
        .with_no_client_auth();

    let connector = TlsConnector::from(Arc::new(tls_config));

    let addr = format!("{}:{}", config.hub_host, config.hub_port);
    let tcp = match TcpStream::connect(&addr).await {
        Ok(s) => s,
        Err(e) => {
            log::error!("Failed to connect to hub at {}: {}", addr, e);
            return EdgeLoopResult::Reconnect;
        }
    };

    let server_name = rustls::pki_types::ServerName::try_from(config.hub_host.clone())
        .unwrap_or_else(|_| rustls::pki_types::ServerName::try_from("remoteingress-hub".to_string()).unwrap());

    let tls_stream = match connector.connect(server_name, tcp).await {
        Ok(s) => s,
        Err(e) => {
            log::error!("TLS handshake failed: {}", e);
            return EdgeLoopResult::Reconnect;
        }
    };

    let (read_half, mut write_half) = tokio::io::split(tls_stream);

    // Send auth line
    let auth_line = format!("EDGE {} {}\n", config.edge_id, config.secret);
    if write_half.write_all(auth_line.as_bytes()).await.is_err() {
        return EdgeLoopResult::Reconnect;
    }

    // Read handshake response line from hub (JSON with initial config)
    let mut buf_reader = BufReader::new(read_half);
    let mut handshake_line = String::new();
    match buf_reader.read_line(&mut handshake_line).await {
        Ok(0) => {
            log::error!("Hub rejected connection (EOF before handshake)");
            return EdgeLoopResult::Reconnect;
        }
        Ok(_) => {}
        Err(e) => {
            log::error!("Failed to read handshake response: {}", e);
            return EdgeLoopResult::Reconnect;
        }
    }

    let handshake: HandshakeConfig = match serde_json::from_str(handshake_line.trim()) {
        Ok(h) => h,
        Err(e) => {
            log::error!("Invalid handshake response: {}", e);
            return EdgeLoopResult::Reconnect;
        }
    };

    log::info!(
        "Handshake from hub: ports {:?}, stun_interval {}s",
        handshake.listen_ports,
        handshake.stun_interval_secs
    );

    *connected.write().await = true;
    let _ = event_tx.send(EdgeEvent::TunnelConnected);
    log::info!("Connected to hub at {}", addr);

    // Store initial ports and emit event
    *listen_ports.write().await = handshake.listen_ports.clone();
    let _ = event_tx.send(EdgeEvent::PortsAssigned {
        listen_ports: handshake.listen_ports.clone(),
    });

    // Start STUN discovery
    let stun_interval = handshake.stun_interval_secs;
    let public_ip_clone = public_ip.clone();
    let event_tx_clone = event_tx.clone();
    let stun_handle = tokio::spawn(async move {
        loop {
            if let Some(ip) = crate::stun::discover_public_ip().await {
                let mut pip = public_ip_clone.write().await;
                let changed = pip.as_ref() != Some(&ip);
                *pip = Some(ip.clone());
                if changed {
                    let _ = event_tx_clone.send(EdgeEvent::PublicIpDiscovered { ip });
                }
            }
            tokio::time::sleep(std::time::Duration::from_secs(stun_interval)).await;
        }
    });

    // Client socket map: stream_id -> sender for writing data back to client
    let client_writers: Arc<Mutex<HashMap<u32, mpsc::Sender<Vec<u8>>>>> =
        Arc::new(Mutex::new(HashMap::new()));

    // Shared tunnel writer
    let tunnel_writer = Arc::new(Mutex::new(write_half));

    // Start TCP listeners for initial ports (hot-reloadable)
    let mut port_listeners: HashMap<u16, JoinHandle<()>> = HashMap::new();
    apply_port_config(
        &handshake.listen_ports,
        &mut port_listeners,
        &tunnel_writer,
        &client_writers,
        active_streams,
        next_stream_id,
        &config.edge_id,
    );

    // Read frames from hub
    let mut frame_reader = FrameReader::new(buf_reader);
    let result = loop {
        tokio::select! {
            frame_result = frame_reader.next_frame() => {
                match frame_result {
                    Ok(Some(frame)) => {
                        match frame.frame_type {
                            FRAME_DATA_BACK => {
                                let writers = client_writers.lock().await;
                                if let Some(tx) = writers.get(&frame.stream_id) {
                                    let _ = tx.send(frame.payload).await;
                                }
                            }
                            FRAME_CLOSE_BACK => {
                                let mut writers = client_writers.lock().await;
                                writers.remove(&frame.stream_id);
                            }
                            FRAME_CONFIG => {
                                if let Ok(update) = serde_json::from_slice::<ConfigUpdate>(&frame.payload) {
                                    log::info!("Config update from hub: ports {:?}", update.listen_ports);
                                    *listen_ports.write().await = update.listen_ports.clone();
                                    let _ = event_tx.send(EdgeEvent::PortsUpdated {
                                        listen_ports: update.listen_ports.clone(),
                                    });
                                    apply_port_config(
                                        &update.listen_ports,
                                        &mut port_listeners,
                                        &tunnel_writer,
                                        &client_writers,
                                        active_streams,
                                        next_stream_id,
                                        &config.edge_id,
                                    );
                                }
                            }
                            _ => {
                                log::warn!("Unexpected frame type {} from hub", frame.frame_type);
                            }
                        }
                    }
                    Ok(None) => {
                        log::info!("Hub disconnected (EOF)");
                        break EdgeLoopResult::Reconnect;
                    }
                    Err(e) => {
                        log::error!("Hub frame error: {}", e);
                        break EdgeLoopResult::Reconnect;
                    }
                }
            }
            _ = shutdown_rx.recv() => {
                break EdgeLoopResult::Shutdown;
            }
        }
    };

    // Cleanup
    stun_handle.abort();
    for (_, h) in port_listeners.drain() {
        h.abort();
    }

    result
}

/// Apply a new port configuration: spawn listeners for added ports, abort removed ports.
fn apply_port_config(
    new_ports: &[u16],
    port_listeners: &mut HashMap<u16, JoinHandle<()>>,
    tunnel_writer: &Arc<Mutex<tokio::io::WriteHalf<tokio_rustls::client::TlsStream<TcpStream>>>>,
    client_writers: &Arc<Mutex<HashMap<u32, mpsc::Sender<Vec<u8>>>>>,
    active_streams: &Arc<AtomicU32>,
    next_stream_id: &Arc<AtomicU32>,
    edge_id: &str,
) {
    let new_set: std::collections::HashSet<u16> = new_ports.iter().copied().collect();
    let old_set: std::collections::HashSet<u16> = port_listeners.keys().copied().collect();

    // Remove ports no longer needed
    for &port in old_set.difference(&new_set) {
        if let Some(handle) = port_listeners.remove(&port) {
            log::info!("Stopping listener on port {}", port);
            handle.abort();
        }
    }

    // Add new ports
    for &port in new_set.difference(&old_set) {
        let tunnel_writer = tunnel_writer.clone();
        let client_writers = client_writers.clone();
        let active_streams = active_streams.clone();
        let next_stream_id = next_stream_id.clone();
        let edge_id = edge_id.to_string();

        let handle = tokio::spawn(async move {
            let listener = match TcpListener::bind(("0.0.0.0", port)).await {
                Ok(l) => l,
                Err(e) => {
                    log::error!("Failed to bind port {}: {}", port, e);
                    return;
                }
            };
            log::info!("Listening on port {}", port);

            loop {
                match listener.accept().await {
                    Ok((client_stream, client_addr)) => {
                        let stream_id = next_stream_id.fetch_add(1, Ordering::Relaxed);
                        let tunnel_writer = tunnel_writer.clone();
                        let client_writers = client_writers.clone();
                        let active_streams = active_streams.clone();
                        let edge_id = edge_id.clone();

                        active_streams.fetch_add(1, Ordering::Relaxed);

                        tokio::spawn(async move {
                            handle_client_connection(
                                client_stream,
                                client_addr,
                                stream_id,
                                port,
                                &edge_id,
                                tunnel_writer,
                                client_writers,
                            )
                            .await;
                            active_streams.fetch_sub(1, Ordering::Relaxed);
                        });
                    }
                    Err(e) => {
                        log::error!("Accept error on port {}: {}", port, e);
                    }
                }
            }
        });
        port_listeners.insert(port, handle);
    }
}

async fn handle_client_connection(
    client_stream: TcpStream,
    client_addr: std::net::SocketAddr,
    stream_id: u32,
    dest_port: u16,
    edge_id: &str,
    tunnel_writer: Arc<Mutex<tokio::io::WriteHalf<tokio_rustls::client::TlsStream<TcpStream>>>>,
    client_writers: Arc<Mutex<HashMap<u32, mpsc::Sender<Vec<u8>>>>>,
) {
    let client_ip = client_addr.ip().to_string();
    let client_port = client_addr.port();

    // Determine edge IP (use 0.0.0.0 as placeholder — hub doesn't use it for routing)
    let edge_ip = "0.0.0.0";

    // Send OPEN frame with PROXY v1 header
    let proxy_header = build_proxy_v1_header(&client_ip, edge_ip, client_port, dest_port);
    let open_frame = encode_frame(stream_id, FRAME_OPEN, proxy_header.as_bytes());
    {
        let mut w = tunnel_writer.lock().await;
        if w.write_all(&open_frame).await.is_err() {
            return;
        }
    }

    // Set up channel for data coming back from hub
    let (back_tx, mut back_rx) = mpsc::channel::<Vec<u8>>(256);
    {
        let mut writers = client_writers.lock().await;
        writers.insert(stream_id, back_tx);
    }

    let (mut client_read, mut client_write) = client_stream.into_split();

    // Task: hub -> client
    let hub_to_client = tokio::spawn(async move {
        while let Some(data) = back_rx.recv().await {
            if client_write.write_all(&data).await.is_err() {
                break;
            }
        }
        let _ = client_write.shutdown().await;
    });

    // Task: client -> hub
    let mut buf = vec![0u8; 32768];
    loop {
        match client_read.read(&mut buf).await {
            Ok(0) => break,
            Ok(n) => {
                let data_frame = encode_frame(stream_id, FRAME_DATA, &buf[..n]);
                let mut w = tunnel_writer.lock().await;
                if w.write_all(&data_frame).await.is_err() {
                    break;
                }
            }
            Err(_) => break,
        }
    }

    // Send CLOSE frame
    let close_frame = encode_frame(stream_id, FRAME_CLOSE, &[]);
    {
        let mut w = tunnel_writer.lock().await;
        let _ = w.write_all(&close_frame).await;
    }

    // Cleanup
    {
        let mut writers = client_writers.lock().await;
        writers.remove(&stream_id);
    }
    hub_to_client.abort();
    let _ = edge_id; // used for logging context
}

/// TLS certificate verifier that accepts any certificate (auth is via shared secret).
#[derive(Debug)]
struct NoCertVerifier;

impl rustls::client::danger::ServerCertVerifier for NoCertVerifier {
    fn verify_server_cert(
        &self,
        _end_entity: &rustls::pki_types::CertificateDer<'_>,
        _intermediates: &[rustls::pki_types::CertificateDer<'_>],
        _server_name: &rustls::pki_types::ServerName<'_>,
        _ocsp_response: &[u8],
        _now: rustls::pki_types::UnixTime,
    ) -> Result<rustls::client::danger::ServerCertVerified, rustls::Error> {
        Ok(rustls::client::danger::ServerCertVerified::assertion())
    }

    fn verify_tls12_signature(
        &self,
        _message: &[u8],
        _cert: &rustls::pki_types::CertificateDer<'_>,
        _dss: &rustls::DigitallySignedStruct,
    ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
        Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
    }

    fn verify_tls13_signature(
        &self,
        _message: &[u8],
        _cert: &rustls::pki_types::CertificateDer<'_>,
        _dss: &rustls::DigitallySignedStruct,
    ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> {
        Ok(rustls::client::danger::HandshakeSignatureValid::assertion())
    }

    fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> {
        vec![
            rustls::SignatureScheme::RSA_PKCS1_SHA256,
            rustls::SignatureScheme::RSA_PKCS1_SHA384,
            rustls::SignatureScheme::RSA_PKCS1_SHA512,
            rustls::SignatureScheme::ECDSA_NISTP256_SHA256,
            rustls::SignatureScheme::ECDSA_NISTP384_SHA384,
            rustls::SignatureScheme::ECDSA_NISTP521_SHA512,
            rustls::SignatureScheme::RSA_PSS_SHA256,
            rustls::SignatureScheme::RSA_PSS_SHA384,
            rustls::SignatureScheme::RSA_PSS_SHA512,
            rustls::SignatureScheme::ED25519,
            rustls::SignatureScheme::ED448,
        ]
    }
}