Add config options to AttestationProvider and Handshaker

Bug: 338559159
Change-Id: I572120c9c32fe17c4965bdd8f3426abe002205d7

oak_session/src/attestation.rs

@@ -17,13 +17,13 @@
 //! This module provides an implementation of the Attestation Provider, which
 //! handles remote attestation between two parties.
 
-use alloc::vec::Vec;
-
 use oak_proto_rust::oak::{
     attestation::v1::{AttestationResults, Endorsements, Evidence},
     session::v1::{AttestRequest, AttestResponse, EndorsedEvidence},
 };
 
+use crate::{config::AttestationProviderConfig, ProtocolEngine};
+
 pub trait Attester {
     fn get_endorsed_evidence(&self) -> anyhow::Result<EndorsedEvidence>;
 }
@@ -36,45 +36,52 @@ pub trait AttestationVerifier {
     ) -> anyhow::Result<AttestationResults>;
 }
 
-#[allow(dead_code)]
-struct AttestationProvider<'a> {
-    self_attesters: Vec<&'a dyn Attester>,
-    peer_verifiers: Vec<&'a dyn AttestationVerifier>,
+/// Configuration of the attestation behavior that the AttestationProtiver will
+/// perform between two parties: Client and Server.
+///
+/// When configuring the Client: "Self" is the Client and "Peer" is the Server.
+/// When configuring the Server: "Self" is the Server and "Peer" is the Client.
+pub enum AttestationType {
+    /// Both parties attest each other.
+    Bidirectional,
+    /// "Self" attests itself to the "Peer".
+    SelfUnidirectional,
+    /// "Peer" attests itself to the "Self".
+    PeerUnidirectional,
 }
 
