deepslate 0.2.0

//! A high-performance Minecraft server proxy written in Rust.
//!
//! This crate can be used as a library to build a custom proxy binary with
//! compile-time plugins, or as a dependency of the default `deepslate` binary.
//!
//! # Usage as a library
//!
//! ```rust,no_run
//! use deepslate::{Proxy, ServerId};
//! use deepslate::event::*;
//! use deepslate::event::events::*;
//!
//! const LOBBY: ServerId = ServerId::new("lobby", "127.0.0.1:25566");
//!
//! struct LobbyPlugin;
//!
//! impl Plugin for LobbyPlugin {
//!     fn register(&self, events: &mut EventManager) {
//!         events.subscribe::<ChooseServerEvent>(PostOrder::NORMAL, |event| {
//!             event.set_result(LOBBY.into());
//!         });
//!
//!         events.observe::<LoginEvent>(Phase::Post, |event| {
//!             println!("Player logged in: {}", event.player.profile.name);
//!         });
//!     }
//! }
//!
//! #[tokio::main]
//! async fn main() {
//!     let proxy = Proxy::builder()
//!         .forwarding_secret("your-secret")
//!         .server(&LOBBY)
//!         .try_servers([&LOBBY])
//!         .plugin(LobbyPlugin)
//!         .build()
//!         .expect("failed to build proxy");
//!     proxy.run().await.expect("proxy error");
//! }
//! ```

/// Authentication and player profiles.
pub mod auth;
/// Proxy configuration.
pub mod config;
/// Connection handling and networking.
pub mod connection;
/// Cryptographic primitives.
pub mod crypto;
/// Event system and lifecycle events.
pub mod event;
/// gRPC control plane service.
#[cfg(feature = "grpc")]
pub mod rpc;
/// Server registry and backend management.
pub mod server;
/// Server status (MOTD) handling.
pub mod status;

use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;

use tokio::net::TcpListener;
use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;
use tracing::{Instrument, error, info, info_span};

pub use config::Config;
pub use event::{EventManager, Phase, PlayerInfo, Plugin, PostOrder};
pub use server::{Server, ServerId, ServerRegistry};

use config::ConfigError;

