address pr comments

atoma-bin/atoma_node.rs

@@ -1,7 +1,6 @@
 use std::{path::Path, str::FromStr, sync::Arc};
 
 use anyhow::{Context, Result};
-#[cfg(feature = "tdx")]
 use atoma_confidential::AtomaConfidentialComputeService;
 use atoma_daemon::{AtomaDaemonConfig, DaemonState};
 use atoma_service::{
@@ -71,8 +70,10 @@ struct Config {
     /// Configuration for the state manager component.
     state: AtomaStateManagerConfig,
 
+    /// Configuration for the daemon component.
     daemon: AtomaDaemonConfig,
 
+    /// Configuration for the proxy component.
     proxy: ProxyConfig,
 }
 
@@ -229,8 +230,7 @@ async fn main() -> Result<()> {
     let (compute_shared_secret_sender, _compute_shared_secret_receiver) =
         tokio::sync::mpsc::unbounded_channel();
 
-    #[cfg(feature = "tdx")]
-    spawn_with_shutdown(
+    let confidential_compute_service_handle = spawn_with_shutdown(
         AtomaConfidentialComputeService::start_confidential_compute_service(
             client.clone(),
             _subscriber_confidential_compute_receiver,
@@ -379,18 +379,27 @@ async fn main() -> Result<()> {
     });
 
     // Wait for shutdown signal and handle cleanup
-    let (subscriber_result, state_manager_result, server_result, daemon_result, _) = try_join!(
+    let (
+        subscriber_result,
+        state_manager_result,
+        server_result,
+        daemon_result,
+        confidential_compute_service_result,
+        _,
+    ) = try_join!(
         subscriber_handle,
         state_manager_handle,
         service_handle,
         daemon_handle,
+        confidential_compute_service_handle,
         ctrl_c
     )?;
     handle_tasks_results(
         subscriber_result,
         state_manager_result,
         server_result,
         daemon_result,
+        confidential_compute_service_result,
     )?;
 
     info!(
@@ -488,6 +497,7 @@ fn handle_tasks_results(
     state_manager_result: Result<()>,
     server_result: Result<()>,
     daemon_result: Result<()>,
+    confidential_compute_service_result: Result<()>,
 ) -> Result<()> {
     let result_handler = |result: Result<()>, message: &str| {
         if let Err(e) = result {
@@ -505,5 +515,9 @@ fn handle_tasks_results(
     result_handler(state_manager_result, "State manager terminated abruptly")?;
     result_handler(server_result, "Server terminated abruptly")?;
     result_handler(daemon_result, "Daemon terminated abruptly")?;
+    result_handler(
+        confidential_compute_service_result,
+        "Confidential compute service terminated abruptly",
+    )?;
     Ok(())
 }

atoma-confidential/src/service.rs

@@ -45,7 +45,7 @@ type ServiceSharedSecretRequest = (
 /// - Graceful shutdown handling
 pub struct AtomaConfidentialComputeService {
     /// Client for interacting with the Sui blockchain to submit attestations and transactions
-    sui_client: Arc<RwLock<AtomaSuiClient>>,
+    _sui_client: Arc<RwLock<AtomaSuiClient>>,
     /// Current key rotation counter
     key_rotation_counter: Option<u64>,
     /// Manages TDX key operations including key rotation and attestation generation
@@ -65,7 +65,7 @@ pub struct AtomaConfidentialComputeService {
 impl AtomaConfidentialComputeService {
     /// Constructor
     pub fn new(
-        sui_client: Arc<RwLock<AtomaSuiClient>>,
+        _sui_client: Arc<RwLock<AtomaSuiClient>>,
         event_receiver: UnboundedReceiver<AtomaEvent>,
         service_decryption_receiver: UnboundedReceiver<ServiceDecryptionRequest>,
         service_encryption_receiver: UnboundedReceiver<ServiceEncryptionRequest>,
@@ -74,7 +74,7 @@ impl AtomaConfidentialComputeService {
     ) -> Result<Self> {
         let key_manager = X25519KeyPairManager::new()?;
         Ok(Self {
-            sui_client,
+            _sui_client,
             key_rotation_counter: None,
             key_manager,
             event_receiver,
@@ -248,50 +248,48 @@ impl AtomaConfidentialComputeService {
     #[instrument(level = "debug", skip_all)]
     async fn submit_node_key_rotation_tdx_attestation(&mut self) -> Result<()> {
         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")]
-        let tdx_quote_bytes = {
+        {
+            let public_key = self.key_manager.get_public_key();
+            let public_key_bytes = public_key.to_bytes();
             let tdx_quote = get_compute_data_attestation(&public_key_bytes)?;
-            tdx_quote.to_bytes()
-        };
-        // NOTE: This is to make it easier to compile the code without TDX support and debug
+            let tdx_quote_bytes = tdx_quote.to_bytes();
+            match self
+                .sui_client
+                .write()
+                .await
+                .submit_key_rotation_remote_attestation(
+                    public_key_bytes,
+                    tdx_quote_bytes,
+                    None,
+                    None,
+                    None,
+                )
+                .await
+            {
+                Ok((digest, key_rotation_counter)) => {
+                    tracing::info!(
+                        target = "atoma-tdx-service",
+                        digest = digest,
+                        key_rotation_counter = key_rotation_counter,
+                        "Submitted node key rotation attestation successfully"
+                    );
+                    self.key_rotation_counter = Some(key_rotation_counter);
+                    Ok(())
+                }
+                Err(e) => {
+                    tracing::error!(
+                        target = "atoma-tdx-service",
+                        error = %e,
+                        "Failed to submit node key rotation attestation"
+                    );
+                    Err(AtomaConfidentialComputeError::SuiClientError(e))
+                }
+            }
+        }
         #[cfg(not(feature = "tdx"))]
-        let tdx_quote_bytes = {
-            const TDX_QUOTE_V4_SIZE: usize = 512;
-            vec![0u8; TDX_QUOTE_V4_SIZE]
-        };
-        match self
-            .sui_client
-            .write()
-            .await
-            .submit_key_rotation_remote_attestation(
-                public_key_bytes,
-                tdx_quote_bytes,
-                None,
-                None,
-                None,
-            )
-            .await
         {
-            Ok((digest, key_rotation_counter)) => {
-                tracing::info!(
-                    target = "atoma-tdx-service",
-                    digest = digest,
-                    key_rotation_counter = key_rotation_counter,
-                    "Submitted node key rotation attestation successfully"
-                );
-                self.key_rotation_counter = Some(key_rotation_counter);
-                Ok(())
-            }
-            Err(e) => {
-                tracing::error!(
-                    target = "atoma-tdx-service",
-                    error = %e,
-                    "Failed to submit node key rotation attestation"
-                );
-                Err(AtomaConfidentialComputeError::SuiClientError(e))
-            }
+            Ok(())
         }
     }