tcproxy 0.1.1

A TCP proxy for PostgreSQL connections with SSH tunnel support and runtime target switching
Documentation
mod cli;
mod config;
mod connection_manager;
mod logging;
mod proxy;

use anyhow::{Context, Result};
use cli::{Cli, Commands};
use config::Config;
use std::path::Path;
use tracing::{debug, error, info, warn};

#[tokio::main]
async fn main() -> Result<()> {
    let args = Cli::parse_args();

    if let Err(e) = logging::init_logging(&args.log_level, args.json_logs) {
        eprintln!("Failed to initialize logging: {}", e);
        std::process::exit(1);
    }

    info!(
        version = env!("CARGO_PKG_VERSION"),
        log_level = %args.log_level,
        json_logs = args.json_logs,
        config_path = %args.config.display(),
        "Starting tcproxy"
    );

    std::panic::set_hook(Box::new(|panic_info| {
        error!(
            panic_message = %panic_info,
            "Application panicked - this indicates a serious bug"
        );
    }));

    let result = match args.command {
        Some(Commands::Start { target, port, host }) => {
            start_proxy(&args.config, &target, port, host)
                .await
                .with_context(|| format!("Failed to start proxy with target '{}'", target))
        }
        Some(Commands::InitConfig { output, force }) => init_config(&output, force)
            .with_context(|| format!("Failed to initialize config at '{}'", output.display())),
        Some(Commands::ValidateConfig) => validate_config(&args.config)
            .with_context(|| format!("Failed to validate config at '{}'", args.config.display())),
        Some(Commands::ShowConfig) => show_config(&args.config)
            .with_context(|| format!("Failed to show config at '{}'", args.config.display())),
        Some(Commands::ListTargets) => list_targets(&args.config)
            .with_context(|| format!("Failed to list targets from '{}'", args.config.display())),
        Some(Commands::HealthCheck { target }) => health_check(&args.config, target.as_deref())
            .await
            .with_context(|| "Failed to perform health check"),
        None => {
            error!("No command provided");
            println!("Error: No command provided. Use 'tcproxy --help' for usage information.");
            println!(
                "To start the proxy, you must specify a target: 'tcproxy start --target <TARGET_NAME>'"
            );
            std::process::exit(1);
        }
    };

    if let Err(e) = result {
        error!(
            error = %e,
            error_chain = ?e.chain().collect::<Vec<_>>(),
            "Application failed"
        );

        eprintln!("Error: {}", e);

        let mut source = e.source();
        while let Some(err) = source {
            eprintln!("  Caused by: {}", err);
            source = err.source();
        }

        std::process::exit(1);
    }

    info!("Application completed successfully");
    Ok(())
}

async fn start_proxy(
    config_path: &Path,
    target_name: &str,
    port_override: Option<u16>,
    host_override: Option<String>,
) -> Result<()> {
    debug!(
        config_path = %config_path.display(),
        target_name = target_name,
        port_override = ?port_override,
        host_override = ?host_override,
        "Loading configuration and starting proxy"
    );

    let mut config = load_config(config_path).with_context(|| {
        format!(
            "Failed to load configuration from '{}'",
            config_path.display()
        )
    })?;

    if let Some(port) = port_override {
        info!(
            old_port = config.proxy.listen_port,
            new_port = port,
            "Overriding listen port from CLI"
        );
        config.proxy.listen_port = port;
    }
    if let Some(host) = host_override.clone() {
        info!(
            old_host = %config.proxy.listen_host,
            new_host = %host,
            "Overriding listen host from CLI"
        );
        config.proxy.listen_host = host;
    }

    let target = config
        .get_target(target_name)
        .with_context(|| format!("Target '{}' not found in configuration", target_name))?;

    info!(
        listen_host = %config.proxy.listen_host,
        listen_port = config.proxy.listen_port,
        target_name = target_name,
        target_host = %target.host,
        target_port = target.port,
        "Starting proxy server"
    );

    if let Some(ssh) = &target.ssh {
        if ssh.enabled {
            info!(
                ssh_host = %ssh.host.as_ref().unwrap_or(&"<not configured>".to_string()),
                ssh_user = %ssh.user.as_ref().unwrap_or(&"<not configured>".to_string()),
                ssh_port = ssh.port.unwrap_or(22),
                auto_reconnect = ssh.auto_reconnect,
                max_reconnect_attempts = ssh.max_reconnect_attempts,
                "SSH tunnel enabled"
            );
        } else {
            info!("SSH tunnel configured but disabled");
        }
    } else {
        info!("Direct connection (no SSH tunnel)");
    }

    let proxy_server = proxy::ProxyServer::new(
        config.proxy.listen_host.clone(),
        config.proxy.listen_port,
        target_name.to_string(),
        target.clone(),
        config.connection_management.clone(),
    );

    let shutdown_signal = async {
        match tokio::signal::ctrl_c().await {
            Ok(()) => {
                info!("Received CTRL+C signal, initiating graceful shutdown");
            }
            Err(e) => {
                error!(error = %e, "Failed to listen for CTRL+C signal");
            }
        }
    };

    info!("Proxy server initialized, starting main loop");

    tokio::select! {
        result = proxy_server.run() => {
            match result {
                Ok(()) => {
                    info!("Proxy server completed successfully");
                }
                Err(e) => {
                    error!(
                        error = %e,
                        error_chain = ?e.chain().collect::<Vec<_>>(),
                        "Proxy server failed"
                    );
                    return Err(e).with_context(|| "Proxy server encountered a fatal error");
                }
            }
        }
        _ = shutdown_signal => {
            info!("Graceful shutdown initiated");
        }
    }

    info!("Proxy server shutdown completed");
    Ok(())
}

