herodb/src/main.rs

// #![allow(unused_imports)]

use std::path::PathBuf;
use tokio::net::TcpListener;

use herodb::server;
use herodb::rpc_server;

use clap::Parser;

/// Simple program to greet a person
#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct Args {
    /// The directory of Redis DB file
    #[arg(long)]
    dir: PathBuf,

    /// The port of the Redis server, default is 6379 if not specified
    #[arg(long)]
    port: Option<u16>,

    /// Enable debug mode
    #[arg(long)]
    debug: bool,

    /// Master encryption key for encrypted databases (deprecated; ignored for data DBs)
    #[arg(long)]
    encryption_key: Option<String>,

    /// Encrypt the database (deprecated; ignored for data DBs)
    #[arg(long)]
    encrypt: bool,

    /// Enable RPC management server
    #[arg(long)]
    enable_rpc: bool,

    /// RPC server port (default: 8080)
    #[arg(long, default_value = "8080")]
    rpc_port: u16,

    /// Enable RPC over Unix Domain Socket (IPC)
    #[arg(long)]
    enable_rpc_ipc: bool,

    /// RPC IPC socket path (Unix Domain Socket)
    #[arg(long, default_value = "/tmp/herodb.ipc")]
    rpc_ipc_path: String,

    /// Use the sled backend
    #[arg(long)]
    sled: bool,

    /// Admin secret used to encrypt DB 0 and authorize admin access (required)
    #[arg(long)]
    admin_secret: String,
}

#[tokio::main]
async fn main() {
    // parse args
    let args = Args::parse();

    // bind port
    let port = args.port.unwrap_or(6379);
    println!("will listen on port: {}", port);
    let listener = TcpListener::bind(format!("127.0.0.1:{}", port))
        .await
        .unwrap();

    // deprecation warnings for legacy flags
    if args.encrypt || args.encryption_key.is_some() {
        eprintln!("warning: --encrypt and --encryption-key are deprecated and ignored for data DBs. Admin DB 0 is always encrypted with --admin-secret.");
    }
    // basic validation for admin secret
    if args.admin_secret.trim().is_empty() {
        eprintln!("error: --admin-secret must not be empty");
        std::process::exit(2);
    }

    // new DB option
    let option = herodb::options::DBOption {
        dir: args.dir.clone(),
        port,
        debug: args.debug,
        encryption_key: args.encryption_key,
        encrypt: args.encrypt,
        backend: if args.sled {
            herodb::options::BackendType::Sled
        } else {
            herodb::options::BackendType::Redb
        },
        admin_secret: args.admin_secret.clone(),
    };

    let backend = option.backend.clone();

    // Bootstrap admin DB 0 before opening any server storage
    if let Err(e) = herodb::admin_meta::ensure_bootstrap(&args.dir, backend.clone(), &args.admin_secret) {
        eprintln!("Failed to bootstrap admin DB 0: {}", e.0);
        std::process::exit(2);
    }

    // new server
    let server = server::Server::new(option).await;

    // Add a small delay to ensure the port is ready
    tokio::time::sleep(std::time::Duration::from_millis(100)).await;

    // Start RPC server if enabled
    let _rpc_handle = if args.enable_rpc {
        let rpc_addr = format!("127.0.0.1:{}", args.rpc_port).parse().unwrap();
        let base_dir = args.dir.clone();

        match rpc_server::start_rpc_server(rpc_addr, base_dir, backend.clone(), args.admin_secret.clone()).await {
            Ok(handle) => {
                println!("RPC management server started on port {}", args.rpc_port);
                Some(handle)
            }
            Err(e) => {
                eprintln!("Failed to start RPC server: {}", e);
                None
            }
        }
    } else {
        None
    };

    // Start IPC (Unix socket) RPC server if enabled
    let _rpc_ipc_handle = if args.enable_rpc_ipc {
        let base_dir = args.dir.clone();
        let ipc_path = args.rpc_ipc_path.clone();

        // Remove stale socket if present
        if std::path::Path::new(&ipc_path).exists() {
            let _ = std::fs::remove_file(&ipc_path);
        }

        match rpc_server::start_rpc_ipc_server(ipc_path.clone(), base_dir, backend.clone(), args.admin_secret.clone()).await {
            Ok(handle) => {
                println!("RPC IPC server started at {}", ipc_path);
                Some(handle)
            }
            Err(e) => {
                eprintln!("Failed to start RPC IPC server: {}", e);
                None
            }
        }
    } else {
        None
    };

    // accept new connections
    loop {
        let stream = listener.accept().await;
        match stream {
            Ok((stream, _)) => {
                println!("accepted new connection");

                let mut sc = server.clone();
                tokio::spawn(async move {
                    if let Err(e) = sc.handle(stream).await {
                        println!("error: {:?}, will close the connection. Bye", e);
                    }
                });
            }
            Err(e) => {
                println!("error: {}", e);
            }
        }
    }
}