Adding config file locations and yaml support

src/main.rs

@@ -3,6 +3,8 @@ mod shutdown;
 mod tcp;
 
 use serde::Deserialize;
+use simple_error::bail;
+use std::error::Error;
 use std::fs::File;
 use std::path::Path;
 use tokio::sync::broadcast;
@@ -19,14 +21,48 @@ struct Config {
     // udp: Vec<Target>,
 }
 
+fn load_yaml(path: &Path) -> Result<Config, Box<dyn Error>> {
+    let file = File::open(path)?;
+    let config: Config = serde_yaml::from_reader(file).expect("Failed to parse!"); //TODO: Print path
+
+    return Ok(config);
+}
+
+fn load_json(path: &Path) -> Result<Config, Box<dyn Error>> {
+    let file = File::open(path)?;
+    let config: Config = serde_json::from_reader(file).expect("Failed to parse!"); //TODO: Print path
+
+    return Ok(config);
+}
+
+fn load_config() -> Result<Config, Box<dyn Error>> {
+    for path in [
+        "config.yaml",
+        "config.json",
+        "/etc/rustocat.yaml",
+        "/etc/rustocat.json",
+    ]
+    .iter()
+    {
+        // if(p)
+        let config = if path.ends_with(".yaml") {
+            load_yaml(Path::new(path))
+        } else {
+            load_json(Path::new(path))
+        };
+        if config.is_ok() {
+            return config;
+        }
+    }
+    bail!("No config file found");
+}
+
 #[tokio::main]
 async fn main() {
-    let json_file_path = Path::new("./config.json");
-    let file = File::open(json_file_path).expect("file not found");
+    let config = load_config().expect("config not found");
 
     let (notify_shutdown, _) = broadcast::channel(1);
 
-    let config: Config = serde_json::from_reader(file).expect("json parse error");
     for target in config.tcp {
         let listener = listener::Listener {
             shutdown: shutdown::Shutdown::new(notify_shutdown.subscribe()),

src/shutdown.rs

@@ -28,9 +28,9 @@ impl Shutdown {
    }
 
    /// Returns `true` if the shutdown signal has been received.
-   pub(crate) fn is_shutdown(&self) -> bool {
-      self.shutdown
-   }
+   // pub(crate) fn is_shutdown(&self) -> bool {
+   //    self.shutdown
+   // }
 
    /// Receive the shutdown notice, waiting if necessary.
    pub(crate) async fn recv(&mut self) {