Create server encryptor from context (project-oak#4281)

Fixes project-oak#4258

oak_crypto/src/encryptor.rs

@@ -23,7 +23,7 @@ use crate::{
  proto::oak::crypto::v1::{AeadEncryptedMessage, EncryptedRequest, EncryptedResponse},
 };
 use alloc::{sync::Arc, vec::Vec};
-use anyhow::{anyhow, Context};
+use anyhow::Context;
 
 /// Info string used by Hybrid Public Key Encryption;
 pub(crate) const OAK_HPKE_INFO: &[u8] = b"Oak Hybrid Public Key Encryption v1";
@@ -146,22 +146,26 @@ impl ClientEncryptor {
 /// Encryptor object for decrypting client requests that are received by the server and encrypting
 /// server responses that will be sent back to the client. Each Encryptor object corresponds to a
 /// single crypto session between the client and the server.
-/// Encryptor state is initialized after receiving an initial request message containing client's
-/// encapsulated public key.
 ///
 /// Sequence numbers for requests and responses are incremented separately, meaning that there could
 /// be multiple responses per request and multiple requests per response.
 pub struct ServerEncryptor {
- recipient_context_generator: Arc<dyn RecipientContextGenerator>,
- recipient_context: Option<RecipientContext>,
+ recipient_context: RecipientContext,
 }
 
 impl ServerEncryptor {
- pub fn new(recipient_context_generator: Arc<dyn RecipientContextGenerator>) -> Self {
- Self {
- recipient_context_generator,
- recipient_context: None,
- }
+ pub fn create(
+ serialized_encapsulated_public_key: &[u8],
+ recipient_context_generator: Arc<dyn RecipientContextGenerator>,
+ ) -> anyhow::Result<Self> {
+ let recipient_context = recipient_context_generator
+ .generate_recipient_context(serialized_encapsulated_public_key)
+ .context("couldn't generate recipient crypto context")?;
+ Ok(Self::new(recipient_context))
+ }
+
+ pub fn new(recipient_context: RecipientContext) -> Self {
+ Self { recipient_context }
  }
 
  /// Decrypts a [`EncryptedRequest`] proto message using AEAD.
@@ -170,35 +174,13 @@ impl ServerEncryptor {
  pub fn decrypt(
  &mut self,
  encrypted_request: &EncryptedRequest,
- ) -> anyhow::Result<(Vec<u8>, Vec<u8>)> {
- match &mut self.recipient_context {
- Some(context) => Self::decrypt_with_context(encrypted_request, context),
- None => {
- let serialized_encapsulated_public_key = encrypted_request
- .serialized_encapsulated_public_key
- .as_ref()
- .context("initial request message doesn't contain encapsulated public key")?;
- let mut recipient_context = self
- .recipient_context_generator
- .generate_recipient_context(serialized_encapsulated_public_key)
- .context("couldn't generate recipient crypto context")?;
- let (plaintext, associated_data) =
- Self::decrypt_with_context(encrypted_request, &mut recipient_context)?;
- self.recipient_context = Some(recipient_context);
- Ok((plaintext, associated_data))
- }
- }
- }
-
- fn decrypt_with_context(
- encrypted_request: &EncryptedRequest,
- context: &mut RecipientContext,
  ) -> anyhow::Result<(Vec<u8>, Vec<u8>)> {
  let encrypted_message = encrypted_request
  .encrypted_message
  .as_ref()
  .context("request doesn't contain encrypted message")?;
- let plaintext = context
+ let plaintext = self
+ .recipient_context
  .open(
  &encrypted_message.ciphertext,
  &encrypted_message.associated_data,
@@ -215,22 +197,16 @@ impl ServerEncryptor {
  plaintext: &[u8],
  associated_data: &[u8],
  ) -> anyhow::Result<EncryptedResponse> {
- match &mut self.recipient_context {
- Some(context) => {
- let ciphertext = context
- .seal(plaintext, associated_data)
- .context("couldn't encrypt response")?;
- let response = EncryptedResponse {
- encrypted_message: Some(AeadEncryptedMessage {
- ciphertext,
- associated_data: associated_data.to_vec(),
- }),
- };
- Ok(response)
- }
- None => Err(anyhow!(
- "couldn't encrypt response because crypto context is not initialized"
- )),
- }
+ let ciphertext = self
+ .recipient_context
+ .seal(plaintext, associated_data)
+ .context("couldn't encrypt response")?;
+ let response = EncryptedResponse {
+ encrypted_message: Some(AeadEncryptedMessage {
+ ciphertext,
+ associated_data: associated_data.to_vec(),
+ }),
+ };
+ Ok(response)
  }
 }

oak_crypto/src/tests.rs

@@ -114,7 +114,7 @@ fn test_encryptor() {
 
  let mut client_encryptor = ClientEncryptor::create(&serialized_server_public_key)
  .expect("couldn't create client encryptor");
- let mut server_encryptor = ServerEncryptor::new(key_provider);
+ let mut server_encryptor = None;
 
  for i in 0..TEST_SESSION_SIZE {
  let test_request_message = [TEST_REQUEST_MESSAGE, &[i as u8]].concat();
@@ -134,13 +134,30 @@ fn test_encryptor() {
  .unwrap()
  .ciphertext
  );
+
+ // Initialize server encryptor.
+ if server_encryptor.is_none() {
+ let serialized_encapsulated_public_key = encrypted_request
+ .serialized_encapsulated_public_key
+ .as_ref()
+ .expect("initial request message doesn't contain encapsulated public key");
+ server_encryptor = Some(
+ ServerEncryptor::create(serialized_encapsulated_public_key, key_provider.clone())
+ .expect("couldn't create server encryptor"),
+ );
+ }
+
  let (decrypted_request, request_associated_data) = server_encryptor
+ .as_mut()
+ .expect("server encryptor is not initialized")
  .decrypt(&encrypted_request)
  .expect("server couldn't decrypt request");
  assert_eq!(test_request_message, decrypted_request);
  assert_eq!(test_request_associated_data, request_associated_data);
 
  let encrypted_response = server_encryptor
+ .as_mut()
+ .expect("server encryptor is not initialized")
  .encrypt(&test_response_message, &test_response_associated_data)
  .expect("server couldn't encrypt response");
  // Check that the message was encrypted.

oak_remote_attestation/src/handler.rs

@@ -54,11 +54,22 @@ impl<H: FnOnce(Vec<u8>) -> anyhow::Result<Vec<u8>>> EncryptionHandler<H> {
 
 impl<H: FnOnce(Vec<u8>) -> anyhow::Result<Vec<u8>>> EncryptionHandler<H> {
  pub fn invoke(self, request_body: &[u8]) -> anyhow::Result<Vec<u8>> {
- let mut server_encryptor = ServerEncryptor::new(self.encryption_key_provider.clone());
-
- // Deserialize and decrypt request.
+ // Deserialize request.
  let encrypted_request = EncryptedRequest::decode(request_body)
  .map_err(|error| anyhow!("couldn't deserialize request: {:?}", error))?;
+
+ // Initialize server encryptor.
+ let serialized_encapsulated_public_key = encrypted_request
+ .serialized_encapsulated_public_key
+ .as_ref()
+ .expect("initial request message doesn't contain encapsulated public key");
+ let mut server_encryptor = ServerEncryptor::create(
+ serialized_encapsulated_public_key,
+ self.encryption_key_provider.clone(),
+ )
+ .map_err(|error| anyhow!("couldn't create server encryptor: {:?}", error))?;
+
+ // Decrypt request.
  let (request, _) = server_encryptor
  .decrypt(&encrypted_request)
  .context("couldn't decrypt request")?;