Skip to content

Commit

Permalink
feat: remove file storing for symmetric key and submit rotate key on …
Browse files Browse the repository at this point in the history 
…node start (#289)

* remove file storing for symmetric key and submit rotate key on node initialization

* refactor code

* address pr comments
  • Loading branch information
@jorgeantonio21
jorgeantonio21 authored Dec 18, 2024
1 parent bcffcea commit 54aa801
Show file tree
Hide file tree
Showing 3 changed files with 62 additions and 92 deletions.
17 changes: 8 additions & 9 deletions atoma-bin/atoma_node.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
);

Expand Down
86 changes: 4 additions & 82 deletions atoma-confidential/src/key_management.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<T> = std::result::Result<T, KeyManagementError>;

/// A struct that manages X25519 key pair operations.
Expand All @@ -32,23 +23,9 @@ impl X25519KeyPairManager {
/// Constructor
#[allow(clippy::new_without_default)]
pub fn new() -> Result<Self> {
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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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)]
Expand Down
51 changes: 50 additions & 1 deletion atoma-confidential/src/service.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<RwLock<AtomaSuiClient>>,
event_receiver: UnboundedReceiver<AtomaEvent>,
service_decryption_receiver: UnboundedReceiver<ServiceDecryptionRequest>,
service_encryption_receiver: UnboundedReceiver<ServiceEncryptionRequest>,
service_shared_secret_receiver: UnboundedReceiver<ServiceSharedSecretRequest>,
shutdown_signal: tokio::sync::watch::Receiver<bool>,
) -> 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
Expand Down Expand Up @@ -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() => {
Expand Down Expand Up @@ -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")]
Expand Down

0 comments on commit 54aa801

Please sign in to comment.