handle ecnrypted response json value appropriately

atoma-network · Dec 6, 2024 · 17dbfcd · 17dbfcd
1 parent aab4f80
commit 17dbfcd
Show file tree

Hide file tree

Showing 4 changed files with 56 additions and 46 deletions.
diff --git a/atoma-service/src/handlers/chat_completions.rs b/atoma-service/src/handlers/chat_completions.rs
@@ -318,26 +318,28 @@ async fn handle_non_streaming_response(
     }
 
     // Handle confidential compute encryption response
-    if let Err(e) = handle_confidential_compute_encryption_response(
+    match handle_confidential_compute_encryption_response(
         &state,
-        &mut response_body,
+        response_body,
         client_encryption_metadata,
     )
     .await
     {
-        error!(
-            target = "atoma-service",
-            event = "chat-completions-handler",
-            "Error handling confidential compute encryption response: {}",
-            e
-        );
-        return Err(StatusCode::INTERNAL_SERVER_ERROR);
+        Ok(response_body) => {
+            // Stop the timer before returning the valid response
+            timer.observe_duration();
+            Ok(Json(response_body).into_response())
+        }
+        Err(e) => {
+            error!(
+                target = "atoma-service",
+                event = "chat-completions-handler",
+                "Error handling confidential compute encryption response: {}",
+                e
+            );
+            Err(StatusCode::INTERNAL_SERVER_ERROR)
+        }
     }
-
-    // Stop the timer before returning the response
-    timer.observe_duration();
-
-    Ok(Json(response_body).into_response())
 }
 
 /// Handles streaming chat completion requests by establishing a Server-Sent Events (SSE) connection.

diff --git a/atoma-service/src/handlers/embeddings.rs b/atoma-service/src/handlers/embeddings.rs
@@ -125,24 +125,26 @@ pub async fn embeddings_handler(
     }
 
     // Handle confidential compute encryption response
-    if let Err(e) = handle_confidential_compute_encryption_response(
+    match handle_confidential_compute_encryption_response(
         &state,
-        &mut response_body,
+        response_body,
         client_encryption_metadata,
     )
     .await
     {
-        error!(
-            target = "atoma-service",
-            event = "embeddings-handler",
-            "Error handling confidential compute encryption response: {}",
-            e
-        );
-        return Err(StatusCode::INTERNAL_SERVER_ERROR);
+        Ok(response_body) => {
+            // Stop the timer before returning the valid response
+            timer.observe_duration();
+            Ok(Json(response_body))
+        }
+        Err(e) => {
+            error!(
+                target = "atoma-service",
+                event = "embeddings-handler",
+                "Error handling confidential compute encryption response: {}",
+                e
+            );
+            Err(StatusCode::INTERNAL_SERVER_ERROR)
+        }
     }
-
-    // Stop the timer before returning the response
-    timer.observe_duration();
-
-    Ok(Json(response_body))
 }
diff --git a/atoma-service/src/handlers/image_generations.rs b/atoma-service/src/handlers/image_generations.rs
@@ -121,23 +121,26 @@ pub async fn image_generations_handler(
     }
 
     // Handle confidential compute encryption response
-    if let Err(e) = handle_confidential_compute_encryption_response(
+    match handle_confidential_compute_encryption_response(
         &state,
-        &mut response_body,
+        response_body,
         client_encryption_metadata,
     )
     .await
     {
-        error!(
-            target = "atoma-service",
-            event = "image-generations-handler",
-            "Error handling confidential compute encryption response: {}",
-            e
-        );
-        return Err(StatusCode::INTERNAL_SERVER_ERROR);
+        Ok(response_body) => {
+            // Stop the timer before returning the valid response
+            timer.observe_duration();
+            Ok(Json(response_body))
+        }
+        Err(e) => {
+            error!(
+                target = "atoma-service",
+                event = "image-generations-handler",
+                "Error handling confidential compute encryption response: {}",
+                e
+            );
+            Err(StatusCode::INTERNAL_SERVER_ERROR)
+        }
     }
-
-    timer.observe_duration();
-
-    Ok(Json(response_body))
 }
diff --git a/atoma-service/src/handlers/mod.rs b/atoma-service/src/handlers/mod.rs
@@ -79,9 +79,9 @@ async fn sign_response_and_update_stack_hash(
 )]
 pub(crate) async fn handle_confidential_compute_encryption_response(
     state: &AppState,
-    response_body: &mut Value,
+    response_body: Value,
     client_encryption_metadata: Option<EncryptionMetadata>,
-) -> Result<(), StatusCode> {
+) -> Result<Value, StatusCode> {
     if let Some(EncryptionMetadata {
         proxy_x25519_public_key,
         salt,
@@ -121,8 +121,11 @@ pub(crate) async fn handle_confidential_compute_encryption_response(
                 );
                 StatusCode::INTERNAL_SERVER_ERROR
             })?;
-        response_body["nonce"] = json!(nonce);
-        response_body["ciphertext"] = json!(ciphertext);
+        Ok(json!({
+            "nonce": nonce,
+            "ciphertext": ciphertext,
+        }))
+    } else {
+        Ok(response_body)
     }
-    Ok(())
 }