Skip to content

Commit

Permalink
refactor code
Browse files Browse the repository at this point in the history
  • Loading branch information
jorgeantonio21 committed Dec 18, 2024
1 parent 32e09c0 commit d2d075d
Show file tree
Hide file tree
Showing 2 changed files with 62 additions and 11 deletions.
22 changes: 13 additions & 9 deletions atoma-bin/atoma_node.rs
Original file line number Diff line number Diff line change
Expand Up @@ -227,17 +227,21 @@ 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(),
)?;

let client_clone = client.clone();
let shutdown_receiver_clone = shutdown_receiver.clone();
spawn_with_shutdown(
async move { confidential_compute_service.run().await },
async move {
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,
)
.await
},
shutdown_sender.clone(),
);

Expand Down
51 changes: 49 additions & 2 deletions atoma-confidential/src/service.rs
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,54 @@ impl AtomaConfidentialComputeService {
shutdown_signal,
})
}

/// 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
///
Expand Down Expand Up @@ -136,8 +184,7 @@ impl AtomaConfidentialComputeService {
"Running confidential compute service, with dh public key: {:?}",
self.key_manager.get_public_key().as_bytes()
);
// Submit the first node key rotation attestation, because the node is starting up afresh
self.submit_node_key_rotation_tdx_attestation().await?;

loop {
tokio::select! {
Some((decryption_request, sender)) = self.service_decryption_receiver.recv() => {
Expand Down

0 comments on commit d2d075d

Please sign in to comment.