feat: add encryption to endpoints

atoma-bin/atoma_node.rs

@@ -213,12 +213,14 @@ async fn main() -> Result<()> {
     let client = Arc::new(RwLock::new(
         AtomaSuiClient::new_from_config(args.config_path).await?,
     ));
-
+    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(),
     )?;
 
@@ -280,6 +282,7 @@ async fn main() -> Result<()> {
         stack_retrieve_sender,
         decryption_sender: app_state_decryption_sender,
         encryption_sender: app_state_encryption_sender,
+        compute_shared_secret_sender,
         tokenizers: Arc::new(tokenizers),
         models: Arc::new(config.service.models),
         chat_completions_service_url: config

atoma-confidential/src/key_management.rs

@@ -1,5 +1,5 @@
 use atoma_utils::encryption::{
-    decrypt_cyphertext, encrypt_plaintext, EncryptionError, NONCE_BYTE_SIZE,
+    decrypt_ciphertext, encrypt_plaintext, EncryptionError, NONCE_BYTE_SIZE,
 };
 use thiserror::Error;
 use x25519_dalek::{PublicKey, SharedSecret, StaticSecret};
@@ -126,9 +126,9 @@ impl X25519KeyPairManager {
     /// let salt = vec![/* salt bytes */];
     /// let nonce = vec![/* nonce bytes */];
     ///
-    /// let plaintext = manager.decrypt_cyphertext(public_key, &ciphertext, &salt, &nonce)?;
+    /// let plaintext = manager.decrypt_ciphertext(public_key, &ciphertext, &salt, &nonce)?;
     /// ```
-    pub fn decrypt_cyphertext(
+    pub fn decrypt_ciphertext(
         &self,
         public_key: [u8; 32],
         ciphertext: &[u8],
@@ -137,7 +137,7 @@ impl X25519KeyPairManager {
     ) -> Result<Vec<u8>> {
         let public_key = PublicKey::from(public_key);
         let shared_secret = self.compute_shared_secret(&public_key);
-        Ok(decrypt_cyphertext(shared_secret, ciphertext, salt, nonce)?)
+        Ok(decrypt_ciphertext(&shared_secret, ciphertext, salt, nonce)?)
     }
 
     /// Encrypts plaintext using X25519 key exchange and symmetric encryption.
@@ -176,7 +176,7 @@ impl X25519KeyPairManager {
     ) -> Result<(Vec<u8>, [u8; NONCE_BYTE_SIZE])> {
         let public_key = PublicKey::from(public_key);
         let shared_secret = self.compute_shared_secret(&public_key);
-        Ok(encrypt_plaintext(plaintext, shared_secret, salt)?)
+        Ok(encrypt_plaintext(plaintext, &shared_secret, salt, None)?)
     }
 
     /// Returns the file path where the private key should be stored.

atoma-confidential/src/service.rs

@@ -5,10 +5,12 @@ use crate::{
     types::{
         ConfidentialComputeDecryptionRequest, ConfidentialComputeDecryptionResponse,
         ConfidentialComputeEncryptionRequest, ConfidentialComputeEncryptionResponse,
+        ConfidentialComputeSharedSecretRequest, ConfidentialComputeSharedSecretResponse,
     },
 };
 use atoma_sui::client::AtomaSuiClient;
 use atoma_sui::{client::AtomaSuiClientError, events::AtomaEvent};
+use atoma_utils::constants::NONCE_SIZE;
 use std::sync::Arc;
 use thiserror::Error;
 use tokio::sync::{mpsc::UnboundedReceiver, oneshot, RwLock};
@@ -18,15 +20,22 @@ use x25519_dalek::PublicKey;
 // TODO: How large can the `ServiceData` be ? Is it feasible to use a Flume channel ?
 
 type Result<T> = std::result::Result<T, AtomaConfidentialComputeError>;
+
 type ServiceDecryptionRequest = (
     ConfidentialComputeDecryptionRequest,
     oneshot::Sender<anyhow::Result<ConfidentialComputeDecryptionResponse>>,
 );
+
 type ServiceEncryptionRequest = (
     ConfidentialComputeEncryptionRequest,
     oneshot::Sender<anyhow::Result<ConfidentialComputeEncryptionResponse>>,
 );
 
+type ServiceSharedSecretRequest = (
+    ConfidentialComputeSharedSecretRequest,
+    oneshot::Sender<ConfidentialComputeSharedSecretResponse>,
+);
+
 /// A service that manages Intel's TDX (Trust Domain Extensions) operations and key rotations.
 ///
 /// The `AtomaConfidentialCompute` handles:
@@ -45,6 +54,8 @@ pub struct AtomaConfidentialComputeService {
     service_decryption_receiver: UnboundedReceiver<ServiceDecryptionRequest>,
     /// Channel receiver for incoming Atoma service requests for encryption and processing
     service_encryption_receiver: UnboundedReceiver<ServiceEncryptionRequest>,
+    /// Channel receiver for incoming Atoma service requests for shared secret computation
+    service_shared_secret_receiver: UnboundedReceiver<ServiceSharedSecretRequest>,
     /// Signal receiver for coordinating graceful shutdown of the service
     shutdown_signal: tokio::sync::watch::Receiver<bool>,
 }
@@ -56,6 +67,7 @@ impl AtomaConfidentialComputeService {
         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<Self> {
         let key_manager = X25519KeyPairManager::new()?;
@@ -65,6 +77,7 @@ impl AtomaConfidentialComputeService {
             event_receiver,
             service_decryption_receiver,
             service_encryption_receiver,
+            service_shared_secret_receiver,
             shutdown_signal,
         })
     }
@@ -81,6 +94,21 @@ impl AtomaConfidentialComputeService {
         self.key_manager.get_public_key()
     }
 
+    /// Returns the shared secret between the node and the proxy
+    ///
+    /// This method computes the shared secret using the node's secret key and the proxy's public key
+    /// and returns it as a `x25519_dalek::StaticSecret`
+    ///
+    /// # Returns
+    /// - `x25519_dalek::StaticSecret`: The shared secret between the node and the proxy
+    pub fn compute_shared_secret(
+        &self,
+        proxy_x25519_public_key: &PublicKey,
+    ) -> x25519_dalek::SharedSecret {
+        self.key_manager
+            .compute_shared_secret(proxy_x25519_public_key)
+    }
+
     /// Starts the TDX service event loop that processes Atoma events and handles graceful shutdown.
     ///
     /// This method runs continuously until a shutdown signal is received, processing two types of events:
@@ -116,6 +144,9 @@ impl AtomaConfidentialComputeService {
                 Some((encryption_request, sender)) = self.service_encryption_receiver.recv() => {
                     self.handle_encryption_request(encryption_request, sender)?;
                 }
+                Some((shared_secret_request, sender)) = self.service_shared_secret_receiver.recv() => {
+                    self.handle_shared_secret_request(shared_secret_request, sender)?;
+                }
                 Some(event) = self.event_receiver.recv() => {
                     self.handle_atoma_event(event).await?;
                 }
@@ -262,7 +293,7 @@ impl AtomaConfidentialComputeService {
             ))
         } else {
             self.key_manager
-                .decrypt_cyphertext(proxy_x25519_public_key, &ciphertext, &salt, &nonce)
+                .decrypt_ciphertext(proxy_x25519_public_key, &ciphertext, &salt, &nonce)
                 .map_err(|e| {
                     tracing::error!(
                         target = "atoma-confidential-compute-service",
@@ -342,6 +373,52 @@ impl AtomaConfidentialComputeService {
             .map_err(|_| AtomaConfidentialComputeError::SenderError)
     }
 
+    /// Handles a shared secret request from a client by computing the shared secret and sending it back.
+    ///
+    /// This method performs the following steps:
+    /// 1. Computes the shared secret using the node's secret key and the proxy's public key
+    /// 2. Generates a random nonce
+    /// 3. Sends the shared secret and nonce back through the provided sender channel
+    ///
+    /// # Arguments
+    /// * `shared_secret_request` - The shared secret request containing:
+    ///   - proxy_x25519_public_key: The public key of the proxy requesting the shared secret
+    /// * `sender` - A oneshot sender channel for returning the shared secret response
+    ///
+    /// # Returns
+    /// * `Ok(())` if the request was handled and response sent successfully
+    /// * `Err(AtomaConfidentialComputeError)` if the response cannot be sent
+    ///
+    /// # Errors
+    /// This function can return `AtomaConfidentialComputeError::SenderError` if the response channel is closed
+    #[instrument(
+        level = "debug",
+        name = "handle_shared_secret_request",
+        skip_all,
+        fields(
+            proxy_public_key = ?shared_secret_request.proxy_x25519_public_key,
+            node_public_key = ?self.key_manager.get_public_key().as_bytes()
+        )
+    )]
+    fn handle_shared_secret_request(
+        &mut self,
+        shared_secret_request: ConfidentialComputeSharedSecretRequest,
+        sender: oneshot::Sender<ConfidentialComputeSharedSecretResponse>,
+    ) -> Result<()> {
+        let ConfidentialComputeSharedSecretRequest {
+            proxy_x25519_public_key,
+        } = shared_secret_request;
+        let shared_secret = self.compute_shared_secret(&PublicKey::from(proxy_x25519_public_key));
+        let nonce = rand::random::<[u8; NONCE_SIZE]>();
+        sender
+            .send(ConfidentialComputeSharedSecretResponse {
+                shared_secret,
+                nonce,
+            })
+            .map_err(|_| AtomaConfidentialComputeError::SenderError)?;
+        Ok(())
+    }
+
     /// Processes incoming Atoma events and handles key rotation requests.
     ///
     /// This method handles two types of events:

atoma-confidential/src/tdx.rs

@@ -3,6 +3,7 @@ use dcap_rs::types::quotes::body::QuoteBody;
 use tdx::device::{Device, DeviceOptions, QuoteV4};
 use thiserror::Error;
 
+/// The size of the data to be attested, for a intel TDX quote.
 pub const TDX_REPORT_DATA_SIZE: usize = 64;
 
 /// Generates a TDX attestation report for the given compute data.

atoma-confidential/src/types.rs

@@ -1,9 +1,8 @@
+use atoma_utils::constants::{NONCE_SIZE, SALT_SIZE};
+
 /// Size of a Diffie-Hellman public key in bytes
 pub const DH_PUBLIC_KEY_SIZE: usize = 32;
 
-/// Size of a cryptographic nonce in bytes
-pub const NONCE_SIZE: usize = 12;
-
 /// A request for confidential computation that includes encrypted data and key exchange parameters
 ///
 /// This struct contains all necessary components for secure communication:
@@ -15,9 +14,9 @@ pub struct ConfidentialComputeDecryptionRequest {
     /// The encrypted data to be processed
     pub ciphertext: Vec<u8>,
     /// Cryptographic nonce used in the encryption process
-    pub nonce: Vec<u8>,
+    pub nonce: [u8; NONCE_SIZE],
     /// Salt value used in key derivation
-    pub salt: Vec<u8>,
+    pub salt: [u8; SALT_SIZE],
     /// Public key component for Diffie-Hellman key exchange
     pub proxy_x25519_public_key: [u8; DH_PUBLIC_KEY_SIZE],
     /// Public key component for Diffie-Hellman key exchange
@@ -44,7 +43,7 @@ pub struct ConfidentialComputeEncryptionRequest {
     /// The plaintext data to be encrypted
     pub plaintext: Vec<u8>,
     /// Salt value used in key derivation
-    pub salt: Vec<u8>,
+    pub salt: [u8; SALT_SIZE],
     /// Public key component for Diffie-Hellman key exchange
     pub proxy_x25519_public_key: [u8; DH_PUBLIC_KEY_SIZE],
 }
@@ -59,3 +58,22 @@ pub struct ConfidentialComputeEncryptionResponse {
     /// Cryptographic nonce used in the encryption process
     pub nonce: [u8; NONCE_SIZE],
 }
+
+/// Request for the shared secret from a confidential computation request
+///
+/// This struct contains the public key of the proxy
+pub struct ConfidentialComputeSharedSecretRequest {
+    /// Public key component for Diffie-Hellman key exchange
+    pub proxy_x25519_public_key: [u8; DH_PUBLIC_KEY_SIZE],
+}
+
+/// Response containing the shared secret from a confidential computation request
+///
+/// This struct contains the shared secret after successful processing of a
+/// ConfidentialComputeSharedSecretRequest.
+pub struct ConfidentialComputeSharedSecretResponse {
+    /// The shared secret resulting from the confidential computation
+    pub shared_secret: x25519_dalek::SharedSecret,
+    /// Cryptographic nonce used in the encryption process
+    pub nonce: [u8; NONCE_SIZE],
+}