-impl<'a> AttestationProvider<'a> {
-    pub fn new(
-        self_attesters: Vec<&'a dyn Attester>,
-        peer_verifiers: Vec<&'a dyn AttestationVerifier>,
-    ) -> Self {
-        Self { self_attesters, peer_verifiers }
-    }
+pub trait AttestationProvider {
+    fn get_attestation_results(self) -> Option<AttestationResults>;
 }
 
 /// Client-side Attestation Provider that initiates remote attestation with the
 /// server.
 #[allow(dead_code)]
 pub struct ClientAttestationProvider<'a> {
-    inner: AttestationProvider<'a>,
+    config: AttestationProviderConfig<'a>,
 }
 
 impl<'a> ClientAttestationProvider<'a> {
-    pub fn new(
-        self_attesters: Vec<&'a dyn Attester>,
-        peer_verifiers: Vec<&'a dyn AttestationVerifier>,
-    ) -> Self {
-        Self { inner: AttestationProvider::new(self_attesters, peer_verifiers) }
+    pub fn new(config: AttestationProviderConfig<'a>) -> Self {
+        Self { config }
     }
+}
 
-    pub fn get_request(&self) -> anyhow::Result<AttestRequest> {
+impl<'a> AttestationProvider for ClientAttestationProvider<'a> {
+    fn get_attestation_results(self) -> Option<AttestationResults> {
         core::unimplemented!();
     }
+}
 
-    pub fn put_response(&self, _response: &AttestResponse) -> anyhow::Result<()> {
+impl<'a> ProtocolEngine<AttestResponse, AttestRequest> for ClientAttestationProvider<'a> {
+    fn get_outgoing_message(&mut self) -> anyhow::Result<Option<AttestRequest>> {
         core::unimplemented!();
     }
 
-    pub fn get_attestation_results(self) -> Option<AttestationResults> {
+    fn put_incoming_message(
+        &mut self,
+        _incoming_message: &AttestResponse,
+    ) -> anyhow::Result<Option<()>> {
         core::unimplemented!();
     }
 }
@@ -83,26 +90,30 @@ impl<'a> ClientAttestationProvider<'a> {
 /// request from the client.
 #[allow(dead_code)]
 pub struct ServerAttestationProvider<'a> {
-    inner: AttestationProvider<'a>,
+    config: AttestationProviderConfig<'a>,
 }
 
 impl<'a> ServerAttestationProvider<'a> {
-    pub fn new(
-        self_attesters: Vec<&'a dyn Attester>,
-        peer_verifiers: Vec<&'a dyn AttestationVerifier>,
-    ) -> Self {
-        Self { inner: AttestationProvider::new(self_attesters, peer_verifiers) }
+    pub fn new(config: AttestationProviderConfig<'a>) -> Self {
+        Self { config }
     }
+}
 
-    pub fn put_request(&self, _request: &AttestRequest) -> anyhow::Result<()> {
+impl<'a> AttestationProvider for ServerAttestationProvider<'a> {
+    fn get_attestation_results(self) -> Option<AttestationResults> {
         core::unimplemented!();
     }
+}
 
-    pub fn get_response(&self) -> anyhow::Result<AttestResponse> {
+impl<'a> ProtocolEngine<AttestRequest, AttestResponse> for ServerAttestationProvider<'a> {
+    fn get_outgoing_message(&mut self) -> anyhow::Result<Option<AttestResponse>> {
         core::unimplemented!();
     }
 
-    pub fn get_attestation_results(self) -> Option<AttestationResults> {
+    fn put_incoming_message(
+        &mut self,
+        _incoming_message: &AttestRequest,
+    ) -> anyhow::Result<Option<()>> {
         core::unimplemented!();
     }
 }

oak_session/src/config.rs

@@ -14,40 +14,77 @@
 // limitations under the License.
 //
 
-use crate::attestation::{AttestationVerifier, Attester};
+use alloc::{vec, vec::Vec};
 
-#[derive(Default)]
+use crate::{
+    attestation::{AttestationType, AttestationVerifier, Attester},
+    handshake::{EncryptionKeyHandle, HandshakeType},
+};
+
+#[allow(dead_code)]
 pub struct SessionConfig<'a> {
-    self_attester: Option<&'a dyn Attester>,
-    peer_verifier: Option<&'a dyn AttestationVerifier>,
+    attestation_provider_config: AttestationProviderConfig<'a>,
+    handshaker_config: HandshakerConfig<'a>,
 }
 
 impl<'a> SessionConfig<'a> {
-    pub fn builder() -> SessionConfigBuilder<'a> {
-        SessionConfigBuilder::default()
+    pub fn builder(attestation_type: AttestationType) -> SessionConfigBuilder<'a> {
+        SessionConfigBuilder::new(attestation_type)
     }
 }
 
-#[derive(Default)]
 pub struct SessionConfigBuilder<'a> {
     config: SessionConfig<'a>,
 }
 
 impl<'a> SessionConfigBuilder<'a> {
-    pub fn set_self_attester(mut self, self_attester: &'a dyn Attester) -> Self {
-        if self.config.self_attester.is_none() {
-            self.config.self_attester = Some(self_attester);
+    fn new(attestation_type: AttestationType) -> Self {
+        let handshake_type = match attestation_type {
+            AttestationType::Bidirectional => HandshakeType::NoiseKK,
+            AttestationType::SelfUnidirectional => HandshakeType::NoiseKN,
+            AttestationType::PeerUnidirectional => HandshakeType::NoiseNK,
+        };
+
+        let attestation_provider_config = AttestationProviderConfig {
+            attestation_type,
+            self_attesters: vec![],
+            peer_verifiers: vec![],
+        };
+
+        let handshaker_config = HandshakerConfig {
+            handshake_type,
+            self_static_private_key: None,
+            peer_static_public_key: None,
+        };
+
+        let config = SessionConfig { attestation_provider_config, handshaker_config };
+        Self { config }
+    }
+
+    pub fn add_self_attester(mut self, attester: &'a dyn Attester) -> Self {
+        self.config.attestation_provider_config.self_attesters.push(attester);
+        self
+    }
+
+    pub fn add_peer_verifier(mut self, verifier: &'a dyn AttestationVerifier) -> Self {
+        self.config.attestation_provider_config.peer_verifiers.push(verifier);
+        self
+    }
+
+    pub fn set_self_private_key(mut self, private_key: &'a dyn EncryptionKeyHandle) -> Self {
+        if self.config.handshaker_config.self_static_private_key.is_none() {
+            self.config.handshaker_config.self_static_private_key = Some(private_key);
         } else {
-            panic!("self attester has already been set");
+            panic!("self private key has already been set");
         }
         self
     }
 
-    pub fn set_peer_verifier(mut self, peer_verifier: &'a dyn AttestationVerifier) -> Self {
-        if self.config.peer_verifier.is_none() {
-            self.config.peer_verifier = Some(peer_verifier);
+    pub fn set_peer_static_public_key(mut self, public_key: &[u8]) -> Self {
+        if self.config.handshaker_config.peer_static_public_key.is_none() {
+            self.config.handshaker_config.peer_static_public_key = Some(public_key.to_vec());
         } else {
-            panic!("peer verifier has already been set");
+            panic!("peer public key has already been set");
         }
         self
     }
@@ -56,3 +93,17 @@ impl<'a> SessionConfigBuilder<'a> {
         self.config
     }
 }
+
+#[allow(dead_code)]
+pub struct AttestationProviderConfig<'a> {
+    attestation_type: AttestationType,
+    self_attesters: Vec<&'a dyn Attester>,
+    peer_verifiers: Vec<&'a dyn AttestationVerifier>,
+}
+
+#[allow(dead_code)]
+pub struct HandshakerConfig<'a> {
+    handshake_type: HandshakeType,
+    self_static_private_key: Option<&'a dyn EncryptionKeyHandle>,
+    peer_static_public_key: Option<Vec<u8>>,
+}

oak_session/src/handshake.rs

@@ -17,13 +17,13 @@
 //! This module provides an implementation of the Handshaker, which
 //! handles cryptographic handshake and secure session creation.
 
-use alloc::vec::Vec;
-
 use oak_proto_rust::oak::{
     crypto::v1::SessionKeys,
     session::v1::{HandshakeRequest, HandshakeResponse},
 };
 
+use crate::{config::HandshakerConfig, ProtocolEngine};
+
 pub trait EncryptionKeyHandle {
     fn derive_session_keys(
         &self,
@@ -34,74 +34,73 @@ pub trait EncryptionKeyHandle {
 
 pub enum HandshakeType {
     NoiseKK,
+    NoiseKN,
     NoiseNK,
 }
 
+pub trait Handshaker {
+    fn derive_session_keys(self) -> Option<SessionKeys>;
+}
+
 /// Client-side Handshaker that initiates the crypto handshake with the server.
 #[allow(dead_code)]
 pub struct ClientHandshaker<'a> {
-    handshake_type: HandshakeType,
-    self_static_private_key: Option<&'a dyn EncryptionKeyHandle>,
-    peer_static_public_key: Option<Vec<u8>>,
+    handshaker_config: HandshakerConfig<'a>,
 }
 
 impl<'a> ClientHandshaker<'a> {
-    pub fn new(
-        handshake_type: HandshakeType,
-        self_static_private_key: Option<&'a dyn EncryptionKeyHandle>,
-        peer_static_public_key: Option<&[u8]>,
-    ) -> Self {
-        Self {
-            handshake_type,
-            self_static_private_key,
-            peer_static_public_key: peer_static_public_key.map(|k| k.to_vec()),
-        }
+    pub fn new(handshaker_config: HandshakerConfig<'a>) -> Self {
+        Self { handshaker_config }
     }
+}
 
-    pub fn get_request(&mut self) -> anyhow::Result<HandshakeRequest> {
+impl<'a> Handshaker for ClientHandshaker<'a> {
+    fn derive_session_keys(self) -> Option<SessionKeys> {
         core::unimplemented!();
     }
+}
 
-    pub fn put_response(&mut self, _response: HandshakeResponse) -> anyhow::Result<()> {
+impl<'a> ProtocolEngine<HandshakeResponse, HandshakeRequest> for ClientHandshaker<'a> {
+    fn get_outgoing_message(&mut self) -> anyhow::Result<Option<HandshakeRequest>> {
         core::unimplemented!();
     }
 
-    pub fn derive_session_keys(self) -> Option<SessionKeys> {
+    fn put_incoming_message(
+        &mut self,
+        _incoming_message: &HandshakeResponse,
+    ) -> anyhow::Result<Option<()>> {
         core::unimplemented!();
     }
 }
 
-/// Server-side Attestation Provider that responds to the crypto handshake
-/// request from the client.
+/// Server-side Handshaker that responds to the crypto handshake request from
+/// the client.
 #[allow(dead_code)]
 pub struct ServerHandshaker<'a> {
-    handshake_type: HandshakeType,
-    self_static_private_key: Option<&'a dyn EncryptionKeyHandle>,
-    peer_static_public_key: Option<Vec<u8>>,
+    handshaker_config: HandshakerConfig<'a>,
 }
 
 impl<'a> ServerHandshaker<'a> {
-    pub fn new(
-        handshake_type: HandshakeType,
-        self_static_private_key: Option<&'a dyn EncryptionKeyHandle>,
-        peer_static_public_key: Option<&[u8]>,
-    ) -> Self {
-        Self {
-            handshake_type,
-            self_static_private_key,
-            peer_static_public_key: peer_static_public_key.map(|k| k.to_vec()),
-        }
+    pub fn new(handshaker_config: HandshakerConfig<'a>) -> Self {
+        Self { handshaker_config }
     }
+}
 
-    pub fn put_request(&mut self, _request: HandshakeRequest) -> anyhow::Result<()> {
+impl<'a> Handshaker for ServerHandshaker<'a> {
+    fn derive_session_keys(self) -> Option<SessionKeys> {
         core::unimplemented!();
     }
+}
 
-    pub fn get_response(&mut self) -> anyhow::Result<HandshakeResponse> {
+impl<'a> ProtocolEngine<HandshakeRequest, HandshakeResponse> for ServerHandshaker<'a> {
+    fn get_outgoing_message(&mut self) -> anyhow::Result<Option<HandshakeResponse>> {
         core::unimplemented!();
     }
 
-    pub fn derive_session_keys(self) -> Option<SessionKeys> {
+    fn put_incoming_message(
+        &mut self,
+        _incoming_message: &HandshakeRequest,
+    ) -> anyhow::Result<Option<()>> {
         core::unimplemented!();
     }
 }