improve Config module, supporting ${} syntax

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"

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<Self, ConfigError> {
-        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<DefaultState>) -> 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,

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"] }

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
 

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<Vec<&str>>) -> 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::<String>("config").cloned() {
-        Config::new(path.as_str()).unwrap()
+    let config = if let Some(path) = matches.get_one::<PathBuf>("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<Vec<&str>>) -> 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),
         _ => {