use crate::error::{IdentityError, Result}; use mtp::crypto::{Keyring, PublicKeyBundle}; use mtp::files::{ FileError, load_keyring, load_public_key_bundle, save_keyring, save_public_key_bundle, }; use std::{ fs, path::{Path, PathBuf}, sync::{ Arc, atomic::{AtomicU64, Ordering}, }, }; pub const KEYRING_PATH: &str = "./omega.mk"; pub const PUBLIC_KEY_PATH: &str = "./omega.mpkb"; pub struct OmegaIdentity { keyring: Arc, } static PUBLIC_BUNDLE_TEMP_COUNTER: AtomicU64 = AtomicU64::new(0); impl OmegaIdentity { pub fn load_or_create(passphrase: &[u8]) -> Result { Self::load_or_create_at( Path::new(KEYRING_PATH), Path::new(PUBLIC_KEY_PATH), passphrase, ) } pub(crate) fn load_or_create_at( keyring_path: impl AsRef, public_key_path: impl AsRef, passphrase: &[u8], ) -> Result { let keyring_path = keyring_path.as_ref(); let public_key_path = public_key_path.as_ref(); let keyring = match load_keyring(keyring_path, passphrase) { Ok(keyring) => keyring, Err(FileError::Io(error)) if error.kind() == std::io::ErrorKind::NotFound => { return Self::create_at(keyring_path, public_key_path, passphrase); } Err(error) => { return Err(IdentityError::Storage { path: keyring_path.to_path_buf(), source: error, } .into()); } }; let identity = Self { keyring: Arc::new(keyring), }; match load_public_key_bundle(public_key_path) { Ok(persisted_bundle) => { identity.verify_public_bundle(public_key_path, &persisted_bundle)? } Err(FileError::Io(error)) if error.kind() == std::io::ErrorKind::NotFound => { Self::persist_public_bundle(&identity.public_key_bundle(), public_key_path)?; } Err(error) => { return Err(IdentityError::Storage { path: public_key_path.to_path_buf(), source: error, } .into()); } } Ok(identity) } fn create_at(keyring_path: &Path, public_key_path: &Path, passphrase: &[u8]) -> Result { let keyring = Keyring::generate(); save_keyring(&keyring, keyring_path, passphrase).map_err(|error| { IdentityError::Storage { path: keyring_path.to_path_buf(), source: error, } })?; Self::persist_public_bundle(&keyring.public_key_bundle(), public_key_path)?; Ok(Self { keyring: Arc::new(keyring), }) } fn persist_public_bundle(bundle: &PublicKeyBundle, path: &Path) -> Result<()> { let temporary_path = temporary_path(path)?; let result = save_public_key_bundle(bundle, &temporary_path).map_err(|source| { IdentityError::Storage { path: path.to_path_buf(), source, } }); if let Err(error) = result { let _ = fs::remove_file(&temporary_path); return Err(error.into()); } if let Err(source) = fs::File::open(&temporary_path).and_then(|file| file.sync_all()) { let _ = fs::remove_file(&temporary_path); return Err(IdentityError::Io { path: path.to_path_buf(), source, } .into()); } if let Err(source) = fs::rename(&temporary_path, path) { let _ = fs::remove_file(&temporary_path); return Err(IdentityError::Io { path: path.to_path_buf(), source, } .into()); } Ok(()) } fn verify_public_bundle(&self, path: &Path, persisted_bundle: &PublicKeyBundle) -> Result<()> { let expected = self.keyring.public_key_bundle().try_as_bytes()?; let actual = persisted_bundle.try_as_bytes()?; if expected != actual { return Err(IdentityError::PublicBundleMismatch { path: path.to_path_buf(), } .into()); } Ok(()) } pub fn keyring(&self) -> &Keyring { &self.keyring } pub fn public_key_bundle(&self) -> PublicKeyBundle { self.keyring().public_key_bundle() } pub fn clone_keyring(&self) -> Result { let bytes = self.keyring.try_to_bytes()?; Ok(Keyring::from_bytes(&bytes)?) } #[cfg(test)] pub(crate) fn from_keyring(keyring: Keyring) -> Self { Self { keyring: Arc::new(keyring), } } } fn temporary_path(path: &Path) -> Result { let parent = path.parent().unwrap_or_else(|| Path::new(".")); let file_name = path.file_name().ok_or_else(|| IdentityError::Io { path: path.to_path_buf(), source: std::io::Error::new( std::io::ErrorKind::InvalidInput, "identity path has no file name", ), })?; let mut temporary_name = file_name.to_os_string(); temporary_name.push(format!( ".tmp-{}-{}", std::process::id(), PUBLIC_BUNDLE_TEMP_COUNTER.fetch_add(1, Ordering::Relaxed) )); Ok(parent.join(temporary_name)) } #[cfg(test)] mod tests { use super::*; use std::fs; use std::sync::atomic::{AtomicU64, Ordering}; fn test_directory() -> std::path::PathBuf { static COUNTER: AtomicU64 = AtomicU64::new(0); let id = COUNTER.fetch_add(1, Ordering::Relaxed); let path = std::env::temp_dir().join(format!("omega-identity-test-{}-{id}", std::process::id())); fs::create_dir_all(&path).expect("create test directory"); path } #[test] fn generated_identity_survives_restart_with_separate_file_formats() { let directory = test_directory(); let keyring_path = directory.join("omega.mk"); let public_key_path = directory.join("omega.mpkb"); let passphrase = b"test-passphrase"; let first = OmegaIdentity::load_or_create_at(&keyring_path, &public_key_path, passphrase) .expect("create identity"); let first_bundle = first.public_key_bundle().try_as_bytes().expect("bundle"); let keyring_bytes = fs::read(&keyring_path).expect("read keyring"); let bundle_bytes = fs::read(&public_key_path).expect("read bundle"); assert_eq!(&keyring_bytes[..4], b"MTMK"); assert_eq!(&bundle_bytes[..4], b"MPKB"); let restarted = OmegaIdentity::load_or_create_at(&keyring_path, &public_key_path, passphrase) .expect("reload identity"); assert_eq!( restarted .public_key_bundle() .try_as_bytes() .expect("bundle"), first_bundle ); assert_eq!( load_public_key_bundle(&public_key_path) .expect("load public bundle") .try_as_bytes() .expect("bundle"), first_bundle ); fs::remove_dir_all(directory).expect("remove test directory"); } #[test] fn invalid_existing_keyring_does_not_create_a_new_identity() { let directory = test_directory(); let keyring_path = directory.join("omega.mk"); let public_key_path = directory.join("omega.mpkb"); fs::write(&keyring_path, b"not-a-keyring").expect("write invalid keyring"); let result = OmegaIdentity::load_or_create_at(&keyring_path, &public_key_path, b"test-passphrase"); assert!(result.is_err()); assert!(!public_key_path.exists()); fs::remove_dir_all(directory).expect("remove test directory"); } #[test] fn missing_public_bundle_is_repaired_without_changing_keyring() { let directory = test_directory(); let keyring_path = directory.join("omega.mk"); let public_key_path = directory.join("omega.mpkb"); let passphrase = b"test-passphrase"; OmegaIdentity::load_or_create_at(&keyring_path, &public_key_path, passphrase) .expect("create identity"); let original_keyring = fs::read(&keyring_path).expect("read keyring"); fs::remove_file(&public_key_path).expect("remove bundle"); let repaired = OmegaIdentity::load_or_create_at(&keyring_path, &public_key_path, passphrase) .expect("repair bundle"); assert_eq!( fs::read(&keyring_path).expect("read keyring"), original_keyring ); assert_eq!( repaired.public_key_bundle().try_as_bytes().expect("bundle"), load_public_key_bundle(&public_key_path) .expect("load bundle") .try_as_bytes() .expect("bundle") ); fs::remove_dir_all(directory).expect("remove test directory"); } #[test] fn mismatched_public_bundle_stops_startup() { let directory = test_directory(); let keyring_path = directory.join("omega.mk"); let public_key_path = directory.join("omega.mpkb"); let passphrase = b"test-passphrase"; OmegaIdentity::load_or_create_at(&keyring_path, &public_key_path, passphrase) .expect("create identity"); let other_keyring = Keyring::generate(); save_public_key_bundle(&other_keyring.public_key_bundle(), &public_key_path) .expect("save mismatched bundle"); assert!(matches!( OmegaIdentity::load_or_create_at(&keyring_path, &public_key_path, passphrase), Err(crate::OmegaError::Identity( IdentityError::PublicBundleMismatch { .. } )) )); fs::remove_dir_all(directory).expect("remove test directory"); } #[test] fn wrong_passphrase_does_not_replace_existing_keyring() { let directory = test_directory(); let keyring_path = directory.join("omega.mk"); let public_key_path = directory.join("omega.mpkb"); let passphrase = b"test-passphrase"; OmegaIdentity::load_or_create_at(&keyring_path, &public_key_path, passphrase) .expect("create identity"); let original_keyring = fs::read(&keyring_path).expect("read keyring"); assert!( OmegaIdentity::load_or_create_at(&keyring_path, &public_key_path, b"wrong-passphrase") .is_err() ); assert_eq!( fs::read(&keyring_path).expect("read keyring"), original_keyring ); fs::remove_dir_all(directory).expect("remove test directory"); } }