From 54aa8016325d00bec8726722b3f330c500c762bd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jorge=20Ant=C3=B3nio?= Date: Wed, 18 Dec 2024 10:13:16 +0000 Subject: [PATCH] feat: remove file storing for symmetric key and submit rotate key on node start (#289) * remove file storing for symmetric key and submit rotate key on node initialization * refactor code * address pr comments --- atoma-bin/atoma_node.rs | 17 +++-- atoma-confidential/src/key_management.rs | 86 ++---------------------- atoma-confidential/src/service.rs | 51 +++++++++++++- 3 files changed, 62 insertions(+), 92 deletions(-) diff --git a/atoma-bin/atoma_node.rs b/atoma-bin/atoma_node.rs index e19e5f91..71ff0a94 100644 --- a/atoma-bin/atoma_node.rs +++ b/atoma-bin/atoma_node.rs @@ -227,17 +227,16 @@ async fn main() -> Result<()> { let (compute_shared_secret_sender, compute_shared_secret_receiver) = tokio::sync::mpsc::unbounded_channel(); - let confidential_compute_service = AtomaConfidentialComputeService::new( - client.clone(), - subscriber_confidential_compute_receiver, - app_state_decryption_receiver, - app_state_encryption_receiver, - compute_shared_secret_receiver, - shutdown_receiver.clone(), - )?; spawn_with_shutdown( - async move { confidential_compute_service.run().await }, + AtomaConfidentialComputeService::start_confidential_compute_service( + client.clone(), + subscriber_confidential_compute_receiver, + app_state_decryption_receiver, + app_state_encryption_receiver, + compute_shared_secret_receiver, + shutdown_receiver.clone(), + ), shutdown_sender.clone(), ); diff --git a/atoma-confidential/src/key_management.rs b/atoma-confidential/src/key_management.rs index 3eec2c5f..7a1066e7 100644 --- a/atoma-confidential/src/key_management.rs +++ b/atoma-confidential/src/key_management.rs @@ -4,15 +4,6 @@ use atoma_utils::encryption::{ use thiserror::Error; use x25519_dalek::{PublicKey, SharedSecret, StaticSecret}; -/// The size of the X25519 secret key in bytes. -const DH_SECRET_KEY_SIZE: usize = 32; - -/// The directory where the private key file is stored. -const KEY_FILE_DIR: &str = "keys"; - -/// The name of the private key file. -const KEY_FILE_NAME: &str = "dh_privkey"; - type Result = std::result::Result; /// A struct that manages X25519 key pair operations. @@ -32,23 +23,9 @@ impl X25519KeyPairManager { /// Constructor #[allow(clippy::new_without_default)] pub fn new() -> Result { - let path = Self::get_key_file_path(); - - if path.exists() { - // Read the existing key from the file - let key_bytes = std::fs::read(&path).map_err(KeyManagementError::IoError)?; - let mut key_bytes_array: [u8; DH_SECRET_KEY_SIZE] = [0u8; DH_SECRET_KEY_SIZE]; - key_bytes_array.copy_from_slice(&key_bytes[..DH_SECRET_KEY_SIZE]); - let secret_key = StaticSecret::from(key_bytes_array); - Ok(Self { secret_key }) - } else { - // Generate a new key - let mut rng = rand::thread_rng(); - let secret_key = StaticSecret::random_from_rng(&mut rng); - let this = Self { secret_key }; - this.write_private_key_to_file()?; - Ok(this) - } + let mut rng = rand::thread_rng(); + let secret_key = StaticSecret::random_from_rng(&mut rng); + Ok(Self { secret_key }) } /// Returns a reference to the current X25519 public key. @@ -83,10 +60,9 @@ impl X25519KeyPairManager { /// /// Note: This operation automatically updates the device options to ensure /// the attestation report reflects the new keypair. - pub fn rotate_keys(&mut self) -> Result<()> { + pub fn rotate_keys(&mut self) { let mut rng = rand::thread_rng(); self.secret_key = StaticSecret::random_from_rng(&mut rng); - self.write_private_key_to_file() } /// Computes the shared secret between the current secret key and a given public key. @@ -178,60 +154,6 @@ impl X25519KeyPairManager { let shared_secret = self.compute_shared_secret(&public_key); Ok(encrypt_plaintext(plaintext, &shared_secret, salt, None)?) } - - /// Returns the file path where the private key should be stored. - /// - /// This method constructs a path by: - /// 1. Starting from the current working directory - /// 2. Adding a "keys" subdirectory - /// 3. Adding the "dh_privkey" file name - /// - /// # Returns - /// * `PathBuf` - Path to the private key file: `./keys/dh_privkey` - /// - /// # Note - /// Currently uses a hardcoded path relative to the current directory. - /// This is primarily intended for development/testing purposes. - /// Production environments should use a more secure and configurable location. - fn get_key_file_path() -> std::path::PathBuf { - // Use a more appropriate path, possibly from config - std::env::current_dir() - .unwrap_or_default() - .join(KEY_FILE_DIR) - .join(KEY_FILE_NAME) - } - - /// Writes the current private key to a file at the root directory. - /// - /// # Warning - /// This function is intended for development/testing purposes only. - /// Writing private keys to disk in production is a security risk. - /// - /// # Returns - /// * `Ok(())` if the write was successful - /// * `Err(KeyManagementError)` if the write failed - pub fn write_private_key_to_file(&self) -> Result<()> { - use std::fs::{self, create_dir_all}; - #[cfg(unix)] - use std::os::unix::fs::PermissionsExt; // Unix-specific permissions - - let path = Self::get_key_file_path(); - - // Ensure directory exists - if let Some(parent) = path.parent() { - create_dir_all(parent).map_err(KeyManagementError::IoError)?; - } - - // Write key with restricted permissions - fs::write(&path, self.secret_key.to_bytes()).map_err(KeyManagementError::IoError)?; - - // Set file permissions to owner read/write only (0600) - #[cfg(unix)] - fs::set_permissions(&path, fs::Permissions::from_mode(0o600)) - .map_err(KeyManagementError::IoError)?; - - Ok(()) - } } #[derive(Debug, Error)] diff --git a/atoma-confidential/src/service.rs b/atoma-confidential/src/service.rs index 15ce6b51..97db3765 100644 --- a/atoma-confidential/src/service.rs +++ b/atoma-confidential/src/service.rs @@ -82,6 +82,54 @@ impl AtomaConfidentialComputeService { }) } + /// Initializes and starts the confidential compute service. + /// + /// This method performs the following steps: + /// 1. Creates a new service instance + /// 2. Submits an initial node key rotation attestation + /// 3. Starts the main service event loop + /// + /// # Arguments + /// * `sui_client` - Arc-wrapped RwLock containing the Sui blockchain client + /// * `event_receiver` - Channel receiver for Atoma events + /// * `service_decryption_receiver` - Channel receiver for decryption requests + /// * `service_encryption_receiver` - Channel receiver for encryption requests + /// * `service_shared_secret_receiver` - Channel receiver for shared secret computation requests + /// * `shutdown_signal` - Watch channel receiver for coordinating service shutdown + /// + /// # Returns + /// * `Ok(())` if the service starts and runs successfully + /// * `Err(AtomaConfidentialComputeError)` if initialization, attestation, or running fails + /// + /// # Errors + /// This function can return: + /// * `AtomaConfidentialComputeError::KeyManagementError` if key initialization fails + /// * `AtomaConfidentialComputeError::SuiClientError` if attestation submission fails + #[instrument(level = "info", skip_all)] + pub async fn start_confidential_compute_service( + sui_client: Arc>, + event_receiver: UnboundedReceiver, + service_decryption_receiver: UnboundedReceiver, + service_encryption_receiver: UnboundedReceiver, + service_shared_secret_receiver: UnboundedReceiver, + shutdown_signal: tokio::sync::watch::Receiver, + ) -> Result<()> { + let mut service = Self::new( + sui_client, + event_receiver, + service_decryption_receiver, + service_encryption_receiver, + service_shared_secret_receiver, + shutdown_signal, + )?; + + // NOTE: Submit the first node key rotation attestation, because the node is starting up afresh + service.submit_node_key_rotation_tdx_attestation().await?; + service.run().await?; + + Ok(()) + } + /// Returns the current public key used by the confidential compute service /// /// This method provides access to the X25519 public key that is currently being used @@ -136,6 +184,7 @@ impl AtomaConfidentialComputeService { "Running confidential compute service, with dh public key: {:?}", self.key_manager.get_public_key().as_bytes() ); + loop { tokio::select! { Some((decryption_request, sender)) = self.service_decryption_receiver.recv() => { @@ -195,7 +244,7 @@ impl AtomaConfidentialComputeService { /// - `AtomaConfidentialComputeError::SuiClientError` if the attestation submission to Sui fails #[instrument(level = "debug", skip_all)] async fn submit_node_key_rotation_tdx_attestation(&mut self) -> Result<()> { - self.key_manager.rotate_keys()?; + self.key_manager.rotate_keys(); let public_key = self.key_manager.get_public_key(); let public_key_bytes = public_key.to_bytes(); #[cfg(feature = "tdx")]