Clean up more

Changochen · Changochen · commit cd14087cd59e · 2026-03-30T10:48:22.000-07:00
Change-Id: I419e86af4a04a4cdbc93c0bb897ecbc1612d25dd
diff --git a/oak_private_memory/BUILD b/oak_private_memory/BUILD
@@ -133,18 +133,6 @@ rust_library(
     ],
 )
 
-rust_library(
-    name = "session",
-    srcs = ["src/session.rs"],
-    deps = [
-        "@oak//oak_proto_rust",
-        "@oak//oak_session",
-        "@oak//oak_session/tls/rust:oak_session_tls",
-        "@oak_crates_index//:anyhow",
-        "@oak_crates_index//:prost",
-    ],
-)
-
 rust_library(
     name = "session_binder",
     srcs = ["src/session_binder.rs"],
@@ -176,7 +164,6 @@ rust_library(
         ":external_db_client",
         ":log",
         ":metrics",
-        ":session",
         ":session_binder",
         "//app",
         "//database",
diff --git a/oak_private_memory/app/BUILD b/oak_private_memory/app/BUILD
@@ -40,7 +40,6 @@ rust_library(
         "//:external_db_client",
         "//:log",
         "//:metrics",
-        "//:session",
         "//database",
         "//proto:sealed_memory_rust_proto",
         "//proto/grpc:sealed_memory_grpc_proto",
diff --git a/oak_private_memory/app/service.rs b/oak_private_memory/app/service.rs
@@ -30,7 +30,6 @@ use sealed_memory_grpc_proto::oak::private_memory::sealed_memory_service_server:
     SealedMemoryService, SealedMemoryServiceServer,
 };
 use sealed_memory_rust_proto::{oak::private_memory::TlsSessionFrame, prelude::v1::*};
-use session::TlsEncryptedSession;
 use tokio::{net::TcpListener, sync::mpsc};
 use tokio_stream::{Stream, StreamExt, wrappers::TcpListenerStream};
 
@@ -200,7 +199,7 @@ impl OakSessionHandler {
 /// data.
 struct TlsSessionHandler {
     metrics: Arc<metrics::Metrics>,
-    session: TlsEncryptedSession,
+    session: oak_session_tls::OakSessionTls,
     application_handler: SealedMemorySessionHandler,
 }
 
@@ -214,7 +213,7 @@ impl TlsSessionHandler {
     ) -> Self {
         Self {
             metrics: metrics.clone(),
-            session: TlsEncryptedSession::new(tls_session),
+            session: tls_session,
             application_handler: SealedMemorySessionHandler::new(
                 metrics.clone(),
                 persistence_tx.clone(),
@@ -229,10 +228,9 @@ impl TlsSessionHandler {
     pub async fn handle_app_request(&mut self, encrypted_request: &[u8]) -> tonic::Result<Vec<u8>> {
         self.metrics.inc_requests(RequestMetricName::total());
 
-        let decrypted_request = self
-            .session
-            .decrypt(encrypted_request)
-            .into_invalid_argument("failed to decrypt TLS request")?;
+        let decrypted_request = self.session.decrypt(encrypted_request).map_err(|e| {
+            tonic::Status::invalid_argument(format!("failed to decrypt TLS request: {e}"))
+        })?;
 
         if decrypted_request.is_empty() {
             // This can happen if the TLS frame only contained handshake data
@@ -246,10 +244,9 @@ impl TlsSessionHandler {
                 self.metrics.inc_failures(RequestMetricName::total());
                 Err(e)
             }
-            Ok(plaintext_response) => self
-                .session
-                .encrypt(&plaintext_response)
-                .into_internal_error("failed to encrypt TLS response"),
+            Ok(plaintext_response) => self.session.encrypt(&plaintext_response).map_err(|e| {
+                tonic::Status::internal(format!("failed to encrypt TLS response: {e}"))
+            }),
         }
     }
 }
diff --git a/oak_private_memory/src/client.rs b/oak_private_memory/src/client.rs
@@ -202,8 +202,16 @@ impl PrivateMemoryClient {
 
         sealed_memory_response.response.ok_or_else(|| anyhow!("empty response"))
     }
+}
+
+#[async_trait]
+pub trait PrivateMemoryAppClient {
+    async fn invoke(
+        &mut self,
+        request: sealed_memory_request::Request,
+    ) -> Result<sealed_memory_response::Response>;
 
-    pub async fn register_user(
+    async fn register_user(
         &mut self,
         pm_uid: &str,
         kek: &[u8],
@@ -221,11 +229,7 @@ impl PrivateMemoryClient {
         }
     }
 
-    pub async fn key_sync(
-        &mut self,
-        pm_uid: &str,
-        kek: &[u8],
-    ) -> Result<key_sync_response::Status> {
+    async fn key_sync(&mut self, pm_uid: &str, kek: &[u8]) -> Result<key_sync_response::Status> {
         let request =
             KeySyncRequest { pm_uid: pm_uid.to_string(), key_encryption_key: kek.to_vec() };
         let response = self.invoke(sealed_memory_request::Request::KeySyncRequest(request)).await?;
@@ -235,15 +239,15 @@ impl PrivateMemoryClient {
         }
     }
 
-    pub async fn add_memory(&mut self, memory: Memory) -> Result<AddMemoryResponse> {
+    async fn add_memory(&mut self, memory: Memory) -> Result<AddMemoryResponse> {
         let request = AddMemoryRequest { memory: Some(memory) };
         let response =
             self.invoke(sealed_memory_request::Request::AddMemoryRequest(request)).await?;
         expect_response_type!(response, sealed_memory_response::Response::AddMemoryResponse)
     }
 
     #[allow(deprecated)]
-    pub async fn get_memories(
+    async fn get_memories(
         &mut self,
         tag: &str,
         page_size: i32,
@@ -261,7 +265,7 @@ impl PrivateMemoryClient {
         expect_response_type!(response, sealed_memory_response::Response::GetMemoriesResponse)
     }
 
-    pub async fn get_memory_by_id(
+    async fn get_memory_by_id(
         &mut self,
         id: &str,
         result_mask: Option<ResultMask>,
@@ -272,7 +276,7 @@ impl PrivateMemoryClient {
         expect_response_type!(response, sealed_memory_response::Response::GetMemoryByIdResponse)
     }
 
-    pub async fn search_memory(
+    async fn search_memory(
         &mut self,
         query: SearchMemoryQuery,
         page_size: i32,
@@ -292,21 +296,21 @@ impl PrivateMemoryClient {
         expect_response_type!(response, sealed_memory_response::Response::SearchMemoryResponse)
     }
 
-    pub async fn delete_memory(&mut self, ids: Vec<String>) -> Result<DeleteMemoryResponse> {
+    async fn delete_memory(&mut self, ids: Vec<String>) -> Result<DeleteMemoryResponse> {
         let request = DeleteMemoryRequest { ids };
         let response =
             self.invoke(sealed_memory_request::Request::DeleteMemoryRequest(request)).await?;
         expect_response_type!(response, sealed_memory_response::Response::DeleteMemoryResponse)
     }
 
-    pub async fn reset_memory(&mut self) -> Result<ResetMemoryResponse> {
+    async fn reset_memory(&mut self) -> Result<ResetMemoryResponse> {
         let request = ResetMemoryRequest::default();
         let response =
             self.invoke(sealed_memory_request::Request::ResetMemoryRequest(request)).await?;
         expect_response_type!(response, sealed_memory_response::Response::ResetMemoryResponse)
     }
 
-    pub async fn get_memories_by_id(
+    async fn get_memories_by_id(
         &mut self,
         ids: Vec<String>,
         result_mask: Option<ResultMask>,
@@ -317,7 +321,7 @@ impl PrivateMemoryClient {
         expect_response_type!(response, sealed_memory_response::Response::GetMemoriesByIdResponse)
     }
 
-    pub async fn get_database_metrics(&mut self) -> Result<GetDatabaseMetricsResponse> {
+    async fn get_database_metrics(&mut self) -> Result<GetDatabaseMetricsResponse> {
         let request = GetDatabaseMetricsRequest::default();
         let response =
             self.invoke(sealed_memory_request::Request::GetDatabaseMetricsRequest(request)).await?;
@@ -328,6 +332,15 @@ impl PrivateMemoryClient {
     }
 }
 
+#[async_trait]
+impl PrivateMemoryAppClient for PrivateMemoryClient {
+    async fn invoke(
+        &mut self,
+        request: sealed_memory_request::Request,
+    ) -> Result<sealed_memory_response::Response> {
+        self.invoke(request).await
+    }
+}
 // ---------------------------------------------------------------------------
 // TLS Client Support
 // ---------------------------------------------------------------------------
@@ -460,109 +473,14 @@ impl PrivateMemoryTlsClient {
 
         sealed_memory_response.response.ok_or_else(|| anyhow!("empty response"))
     }
+}
 
-    pub async fn register_user(
-        &mut self,
-        pm_uid: &str,
-        kek: &[u8],
-    ) -> Result<user_registration_response::Status> {
-        let request = UserRegistrationRequest {
-            pm_uid: pm_uid.to_string(),
-            key_encryption_key: kek.to_vec(),
-            boot_strap_info: Some(KeyDerivationInfo::default()),
-        };
-        let response =
-            self.invoke(sealed_memory_request::Request::UserRegistrationRequest(request)).await?;
-        match response {
-            sealed_memory_response::Response::UserRegistrationResponse(resp) => Ok(resp.status()),
-            _ => Err(anyhow!("unexpected response type for user registration")),
-        }
-    }
-
-    pub async fn key_sync(
-        &mut self,
-        pm_uid: &str,
-        kek: &[u8],
-    ) -> Result<key_sync_response::Status> {
-        let request =
-            KeySyncRequest { pm_uid: pm_uid.to_string(), key_encryption_key: kek.to_vec() };
-        let response = self.invoke(sealed_memory_request::Request::KeySyncRequest(request)).await?;
-        match response {
-            sealed_memory_response::Response::KeySyncResponse(resp) => Ok(resp.status()),
-            _ => Err(anyhow!("unexpected response type for key sync")),
-        }
-    }
-
-    pub async fn add_memory(&mut self, memory: Memory) -> Result<AddMemoryResponse> {
-        let request = AddMemoryRequest { memory: Some(memory) };
-        let response =
-            self.invoke(sealed_memory_request::Request::AddMemoryRequest(request)).await?;
-        expect_response_type!(response, sealed_memory_response::Response::AddMemoryResponse)
-    }
-
-    pub async fn search_memory(
-        &mut self,
-        query: SearchMemoryQuery,
-        page_size: i32,
-        result_mask: Option<ResultMask>,
-        page_token: &str,
-        keep_all_llm_views: bool,
-    ) -> Result<SearchMemoryResponse> {
-        let request = SearchMemoryRequest {
-            query: Some(query),
-            page_size,
-            result_mask,
-            page_token: page_token.to_string(),
-            keep_all_llm_views,
-        };
-        let response =
-            self.invoke(sealed_memory_request::Request::SearchMemoryRequest(request)).await?;
-        expect_response_type!(response, sealed_memory_response::Response::SearchMemoryResponse)
-    }
-
-    pub async fn delete_memory(&mut self, ids: Vec<String>) -> Result<DeleteMemoryResponse> {
-        let request = DeleteMemoryRequest { ids };
-        let response =
-            self.invoke(sealed_memory_request::Request::DeleteMemoryRequest(request)).await?;
-        expect_response_type!(response, sealed_memory_response::Response::DeleteMemoryResponse)
-    }
-
-    pub async fn reset_memory(&mut self) -> Result<ResetMemoryResponse> {
-        let request = ResetMemoryRequest::default();
-        let response =
-            self.invoke(sealed_memory_request::Request::ResetMemoryRequest(request)).await?;
-        expect_response_type!(response, sealed_memory_response::Response::ResetMemoryResponse)
-    }
-
-    pub async fn get_memory_by_id(
-        &mut self,
-        id: &str,
-        result_mask: Option<ResultMask>,
-    ) -> Result<GetMemoryByIdResponse> {
-        let request = GetMemoryByIdRequest { id: id.to_string(), result_mask };
-        let response =
-            self.invoke(sealed_memory_request::Request::GetMemoryByIdRequest(request)).await?;
-        expect_response_type!(response, sealed_memory_response::Response::GetMemoryByIdResponse)
-    }
-
-    pub async fn get_memories_by_id(
+#[async_trait]
+impl PrivateMemoryAppClient for PrivateMemoryTlsClient {
+    async fn invoke(
         &mut self,
-        ids: Vec<String>,
-        result_mask: Option<ResultMask>,
-    ) -> Result<GetMemoriesByIdResponse> {
-        let request = GetMemoriesByIdRequest { ids, result_mask };
-        let response =
-            self.invoke(sealed_memory_request::Request::GetMemoriesByIdRequest(request)).await?;
-        expect_response_type!(response, sealed_memory_response::Response::GetMemoriesByIdResponse)
-    }
-
-    pub async fn get_database_metrics(&mut self) -> Result<GetDatabaseMetricsResponse> {
-        let request = GetDatabaseMetricsRequest::default();
-        let response =
-            self.invoke(sealed_memory_request::Request::GetDatabaseMetricsRequest(request)).await?;
-        expect_response_type!(
-            response,
-            sealed_memory_response::Response::GetDatabaseMetricsResponse
-        )
+        request: sealed_memory_request::Request,
+    ) -> Result<sealed_memory_response::Response> {
+        self.invoke(request).await
     }
 }
diff --git a/oak_private_memory/src/lib.rs b/oak_private_memory/src/lib.rs
@@ -18,5 +18,4 @@ pub use encryption;
 pub use external_db_client;
 pub use log;
 pub use metrics;
-pub use session;
 pub use session_binder;
diff --git a/oak_private_memory/src/session.rs b/oak_private_memory/src/session.rs
diff --git a/oak_private_memory/test/client_test.rs b/oak_private_memory/test/client_test.rs
@@ -20,7 +20,7 @@ use attestation_testing::{
     DUMMY_ATTESTATION_ID, DummySessionBindingVerifierProvider, RejectingVerifier,
     dummy_client_session_config,
 };
-use client::PrivateMemoryClient;
+use client::{PrivateMemoryAppClient, PrivateMemoryClient};
 use oak_session::{attestation::AttestationType, config::SessionConfig, handshake::HandshakeType};
 use private_memory_test_utils::{start_server, system_time_to_timestamp};
 use sealed_memory_rust_proto::{
diff --git a/oak_private_memory/test/e2e_test.rs b/oak_private_memory/test/e2e_test.rs
diff --git a/oak_private_memory/test/tls_client_test.rs b/oak_private_memory/test/tls_client_test.rs
