v0.25.1

src/stream/model.rs

@@ -0,0 +1,399 @@
+use anyhow::{ensure, Context, Result};
+use bytes::Bytes;
+use serde::{Deserialize, Serialize};
+use std::{
+    collections::{HashMap, HashSet},
+    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
+    sync::{
+        atomic::{AtomicU64, Ordering},
+        Arc,
+    },
+    time::Duration,
+};
+use tokio::{
+    io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt},
+    net::{TcpListener, TcpStream, UdpSocket},
+    sync::{mpsc, Mutex, Semaphore},
+    task::JoinSet,
+    time::Instant,
+};
+use tokio_util::sync::CancellationToken;
+
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum Protocol {
+    #[default]
+    Tcp,
+    Udp,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(default, deny_unknown_fields)]
+pub struct Host {
+    pub id: String,
+    pub name: String,
+    pub enabled: bool,
+    pub protocol: Protocol,
+    pub listen: SocketAddr,
+    pub upstream: SocketAddr,
+    /// TCP connect timeout. Kept in the common schema so TCP/UDP rows stay simple.
+    pub connect_timeout_seconds: u64,
+    /// Inactivity timeout for a TCP relay or a per-client UDP session.
+    pub idle_timeout_seconds: u64,
+    /// Maximum concurrent TCP connections or UDP client sessions.
+    pub max_connections: usize,
+}
+impl Default for Host {
+    fn default() -> Self {
+        Self {
+            id: String::new(),
+            name: String::new(),
+            enabled: true,
+            protocol: Protocol::Tcp,
+            listen: "127.0.0.1:10000".parse().expect("static stream listen"),
+            upstream: "127.0.0.1:10001".parse().expect("static stream upstream"),
+            connect_timeout_seconds: 5,
+            idle_timeout_seconds: 300,
+            max_connections: 1024,
+        }
+    }
+}
+
+/// Fixed-cardinality telemetry for one configured L4 stream. The registry is
+/// created from the active startup configuration, so stream IDs cannot be
+/// attacker-controlled or grow without a configuration/restart boundary.
+#[derive(Default, Debug)]
+pub struct Metrics {
+    pub sessions_total: AtomicU64,
+    pub sessions_completed: AtomicU64,
+    pub active: AtomicU64,
+    pub rejected: AtomicU64,
+    pub errors: AtomicU64,
+    pub bytes_in: AtomicU64,
+    pub bytes_out: AtomicU64,
+    pub packets_in: AtomicU64,
+    pub packets_out: AtomicU64,
+}
+
+#[derive(Debug, Clone, Serialize)]
+pub struct MetricsView {
+    pub sessions_total: u64,
+    pub sessions_completed: u64,
+    pub active: u64,
+    pub rejected: u64,
+    pub errors: u64,
+    pub bytes_in: u64,
+    pub bytes_out: u64,
+    pub packets_in: u64,
+    pub packets_out: u64,
+}
+
+impl Metrics {
+    pub fn view(&self) -> MetricsView {
+        let load = |value: &AtomicU64| value.load(Ordering::Relaxed);
+        MetricsView {
+            sessions_total: load(&self.sessions_total),
+            sessions_completed: load(&self.sessions_completed),
+            active: load(&self.active),
+            rejected: load(&self.rejected),
+            errors: load(&self.errors),
+            bytes_in: load(&self.bytes_in),
+            bytes_out: load(&self.bytes_out),
+            packets_in: load(&self.packets_in),
+            packets_out: load(&self.packets_out),
+        }
+    }
+}
+
+struct ActiveSession {
+    metrics: Arc<Metrics>,
+}
+
+impl ActiveSession {
+    fn begin(metrics: Arc<Metrics>) -> Self {
+        metrics.active.fetch_add(1, Ordering::Relaxed);
+        Self { metrics }
+    }
+}
+
+impl Drop for ActiveSession {
+    fn drop(&mut self) {
+        self.metrics.active.fetch_sub(1, Ordering::Relaxed);
+        self.metrics
+            .sessions_completed
+            .fetch_add(1, Ordering::Relaxed);
+    }
+}
+
+pub fn validate(
+    hosts: &mut [Host],
+    http_listeners: &[crate::vhost::Listener],
+    admin: SocketAddr,
+    prometheus: Option<SocketAddr>,
+) -> Result<()> {
+    ensure!(hosts.len() <= 128, "at most 128 stream hosts");
+    let mut ids = HashSet::new();
+    for host in hosts.iter_mut() {
+        ensure!(
+            crate::config::valid_id(&host.id) && ids.insert(host.id.clone()),
+            "invalid/duplicate stream host id"
+        );
+        host.name = host.name.trim().to_owned();
+        ensure!(
+            host.name.len() <= 96 && !host.name.chars().any(char::is_control),
+            "stream host name must be at most 96 bytes without controls"
+        );
+        ensure!(
+            host.listen.port() > 0 && host.upstream.port() > 0,
+            "stream listener/upstream ports must be nonzero"
+        );
+        ensure!(
+            !host.listen.ip().is_multicast() && !host.upstream.ip().is_multicast(),
+            "stream listener/upstream must not use multicast addresses"
+        );
+        ensure!(
+            !host.upstream.ip().is_unspecified(),
+            "stream upstream must use a concrete IP address"
+        );
+        ensure!(
+            host.connect_timeout_seconds >= 1 && host.connect_timeout_seconds <= 300,
+            "stream connect timeout must be 1..300 seconds"
+        );
+        ensure!(
+            host.idle_timeout_seconds >= 1 && host.idle_timeout_seconds <= 86400,
+            "stream idle timeout must be 1..86400 seconds"
+        );
+        ensure!(
+            (1..=65535).contains(&host.max_connections),
+            "stream connection/session limit must be 1..65535"
+        );
+    }
+
+    for (index, host) in hosts.iter().enumerate() {
+        if !host.enabled {
+            continue;
+        }
+        for other in hosts.iter().skip(index + 1).filter(|item| item.enabled) {
+            if host.protocol == other.protocol {
+                ensure!(
+                    !crate::vhost::socket_overlap(host.listen, other.listen),
+                    "stream listeners {} and {} overlap",
+                    host.id,
+                    other.id
+                );
+            }
+        }
+        match host.protocol {
+            Protocol::Tcp => {
+                ensure!(
+                    !crate::vhost::socket_overlap(host.listen, admin),
+                    "TCP stream {} overlaps administration listener",
+                    host.id
+                );
+                if let Some(metrics) = prometheus {
+                    ensure!(
+                        !crate::vhost::socket_overlap(host.listen, metrics),
+                        "TCP stream {} overlaps Prometheus listener",
+                        host.id
+                    );
+                }
+                for listener in http_listeners {
+                    ensure!(
+                        !crate::vhost::socket_overlap(host.listen, listener.address),
+                        "TCP stream {} overlaps HTTP/TLS listener",
+                        host.id
+                    );
+                }
+            }
+            Protocol::Udp => {
+                for listener in http_listeners
+                    .iter()
+                    .filter(|listener| listener.http3.is_some())
+                {
+                    ensure!(
+                        !crate::vhost::socket_overlap(host.listen, listener.address),
+                        "UDP stream {} overlaps HTTP/3 listener",
+                        host.id
+                    );
+                }
+            }
+        }
+    }
+
+    // Reject direct loops into another RustProxy listener of the same transport.
+    for host in hosts.iter().filter(|item| item.enabled) {
+        match host.protocol {
+            Protocol::Tcp => {
+                ensure!(
+                    !crate::vhost::socket_overlap(host.upstream, admin)
+                        && prometheus
+                            .is_none_or(|addr| !crate::vhost::socket_overlap(host.upstream, addr))
+                        && !http_listeners
+                            .iter()
+                            .any(|listener| crate::vhost::socket_overlap(
+                                host.upstream,
+                                listener.address
+                            ))
+                        && !hosts
+                            .iter()
+                            .filter(|other| other.enabled && other.protocol == Protocol::Tcp)
+                            .any(|other| crate::vhost::socket_overlap(host.upstream, other.listen)),
+                    "TCP stream {} upstream points at a RustProxy listener",
+                    host.id
+                );
+            }
+            Protocol::Udp => {
+                ensure!(
+                    !http_listeners
+                        .iter()
+                        .filter(|listener| listener.http3.is_some())
+                        .any(|listener| crate::vhost::socket_overlap(
+                            host.upstream,
+                            listener.address
+                        ))
+                        && !hosts
+                            .iter()
+                            .filter(|other| other.enabled && other.protocol == Protocol::Udp)
+                            .any(|other| crate::vhost::socket_overlap(host.upstream, other.listen)),
+                    "UDP stream {} upstream points at a RustProxy UDP listener",
+                    host.id
+                );
+            }
+        }
+    }
+    Ok(())
+}
+
+pub async fn serve_tcp(
+    host: Host,
+    listener: TcpListener,
+    shutdown: CancellationToken,
+    grace: Duration,
+    metrics: Arc<Metrics>,
+) -> Result<()> {
+    let slots = Arc::new(Semaphore::new(host.max_connections));
+    let mut tasks = JoinSet::new();
+    loop {
+        tokio::select! {
+            biased;
+            _ = shutdown.cancelled() => break,
+            _ = tasks.join_next(), if !tasks.is_empty() => {},
+            accepted = listener.accept() => {
+                let (client, peer) = accepted.context("stream TCP accept failed")?;
+                let Ok(permit) = slots.clone().try_acquire_owned() else {
+                    metrics.rejected.fetch_add(1, Ordering::Relaxed);
+                    drop(client);
+                    continue;
+                };
+                metrics.sessions_total.fetch_add(1, Ordering::Relaxed);
+                let cfg = host.clone();
+                let stop = shutdown.clone();
+                let stream_metrics = metrics.clone();
+                let active = ActiveSession::begin(stream_metrics.clone());
+                tasks.spawn(async move {
+                    let _active = active;
+                    let _permit = permit;
+                    let upstream = match tokio::time::timeout(
+                        Duration::from_secs(cfg.connect_timeout_seconds),
+                        TcpStream::connect(cfg.upstream),
+                    ).await {
+                        Ok(Ok(stream)) => stream,
+                        Ok(Err(error)) => {
+                            stream_metrics.errors.fetch_add(1, Ordering::Relaxed);
+                            tracing::debug!(stream=%cfg.id, %peer, %error, "TCP stream upstream connect failed");
+                            return;
+                        },
+                        Err(_) => {
+                            stream_metrics.errors.fetch_add(1, Ordering::Relaxed);
+                            tracing::debug!(stream=%cfg.id, %peer, "TCP stream upstream connect timed out");
+                            return;
+                        },
+                    };
+                    if let Err(error) = relay_tcp(client, upstream, stop, Duration::from_secs(cfg.idle_timeout_seconds), &stream_metrics).await {
+                        stream_metrics.errors.fetch_add(1, Ordering::Relaxed);
+                        tracing::debug!(stream=%cfg.id, %peer, %error, "TCP stream relay ended");
+                    }
+                });
+            }
+        }
+    }
+    let _ = tokio::time::timeout(grace, async { while tasks.join_next().await.is_some() {} }).await;
+    tasks.abort_all();
+    while tasks.join_next().await.is_some() {}
+    Ok(())
+}
+
+async fn copy_with_activity<R, W>(
+    mut reader: R,
+    mut writer: W,
+    activity: tokio::sync::watch::Sender<Instant>,
+    transferred: &AtomicU64,
+) -> std::io::Result<u64>
+where
+    R: AsyncRead + Unpin,
+    W: AsyncWrite + Unpin,
+{
+    let mut total = 0u64;
+    let mut buffer = [0u8; 32 * 1024];
+    loop {
+        let read = reader.read(&mut buffer).await?;
+        if read == 0 {
+            writer.shutdown().await?;
+            return Ok(total);
+        }
+        writer.write_all(&buffer[..read]).await?;
+        total = total.saturating_add(read as u64);
+        transferred.fetch_add(read as u64, Ordering::Relaxed);
+        activity.send_replace(Instant::now());
+    }
+}
+
+async fn relay_tcp(
+    client: TcpStream,
+    upstream: TcpStream,
+    shutdown: CancellationToken,
+    idle: Duration,
+    metrics: &Metrics,
+) -> Result<()> {
+    let (client_read, client_write) = client.into_split();
+    let (upstream_read, upstream_write) = upstream.into_split();
+    let (activity, mut activity_rx) = tokio::sync::watch::channel(Instant::now());
+    let c2u = copy_with_activity(
+        client_read,
+        upstream_write,
+        activity.clone(),
+        &metrics.bytes_in,
+    );
+    let u2c = copy_with_activity(
+        upstream_read,
+        client_write,
+        activity.clone(),
+        &metrics.bytes_out,
+    );
+    drop(activity);
+    tokio::pin!(c2u);
+    tokio::pin!(u2c);
+    let mut c2u_done = false;
+    let mut u2c_done = false;
+    let deadline = tokio::time::sleep_until(Instant::now() + idle);
+    tokio::pin!(deadline);
+    loop {
+        tokio::select! {
+            _ = shutdown.cancelled() => break,
+            result = &mut c2u, if !c2u_done => { result?; c2u_done = true; if u2c_done { break; } },
+            result = &mut u2c, if !u2c_done => { result?; u2c_done = true; if c2u_done { break; } },
+            changed = activity_rx.changed() => {
+                if changed.is_err() { break; }
+                deadline.as_mut().reset(*activity_rx.borrow() + idle);
+            },
+            _ = &mut deadline => break,
+        }
+    }
+    Ok(())
+}
+
+#[derive(Clone)]
+struct UdpSession {
+    id: u64,
+    tx: mpsc::Sender<Bytes>,
+}