Create separate files for psks.

6 files changed, 136 insertions, 93 deletions

diff --git a/src/config.rs b/src/config.rs
index 08db0eb..95e80a8 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -32,7 +32,7 @@ pub struct GlobalConfig {
     pub min_keepalive: u32,
     #[serde(default = "default_max_keepalive")]
     pub max_keepalive: u32,
-    #[serde(rename = "peer")]
+    #[serde(default, rename = "peer")]
     pub peers: HashMap<Key, Peer>,
 }
 
@@ -82,7 +82,7 @@ impl Default for UpdaterConfig {
 #[serde(deny_unknown_fields)]
 #[derive(serde_derive::Serialize, serde_derive::Deserialize, Default, Clone, Debug)]
 pub struct Config {
-    pub state_directory: Option<PathBuf>,
+    pub runtime_directory: Option<PathBuf>,
 
     #[serde(flatten)]
     pub global: GlobalConfig,
diff --git a/src/fileutil.rs b/src/fileutil.rs
new file mode 100644
index 0000000..0bbafc0
--- /dev/null
+++ b/src/fileutil.rs
@@ -0,0 +1,97 @@
+use std::ffi::OsString;
+#[cfg(unix)]
+use std::os::unix::fs::OpenOptionsExt;
+use std::path::{Path, PathBuf};
+use std::{fs, io, mem};
+
+#[repr(transparent)]
+pub struct Temp {
+    path: PathBuf,
+}
+
+impl Drop for Temp {
+    fn drop(&mut self) {
+        if self.path.as_os_str().is_empty() {
+            return;
+        }
+        fs::remove_file(&self.path).expect("Failed to clean up temporary file");
+    }
+}
+
+impl Temp {
+    #[inline]
+    pub fn path(&self) -> &Path {
+        &*self.path
+    }
+
+    #[inline]
+    pub fn leave(mut self) -> PathBuf {
+        mem::replace(&mut self.path, PathBuf::new())
+    }
+
+    #[inline]
+    pub fn rename_to(self, to: impl AsRef<Path>) -> io::Result<()> {
+        fs::rename(self.leave(), to)
+    }
+}
+
+pub struct Writer {
+    inner: Temp,
+    file: fs::File,
+}
+
+impl Writer {
+    pub fn new(path: PathBuf) -> io::Result<Self> {
+        let mut file = fs::OpenOptions::new();
+        file.create_new(true);
+        file.append(true);
+        #[cfg(unix)]
+        file.mode(0o0600);
+        let file = file.open(&path)?;
+
+        Ok(Self {
+            inner: Temp { path },
+            file,
+        })
+    }
+
+    #[inline]
+    pub fn file(&mut self) -> &mut fs::File {
+        &mut self.file
+    }
+
+    #[inline]
+    pub fn sync_done(self) -> io::Result<Temp> {
+        self.file.sync_data()?;
+        Ok(self.inner)
+    }
+
+    #[inline]
+    pub fn done(self) -> Temp {
+        self.inner
+    }
+}
+
+pub fn update(path: &Path, data: &[u8]) -> io::Result<()> {
+    let mut tmp = OsString::from(path);
+    tmp.push(".tmp");
+    let mut tmp = Writer::new(PathBuf::from(tmp))?;
+    io::Write::write_all(tmp.file(), data)?;
+    tmp.sync_done()?.rename_to(path)
+}
+
+pub fn load(path: &Path) -> io::Result<Option<Vec<u8>>> {
+    let mut file = match fs::File::open(&path) {
+        Ok(file) => file,
+        Err(e) => {
+            if e.kind() == io::ErrorKind::NotFound {
+                return Ok(None);
+            }
+            return Err(e);
+        }
+    };
+
+    let mut data = Vec::new();
+    io::Read::read_to_end(&mut file, &mut data)?;
+    Ok(Some(data))
+}
diff --git a/src/main.rs b/src/main.rs
index b6b5ce1..731357f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -17,6 +17,7 @@ use std::{fs, io};
 use toml;
 
 mod config;
+mod fileutil;
 mod manager;
 mod model;
 mod proto;
@@ -188,7 +189,7 @@ fn run_with_cmdline(argv0: &str, args: &mut impl Iterator<Item = OsString>) -> i
 
 fn run_daemon(ifname: OsString, mut config: config::Config) -> i32 {
     maybe_get_var(&mut config.updater.cache_directory, "CACHE_DIRECTORY");
-    maybe_get_var(&mut config.state_directory, "RUNTIME_DIRECTORY");
+    maybe_get_var(&mut config.runtime_directory, "RUNTIME_DIRECTORY");
 
     let mut m = match manager::Manager::new(ifname, config) {
         Ok(m) => m,
diff --git a/src/manager/mod.rs b/src/manager/mod.rs
index aa8b792..9afce72 100644
--- a/src/manager/mod.rs
+++ b/src/manager/mod.rs
@@ -2,55 +2,11 @@
 //
 // See COPYING.
 
-use crate::{config, model, proto, wg};
+use crate::{config, fileutil, model, proto, wg};
 use std::ffi::OsString;
-#[cfg(unix)]
-use std::os::unix::fs::OpenOptionsExt;
-use std::path::{Path, PathBuf};
+use std::io;
+use std::path::PathBuf;
 use std::time::{Duration, Instant, SystemTime};
-use std::{fs, io};
-
-fn update_file(path: &Path, data: &[u8]) -> io::Result<()> {
-    let mut tmp_path = OsString::from(path);
-    tmp_path.push(".tmp");
-    let tmp_path = PathBuf::from(tmp_path);
-
-    let mut file = {
-        let mut file = fs::OpenOptions::new();
-        file.append(true);
-        file.create_new(true);
-        #[cfg(unix)]
-        file.mode(0o0600);
-        file.open(&tmp_path)?
-    };
-
-    let r = io::Write::write_all(&mut file, data)
-        .and_then(|_| file.sync_data())
-        .and_then(|_| fs::rename(&tmp_path, &path));
-
-    if r.is_err() {
-        fs::remove_file(&tmp_path).unwrap_or_else(|e2| {
-            eprintln!("<3>Failed to clean up [{}]: {}", tmp_path.display(), e2);
-        });
-    }
-    r
-}
-
-fn load_file(path: &Path) -> io::Result<Option<Vec<u8>>> {
-    let mut file = match fs::File::open(&path) {
-        Ok(file) => file,
-        Err(e) => {
-            if e.kind() == io::ErrorKind::NotFound {
-                return Ok(None);
-            }
-            return Err(e);
-        }
-    };
-
-    let mut data = Vec::new();
-    io::Read::read_to_end(&mut file, &mut data)?;
-    Ok(Some(data))
-}
 
 struct Source {
     name: String,
@@ -70,18 +26,23 @@ pub struct Manager {
     global_config: config::GlobalConfig,
     sources: Vec<Source>,
     current: model::Config,
-    state_directory: Option<PathBuf>,
+    state_path: PathBuf,
     updater: updater::Updater,
 }
 
 impl Manager {
     pub fn new(ifname: OsString, c: config::Config) -> io::Result<Self> {
+        let runtime_directory = c.runtime_directory.ok_or_else(|| {
+            io::Error::new(io::ErrorKind::InvalidInput, "Runtime directory required")
+        })?;
+        let mut state_path = runtime_directory.clone();
+        state_path.push("state.json");
         let mut m = Self {
-            dev: wg::Device::open(ifname)?,
+            dev: wg::Device::open(ifname, runtime_directory)?,
             global_config: c.global,
             sources: vec![],
             current: model::Config::empty(),
-            state_directory: c.state_directory,
+            state_path,
             updater: updater::Updater::new(c.updater),
         };
 
@@ -94,19 +55,8 @@ impl Manager {
         Ok(m)
     }
 
-    fn state_path(&self) -> Option<PathBuf> {
-        let mut path = self.state_directory.as_ref()?.clone();
-        path.push("state.json");
-        Some(path)
-    }
-
     fn current_load(&mut self) -> bool {
-        let path = match self.state_path() {
-            Some(v) => v,
-            None => return false,
-        };
-
-        let data = match load_file(&path) {
+        let data = match fileutil::load(&self.state_path) {
             Ok(Some(data)) => data,
             Ok(None) => {
                 return false;
@@ -131,13 +81,8 @@ impl Manager {
     }
 
     fn current_update(&mut self, c: &model::Config) {
-        let path = match self.state_path() {
-            Some(v) => v,
-            None => return,
-        };
-
         let data = serde_json::to_vec(c).unwrap();
-        match update_file(&path, &data) {
+        match fileutil::update(&self.state_path, &data) {
             Ok(()) => {}
             Err(e) => {
                 eprintln!("<3>Failed to persist interface state: {}", e);
diff --git a/src/manager/updater.rs b/src/manager/updater.rs
index 14d55fc..67ed953 100644
--- a/src/manager/updater.rs
+++ b/src/manager/updater.rs
@@ -2,8 +2,8 @@
 //
 // See COPYING.
 
-use super::{load_file, update_file, Source};
-use crate::{config, proto};
+use super::Source;
+use crate::{config, fileutil, proto};
 use std::ffi::{OsStr, OsString};
 use std::path::PathBuf;
 use std::time::{Duration, Instant};
@@ -31,7 +31,7 @@ impl Updater {
         };
 
         let data = serde_json::to_vec(&src.data).unwrap();
-        match update_file(&path, &data) {
+        match fileutil::update(&path, &data) {
             Ok(()) => {}
             Err(e) => {
                 eprintln!("<4>Failed to cache [{}]: {}", &src.name, e);
@@ -40,12 +40,12 @@ impl Updater {
     }
 
     pub fn cache_load(&self, src: &mut Source) -> bool {
-        let path = match  self.cache_path(src) {
+        let path = match self.cache_path(src) {
             Some(v) => v,
             None => return false,
         };
 
-        let data = match load_file(&path) {
+        let data = match fileutil::load(&path) {
             Ok(Some(data)) => data,
             Ok(None) => {
                 return false;
diff --git a/src/wg.rs b/src/wg.rs
index 7565a02..004e6d8 100644
--- a/src/wg.rs
+++ b/src/wg.rs
@@ -2,19 +2,21 @@
 //
 // See COPYING.
 
-use crate::model;
+use crate::{fileutil, model};
 use std::ffi::{OsStr, OsString};
+use std::path::PathBuf;
 use std::process::{Command, Stdio};
-use std::{env, io};
+use std::{env, io, mem};
 
 pub struct Device {
     ifname: OsString,
+    tmpdir: PathBuf,
 }
 
 impl Device {
     #[inline]
-    pub fn open(ifname: OsString) -> io::Result<Self> {
-        let dev = Self { ifname };
+    pub fn open(ifname: OsString, tmpdir: PathBuf) -> io::Result<Self> {
+        let dev = Self { ifname, tmpdir };
         let _ = dev.get_public_key()?;
         Ok(dev)
     }
@@ -55,11 +57,10 @@ impl Device {
 
     pub fn apply_diff(&mut self, old: &model::Config, new: &model::Config) -> io::Result<()> {
         let mut proc = Self::wg_command();
-        proc.stdin(Stdio::piped());
         proc.arg("set");
         proc.arg(&self.ifname);
 
-        let mut psks = String::new();
+        let mut tmps = vec![];
 
         for (pubkey, conf) in &new.peers {
             let old_endpoint;
@@ -87,11 +88,16 @@ impl Device {
 
             if let Some(psk) = &conf.psk {
                 proc.arg("preshared-key");
-                proc.arg("/dev/stdin");
+                let mut tmp = self.tmpdir.clone();
+                tmp.push(format!("tmp-{}", tmps.len()));
+                let mut tmp = fileutil::Writer::new(tmp)?;
                 {
-                    use std::fmt::Write;
-                    writeln!(&mut psks, "{}", psk).unwrap();
+                    use io::Write;
+                    writeln!(tmp.file(), "{}", psk)?;
                 }
+                let tmp = tmp.done();
+                proc.arg(tmp.path());
+                tmps.push(tmp);
             }
 
             let mut ips = String::new();
@@ -124,14 +130,8 @@ impl Device {
             proc.arg("remove");
         }
 
-        let mut proc = proc.spawn()?;
-        {
-            use std::io::Write;
-            let stdin = proc.stdin.as_mut().unwrap();
-            write!(stdin, "{}", psks)?;
-        }
-
-        let r = proc.wait()?;
+        let r = proc.status()?;
+        mem::drop(tmps);
         if !r.success() {
             return Err(io::Error::new(io::ErrorKind::Other, "Child process failed"));
         }