merge: custom db ids for entities

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

@@ -1,29 +1,48 @@
 use crate::db::{self, DatabasePool, schema};
+use arbiter_crypto::safecell::{SafeCell, SafeCellHandle as _};
 use arbiter_proto::{BOOTSTRAP_PATH, home_path};
 
 use diesel::QueryDsl;
 use diesel_async::RunQueryDsl;
 use kameo::{Actor, messages};
-use rand::{RngExt, distr::Alphanumeric, make_rng, rngs::StdRng};
+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;
 
 const TOKEN_LENGTH: usize = 64;
 
-pub async fn generate_token() -> Result<String, std::io::Error> {
-    let rng: StdRng = make_rng();
+async fn write_token_file(path: &Path, content: &str) -> Result<(), std::io::Error> {
+    tokio::fs::write(path, content.as_bytes()).await?;
 
-    let token = rng.sample_iter(Alphanumeric).take(TOKEN_LENGTH).fold(
-        String::default(),
-        |mut accum, char| {
-            accum += char.to_string().as_str();
-            accum
-        },
-    );
+    #[cfg(unix)]
+    {
+        use std::os::unix::fs::PermissionsExt as _;
+        tokio::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600)).await?;
+    }
 
-    tokio::fs::write(home_path()?.join(BOOTSTRAP_PATH), token.as_str()).await?;
+    Ok(())
+}
 
-    Ok(token)
+async fn generate_token(path: &Path) -> Result<SafeCell<[u8; TOKEN_LENGTH]>, std::io::Error> {
+    let mut cell = SafeCell::new([0u8; TOKEN_LENGTH]);
+    {
+        let mut buf = cell.write();
+        for (slot, b) in buf
+            .iter_mut()
+            .zip(UnwrapErr(SysRng).sample_iter(Alphanumeric))
+        {
+            *slot = b;
+        }
+    }
+
+    let token_str = cell.read_inline(|buf| String::from_utf8_lossy(buf.as_ref()).into_owned());
+
+    write_token_file(path, &token_str).await?;
+
+    Ok(cell)
 }
 
 #[derive(Error, Debug)]
@@ -40,7 +59,8 @@ pub enum Error {
 
 #[derive(Actor)]
 pub struct Bootstrapper {
-    token: Option<String>,
+    token: Option<SafeCell<[u8; TOKEN_LENGTH]>>,
+    token_path: Option<PathBuf>,
 }
 
 impl Bootstrapper {
@@ -54,34 +74,37 @@ impl Bootstrapper {
                 .await?
         };
 
-        let token = if row_count == 0 {
-            let token = generate_token().await?;
-            Some(token)
+        let (token, token_path) = if row_count == 0 {
+            let path = home_path()?.join(BOOTSTRAP_PATH);
+            let token = generate_token(&path).await?;
+            (Some(token), Some(path))
         } else {
-            None
+            (None, None)
         };
 
-        Ok(Self { token })
+        Ok(Self { token, token_path })
+    }
+}
+
+impl Bootstrapper {
+    fn is_correct_token(&mut self, token: &[u8]) -> bool {
+        self.token.as_mut().is_some_and(|expected| {
+            expected.read_inline(|exp| bool::from(exp.as_ref().ct_eq(token)))
+        })
     }
 }
 
 #[messages]
 impl Bootstrapper {
     #[message]
-    pub fn is_correct_token(&self, token: String) -> bool {
-        self.token.as_ref().is_some_and(|expected| {
-            let expected_bytes = expected.as_bytes();
-            let token_bytes = token.as_bytes();
-
-            let choice = expected_bytes.ct_eq(token_bytes);
-            bool::from(choice)
-        })
-    }
-
-    #[message]
-    pub fn consume_token(&mut self, token: String) -> bool {
-        if self.is_correct_token(token) {
+    pub async fn consume_token(&mut self, token: Vec<u8>) -> bool {
+        if self.is_correct_token(&token) {
             self.token = None;
+            if let Some(path) = self.token_path.take()
+                && 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
@@ -92,7 +115,9 @@ impl Bootstrapper {
 #[messages]
 impl Bootstrapper {
     #[message]
-    pub fn get_token(&self) -> Option<String> {
-        self.token.clone()
+    pub fn get_token(&mut self) -> Option<String> {
+        self.token
+            .as_mut()
+            .map(|cell| cell.read_inline(|buf| String::from_utf8_lossy(buf.as_ref()).into_owned()))
     }
 }