dpe-rs/src/noise.rs - open-dice - Git at Google

 // Copyright 2024 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License"); you may not
 // use this file except in compliance with the License. You may obtain a copy of
 // the License at
 //
 //     https://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 // License for the specific language governing permissions and limitations under
 // the License.

 //! An encrypted session implementation which uses
 //! Noise_NK_X25519_AESGCM_SHA512 and Noise_NNpsk0_X25519_AESGCM_SHA512.

 use crate::crypto::{
     Commit, Counter, DhPrivateKey, DhPublicKey, HandshakeMessage,
     HandshakePayload, Hash, SessionCrypto,
 };
 use crate::error::{DpeResult, ErrCode};
 use crate::memory::Message;
 use core::marker::PhantomData;
 use log::{debug, error};
 use noise_protocol::{HandshakeStateBuilder, Hash as NoiseHash, U8Array};

 impl From<noise_protocol::Error> for ErrCode {
     fn from(_err: noise_protocol::Error) -> Self {
         ErrCode::InvalidArgument
     }
 }

 impl<NoiseHash> From<&NoiseHash> for Hash
 where
     NoiseHash: U8Array,
 {
     fn from(value: &NoiseHash) -> Self {
         // The Noise hash size may not match HASH_SIZE.
         Hash::from_slice_infallible(value.as_slice())
     }
 }

 /// A cipher state type that can be used as a
 /// [`SessionCipherState`](crate::crypto::SessionCrypto::SessionCipherState).
 pub struct NoiseCipherState<C: noise_protocol::Cipher> {
     k: C::Key,
     n: u64,
     n_staged: u64,
 }

 impl<C: noise_protocol::Cipher> Clone for NoiseCipherState<C> {
     fn clone(&self) -> Self {
         Self { k: self.k.clone(), n: self.n, n_staged: self.n_staged }
     }
 }

 impl<C: noise_protocol::Cipher> Default for NoiseCipherState<C> {
     fn default() -> Self {
         Self { k: C::Key::new(), n: 0, n_staged: 0 }
     }
 }

 impl<C: noise_protocol::Cipher> core::fmt::Debug for NoiseCipherState<C> {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         write!(f, "k: redacted, n: {}", self.n)?;
         Ok(())
     }
 }

 impl<C: noise_protocol::Cipher> core::hash::Hash for NoiseCipherState<C> {
     fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
         self.k.as_slice().hash(state);
         self.n.hash(state);
         self.n_staged.hash(state);
     }
 }

 #[cfg(test)]
 impl<C: noise_protocol::Cipher> PartialEq for NoiseCipherState<C> {
     fn eq(&self, other: &Self) -> bool {
         self.k.as_slice() == other.k.as_slice()
             && self.n == other.n
             && self.n_staged == other.n_staged
     }
 }

 #[cfg(test)]
 impl<C: noise_protocol::Cipher> Eq for NoiseCipherState<C> {}

 impl<C: noise_protocol::Cipher> Counter for NoiseCipherState<C> {
     fn n(&self) -> u64 {
         self.n
     }
     fn set_n(&mut self, n: u64) {
         self.n = n;
     }
 }

 impl<C: noise_protocol::Cipher> Commit for NoiseCipherState<C> {
     // Called when an encrypted message is finalized to commit the new cipher
     // state.
     fn commit(&mut self) {
         self.n = self.n_staged;
     }
 }

 impl<C: noise_protocol::Cipher> From<&noise_protocol::CipherState<C>>
     for NoiseCipherState<C>
 {
     fn from(cs: &noise_protocol::CipherState<C>) -> Self {
         let (key, counter) = cs.clone().extract();
         NoiseCipherState { k: key, n: counter, n_staged: counter }
     }
 }

 /// Returns the public key corresponding to a given `dh_private_key`.
 pub fn get_dh_public_key<D: noise_protocol::DH>(
     dh_private_key: &DhPrivateKey,
 ) -> DpeResult<DhPublicKey> {
     DhPublicKey::from_slice(
         D::pubkey(&D::Key::from_slice(dh_private_key.as_slice())).as_slice(),
     )
 }

 /// A trait representing [`NoiseSessionCrypto`] dependencies.
 pub trait NoiseCryptoDeps {
     /// Cipher type
     type Cipher: noise_protocol::Cipher;
     /// DH type
     type DH: noise_protocol::DH;
     /// Hash type
     type Hash: noise_protocol::Hash;
 }

 /// A Noise implementation of the [`SessionCrypto`] trait.
 pub struct NoiseSessionCrypto<D: NoiseCryptoDeps> {
     #[allow(dead_code)]
     phantom: PhantomData<D>,
 }

 impl<D> Clone for NoiseSessionCrypto<D>
 where
     D: NoiseCryptoDeps,
 {
     fn clone(&self) -> Self {
         Self { phantom: Default::default() }
     }
 }

 impl<D> Default for NoiseSessionCrypto<D>
 where
     D: NoiseCryptoDeps,
 {
     fn default() -> Self {
         Self { phantom: Default::default() }
     }
 }

 impl<D> core::fmt::Debug for NoiseSessionCrypto<D>
 where
     D: NoiseCryptoDeps,
 {
     fn fmt(&self, _: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         Ok(())
     }
 }

 impl<D> core::hash::Hash for NoiseSessionCrypto<D>
 where
     D: NoiseCryptoDeps,
 {
     fn hash<Hr: core::hash::Hasher>(&self, _: &mut Hr) {}
 }

 impl<D> PartialEq for NoiseSessionCrypto<D>
 where
     D: NoiseCryptoDeps,
 {
     fn eq(&self, _: &Self) -> bool {
         true
     }
 }

 impl<D> Eq for NoiseSessionCrypto<D> where D: NoiseCryptoDeps {}

 impl<D> SessionCrypto for NoiseSessionCrypto<D>
 where
     D: NoiseCryptoDeps,
 {
     type SessionCipherState = NoiseCipherState<D::Cipher>;

     /// Implements the responder role of a Noise_NK handshake.
     fn new_session_handshake(
         static_dh_key: &DhPrivateKey,
         initiator_handshake: &HandshakeMessage,
         payload: &HandshakePayload,
         responder_handshake: &mut HandshakeMessage,
         decrypt_cipher_state: &mut NoiseCipherState<D::Cipher>,
         encrypt_cipher_state: &mut NoiseCipherState<D::Cipher>,
         psk_seed: &mut Hash,
     ) -> DpeResult<()> {
         #[allow(unused_results)]
         let mut handshake: noise_protocol::HandshakeState<
             D::DH,
             D::Cipher,
             D::Hash,
         > = {
             let mut builder = HandshakeStateBuilder::new();
             builder.set_pattern(noise_protocol::patterns::noise_nk());
             builder.set_is_initiator(false);
             builder.set_prologue(&[]);
             builder.set_s(<D::DH as noise_protocol::DH>::Key::from_slice(
                 static_dh_key.as_slice(),
             ));
             builder.build_handshake_state()
         };
         handshake.read_message(initiator_handshake.as_slice(), &mut [])?;
         handshake.write_message(
             payload.as_slice(),
             responder_handshake.as_mut_sized(
                 handshake.get_next_message_overhead() + payload.len(),
             )?,
         )?;
         assert!(handshake.completed());
         let ciphers = handshake.get_ciphers();
         *decrypt_cipher_state = (&ciphers.0).into();
         *encrypt_cipher_state = (&ciphers.1).into();
         debug!("get_hash");
         *psk_seed = Hash::from_slice(handshake.get_hash())?;
         Ok(())
     }

     /// Implements the responder role of a Noise_NNpsk0 handshake.
     fn derive_session_handshake(
         psk: &Hash,
         initiator_handshake: &HandshakeMessage,
         payload: &HandshakePayload,
         responder_handshake: &mut HandshakeMessage,
         decrypt_cipher_state: &mut NoiseCipherState<D::Cipher>,
         encrypt_cipher_state: &mut NoiseCipherState<D::Cipher>,
         psk_seed: &mut Hash,
     ) -> DpeResult<()> {
         #[allow(unused_results)]
         let mut handshake: noise_protocol::HandshakeState<
             D::DH,
             D::Cipher,
             D::Hash,
         > = {
             let mut builder = HandshakeStateBuilder::new();
             builder.set_pattern(noise_protocol::patterns::noise_nn_psk0());
             builder.set_is_initiator(false);
             builder.set_prologue(&[]);
             builder.build_handshake_state()
         };
         handshake
             .push_psk(psk.as_slice().get(..32).ok_or(ErrCode::InternalError)?);
         handshake.read_message(initiator_handshake.as_slice(), &mut [])?;
         handshake.write_message(
             payload.as_slice(),
             responder_handshake.as_mut_sized(
                 handshake.get_next_message_overhead() + payload.len(),
             )?,
         )?;
         let ciphers = handshake.get_ciphers();
         *decrypt_cipher_state = (&ciphers.0).into();
         *encrypt_cipher_state = (&ciphers.1).into();
         *psk_seed = Hash::from_slice(handshake.get_hash())?;
         Ok(())
     }

     /// Encrypts a Noise transport message in place.
     fn session_encrypt(
         cipher_state: &mut NoiseCipherState<D::Cipher>,
         in_place_buffer: &mut Message,
     ) -> DpeResult<()> {
         let mut cs = noise_protocol::CipherState::<D::Cipher>::new(
             cipher_state.k.as_slice(),
             cipher_state.n,
         );
         let plaintext_len = in_place_buffer.len();
         let _ = cs.encrypt_in_place(
             in_place_buffer.as_mut_sized(
                 plaintext_len
                     + <D::Cipher as noise_protocol::Cipher>::tag_len(),
             )?,
             plaintext_len,
         );
         // Encrypting a message is usually not the final step in preparing
         // the message for transport. If a subsequent step fails, it is
         // better for 'n' to remain unchanged so we don't get out of sync.
         (_, cipher_state.n_staged) = cs.extract();
         Ok(())
     }

     /// Decrypts a Noise transport message in place.
     fn session_decrypt(
         cipher_state: &mut NoiseCipherState<D::Cipher>,
         in_place_buffer: &mut Message,
     ) -> DpeResult<()> {
         let mut cs = noise_protocol::CipherState::<D::Cipher>::new(
             cipher_state.k.as_slice(),
             cipher_state.n,
         );
         let ciphertext_len = in_place_buffer.len();
         let plaintext_len = match cs
             .decrypt_in_place(in_place_buffer.vec.as_mut(), ciphertext_len)
         {
             Ok(length) => length,
             _ => {
                 error!("Session decrypt failed");
                 return Err(ErrCode::InvalidCommand);
             }
         };
         in_place_buffer.vec.truncate(plaintext_len);
         (_, cipher_state.n) = cs.extract();
         Ok(())
     }

     /// Derives a responder-side PSK.
     fn derive_psk_from_session(
         psk_seed: &Hash,
         decrypt_cipher_state: &NoiseCipherState<D::Cipher>,
         encrypt_cipher_state: &NoiseCipherState<D::Cipher>,
     ) -> DpeResult<Hash> {
         let mut hasher: D::Hash = Default::default();
         hasher.input(psk_seed.as_slice());
         // Use the decrypt state as it was before we decrypted the current
         // command message. This allows clients to compute the PSK using
         // the cipher states as they are before the client sends the
         // command.
         hasher.input(&(decrypt_cipher_state.n() - 1).to_le_bytes());
         hasher.input(&encrypt_cipher_state.n().to_le_bytes());
         Ok((&hasher.result()).into())
     }
 }

 /// A SessionClient implements the initiator side of an encrypted session. A
 /// DPE does not use this itself, it is useful for clients and testing.
 pub struct SessionClient<D>
 where
     D: NoiseCryptoDeps,
 {
     handshake_state:
         Option<noise_protocol::HandshakeState<D::DH, D::Cipher, D::Hash>>,
     /// Cipher state for encrypting messages to a DPE.
     pub encrypt_cipher_state: NoiseCipherState<D::Cipher>,
     /// Cipher state for decrypting messages from a DPE.
     pub decrypt_cipher_state: NoiseCipherState<D::Cipher>,
     /// PSK seed for deriving sessions. See [`derive_psk`].
     ///
     /// [`derive_psk`]: #method.derive_psk
     pub psk_seed: Hash,
 }

 impl<D> Clone for SessionClient<D>
 where
     D: NoiseCryptoDeps,
 {
     fn clone(&self) -> Self {
         Self {
             handshake_state: self.handshake_state.clone(),
             encrypt_cipher_state: self.encrypt_cipher_state.clone(),
             decrypt_cipher_state: self.decrypt_cipher_state.clone(),
             psk_seed: self.psk_seed.clone(),
         }
     }
 }

 impl<D> Default for SessionClient<D>
 where
     D: NoiseCryptoDeps,
 {
     fn default() -> Self {
         Self::new()
     }
 }

 impl<D> core::fmt::Debug for SessionClient<D>
 where
     D: NoiseCryptoDeps,
 {
     fn fmt(&self, _: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         Ok(())
     }
 }

 impl<D> SessionClient<D>
 where
     D: NoiseCryptoDeps,
 {
     /// Creates a new SessionClient instance. Set up by starting and finishing a
     /// handshake.
     pub fn new() -> Self {
         Self {
             handshake_state: Default::default(),
             encrypt_cipher_state: Default::default(),
             decrypt_cipher_state: Default::default(),
             psk_seed: Default::default(),
         }
     }

     /// Starts a handshake using a known `public_key` and returns a message that
     /// works with the DPE OpenSession command.
     pub fn start_handshake_with_known_public_key(
         &mut self,
         public_key: &DhPublicKey,
     ) -> DpeResult<HandshakeMessage> {
         #[allow(unused_results)]
         let mut handshake_state = {
             let mut builder = HandshakeStateBuilder::new();
             builder.set_pattern(noise_protocol::patterns::noise_nk());
             builder.set_is_initiator(true);
             builder.set_prologue(&[]);
             builder.set_rs(<D::DH as noise_protocol::DH>::Pubkey::from_slice(
                 public_key.as_slice(),
             ));
             builder.build_handshake_state()
         };
         let mut message = HandshakeMessage::new();
         handshake_state.write_message(
             &[],
             message
                 .as_mut_sized(handshake_state.get_next_message_overhead())?,
         )?;
         self.handshake_state = Some(handshake_state);
         Ok(message)
     }

     /// Starts a handshake using a `psk` and returns a message that works with
     /// the DPE DeriveContext command. Use [`derive_psk`] to obtain this value
     /// from an existing session.
     ///
     /// [`derive_psk`]: #method.derive_psk
     pub fn start_handshake_with_psk(
         &mut self,
         psk: &Hash,
     ) -> DpeResult<HandshakeMessage> {
         #[allow(unused_results)]
         let mut handshake_state = {
             let mut builder = HandshakeStateBuilder::new();
             builder.set_pattern(noise_protocol::patterns::noise_nn_psk0());
             builder.set_is_initiator(true);
             builder.set_prologue(&[]);
             builder.build_handshake_state()
         };
         handshake_state
             .push_psk(psk.as_slice().get(..32).ok_or(ErrCode::InternalError)?);
         let mut message = HandshakeMessage::new();
         handshake_state.write_message(
             &[],
             message
                 .as_mut_sized(handshake_state.get_next_message_overhead())?,
         )?;
         self.handshake_state = Some(handshake_state);
         Ok(message)
     }

     /// Finishes a handshake started using one of the start_handshake_* methods.
     /// On success, returns the handshake payload from the responder and sets up
     /// internal state for subsequent calls to encrypt and decrypt.
     pub fn finish_handshake(
         &mut self,
         responder_handshake: &HandshakeMessage,
     ) -> DpeResult<HandshakePayload> {
         match self.handshake_state {
             None => Err(ErrCode::InvalidArgument),
             Some(ref mut handshake) => {
                 let mut payload = HandshakePayload::new();
                 handshake.read_message(
                     responder_handshake.as_slice(),
                     payload.as_mut_sized(
                         responder_handshake.len()
                             - handshake.get_next_message_overhead(),
                     )?,
                 )?;
                 let ciphers = handshake.get_ciphers();
                 self.encrypt_cipher_state = (&ciphers.0).into();
                 self.decrypt_cipher_state = (&ciphers.1).into();
                 self.psk_seed = Hash::from_slice(handshake.get_hash())?;
                 Ok(payload)
             }
         }
     }

     /// Derives a PSK from the current session.
     pub fn derive_psk(&self) -> Hash {
         // Note this is from a client perspective so the counters are hashed
         // encrypt first and unmodified from their current state. A DPE will
         // reverse the order and decrement the first counter in order to derive
         // the same value (see derive_psk_from_session).
         let mut hasher: D::Hash = Default::default();
         hasher.input(self.psk_seed.as_slice());
         hasher.input(&self.encrypt_cipher_state.n().to_le_bytes());
         hasher.input(&self.decrypt_cipher_state.n().to_le_bytes());
         (&hasher.result()).into()
     }

     /// Encrypts a message to send to a DPE and commits cipher state changes.
     pub fn encrypt(&mut self, in_place_buffer: &mut Message) -> DpeResult<()> {
         NoiseSessionCrypto::<D>::session_encrypt(
             &mut self.encrypt_cipher_state,
             in_place_buffer,
         )?;
         self.encrypt_cipher_state.commit();
         Ok(())
     }

     /// Decrypts a message from a DPE.
     pub fn decrypt(&mut self, in_place_buffer: &mut Message) -> DpeResult<()> {
         NoiseSessionCrypto::<D>::session_decrypt(
             &mut self.decrypt_cipher_state,
             in_place_buffer,
         )
     }
 }

 #[cfg(test)]
 mod tests {
     use super::*;

     struct DepsForTesting {}
     impl NoiseCryptoDeps for DepsForTesting {
         type Cipher = noise_rust_crypto::Aes256Gcm;
         type DH = noise_rust_crypto::X25519;
         type Hash = noise_rust_crypto::Sha512;
     }

     type SessionCryptoForTesting = NoiseSessionCrypto<DepsForTesting>;

     type SessionClientForTesting = SessionClient<DepsForTesting>;

     type CipherStateForTesting = NoiseCipherState<noise_rust_crypto::Aes256Gcm>;

     #[test]
     fn end_to_end_session() {
         let mut client = SessionClientForTesting::new();
         let dh_key: DhPrivateKey = Default::default();
         let dh_public_key = get_dh_public_key::<
             <DepsForTesting as NoiseCryptoDeps>::DH,
         >(&dh_key)
         .unwrap();
         let handshake1 = client
             .start_handshake_with_known_public_key(&dh_public_key)
             .unwrap();
         let mut dpe_decrypt_cs: CipherStateForTesting = Default::default();
         let mut dpe_encrypt_cs: CipherStateForTesting = Default::default();
         let mut psk_seed = Default::default();
         let mut handshake2 = Default::default();
         let payload = HandshakePayload::from_slice("pay".as_bytes()).unwrap();
         SessionCryptoForTesting::new_session_handshake(
             &dh_key,
             &handshake1,
             &payload,
             &mut handshake2,
             &mut dpe_decrypt_cs,
             &mut dpe_encrypt_cs,
             &mut psk_seed,
         )
         .unwrap();
         assert_eq!(payload, client.finish_handshake(&handshake2).unwrap());

         // Check that the session works.
         let mut buffer = Message::from_slice("message".as_bytes()).unwrap();
         client.encrypt(&mut buffer).unwrap();
         SessionCryptoForTesting::session_decrypt(
             &mut dpe_decrypt_cs,
             &mut buffer,
         )
         .unwrap();
         assert_eq!("message".as_bytes(), buffer.as_slice());
         SessionCryptoForTesting::session_encrypt(
             &mut dpe_encrypt_cs,
             &mut buffer,
         )
         .unwrap();
         dpe_encrypt_cs.commit();
         client.decrypt(&mut buffer).unwrap();
         assert_eq!("message".as_bytes(), buffer.as_slice());

         // Do it again to check session state still works.
         client.encrypt(&mut buffer).unwrap();
         SessionCryptoForTesting::session_decrypt(
             &mut dpe_decrypt_cs,
             &mut buffer,
         )
         .unwrap();
         assert_eq!("message".as_bytes(), buffer.as_slice());
         SessionCryptoForTesting::session_encrypt(
             &mut dpe_encrypt_cs,
             &mut buffer,
         )
         .unwrap();
         dpe_encrypt_cs.commit();
         client.decrypt(&mut buffer).unwrap();
         assert_eq!("message".as_bytes(), buffer.as_slice());
     }

     #[test]
     fn derived_session() {
         // Set up a session from which to derive.
         let mut client = SessionClientForTesting::new();
         let dh_key: DhPrivateKey = Default::default();
         let dh_public_key = get_dh_public_key::<
             <DepsForTesting as NoiseCryptoDeps>::DH,
         >(&dh_key)
         .unwrap();
         let handshake1 = client
             .start_handshake_with_known_public_key(&dh_public_key)
             .unwrap();
         let mut dpe_decrypt_cs = Default::default();
         let mut dpe_encrypt_cs = Default::default();
         let mut psk_seed = Default::default();
         let mut handshake2 = Default::default();
         let payload = HandshakePayload::from_slice("pay".as_bytes()).unwrap();
         SessionCryptoForTesting::new_session_handshake(
             &dh_key,
             &handshake1,
             &payload,
             &mut handshake2,
             &mut dpe_decrypt_cs,
             &mut dpe_encrypt_cs,
             &mut psk_seed,
         )
         .unwrap();
         assert_eq!(payload, client.finish_handshake(&handshake2).unwrap());

         // Derive a second session.
         let mut client2 = SessionClientForTesting::new();
         let client_psk = client.derive_psk();
         // Simulate the session state after command decryption on the DPE side
         // as expected by the DPE PSK logic.
         let mut buffer = Message::from_slice("message".as_bytes()).unwrap();
         client.encrypt(&mut buffer).unwrap();
         SessionCryptoForTesting::session_decrypt(
             &mut dpe_decrypt_cs,
             &mut buffer,
         )
         .unwrap();
         let dpe_psk = SessionCryptoForTesting::derive_psk_from_session(
             &psk_seed,
             &dpe_decrypt_cs,
             &dpe_encrypt_cs,
         )
         .unwrap();
         let handshake1 = client2.start_handshake_with_psk(&client_psk).unwrap();
         let mut dpe_decrypt_cs2 = Default::default();
         let mut dpe_encrypt_cs2 = Default::default();
         let mut psk_seed2 = Default::default();
         SessionCryptoForTesting::derive_session_handshake(
             &dpe_psk,
             &handshake1,
             &payload,
             &mut handshake2,
             &mut dpe_decrypt_cs2,
             &mut dpe_encrypt_cs2,
             &mut psk_seed2,
         )
         .unwrap();
         assert_eq!(payload, client2.finish_handshake(&handshake2).unwrap());

         // Check that the second session works.
         let mut buffer = Message::from_slice("message".as_bytes()).unwrap();
         client2.encrypt(&mut buffer).unwrap();
         SessionCryptoForTesting::session_decrypt(
             &mut dpe_decrypt_cs2,
             &mut buffer,
         )
         .unwrap();
         assert_eq!("message".as_bytes(), buffer.as_slice());
         SessionCryptoForTesting::session_encrypt(
             &mut dpe_encrypt_cs2,
             &mut buffer,
         )
         .unwrap();
         dpe_encrypt_cs2.commit();
         client2.decrypt(&mut buffer).unwrap();
         assert_eq!("message".as_bytes(), buffer.as_slice());

         // Check that the first session also still works.
         let mut buffer = Message::from_slice("message".as_bytes()).unwrap();
         client.encrypt(&mut buffer).unwrap();
         SessionCryptoForTesting::session_decrypt(
             &mut dpe_decrypt_cs,
             &mut buffer,
         )
         .unwrap();
         assert_eq!("message".as_bytes(), buffer.as_slice());
         SessionCryptoForTesting::session_encrypt(
             &mut dpe_encrypt_cs,
             &mut buffer,
         )
         .unwrap();
         dpe_encrypt_cs.commit();
         client.decrypt(&mut buffer).unwrap();
         assert_eq!("message".as_bytes(), buffer.as_slice());
     }
 }
	// Copyright 2024 Google LLC
	//
	// Licensed under the Apache License, Version 2.0 (the "License"); you may not
	// use this file except in compliance with the License. You may obtain a copy of
	// the License at
	//
	// https://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
	// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
	// License for the specific language governing permissions and limitations under
	// the License.

	//! An encrypted session implementation which uses
	//! Noise_NK_X25519_AESGCM_SHA512 and Noise_NNpsk0_X25519_AESGCM_SHA512.

	use crate::crypto::{
	Commit, Counter, DhPrivateKey, DhPublicKey, HandshakeMessage,
	HandshakePayload, Hash, SessionCrypto,
	};
	use crate::error::{DpeResult, ErrCode};
	use crate::memory::Message;
	use core::marker::PhantomData;
	use log::{debug, error};
	use noise_protocol::{HandshakeStateBuilder, Hash as NoiseHash, U8Array};

	impl From<noise_protocol::Error> for ErrCode {
	fn from(_err: noise_protocol::Error) -> Self {
	ErrCode::InvalidArgument
	}
	}

	impl<NoiseHash> From<&NoiseHash> for Hash
	where
	NoiseHash: U8Array,
	{
	fn from(value: &NoiseHash) -> Self {
	// The Noise hash size may not match HASH_SIZE.
	Hash::from_slice_infallible(value.as_slice())
	}
	}

	/// A cipher state type that can be used as a
	/// [`SessionCipherState`](crate::crypto::SessionCrypto::SessionCipherState).
	pub struct NoiseCipherState<C: noise_protocol::Cipher> {
	k: C::Key,
	n: u64,
	n_staged: u64,
	}

	impl<C: noise_protocol::Cipher> Clone for NoiseCipherState<C> {
	fn clone(&self) -> Self {
	Self { k: self.k.clone(), n: self.n, n_staged: self.n_staged }
	}
	}

	impl<C: noise_protocol::Cipher> Default for NoiseCipherState<C> {
	fn default() -> Self {
	Self { k: C::Key::new(), n: 0, n_staged: 0 }
	}
	}

	impl<C: noise_protocol::Cipher> core::fmt::Debug for NoiseCipherState<C> {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
	write!(f, "k: redacted, n: {}", self.n)?;
	Ok(())
	}
	}

	impl<C: noise_protocol::Cipher> core::hash::Hash for NoiseCipherState<C> {
	fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
	self.k.as_slice().hash(state);
	self.n.hash(state);
	self.n_staged.hash(state);
	}
	}

	#[cfg(test)]
	impl<C: noise_protocol::Cipher> PartialEq for NoiseCipherState<C> {
	fn eq(&self, other: &Self) -> bool {
	self.k.as_slice() == other.k.as_slice()
	&& self.n == other.n
	&& self.n_staged == other.n_staged
	}
	}

	#[cfg(test)]
	impl<C: noise_protocol::Cipher> Eq for NoiseCipherState<C> {}

	impl<C: noise_protocol::Cipher> Counter for NoiseCipherState<C> {
	fn n(&self) -> u64 {
	self.n
	}
	fn set_n(&mut self, n: u64) {
	self.n = n;
	}
	}

	impl<C: noise_protocol::Cipher> Commit for NoiseCipherState<C> {
	// Called when an encrypted message is finalized to commit the new cipher
	// state.
	fn commit(&mut self) {
	self.n = self.n_staged;
	}
	}

	impl<C: noise_protocol::Cipher> From<&noise_protocol::CipherState<C>>
	for NoiseCipherState<C>
	{
	fn from(cs: &noise_protocol::CipherState<C>) -> Self {
	let (key, counter) = cs.clone().extract();
	NoiseCipherState { k: key, n: counter, n_staged: counter }
	}
	}

	/// Returns the public key corresponding to a given `dh_private_key`.
	pub fn get_dh_public_key<D: noise_protocol::DH>(
	dh_private_key: &DhPrivateKey,
	) -> DpeResult<DhPublicKey> {
	DhPublicKey::from_slice(
	D::pubkey(&D::Key::from_slice(dh_private_key.as_slice())).as_slice(),
	)
	}

	/// A trait representing [`NoiseSessionCrypto`] dependencies.
	pub trait NoiseCryptoDeps {
	/// Cipher type
	type Cipher: noise_protocol::Cipher;
	/// DH type
	type DH: noise_protocol::DH;
	/// Hash type
	type Hash: noise_protocol::Hash;
	}

	/// A Noise implementation of the [`SessionCrypto`] trait.
	pub struct NoiseSessionCrypto<D: NoiseCryptoDeps> {
	#[allow(dead_code)]
	phantom: PhantomData<D>,
	}

	impl<D> Clone for NoiseSessionCrypto<D>
	where
	D: NoiseCryptoDeps,
	{
	fn clone(&self) -> Self {
	Self { phantom: Default::default() }
	}
	}

	impl<D> Default for NoiseSessionCrypto<D>
	where
	D: NoiseCryptoDeps,
	{
	fn default() -> Self {
	Self { phantom: Default::default() }
	}
	}

	impl<D> core::fmt::Debug for NoiseSessionCrypto<D>
	where
	D: NoiseCryptoDeps,
	{
	fn fmt(&self, _: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
	Ok(())
	}
	}

	impl<D> core::hash::Hash for NoiseSessionCrypto<D>
	where
	D: NoiseCryptoDeps,
	{
	fn hash<Hr: core::hash::Hasher>(&self, _: &mut Hr) {}
	}

	impl<D> PartialEq for NoiseSessionCrypto<D>
	where
	D: NoiseCryptoDeps,
	{
	fn eq(&self, _: &Self) -> bool {
	true
	}
	}

	impl<D> Eq for NoiseSessionCrypto<D> where D: NoiseCryptoDeps {}

	impl<D> SessionCrypto for NoiseSessionCrypto<D>
	where
	D: NoiseCryptoDeps,
	{
	type SessionCipherState = NoiseCipherState<D::Cipher>;

	/// Implements the responder role of a Noise_NK handshake.
	fn new_session_handshake(
	static_dh_key: &DhPrivateKey,
	initiator_handshake: &HandshakeMessage,
	payload: &HandshakePayload,
	responder_handshake: &mut HandshakeMessage,
	decrypt_cipher_state: &mut NoiseCipherState<D::Cipher>,
	encrypt_cipher_state: &mut NoiseCipherState<D::Cipher>,
	psk_seed: &mut Hash,
	) -> DpeResult<()> {
	#[allow(unused_results)]
	let mut handshake: noise_protocol::HandshakeState<
	D::DH,
	D::Cipher,
	D::Hash,
	> = {
	let mut builder = HandshakeStateBuilder::new();
	builder.set_pattern(noise_protocol::patterns::noise_nk());
	builder.set_is_initiator(false);
	builder.set_prologue(&[]);
	builder.set_s(<D::DH as noise_protocol::DH>::Key::from_slice(
	static_dh_key.as_slice(),
	));
	builder.build_handshake_state()
	};
	handshake.read_message(initiator_handshake.as_slice(), &mut [])?;
	handshake.write_message(
	payload.as_slice(),
	responder_handshake.as_mut_sized(
	handshake.get_next_message_overhead() + payload.len(),
	)?,
	)?;
	assert!(handshake.completed());
	let ciphers = handshake.get_ciphers();
	*decrypt_cipher_state = (&ciphers.0).into();
	*encrypt_cipher_state = (&ciphers.1).into();
	debug!("get_hash");
	*psk_seed = Hash::from_slice(handshake.get_hash())?;
	Ok(())
	}

	/// Implements the responder role of a Noise_NNpsk0 handshake.
	fn derive_session_handshake(
	psk: &Hash,
	initiator_handshake: &HandshakeMessage,
	payload: &HandshakePayload,
	responder_handshake: &mut HandshakeMessage,
	decrypt_cipher_state: &mut NoiseCipherState<D::Cipher>,
	encrypt_cipher_state: &mut NoiseCipherState<D::Cipher>,
	psk_seed: &mut Hash,
	) -> DpeResult<()> {
	#[allow(unused_results)]
	let mut handshake: noise_protocol::HandshakeState<
	D::DH,
	D::Cipher,
	D::Hash,
	> = {
	let mut builder = HandshakeStateBuilder::new();
	builder.set_pattern(noise_protocol::patterns::noise_nn_psk0());
	builder.set_is_initiator(false);
	builder.set_prologue(&[]);
	builder.build_handshake_state()
	};
	handshake
	.push_psk(psk.as_slice().get(..32).ok_or(ErrCode::InternalError)?);
	handshake.read_message(initiator_handshake.as_slice(), &mut [])?;
	handshake.write_message(
	payload.as_slice(),
	responder_handshake.as_mut_sized(
	handshake.get_next_message_overhead() + payload.len(),
	)?,
	)?;
	let ciphers = handshake.get_ciphers();
	*decrypt_cipher_state = (&ciphers.0).into();
	*encrypt_cipher_state = (&ciphers.1).into();
	*psk_seed = Hash::from_slice(handshake.get_hash())?;
	Ok(())
	}

	/// Encrypts a Noise transport message in place.
	fn session_encrypt(
	cipher_state: &mut NoiseCipherState<D::Cipher>,
	in_place_buffer: &mut Message,
	) -> DpeResult<()> {
	let mut cs = noise_protocol::CipherState::<D::Cipher>::new(
	cipher_state.k.as_slice(),
	cipher_state.n,
	);
	let plaintext_len = in_place_buffer.len();
	let _ = cs.encrypt_in_place(
	in_place_buffer.as_mut_sized(
	plaintext_len
	+ <D::Cipher as noise_protocol::Cipher>::tag_len(),
	)?,
	plaintext_len,
	);
	// Encrypting a message is usually not the final step in preparing
	// the message for transport. If a subsequent step fails, it is
	// better for 'n' to remain unchanged so we don't get out of sync.
	(_, cipher_state.n_staged) = cs.extract();
	Ok(())
	}

	/// Decrypts a Noise transport message in place.
	fn session_decrypt(
	cipher_state: &mut NoiseCipherState<D::Cipher>,
	in_place_buffer: &mut Message,
	) -> DpeResult<()> {
	let mut cs = noise_protocol::CipherState::<D::Cipher>::new(
	cipher_state.k.as_slice(),
	cipher_state.n,
	);
	let ciphertext_len = in_place_buffer.len();
	let plaintext_len = match cs
	.decrypt_in_place(in_place_buffer.vec.as_mut(), ciphertext_len)
	{
	Ok(length) => length,
	_ => {
	error!("Session decrypt failed");
	return Err(ErrCode::InvalidCommand);
	}
	};
	in_place_buffer.vec.truncate(plaintext_len);
	(_, cipher_state.n) = cs.extract();
	Ok(())
	}

	/// Derives a responder-side PSK.
	fn derive_psk_from_session(
	psk_seed: &Hash,
	decrypt_cipher_state: &NoiseCipherState<D::Cipher>,
	encrypt_cipher_state: &NoiseCipherState<D::Cipher>,
	) -> DpeResult<Hash> {
	let mut hasher: D::Hash = Default::default();
	hasher.input(psk_seed.as_slice());
	// Use the decrypt state as it was before we decrypted the current
	// command message. This allows clients to compute the PSK using
	// the cipher states as they are before the client sends the
	// command.
	hasher.input(&(decrypt_cipher_state.n() - 1).to_le_bytes());
	hasher.input(&encrypt_cipher_state.n().to_le_bytes());
	Ok((&hasher.result()).into())
	}
	}

	/// A SessionClient implements the initiator side of an encrypted session. A
	/// DPE does not use this itself, it is useful for clients and testing.
	pub struct SessionClient<D>
	where
	D: NoiseCryptoDeps,
	{
	handshake_state:
	Option<noise_protocol::HandshakeState<D::DH, D::Cipher, D::Hash>>,
	/// Cipher state for encrypting messages to a DPE.
	pub encrypt_cipher_state: NoiseCipherState<D::Cipher>,
	/// Cipher state for decrypting messages from a DPE.
	pub decrypt_cipher_state: NoiseCipherState<D::Cipher>,
	/// PSK seed for deriving sessions. See [`derive_psk`].
	///
	/// [`derive_psk`]: #method.derive_psk
	pub psk_seed: Hash,
	}

	impl<D> Clone for SessionClient<D>
	where
	D: NoiseCryptoDeps,
	{
	fn clone(&self) -> Self {
	Self {
	handshake_state: self.handshake_state.clone(),
	encrypt_cipher_state: self.encrypt_cipher_state.clone(),
	decrypt_cipher_state: self.decrypt_cipher_state.clone(),
	psk_seed: self.psk_seed.clone(),
	}
	}
	}

	impl<D> Default for SessionClient<D>
	where
	D: NoiseCryptoDeps,
	{
	fn default() -> Self {
	Self::new()
	}
	}

	impl<D> core::fmt::Debug for SessionClient<D>
	where
	D: NoiseCryptoDeps,
	{
	fn fmt(&self, _: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
	Ok(())
	}
	}

	impl<D> SessionClient<D>
	where
	D: NoiseCryptoDeps,
	{
	/// Creates a new SessionClient instance. Set up by starting and finishing a
	/// handshake.
	pub fn new() -> Self {
	Self {
	handshake_state: Default::default(),
	encrypt_cipher_state: Default::default(),
	decrypt_cipher_state: Default::default(),
	psk_seed: Default::default(),
	}
	}

	/// Starts a handshake using a known `public_key` and returns a message that
	/// works with the DPE OpenSession command.
	pub fn start_handshake_with_known_public_key(
	&mut self,
	public_key: &DhPublicKey,
	) -> DpeResult<HandshakeMessage> {
	#[allow(unused_results)]
	let mut handshake_state = {
	let mut builder = HandshakeStateBuilder::new();
	builder.set_pattern(noise_protocol::patterns::noise_nk());
	builder.set_is_initiator(true);
	builder.set_prologue(&[]);
	builder.set_rs(<D::DH as noise_protocol::DH>::Pubkey::from_slice(
	public_key.as_slice(),
	));
	builder.build_handshake_state()
	};
	let mut message = HandshakeMessage::new();
	handshake_state.write_message(
	&[],
	message
	.as_mut_sized(handshake_state.get_next_message_overhead())?,
	)?;
	self.handshake_state = Some(handshake_state);
	Ok(message)
	}

	/// Starts a handshake using a `psk` and returns a message that works with
	/// the DPE DeriveContext command. Use [`derive_psk`] to obtain this value
	/// from an existing session.
	///
	/// [`derive_psk`]: #method.derive_psk
	pub fn start_handshake_with_psk(
	&mut self,
	psk: &Hash,
	) -> DpeResult<HandshakeMessage> {
	#[allow(unused_results)]
	let mut handshake_state = {
	let mut builder = HandshakeStateBuilder::new();
	builder.set_pattern(noise_protocol::patterns::noise_nn_psk0());
	builder.set_is_initiator(true);
	builder.set_prologue(&[]);
	builder.build_handshake_state()
	};
	handshake_state
	.push_psk(psk.as_slice().get(..32).ok_or(ErrCode::InternalError)?);
	let mut message = HandshakeMessage::new();
	handshake_state.write_message(
	&[],
	message
	.as_mut_sized(handshake_state.get_next_message_overhead())?,
	)?;
	self.handshake_state = Some(handshake_state);
	Ok(message)
	}

	/// Finishes a handshake started using one of the start_handshake_* methods.
	/// On success, returns the handshake payload from the responder and sets up
	/// internal state for subsequent calls to encrypt and decrypt.
	pub fn finish_handshake(
	&mut self,
	responder_handshake: &HandshakeMessage,
	) -> DpeResult<HandshakePayload> {
	match self.handshake_state {
	None => Err(ErrCode::InvalidArgument),
	Some(ref mut handshake) => {
	let mut payload = HandshakePayload::new();
	handshake.read_message(
	responder_handshake.as_slice(),
	payload.as_mut_sized(
	responder_handshake.len()
	- handshake.get_next_message_overhead(),
	)?,
	)?;
	let ciphers = handshake.get_ciphers();
	self.encrypt_cipher_state = (&ciphers.0).into();
	self.decrypt_cipher_state = (&ciphers.1).into();
	self.psk_seed = Hash::from_slice(handshake.get_hash())?;
	Ok(payload)
	}
	}
	}

	/// Derives a PSK from the current session.
	pub fn derive_psk(&self) -> Hash {
	// Note this is from a client perspective so the counters are hashed
	// encrypt first and unmodified from their current state. A DPE will
	// reverse the order and decrement the first counter in order to derive
	// the same value (see derive_psk_from_session).
	let mut hasher: D::Hash = Default::default();
	hasher.input(self.psk_seed.as_slice());
	hasher.input(&self.encrypt_cipher_state.n().to_le_bytes());
	hasher.input(&self.decrypt_cipher_state.n().to_le_bytes());
	(&hasher.result()).into()
	}

	/// Encrypts a message to send to a DPE and commits cipher state changes.
	pub fn encrypt(&mut self, in_place_buffer: &mut Message) -> DpeResult<()> {
	NoiseSessionCrypto::<D>::session_encrypt(
	&mut self.encrypt_cipher_state,
	in_place_buffer,
	)?;
	self.encrypt_cipher_state.commit();
	Ok(())
	}

	/// Decrypts a message from a DPE.
	pub fn decrypt(&mut self, in_place_buffer: &mut Message) -> DpeResult<()> {
	NoiseSessionCrypto::<D>::session_decrypt(
	&mut self.decrypt_cipher_state,
	in_place_buffer,
	)
	}
	}

	#[cfg(test)]
	mod tests {
	use super::*;

	struct DepsForTesting {}
	impl NoiseCryptoDeps for DepsForTesting {
	type Cipher = noise_rust_crypto::Aes256Gcm;
	type DH = noise_rust_crypto::X25519;
	type Hash = noise_rust_crypto::Sha512;
	}

	type SessionCryptoForTesting = NoiseSessionCrypto<DepsForTesting>;

	type SessionClientForTesting = SessionClient<DepsForTesting>;

	type CipherStateForTesting = NoiseCipherState<noise_rust_crypto::Aes256Gcm>;

	#[test]
	fn end_to_end_session() {
	let mut client = SessionClientForTesting::new();
	let dh_key: DhPrivateKey = Default::default();
	let dh_public_key = get_dh_public_key::<
	<DepsForTesting as NoiseCryptoDeps>::DH,
	>(&dh_key)
	.unwrap();
	let handshake1 = client
	.start_handshake_with_known_public_key(&dh_public_key)
	.unwrap();
	let mut dpe_decrypt_cs: CipherStateForTesting = Default::default();
	let mut dpe_encrypt_cs: CipherStateForTesting = Default::default();
	let mut psk_seed = Default::default();
	let mut handshake2 = Default::default();
	let payload = HandshakePayload::from_slice("pay".as_bytes()).unwrap();
	SessionCryptoForTesting::new_session_handshake(
	&dh_key,
	&handshake1,
	&payload,
	&mut handshake2,
	&mut dpe_decrypt_cs,
	&mut dpe_encrypt_cs,
	&mut psk_seed,
	)
	.unwrap();
	assert_eq!(payload, client.finish_handshake(&handshake2).unwrap());

	// Check that the session works.
	let mut buffer = Message::from_slice("message".as_bytes()).unwrap();
	client.encrypt(&mut buffer).unwrap();
	SessionCryptoForTesting::session_decrypt(
	&mut dpe_decrypt_cs,
	&mut buffer,
	)
	.unwrap();
	assert_eq!("message".as_bytes(), buffer.as_slice());
	SessionCryptoForTesting::session_encrypt(
	&mut dpe_encrypt_cs,
	&mut buffer,
	)
	.unwrap();
	dpe_encrypt_cs.commit();
	client.decrypt(&mut buffer).unwrap();
	assert_eq!("message".as_bytes(), buffer.as_slice());

	// Do it again to check session state still works.
	client.encrypt(&mut buffer).unwrap();
	SessionCryptoForTesting::session_decrypt(
	&mut dpe_decrypt_cs,
	&mut buffer,
	)
	.unwrap();
	assert_eq!("message".as_bytes(), buffer.as_slice());
	SessionCryptoForTesting::session_encrypt(
	&mut dpe_encrypt_cs,
	&mut buffer,
	)
	.unwrap();
	dpe_encrypt_cs.commit();
	client.decrypt(&mut buffer).unwrap();
	assert_eq!("message".as_bytes(), buffer.as_slice());
	}

	#[test]
	fn derived_session() {
	// Set up a session from which to derive.
	let mut client = SessionClientForTesting::new();
	let dh_key: DhPrivateKey = Default::default();
	let dh_public_key = get_dh_public_key::<
	<DepsForTesting as NoiseCryptoDeps>::DH,
	>(&dh_key)
	.unwrap();
	let handshake1 = client
	.start_handshake_with_known_public_key(&dh_public_key)
	.unwrap();
	let mut dpe_decrypt_cs = Default::default();
	let mut dpe_encrypt_cs = Default::default();
	let mut psk_seed = Default::default();
	let mut handshake2 = Default::default();
	let payload = HandshakePayload::from_slice("pay".as_bytes()).unwrap();
	SessionCryptoForTesting::new_session_handshake(
	&dh_key,
	&handshake1,
	&payload,
	&mut handshake2,
	&mut dpe_decrypt_cs,
	&mut dpe_encrypt_cs,
	&mut psk_seed,
	)
	.unwrap();
	assert_eq!(payload, client.finish_handshake(&handshake2).unwrap());

	// Derive a second session.
	let mut client2 = SessionClientForTesting::new();
	let client_psk = client.derive_psk();
	// Simulate the session state after command decryption on the DPE side
	// as expected by the DPE PSK logic.
	let mut buffer = Message::from_slice("message".as_bytes()).unwrap();
	client.encrypt(&mut buffer).unwrap();
	SessionCryptoForTesting::session_decrypt(
	&mut dpe_decrypt_cs,
	&mut buffer,
	)
	.unwrap();
	let dpe_psk = SessionCryptoForTesting::derive_psk_from_session(
	&psk_seed,
	&dpe_decrypt_cs,
	&dpe_encrypt_cs,
	)
	.unwrap();
	let handshake1 = client2.start_handshake_with_psk(&client_psk).unwrap();
	let mut dpe_decrypt_cs2 = Default::default();
	let mut dpe_encrypt_cs2 = Default::default();
	let mut psk_seed2 = Default::default();
	SessionCryptoForTesting::derive_session_handshake(
	&dpe_psk,
	&handshake1,
	&payload,
	&mut handshake2,
	&mut dpe_decrypt_cs2,
	&mut dpe_encrypt_cs2,
	&mut psk_seed2,
	)
	.unwrap();
	assert_eq!(payload, client2.finish_handshake(&handshake2).unwrap());

	// Check that the second session works.
	let mut buffer = Message::from_slice("message".as_bytes()).unwrap();
	client2.encrypt(&mut buffer).unwrap();
	SessionCryptoForTesting::session_decrypt(
	&mut dpe_decrypt_cs2,
	&mut buffer,
	)
	.unwrap();
	assert_eq!("message".as_bytes(), buffer.as_slice());
	SessionCryptoForTesting::session_encrypt(
	&mut dpe_encrypt_cs2,
	&mut buffer,
	)
	.unwrap();
	dpe_encrypt_cs2.commit();
	client2.decrypt(&mut buffer).unwrap();
	assert_eq!("message".as_bytes(), buffer.as_slice());

	// Check that the first session also still works.
	let mut buffer = Message::from_slice("message".as_bytes()).unwrap();
	client.encrypt(&mut buffer).unwrap();
	SessionCryptoForTesting::session_decrypt(
	&mut dpe_decrypt_cs,
	&mut buffer,
	)
	.unwrap();
	assert_eq!("message".as_bytes(), buffer.as_slice());
	SessionCryptoForTesting::session_encrypt(
	&mut dpe_encrypt_cs,
	&mut buffer,
	)
	.unwrap();
	dpe_encrypt_cs.commit();
	client.decrypt(&mut buffer).unwrap();
	assert_eq!("message".as_bytes(), buffer.as_slice());
	}
	}