server/crates/arbiter-server/src/actors/bootstrap.rs

@@ -6,22 +6,36 @@ use diesel_async::RunQueryDsl;
 use kameo::{Actor, messages};
 use rand::{RngExt, distr::Alphanumeric, rngs::SysRng};
 use rand_core::UnwrapErr;
+use std::path::{Path, PathBuf};
 use subtle::ConstantTimeEq as _;
 use thiserror::Error;
+use tracing::warn;
 use zeroize::Zeroizing;
 
 const TOKEN_LENGTH: usize = 64;
 
-pub async fn generate_token() -> Result<String, std::io::Error> {
-    let token = UnwrapErr(SysRng).sample_iter(Alphanumeric).take(TOKEN_LENGTH).fold(
-        String::default(),
-        |mut accum, char| {
+async fn write_token_file(path: &Path, content: &str) -> Result<(), std::io::Error> {
+    tokio::fs::write(path, content.as_bytes()).await?;
+
+    #[cfg(unix)]
+    {
+        use std::os::unix::fs::PermissionsExt as _;
+        tokio::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600)).await?;
+    }
+
+    Ok(())
+}
+
+async fn generate_token(path: &Path) -> Result<String, std::io::Error> {
+    let token = UnwrapErr(SysRng)
+        .sample_iter(Alphanumeric)
+        .take(TOKEN_LENGTH)
+        .fold(String::default(), |mut accum, char| {
             accum += char.to_string().as_str();
             accum
-        },
-    );
+        });
 
-    tokio::fs::write(home_path()?.join(BOOTSTRAP_PATH), token.as_str()).await?;
+    write_token_file(path, &token).await?;
 
     Ok(token)
 }
@@ -41,6 +55,7 @@ pub enum Error {
 #[derive(Actor)]
 pub struct Bootstrapper {
     token: Option<Zeroizing<String>>,
+    token_path: Option<PathBuf>,
 }
 
 impl Bootstrapper {
@@ -54,14 +69,15 @@ impl Bootstrapper {
                 .await?
         };
 
-        let token = if row_count == 0 {
-            let token = generate_token().await?;
-            Some(Zeroizing::new(token))
+        let (token, token_path) = if row_count == 0 {
+            let path = home_path()?.join(BOOTSTRAP_PATH);
+            let token = generate_token(&path).await?;
+            (Some(Zeroizing::new(token)), Some(path))
         } else {
-            None
+            (None, None)
         };
 
-        Ok(Self { token })
+        Ok(Self { token, token_path })
     }
 }
 
@@ -80,9 +96,14 @@ impl Bootstrapper {
 #[messages]
 impl Bootstrapper {
     #[message]
-    pub fn consume_token(&mut self, token: Zeroizing<String>) -> bool {
+    pub async fn consume_token(&mut self, token: Zeroizing<String>) -> bool {
         if self.is_correct_token(&token) {
             self.token = None;
+            if let Some(path) = self.token_path.take() {
+                if let Err(e) = tokio::fs::remove_file(&path).await {
+                    warn!(error = ?e, path = ?path, "Failed to delete bootstrap token file after consumption");
+                }
+            }
             true
         } else {
             false