diff --git a/common/Cargo.toml b/common/Cargo.toml index d8343cd1d..b72dcae6f 100644 --- a/common/Cargo.toml +++ b/common/Cargo.toml @@ -17,3 +17,4 @@ clap = { workspace = true, features = ["derive"] } idgenerator = { workspace = true } serde = { workspace = true } config = { workspace = true } +envsubst = "0.2.1" diff --git a/common/src/config.rs b/common/src/config.rs index 319245864..223572ce6 100644 --- a/common/src/config.rs +++ b/common/src/config.rs @@ -1,10 +1,16 @@ +use std::cell::RefCell; +use std::collections::HashMap; use c::{ConfigError, FileFormat}; use config as c; use serde::{Deserialize, Serialize}; use std::path::PathBuf; +use std::rc::Rc; +use config::builder::DefaultState; +use config::{Source, ValueKind}; #[derive(Serialize, Deserialize, Debug, Default, Clone)] pub struct Config { + pub base_dir: PathBuf, pub log: LogConfig, pub database: DbConfig, pub ssh: SshConfig, @@ -17,8 +23,11 @@ pub struct Config { impl Config { pub fn new(path: &str) -> Result { - let builder = c::Config::builder().add_source(c::File::new(path, FileFormat::Toml)); - let config = builder.build().unwrap(); + let builder = c::Config::builder() + .add_source(c::File::new(path, FileFormat::Toml)) + .add_source(c::Environment::with_prefix("mega")); // e.g. MEGA_BASE_DIR == base_dir + // support ${} variable substitution + let config = variable_placeholder_substitute(builder); Config::from_config(config) } @@ -29,6 +38,77 @@ impl Config { } } +/// supports braces-delimited variables (i.e. ${foo}) in config. +/// ### Example: +/// ```toml +/// base_dir = "/tmp/.mega" +/// [log] +/// log_path = "${base_dir}/logs" +/// ``` +/// ### Limitations: +/// - only support `String` type. +/// - vars apply from up to down +fn variable_placeholder_substitute(mut builder: c::ConfigBuilder) -> c::Config { + // `Config::set` is deprecated, use `ConfigBuilder::set_override` instead + let config = builder.clone().build().unwrap(); // initial config + let mut vars = HashMap::new(); + // top-level variables + for (k, mut v) in config.collect().unwrap() { // a copy + if let ValueKind::String(str) = &v.kind { + if envsubst::is_templated(str) { + let new_str = envsubst::substitute(str, &vars).unwrap(); + v.kind = ValueKind::String(new_str.clone()); + builder = builder.set_override(&k, v).unwrap(); + vars.insert(k, new_str); + } else { + vars.insert(k, str.clone()); + } + } + } + // second-level or nested variables + // extract all config k-v + let map = Rc::new(RefCell::new(HashMap::new())); + for (k, v) in config.collect().unwrap() { + if let ValueKind::Table(_) = v.kind { + let map_c = map.clone(); + traverse_config(&k, &v, &move |key: &str, value: &c::Value| { + if let ValueKind::String(_) = value.kind { + map_c.borrow_mut().insert(key.to_string(), value.clone()); + } + }); + } + } + + // do substitution: ${} -> real value + for (k, mut v) in Rc::try_unwrap(map).unwrap().into_inner() { + let mut str = v.clone().into_string().unwrap(); + if envsubst::is_templated(&str) { + let new_str = envsubst::substitute(&str, &vars).unwrap(); + println!("{}: {} -> {}", k, str, &new_str); + v.kind = ValueKind::String(new_str.clone()); + builder = builder.set_override(&k, v).unwrap(); + str = new_str; + } + vars.insert(k, str); + } + + builder.build().unwrap() +} + +/// visitor pattern: traverse each config & execute the closure `f` +fn traverse_config(key: &str, value: &c::Value, f: &impl Fn(&str, &c::Value)) { + match &value.kind { + ValueKind::Table(table) => { + for (k, v) in table.iter() { + // join keys by '.' + let new_key = if key.is_empty() { k.clone() } else { format!("{}.{}", key, k) }; + traverse_config(&new_key, v, f); + } + } + _ => f(key, value), + } +} + #[derive(Serialize, Deserialize, Debug, Clone)] pub struct LogConfig { pub log_path: PathBuf, diff --git a/mega/Cargo.toml b/mega/Cargo.toml index 97a004acd..6cb33524f 100644 --- a/mega/Cargo.toml +++ b/mega/Cargo.toml @@ -29,6 +29,7 @@ rand = { workspace = true } smallvec = { workspace = true } config = { workspace = true } shadow-rs = { workspace = true } +ctrlc = "3.4.4" [dev-dependencies] reqwest = { workspace = true , features = ["stream", "json"] } diff --git a/mega/config.toml b/mega/config.toml index e30d25f86..e8417e79d 100644 --- a/mega/config.toml +++ b/mega/config.toml @@ -1,8 +1,12 @@ +# the directory where the data files is located, such as logs, database, etc. +# can be overrided by environment variable `MEGA_BASE_DIR` +base_dir = "/tmp/.mega" + # Filling the following environment variables with values you set ## Logging Configuration [log] # The path which log file is saved -log_path = "/tmp/.mega/logs" +log_path = "${base_dir}/logs" # log level level = "debug" @@ -17,7 +21,7 @@ print_std = true db_type = "sqlite" # used for sqlite -db_path = "/tmp/.mega/mega.db" +db_path = "${base_dir}/mega.db" # database connection url db_url = "postgres://mega:mega@localhost:5432/mega" @@ -33,7 +37,7 @@ sqlx_logging = false [ssh] -ssh_key_path = "/tmp/.mega/ssh" +ssh_key_path = "${base_dir}/ssh" [storage] # raw object stroage type, can be `local` or `remote` @@ -43,9 +47,9 @@ raw_obj_storage_type = "LOCAL" big_obj_threshold = 1024 # set the local path of the project storage -raw_obj_local_path = "/tmp/.mega/objects" +raw_obj_local_path = "${base_dir}/objects" -lfs_obj_local_path = "/tmp/.mega/lfs" +lfs_obj_local_path = "${base_dir}/lfs" obs_access_key = "" obs_secret_key = "" @@ -68,7 +72,7 @@ import_dir = "/third-part" pack_decode_mem_size = 4 # The location where the object stored when the memory used by decode exceeds the limit -pack_decode_cache_path = "/tmp/.mega/cache" +pack_decode_cache_path = "${base_dir}/cache" clean_cache_after_decode = true diff --git a/mega/src/cli.rs b/mega/src/cli.rs index fd496a4fe..0cffcc849 100644 --- a/mega/src/cli.rs +++ b/mega/src/cli.rs @@ -1,7 +1,7 @@ //! Cli module is responsible for parsing command line arguments and executing the appropriate. use std::env; - +use std::path::PathBuf; use clap::{Arg, ArgMatches, Command}; use tracing_subscriber::fmt::writer::MakeWriterExt; @@ -31,8 +31,8 @@ pub fn parse(args: Option>) -> MegaResult { // Get the path to the config file in the current directory let config_path = current_dir.join("config.toml"); - let config = if let Some(path) = matches.get_one::("config").cloned() { - Config::new(path.as_str()).unwrap() + let config = if let Some(path) = matches.get_one::("config").cloned() { + Config::new(path.to_str().unwrap()).unwrap() } else if config_path.exists() { Config::new(config_path.to_str().unwrap()).unwrap() } else { @@ -42,6 +42,11 @@ pub fn parse(args: Option>) -> MegaResult { init_log(&config.log); + ctrlc::set_handler(move || { + tracing::info!("Received Ctrl-C signal, exiting..."); + std::process::exit(0); + }).unwrap(); + let (cmd, subcommand_args) = match matches.subcommand() { Some((cmd, args)) => (cmd, args), _ => {