/// Error type for public proxy operations.
///
/// Returned by [`ProxyBuilder::build`] and [`Proxy::run`] to allow callers to
/// match on specific failure modes rather than relying on string messages.
#[derive(Debug, thiserror::Error)]
pub enum ProxyError {
    /// Configuration loading or validation failed.
    #[error("configuration error: {0}")]
    Config(#[from] ConfigError),

    /// RSA key pair generation failed.
    #[error("RSA key generation failed")]
    KeyGeneration(#[source] Box<dyn std::error::Error + Send + Sync>),

    /// The HTTP client could not be initialised.
    #[error("HTTP client initialization failed")]
    HttpClient(#[from] reqwest::Error),

    /// A server with the given ID was already registered.
    #[error("duplicate server ID: {0}")]
    DuplicateServer(String),

    /// The TCP listener could not bind to the configured address.
    #[error("failed to bind listener")]
    Bind(#[source] std::io::Error),
}

use connection::ConnectionError;
use connection::MinecraftConnection;
use crypto::ServerKeyPair;
#[cfg(feature = "grpc")]
use rpc::DeepslateService;
#[cfg(feature = "grpc")]
use rpc::proto::deepslate_server::DeepslateServer;

/// The Deepslate proxy server.
pub struct Proxy {
    config: Arc<Config>,
    registry: Arc<ServerRegistry>,
    key_pair: Arc<ServerKeyPair>,
    event_manager: Arc<EventManager>,
    http_client: reqwest::Client,
}

impl Proxy {
    /// Create a new proxy builder.
    #[must_use]
    pub fn builder() -> ProxyBuilder {
        ProxyBuilder {
            config: None,
            config_overrides: ConfigOverrides::default(),
            bootstrap_servers: Vec::new(),
            try_servers: None,
            forced_hosts: Vec::new(),
            plugins: Vec::new(),
        }
    }

    /// Get a reference to the proxy configuration.
    #[must_use]
    pub const fn config(&self) -> &Arc<Config> {
        &self.config
    }

    /// Get a reference to the server registry for external management.
    #[must_use]
    pub const fn registry(&self) -> &Arc<ServerRegistry> {
        &self.registry
    }

    /// Get a reference to the event manager.
    #[must_use]
    pub const fn event_manager(&self) -> &Arc<EventManager> {
        &self.event_manager
    }

    /// Run the proxy server.
    ///
    /// This starts the Minecraft listener and, when the `grpc` feature is
    /// enabled, the gRPC control plane. The proxy runs until a shutdown signal
    /// is received (SIGINT/SIGTERM).
    ///
    /// # Errors
    ///
    /// Returns [`ProxyError::Bind`] if the TCP listener cannot bind to the
    /// configured address.
    ///
    /// # Panics
    ///
    /// Panics if the gRPC reflection service fails to build (should not happen
    /// with a valid compiled proto descriptor).
    #[allow(clippy::too_many_lines)]
    pub async fn run(&self) -> Result<(), ProxyError> {
        // Bind the proxy listener first (fail fast on port conflicts)
        let listener = TcpListener::bind(self.config.listen_addr)
            .await
            .map_err(ProxyError::Bind)?;
        info!(addr = %self.config.listen_addr, "proxy listening");

        // Start gRPC control plane
        #[cfg(feature = "grpc")]
        let grpc_addr = self.config.grpc_addr;
        #[cfg(feature = "grpc")]
        let grpc_registry = Arc::clone(&self.registry);
        #[cfg(feature = "grpc")]
        let mut grpc_handle = tokio::spawn(async move {
            let grpc_service = DeepslateService::new(grpc_registry);

            #[cfg(feature = "grpc-reflection")]
            let reflection_service = tonic_reflection::server::Builder::configure()
                .register_encoded_file_descriptor_set(rpc::proto::FILE_DESCRIPTOR_SET)
                .build_v1()
                .expect("failed to build reflection service");

            info!(addr = %grpc_addr, "gRPC control plane listening");
            let builder =
                tonic::transport::Server::builder().add_service(DeepslateServer::new(grpc_service));

            #[cfg(feature = "grpc-reflection")]
            let builder = builder.add_service(reflection_service);

            builder.serve(grpc_addr).await
        });

        // Shared shutdown signal for all connection tasks
        let shutdown = CancellationToken::new();

        // Track active connection tasks for graceful drain
        let mut connections = JoinSet::new();
        let session_counter = AtomicUsize::new(0);

        // Accept loop — runs until a shutdown signal or gRPC exit
        loop {
            // When the `grpc` feature is disabled, this future never resolves so
            // the select loop is not affected.
            #[cfg(feature = "grpc")]
            let grpc_fut = &mut grpc_handle;
            #[cfg(not(feature = "grpc"))]
            let grpc_fut = std::future::pending::<Result<(), String>>();

            tokio::select! {
                result = listener.accept() => {
                    match result {
                        Ok((stream, addr)) => {
                            if let Err(e) = stream.set_nodelay(true) {
                                tracing::warn!(error = %e, "failed to set TCP_NODELAY on client stream");
                            }
                            let session_id = session_counter.fetch_add(1, Ordering::Relaxed);
                            let config = Arc::clone(&self.config);
                            let key_pair = Arc::clone(&self.key_pair);
                            let registry = Arc::clone(&self.registry);
                            let http_client = self.http_client.clone();
                            let event_manager = Arc::clone(&self.event_manager);
                            let conn_shutdown = shutdown.clone();

                            connections.spawn(
                                async move {
                                    let conn = MinecraftConnection::new(
                                        stream,
                                        libdeflater::CompressionLvl::new(config.compression_level)
                                            .expect("validated in Config::validate"),
                                        Duration::from_millis(config.read_timeout_ms),
                                    );
                                    if let Err(e) = Box::pin(connection::client::handle_client(
                                        conn,
                                        addr,
                                        config,
                                        key_pair,
                                        registry,
                                        http_client,
                                        event_manager,
                                        conn_shutdown,
                                    ))
                                    .await
                                    {
                                        match &e {
                                            ConnectionError::Io(io_err) => {
                                                tracing::debug!(error = %io_err, "client I/O error");
                                            }
                                            ConnectionError::Protocol(proto_err) => {
                                                tracing::debug!(error = %proto_err, "protocol error");
                                            }
                                            ConnectionError::Auth(auth_err) => {
                                                tracing::debug!(error = %auth_err, "authentication error");
                                            }
                                            ConnectionError::BackendFailed { reason } => {
                                                tracing::debug!(reason, "backend connection failed");
                                            }
                                            ConnectionError::Timeout => {
                                                tracing::debug!("connection timed out");
                                            }
                                        }
                                    }
                                }
                                .instrument(info_span!(
                                    "conn",
                                    sid = session_id,
                                    ip = %addr.ip(),
                                    port = addr.port()
                                )),
                            );
                        }
                        Err(e) => {
                            error!(error = %e, "failed to accept connection");
                        }
                    }
                }

                // Reap completed connection tasks to avoid unbounded growth
                Some(result) = connections.join_next(), if !connections.is_empty() => {
                    if let Err(e) = result {
                        tracing::debug!(error = %e, "connection task panicked");
                    }
                }

                _ = tokio::signal::ctrl_c() => {
                    info!("received shutdown signal");
                    break;
                }

                result = grpc_fut => {
                    if let Err(e) = result {
                        error!(error = %e, "gRPC server error");
                    }
                    break;
                }
            }
        }

        // Begin graceful shutdown
        info!("shutting down");

        // 1. Signal all connection tasks to send disconnect messages
        shutdown.cancel();

        // 2. Wait for active connections to drain (send disconnects and clean up)
        let drain = self.config.shutdown_drain;
        if !connections.is_empty() {
            info!(
                connections = connections.len(),
                drain_ms = u64::try_from(drain.as_millis()).unwrap_or(u64::MAX),
                "draining active connections"
            );
            let _ = tokio::time::timeout(drain, async {
                while connections.join_next().await.is_some() {}
            })
            .await;

            let remaining = connections.len();
            if remaining > 0 {
                info!(
                    remaining,
                    "drain timeout reached, dropping remaining connections"
                );
                connections.abort_all();
            }
        }

        Ok(())
    }
}

/// Builder for constructing a [`Proxy`] instance.
pub struct ProxyBuilder {
    config: Option<Config>,
    config_overrides: ConfigOverrides,
    bootstrap_servers: Vec<ServerId>,
    try_servers: Option<Vec<&'static str>>,
    forced_hosts: Vec<(String, Vec<&'static str>)>,
    plugins: Vec<Box<dyn Plugin>>,
}

impl ProxyBuilder {
    /// Set the proxy configuration. If not called, configuration is loaded
    /// from environment variables.
    #[must_use]
    pub fn config(mut self, config: Config) -> Self {
        self.config = Some(config);
        self
    }

    /// Set the proxy listen address.
    #[must_use]
    pub const fn listen_addr(mut self, listen_addr: std::net::SocketAddr) -> Self {
        self.config_overrides.listen_addr = Some(listen_addr);
        self
    }

    /// Set the gRPC control plane listen address.
    #[cfg(feature = "grpc")]
    #[must_use]
    pub const fn grpc_addr(mut self, grpc_addr: std::net::SocketAddr) -> Self {
        self.config_overrides.grpc_addr = Some(grpc_addr);
        self
    }

    /// Enable or disable Mojang authentication.
    #[must_use]
    pub const fn online_mode(mut self, online_mode: bool) -> Self {
        self.config_overrides.online_mode = Some(online_mode);
        self
    }

    /// Set the Velocity modern forwarding secret.
    #[must_use]
    pub fn forwarding_secret(mut self, forwarding_secret: impl AsRef<[u8]>) -> Self {
        self.config_overrides.forwarding_secret = Some(forwarding_secret.as_ref().to_vec());
        self
    }

    /// Set the login compression threshold.
    #[must_use]
    pub const fn compression_threshold(mut self, compression_threshold: i32) -> Self {
        self.config_overrides.compression_threshold = Some(compression_threshold);
        self
    }

    /// Set the zlib compression level.
    #[must_use]
    pub const fn compression_level(mut self, compression_level: i32) -> Self {
        self.config_overrides.compression_level = Some(compression_level);
        self
    }

    /// Set the server list MOTD.
    #[must_use]
    pub fn motd(mut self, motd: impl Into<String>) -> Self {
        self.config_overrides.motd = Some(motd.into());
        self
    }

    /// Set the advertised maximum player count.
    #[must_use]
    pub const fn max_players(mut self, max_players: i32) -> Self {
        self.config_overrides.max_players = Some(max_players);
        self
    }

    /// Set the client read timeout in milliseconds.
    #[must_use]
    pub const fn read_timeout_ms(mut self, read_timeout_ms: u64) -> Self {
        self.config_overrides.read_timeout_ms = Some(read_timeout_ms);
        self
    }

    /// Set the initial server try order.
    ///
    /// Accepts an iterable of [`ServerId`] references. The try order determines
    /// which backend server a player is connected to first.
    #[must_use]
    pub fn try_servers<I>(mut self, try_servers: I) -> Self
    where
        I: IntoIterator<Item = &'static ServerId>,
    {
        self.try_servers = Some(try_servers.into_iter().map(|s| s.id).collect());
        self
    }

    /// Add a forced-host entry mapping a hostname to an ordered list of servers.
    ///
    /// When a player connects using the given `hostname`, the proxy tries the
    /// listed servers (in order) instead of the global try list. Only applies
    /// to the initial connection, not server switches.
    ///
    /// The hostname is lowercased for case-insensitive matching.
    ///
    /// Can be called multiple times for different hostnames.
    #[must_use]
    pub fn forced_host<I>(mut self, hostname: impl Into<String>, servers: I) -> Self
    where
        I: IntoIterator<Item = &'static ServerId>,
    {
        let ids: Vec<&'static str> = servers.into_iter().map(|s| s.id).collect();
        self.forced_hosts
            .push((hostname.into().to_lowercase(), ids));
        self
    }

    /// Set the tracing filter string.
    #[must_use]
    pub fn log_level(mut self, log_level: impl Into<String>) -> Self {
        self.config_overrides.log_level = Some(log_level.into());
        self
    }

    /// Enable or disable JSON log output.
    #[must_use]
    pub const fn log_json(mut self, log_json: bool) -> Self {
        self.config_overrides.log_json = Some(log_json);
        self
    }

    /// Set how long to wait for active connections to drain on shutdown.
    #[must_use]
    pub const fn shutdown_drain(mut self, shutdown_drain: std::time::Duration) -> Self {
        self.config_overrides.shutdown_drain = Some(shutdown_drain);
        self
    }

    /// Register a backend server during proxy construction.
    #[must_use]
    pub fn server(mut self, server: &'static ServerId) -> Self {
        self.bootstrap_servers.push(*server);
        self
    }

    /// Register multiple backend servers during proxy construction.
    #[must_use]
    pub fn servers<I>(mut self, servers: I) -> Self
    where
        I: IntoIterator<Item = &'static ServerId>,
    {
        self.bootstrap_servers.extend(servers.into_iter().copied());
        self
    }

    /// Add a plugin. Multiple plugins can be added; each registers its own
    /// event handlers. Plugins are registered in the order they are added.
    #[must_use]
    pub fn plugin(mut self, plugin: impl Plugin) -> Self {
        self.plugins.push(Box::new(plugin));
        self
    }

    /// Build the proxy.
    ///
    /// All registered plugins will have their [`Plugin::register`] method called
    /// to populate the event manager.
    ///
    /// If builder setters are used without calling [`ProxyBuilder::config`], the
    /// remaining configuration fields fall back to [`Config::default`]. If
    /// [`ProxyBuilder::config`] is provided, builder setters override fields from
    /// that config.
    ///
    /// # Errors
    ///
    /// Returns a [`ProxyError`] if configuration loading, RSA key generation,
    /// HTTP client creation, or server registration fails.
    pub fn build(self) -> Result<Proxy, ProxyError> {
        let mut config = match self.config {
            Some(config) => config,
            None if self.config_overrides.has_any() || !self.bootstrap_servers.is_empty() => {
                Config::default()
            }
            None => Config::from_env()?,
        };
        self.config_overrides.apply(&mut config);
        let config = config.validate()?;

        let key_pair = ServerKeyPair::generate().map_err(ProxyError::KeyGeneration)?;
        info!("generated RSA key pair");

        let mut event_manager = EventManager::new();
        for plugin in &self.plugins {
            plugin.register(&mut event_manager);
        }

        let http_client = reqwest::Client::builder()
            .user_agent(concat!("Deepslate/", env!("CARGO_PKG_VERSION")))
            .connect_timeout(Duration::from_secs(5))
            .timeout(Duration::from_secs(10))
            .build()?;

        let registry = ServerRegistry::new();
        for server_id in &self.bootstrap_servers {
            let server = Server::from(server_id);
            if !registry.register(&server) {
                return Err(ProxyError::DuplicateServer(server_id.id.to_string()));
            }
        }

        // Builder try_servers (typed) takes precedence, then config (string-based
        // from env vars / explicit Config struct).
        if let Some(try_servers) = self.try_servers {
            let order: Vec<String> = try_servers.into_iter().map(String::from).collect();
            registry.set_try_order(order);
        } else if !config.try_servers.is_empty() {
            registry.set_try_order(config.try_servers.clone());
        }

        // Builder forced_hosts (typed) takes precedence, then config (string-based
        // from env vars / explicit Config struct).
        if !self.forced_hosts.is_empty() {
            let map: std::collections::HashMap<String, Vec<String>> = self
                .forced_hosts
                .into_iter()
                .map(|(host, ids)| (host, ids.into_iter().map(String::from).collect()))
                .collect();
            registry.set_forced_hosts(map);
        } else if !config.forced_hosts.is_empty() {
            registry.set_forced_hosts(config.forced_hosts.clone());
        }

        Ok(Proxy {
            config: Arc::new(config),
            registry: Arc::new(registry),
            key_pair: Arc::new(key_pair),
            event_manager: Arc::new(event_manager),
            http_client,
        })
    }
}

#[derive(Default)]
struct ConfigOverrides {
    listen_addr: Option<std::net::SocketAddr>,
    #[cfg(feature = "grpc")]
    grpc_addr: Option<std::net::SocketAddr>,
    online_mode: Option<bool>,
    forwarding_secret: Option<Vec<u8>>,
    compression_threshold: Option<i32>,
    compression_level: Option<i32>,
    motd: Option<String>,
    max_players: Option<i32>,
    read_timeout_ms: Option<u64>,
    log_level: Option<String>,
    log_json: Option<bool>,
    shutdown_drain: Option<std::time::Duration>,
}

impl ConfigOverrides {
    const fn has_any(&self) -> bool {
        // NOTE: `grpc_addr` is checked separately below when the feature is on.
        let base = self.listen_addr.is_some()
            || self.online_mode.is_some()
            || self.forwarding_secret.is_some()
            || self.compression_threshold.is_some()
            || self.compression_level.is_some()
            || self.motd.is_some()
            || self.max_players.is_some()
            || self.read_timeout_ms.is_some()
            || self.log_level.is_some()
            || self.log_json.is_some()
            || self.shutdown_drain.is_some();

        #[cfg(feature = "grpc")]
        let base = base || self.grpc_addr.is_some();

        base
    }

    fn apply(self, config: &mut Config) {
        if let Some(listen_addr) = self.listen_addr {
            config.listen_addr = listen_addr;
        }
        #[cfg(feature = "grpc")]
        if let Some(grpc_addr) = self.grpc_addr {
            config.grpc_addr = grpc_addr;
        }
        if let Some(online_mode) = self.online_mode {
            config.online_mode = online_mode;
        }
        if let Some(forwarding_secret) = self.forwarding_secret {
            config.forwarding_secret = forwarding_secret;
        }
        if let Some(compression_threshold) = self.compression_threshold {
            config.compression_threshold = compression_threshold;
        }
        if let Some(compression_level) = self.compression_level {
            config.compression_level = compression_level;
        }
        if let Some(motd) = self.motd {
            config.motd = motd;
        }
        if let Some(max_players) = self.max_players {
            config.max_players = max_players;
        }
        if let Some(read_timeout_ms) = self.read_timeout_ms {
            config.read_timeout_ms = read_timeout_ms;
        }
        if let Some(log_level) = self.log_level {
            config.log_level = log_level;
        }
        if let Some(log_json) = self.log_json {
            config.log_json = log_json;
        }
        if let Some(shutdown_drain) = self.shutdown_drain {
            config.shutdown_drain = shutdown_drain;
        }
    }
}

/// Initialize the tracing subscriber based on configuration.
pub fn init_tracing(config: &Config) {
    use tracing_subscriber::EnvFilter;

    let filter = EnvFilter::try_new(&config.log_level).unwrap_or_else(|_| EnvFilter::new("info"));

    if config.log_json {
        tracing_subscriber::fmt()
            .with_env_filter(filter)
            .json()
            .init();
    } else {
        tracing_subscriber::fmt().with_env_filter(filter).init();
    }
}

#[cfg(test)]
mod tests {
    use std::net::{Ipv4Addr, SocketAddr};

    use super::*;

    #[test]
    fn builder_uses_code_config_without_env() {
        let proxy = Proxy::builder()
            .forwarding_secret("secret")
            .listen_addr(SocketAddr::from((Ipv4Addr::LOCALHOST, 25_565)))
            .motd("Configured in code")
            .build()
            .unwrap();

        assert_eq!(
            proxy.config.listen_addr,
            SocketAddr::from((Ipv4Addr::LOCALHOST, 25_565))
        );
        assert_eq!(proxy.config.motd, "Configured in code");
    }

    #[cfg(feature = "grpc")]
    #[test]
    fn builder_sets_grpc_addr() {
        let proxy = Proxy::builder()
            .forwarding_secret("secret")
            .grpc_addr(SocketAddr::from((Ipv4Addr::LOCALHOST, 25_577)))
            .build()
            .unwrap();

        assert_eq!(
            proxy.config.grpc_addr,
            SocketAddr::from((Ipv4Addr::LOCALHOST, 25_577))
        );
    }

    #[test]
    fn builder_rejects_invalid_manual_config() {
        let result = Proxy::builder()
            .forwarding_secret("secret")
            .compression_level(13)
            .build();

        assert!(result.is_err());
    }

    #[test]
    fn builder_bootstraps_servers_and_try_order() {
        static LOBBY: ServerId = ServerId::new("lobby", "127.0.0.1:25566");
        static SURVIVAL: ServerId = ServerId::new("survival", "127.0.0.1:25567");

        let proxy = Proxy::builder()
            .forwarding_secret("secret")
            .server(&LOBBY)
            .server(&SURVIVAL)
            .try_servers([&SURVIVAL, &LOBBY])
            .build()
            .unwrap();

        assert_eq!(proxy.registry.list().len(), 2);
        assert_eq!(proxy.registry.try_order(), vec!["survival", "lobby"]);
        assert_eq!(proxy.registry.select_initial().unwrap().id, "survival");
    }

    #[test]
    fn builder_overrides_explicit_config_fields() {
        let proxy = Proxy::builder()
            .config(Config {
                forwarding_secret: b"base-secret".to_vec(),
                motd: "Base MOTD".to_string(),
                ..Config::default()
            })
            .forwarding_secret("override-secret")
            .motd("Override MOTD")
            .build()
            .unwrap();

        assert_eq!(proxy.config.forwarding_secret, b"override-secret".to_vec());
        assert_eq!(proxy.config.motd, "Override MOTD");
    }
}