fn init_config(output_path: &Path, force: bool) -> anyhow::Result<()> {
    if output_path.exists() && !force {
        anyhow::bail!(
            "Configuration file already exists at {}. Use --force to overwrite.",
            output_path.display()
        );
    }

    let default_config = Config::default();
    default_config.to_file(&output_path.to_path_buf())?;

    tracing::info!(
        "Generated default configuration at {}",
        output_path.display()
    );
    println!("Configuration file created at: {}", output_path.display());
    println!(
        "Available targets: {}",
        default_config.list_targets().join(", ")
    );
    println!();
    println!("To start the proxy, specify a target:");
    println!("  cargo run -- start --target local");
    println!("  cargo run -- start --target production");
    println!("  cargo run -- start --target development");

    Ok(())
}

fn validate_config(config_path: &Path) -> anyhow::Result<()> {
    let config = load_config(config_path)?;
    config.validate()?;

    tracing::info!("Configuration is valid");
    println!("✓ Configuration is valid");
    println!("Available targets: {}", config.list_targets().join(", "));

    Ok(())
}

fn show_config(config_path: &Path) -> anyhow::Result<()> {
    let config = load_config(config_path)?;
    let yaml = serde_yaml::to_string(&config)?;

    println!("Current configuration:");
    println!("{}", yaml);

    Ok(())
}

fn list_targets(config_path: &Path) -> anyhow::Result<()> {
    let config = load_config(config_path)?;

    println!("Available targets:");
    for target_name in config.list_targets() {
        let target = &config.targets[&target_name];

        println!("{}", target_name);
        println!("    Host: {}:{}", target.host, target.port);

        if let Some(ssh) = &target.ssh {
            if ssh.enabled {
                println!(
                    "    SSH: {}@{}",
                    ssh.user.as_ref().unwrap_or(&"<user>".to_string()),
                    ssh.host.as_ref().unwrap_or(&"<host>".to_string())
                );
            } else {
                println!("    SSH: disabled");
            }
        } else {
            println!("    SSH: not configured");
        }
        println!();
    }

    println!("Usage:");
    println!("  cargo run -- start --target <TARGET_NAME>");

    Ok(())
}

async fn health_check(config_path: &Path, target_name: Option<&str>) -> anyhow::Result<()> {
    let config = load_config(config_path)?;

    let targets_to_check = if let Some(target) = target_name {
        vec![target.to_string()]
    } else {
        config.list_targets()
    };

    println!("Health Check Results:");
    println!("====================");

    for target in targets_to_check {
        match config.get_target(&target) {
            Ok(target_config) => {
                println!("✓ Target '{}' configuration is valid", target);
                println!("  Host: {}:{}", target_config.host, target_config.port);

                if let Some(ssh) = &target_config.ssh {
                    if ssh.enabled {
                        println!(
                            "  SSH: Enabled ({}@{})",
                            ssh.user.as_ref().unwrap_or(&"<user>".to_string()),
                            ssh.host.as_ref().unwrap_or(&"<host>".to_string())
                        );
                        println!("    Auto-reconnect: {}", ssh.auto_reconnect);
                        println!("    Max reconnect attempts: {}", ssh.max_reconnect_attempts);
                        println!(
                            "    Reconnect interval: {}s",
                            ssh.reconnect_interval_seconds
                        );
                    } else {
                        println!("  SSH: Disabled");
                    }
                } else {
                    println!("  SSH: Not configured");
                }

                println!("  Status: Configuration valid (connectivity not tested in this version)");
            }
            Err(e) => {
                println!("✗ Target '{}': {}", target, e);
            }
        }
        println!();
    }

    Ok(())
}

fn load_config(config_path: &Path) -> Result<Config> {
    debug!(
        config_path = %config_path.display(),
        config_exists = config_path.exists(),
        "Loading configuration"
    );

    if !config_path.exists() {
        warn!(
            config_path = %config_path.display(),
            "Configuration file not found, using default configuration"
        );
        let default_config = Config::default();
        debug!(
            targets_count = default_config.targets.len(),
            listen_port = default_config.proxy.listen_port,
            "Using default configuration"
        );
        return Ok(default_config);
    }

    let config = Config::from_file(&config_path.to_path_buf()).with_context(|| {
        format!(
            "Failed to parse configuration file '{}'",
            config_path.display()
        )
    })?;

    config.validate().with_context(|| {
        format!(
            "Configuration validation failed for '{}'",
            config_path.display()
        )
    })?;

    info!(
        config_path = %config_path.display(),
        targets_count = config.targets.len(),
        listen_host = %config.proxy.listen_host,
        listen_port = config.proxy.listen_port,
        available_targets = ?config.list_targets(),
        "Configuration loaded successfully"
    );

    Ok(config)
}