Create EncryptionProvider trait

It improves readability over the Fn-trait based version. Change-Id: I2c760cf97bf778cab2a2a6af6b585b7d7d4af218
project-oak · Sep 20, 2024 · 3f3564d · 3f3564d
1 parent 25ea019
commit 3f3564d
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 10 deletions.
diff --git a/oak_session/src/config.rs b/oak_session/src/config.rs
@@ -49,6 +49,19 @@ pub struct SessionConfigBuilder {
     config: SessionConfig,
 }
 
+pub trait EncryptorProvider: Send {
+    fn provide_encryptor(&self, session_keys: SessionKeys) -> Result<Box<dyn Encryptor>, Error>;
+}
+
+pub struct OrderedChannelEncryptorProvider;
+
+impl EncryptorProvider for OrderedChannelEncryptorProvider {
+    fn provide_encryptor(&self, session_keys: SessionKeys) -> Result<Box<dyn Encryptor>, Error> {
+        TryInto::<OrderedChannelEncryptor>::try_into(session_keys)
+            .map(|v| Box::new(v) as Box<dyn Encryptor>)
+    }
+}
+
 impl SessionConfigBuilder {
     fn new(attestation_type: AttestationType, handshake_type: HandshakeType) -> Self {
         let attestation_provider_config = AttestationProviderConfig {
@@ -65,12 +78,8 @@ impl SessionConfigBuilder {
             peer_attestation_binding_public_key: None,
         };
 
-        let encryptor_config = EncryptorConfig {
-            encryptor_provider: Box::new(|sk| {
-                <SessionKeys as TryInto<OrderedChannelEncryptor>>::try_into(sk)
-                    .map(|v| Box::new(v) as Box<dyn Encryptor>)
-            }),
-        };
+        let encryptor_config =
+            EncryptorConfig { encryptor_provider: Box::new(OrderedChannelEncryptorProvider) };
 
         let config =
             SessionConfig { attestation_provider_config, handshaker_config, encryptor_config };
@@ -111,7 +120,7 @@ impl SessionConfigBuilder {
 
     pub fn set_encryption_provider(
         mut self,
-        encryptor_provider: Box<dyn Send + Fn(SessionKeys) -> Result<Box<dyn Encryptor>, Error>>,
+        encryptor_provider: Box<dyn EncryptorProvider>,
     ) -> Self {
         self.config.encryptor_config.encryptor_provider = encryptor_provider;
         self
@@ -143,5 +152,5 @@ pub struct HandshakerConfig {
 }
 
 pub struct EncryptorConfig {
-    pub encryptor_provider: Box<dyn Send + Fn(SessionKeys) -> Result<Box<dyn Encryptor>, Error>>,
+    pub encryptor_provider: Box<dyn EncryptorProvider>,
 }
diff --git a/oak_session/src/session.rs b/oak_session/src/session.rs
@@ -192,7 +192,9 @@ impl ProtocolEngine<SessionResponse, SessionRequest> for ClientSession {
                         )?;
                     }
                     self.encryptor = Some(
-                        (self.encryptor_config.encryptor_provider)(handshake_result.session_keys)
+                        self.encryptor_config
+                            .encryptor_provider
+                            .provide_encryptor(handshake_result.session_keys)
                             .context("couldn't create an encryptor from the session key")?,
                     )
                 }
@@ -333,7 +335,9 @@ impl ProtocolEngine<SessionRequest, SessionResponse> for ServerSession {
                         )?;
                     }
                     self.encryptor = Some(
-                        (self.encryptor_config.encryptor_provider)(handshake_result.session_keys)
+                        self.encryptor_config
+                            .encryptor_provider
+                            .provide_encryptor(handshake_result.session_keys)
                             .context("couldn't create an encryptor from the session key")?,
                     )
                 }