Rename RecipientContextGenerator into EncryptionKeyHandle in Rust

oak_attestation/src/handler.rs

@@ -20,8 +20,7 @@ use core::future::Future;
 use anyhow::Context;
 use oak_crypto::{
     encryptor::{
-        AsyncRecipientContextGenerator, AsyncServerEncryptor, RecipientContextGenerator,
-        ServerEncryptor,
+        AsyncEncryptionKeyHandle, AsyncServerEncryptor, EncryptionKeyHandle, ServerEncryptor,
     },
     proto::oak::crypto::v1::{EncryptedRequest, EncryptedResponse},
 };
@@ -42,17 +41,14 @@ pub struct PublicKeyInfo {
 /// based on the provided encryption key.
 pub struct EncryptionHandler<H: FnOnce(Vec<u8>) -> Vec<u8>> {
     // TODO(#3442): Use attester to attest to the public key.
-    encryption_key_provider: Arc<dyn RecipientContextGenerator>,
+    encryption_key_handle: Arc<dyn EncryptionKeyHandle>,
     request_handler: H,
 }
 
 impl<H: FnOnce(Vec<u8>) -> Vec<u8>> EncryptionHandler<H> {
-    pub fn create(
-        encryption_key_provider: Arc<dyn RecipientContextGenerator>,
-        request_handler: H,
-    ) -> Self {
+    pub fn create(encryption_key_handle: Arc<dyn EncryptionKeyHandle>, request_handler: H) -> Self {
         Self {
-            encryption_key_provider,
+            encryption_key_handle,
             request_handler,
         }
     }
@@ -67,7 +63,7 @@ impl<H: FnOnce(Vec<u8>) -> Vec<u8>> EncryptionHandler<H> {
             .context("initial request message doesn't contain encapsulated public key")?;
         let mut server_encryptor = ServerEncryptor::create(
             serialized_encapsulated_public_key,
-            self.encryption_key_provider.clone(),
+            self.encryption_key_handle.clone(),
         )
         .context("couldn't create server encryptor")?;
 
@@ -90,27 +86,27 @@ impl<H: FnOnce(Vec<u8>) -> Vec<u8>> EncryptionHandler<H> {
 
 /// Wraps a closure to an underlying function with request encryption and response decryption logic,
 /// based on the provided encryption key.
-/// [`AsyncEncryptionHandler`] can be used when an [`AsyncRecipientContextGenerator`] is needed.
+/// [`AsyncEncryptionHandler`] can be used when an [`AsyncEncryptionKeyHandle`] is needed.
 pub struct AsyncEncryptionHandler<G, H, F>
 where
-    G: AsyncRecipientContextGenerator + Send + Sync,
+    G: AsyncEncryptionKeyHandle + Send + Sync,
     H: FnOnce(Vec<u8>) -> F,
     F: Future<Output = Vec<u8>>,
 {
     // TODO(#3442): Use attester to attest to the public key.
-    recipient_context_generator: Arc<G>,
+    encryption_key_handle: Arc<G>,
     request_handler: H,
 }
 
 impl<G, H, F> AsyncEncryptionHandler<G, H, F>
 where
-    G: AsyncRecipientContextGenerator + Send + Sync,
+    G: AsyncEncryptionKeyHandle + Send + Sync,
     H: FnOnce(Vec<u8>) -> F,
     F: Future<Output = Vec<u8>>,
 {
-    pub fn create(recipient_context_generator: Arc<G>, request_handler: H) -> Self {
+    pub fn create(encryption_key_handle: Arc<G>, request_handler: H) -> Self {
         Self {
-            recipient_context_generator,
+            encryption_key_handle,
             request_handler,
         }
     }
@@ -120,8 +116,7 @@ where
         encrypted_request: &EncryptedRequest,
     ) -> anyhow::Result<EncryptedResponse> {
         // Initialize server encryptor.
-        let mut server_encryptor =
-            AsyncServerEncryptor::new(self.recipient_context_generator.as_ref());
+        let mut server_encryptor = AsyncServerEncryptor::new(self.encryption_key_handle.as_ref());
 
         // Decrypt request.
         let (request, _associated_data) = server_encryptor

oak_containers_orchestrator/src/crypto.rs

@@ -17,7 +17,7 @@ use std::sync::Arc;
 
 use anyhow::{anyhow, Context};
 use hpke::{kem::X25519HkdfSha256, Deserializable, Kem};
-use oak_crypto::encryptor::{EncryptionKeyProvider, RecipientContextGenerator, ServerEncryptor};
+use oak_crypto::encryptor::{EncryptionKeyHandle, EncryptionKeyProvider, ServerEncryptor};
 use tonic::{Request, Response};
 
 use crate::proto::oak::{

oak_containers_sdk/src/crypto.rs

@@ -16,7 +16,7 @@
 use anyhow::Context;
 use async_trait::async_trait;
 use oak_crypto::{
-    encryptor::AsyncRecipientContextGenerator, hpke::RecipientContext,
+    encryptor::AsyncEncryptionKeyHandle, hpke::RecipientContext,
     proto::oak::crypto::v1::SessionKeys,
 };
 use tonic::transport::{Endpoint, Uri};
@@ -71,28 +71,6 @@ impl OrchestratorCryptoClient {
     }
 }
 
-#[async_trait(?Send)]
-pub trait EncryptionKeyHandle {
-    async fn derive_session_keys(
-        &self,
-        encapsulated_public_key: &[u8],
-    ) -> anyhow::Result<RecipientContext>;
-}
-
-#[async_trait(?Send)]
-impl<T> EncryptionKeyHandle for T
-where
-    T: AsyncRecipientContextGenerator,
-{
-    async fn derive_session_keys(
-        &self,
-        encapsulated_public_key: &[u8],
-    ) -> anyhow::Result<RecipientContext> {
-        self.generate_recipient_context(encapsulated_public_key)
-            .await
-    }
-}
-
 pub struct InstanceEncryptionKeyHandle {
     orchestrator_crypto_client: OrchestratorCryptoClient,
 }
@@ -108,7 +86,7 @@ impl InstanceEncryptionKeyHandle {
 }
 
 #[async_trait]
-impl AsyncRecipientContextGenerator for InstanceEncryptionKeyHandle {
+impl AsyncEncryptionKeyHandle for InstanceEncryptionKeyHandle {
     async fn generate_recipient_context(
         &self,
         encapsulated_public_key: &[u8],

oak_containers_sdk/src/lib.rs

@@ -42,5 +42,5 @@ static IGNORED_ENDPOINT_URI: &str = "file://[::]:0";
 const ORCHESTRATOR_IPC_SOCKET: &str = "/oak_utils/orchestrator_ipc";
 
 // Re-export structs so that they are available at the top level of the SDK.
-pub use crypto::{EncryptionKeyHandle, InstanceEncryptionKeyHandle};
+pub use crypto::InstanceEncryptionKeyHandle;
 pub use orchestrator_client::OrchestratorClient;

oak_crypto/src/encryptor.rs

@@ -106,40 +106,14 @@ impl EncryptionKeyProvider {
     }
 }
 
-// This trait just aliases [`RecipientContextGenerator`], while using different naming
-// as defined in the Oak SDK design doc.
-// TODO(#3841): rename the relevant trait and struct in our crypto crates.
-/// Generate [`SessionKeys`] for the provided public key.
 pub trait EncryptionKeyHandle {
-    fn generate_session_keys(&self, encapsulated_public_key: &[u8]) -> anyhow::Result<SessionKeys>;
-}
-
-// Alias this struct in order to conform to the naming outlined in the restricted kernel SDK design
-// doc.
-// TODO(#3841): rename the relevant struct and remove this alias.
-use crate::encryptor::RecipientContext as SessionKeys;
-
-impl<T> RecipientContextGenerator for T
-where
-    T: EncryptionKeyHandle,
-{
-    fn generate_recipient_context(
-        &self,
-        encapsulated_public_key: &[u8],
-    ) -> anyhow::Result<SessionKeys> {
-        self.generate_session_keys(encapsulated_public_key)
-    }
-}
-
-pub trait RecipientContextGenerator {
-    // TODO(#3841): Implement Oak Kernel Crypto API and return corresponding session keys instead.
     fn generate_recipient_context(
         &self,
         encapsulated_public_key: &[u8],
     ) -> anyhow::Result<RecipientContext>;
 }
 
-impl RecipientContextGenerator for EncryptionKeyProvider {
+impl EncryptionKeyHandle for EncryptionKeyProvider {
     fn generate_recipient_context(
         &self,
         encapsulated_public_key: &[u8],
@@ -150,20 +124,20 @@ impl RecipientContextGenerator for EncryptionKeyProvider {
 }
 
 #[async_trait]
-pub trait AsyncRecipientContextGenerator {
+pub trait AsyncEncryptionKeyHandle {
     async fn generate_recipient_context(
         &self,
         encapsulated_public_key: &[u8],
     ) -> anyhow::Result<RecipientContext>;
 }
 
 #[async_trait]
-impl AsyncRecipientContextGenerator for EncryptionKeyProvider {
+impl AsyncEncryptionKeyHandle for EncryptionKeyProvider {
     async fn generate_recipient_context(
         &self,
         encapsulated_public_key: &[u8],
     ) -> anyhow::Result<RecipientContext> {
-        (self as &dyn RecipientContextGenerator).generate_recipient_context(encapsulated_public_key)
+        (self as &dyn EncryptionKeyHandle).generate_recipient_context(encapsulated_public_key)
     }
 }
 
@@ -258,9 +232,9 @@ pub struct ServerEncryptor {
 impl ServerEncryptor {
     pub fn create(
         serialized_encapsulated_public_key: &[u8],
-        recipient_context_generator: Arc<dyn RecipientContextGenerator>,
+        encryption_key_handle: Arc<dyn EncryptionKeyHandle>,
     ) -> anyhow::Result<Self> {
-        let recipient_context = recipient_context_generator
+        let recipient_context = encryption_key_handle
             .generate_recipient_context(serialized_encapsulated_public_key)
             .context("couldn't generate recipient crypto context")?;
         Ok(Self::new(recipient_context))
@@ -331,19 +305,19 @@ impl ServerEncryptor {
 // the Restricted Kernel.
 pub struct AsyncServerEncryptor<'a, G>
 where
-    G: AsyncRecipientContextGenerator + Send + Sync,
+    G: AsyncEncryptionKeyHandle + Send + Sync,
 {
-    recipient_context_generator: &'a G,
+    encryption_key_handle: &'a G,
     inner: Option<ServerEncryptor>,
 }
 
 impl<'a, G> AsyncServerEncryptor<'a, G>
 where
-    G: AsyncRecipientContextGenerator + Send + Sync,
+    G: AsyncEncryptionKeyHandle + Send + Sync,
 {
-    pub fn new(recipient_context_generator: &'a G) -> Self {
+    pub fn new(encryption_key_handle: &'a G) -> Self {
         Self {
-            recipient_context_generator,
+            encryption_key_handle,
             inner: None,
         }
     }
@@ -363,7 +337,7 @@ where
                     .as_ref()
                     .context("initial request message doesn't contain encapsulated public key")?;
                 let recipient_context = self
-                    .recipient_context_generator
+                    .encryption_key_handle
                     .generate_recipient_context(serialized_encapsulated_public_key)
                     .await
                     .context("couldn't generate recipient crypto context")?;

oak_crypto/src/tests.rs

@@ -21,8 +21,8 @@ use async_trait::async_trait;
 
 use crate::{
     encryptor::{
-        AsyncRecipientContextGenerator, AsyncServerEncryptor, ClientEncryptor,
-        EncryptionKeyProvider, ServerEncryptor, OAK_HPKE_INFO,
+        AsyncEncryptionKeyHandle, AsyncServerEncryptor, ClientEncryptor, EncryptionKeyProvider,
+        ServerEncryptor, OAK_HPKE_INFO,
     },
     hpke::{
         aead::{AEAD_ALGORITHM_KEY_SIZE_BYTES, AEAD_NONCE_SIZE_BYTES},
@@ -125,8 +125,8 @@ fn test_hpke() {
 
 #[test]
 fn test_encryptor() {
-    let key_provider = Arc::new(EncryptionKeyProvider::generate());
-    let serialized_server_public_key = key_provider.get_serialized_public_key();
+    let encryption_key = Arc::new(EncryptionKeyProvider::generate());
+    let serialized_server_public_key = encryption_key.get_serialized_public_key();
 
     let mut client_encryptor = ClientEncryptor::create(&serialized_server_public_key)
         .expect("couldn't create client encryptor");
@@ -152,7 +152,7 @@ fn test_encryptor() {
             .as_ref()
             .expect("initial request message doesn't contain encapsulated public key");
         server_encryptor = Some(
-            ServerEncryptor::create(serialized_encapsulated_public_key, key_provider.clone())
+            ServerEncryptor::create(serialized_encapsulated_public_key, encryption_key.clone())
                 .expect("couldn't create server encryptor"),
         );
     }
@@ -186,17 +186,17 @@ fn test_encryptor() {
     assert_eq!(TEST_RESPONSE_ASSOCIATED_DATA, response_associated_data);
 }
 
-struct TestEncryptionKeyProvider {
+struct TestEncryptionKey {
     key_pair: KeyPair,
 }
 
-impl Default for TestEncryptionKeyProvider {
+impl Default for TestEncryptionKey {
     fn default() -> Self {
         Self::new()
     }
 }
 
-impl TestEncryptionKeyProvider {
+impl TestEncryptionKey {
     pub fn new() -> Self {
         Self {
             key_pair: KeyPair::generate(),
@@ -209,7 +209,7 @@ impl TestEncryptionKeyProvider {
 }
 
 #[async_trait]
-impl AsyncRecipientContextGenerator for TestEncryptionKeyProvider {
+impl AsyncEncryptionKeyHandle for TestEncryptionKey {
     async fn generate_recipient_context(
         &self,
         encapsulated_public_key: &[u8],
@@ -221,12 +221,12 @@ impl AsyncRecipientContextGenerator for TestEncryptionKeyProvider {
 
 #[tokio::test]
 async fn test_async_encryptor() {
-    let key_provider = TestEncryptionKeyProvider::new();
-    let serialized_server_public_key = key_provider.get_serialized_public_key();
+    let encryption_key = TestEncryptionKey::new();
+    let serialized_server_public_key = encryption_key.get_serialized_public_key();
 
     let mut client_encryptor = ClientEncryptor::create(&serialized_server_public_key)
         .expect("couldn't create client encryptor");
-    let mut server_encryptor = AsyncServerEncryptor::new(&key_provider);
+    let mut server_encryptor = AsyncServerEncryptor::new(&encryption_key);
 
     let encrypted_request = client_encryptor
         .encrypt(TEST_REQUEST_MESSAGE, TEST_REQUEST_ASSOCIATED_DATA)

oak_functions_containers_app/src/lib.rs

@@ -23,7 +23,7 @@ use std::{
 
 use anyhow::Context;
 use oak_attestation::handler::AsyncEncryptionHandler;
-use oak_crypto::encryptor::AsyncRecipientContextGenerator;
+use oak_crypto::encryptor::AsyncEncryptionKeyHandle;
 use oak_functions_service::{
     instance::OakFunctionsInstance,
     proto::oak::functions::{
@@ -61,13 +61,13 @@ pub mod proto {
 }
 
 // Instance of the OakFunctions service for Oak Containers.
-pub struct OakFunctionsContainersService<G: AsyncRecipientContextGenerator + Send + Sync> {
+pub struct OakFunctionsContainersService<G: AsyncEncryptionKeyHandle + Send + Sync> {
     instance: OnceLock<OakFunctionsInstance>,
     encryption_key_handle: Arc<G>,
     observer: Option<Arc<dyn Observer + Send + Sync>>,
 }
 
-impl<G: AsyncRecipientContextGenerator + Send + Sync> OakFunctionsContainersService<G> {
+impl<G: AsyncEncryptionKeyHandle + Send + Sync> OakFunctionsContainersService<G> {
     pub fn new(
         encryption_key_handle: Arc<G>,
         observer: Option<Arc<dyn Observer + Send + Sync>>,
@@ -110,7 +110,7 @@ fn map_status(status: micro_rpc::Status) -> tonic::Status {
 }
 
 #[tonic::async_trait]
-impl<G: AsyncRecipientContextGenerator + Send + Sync + 'static> OakFunctions
+impl<G: AsyncEncryptionKeyHandle + Send + Sync + 'static> OakFunctions
     for OakFunctionsContainersService<G>
 {
     async fn initialize(
@@ -375,7 +375,7 @@ static GRPC_SUCCESS: http::header::HeaderValue = http::header::HeaderValue::from
 const GRPC_STATUS_HEADER_CODE: &str = "grpc-status";
 
 // Starts up and serves an OakFunctionsContainersService instance from the provided TCP listener.
-pub async fn serve<G: AsyncRecipientContextGenerator + Send + Sync + 'static>(
+pub async fn serve<G: AsyncEncryptionKeyHandle + Send + Sync + 'static>(
     listener: TcpListener,
     encryption_key_handle: Arc<G>,
     meter: Meter,