Skip to content

Commit

Permalink
Replace concatenated bytes in the Noise implementation with a struct
Browse files Browse the repository at this point in the history 
Change-Id: I4a3fb151cf3405f614000fd0112ca16320515bf0
  • Loading branch information
@k-naliuka
k-naliuka committed Sep 16, 2024
1 parent ed30dc3 commit b137c25
Show file tree
Hide file tree
Showing 4 changed files with 112 additions and 122 deletions.
5 changes: 2 additions & 3 deletions oak_crypto/src/identity_key.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ use crate::noise_handshake::{p256_scalar_mult, P256Scalar};

pub trait IdentityKeyHandle {
fn get_public_key(&self) -> anyhow::Result<Vec<u8>>;
fn derive_dh_secret(&self, peer_public_key: Vec<u8>) -> anyhow::Result<Vec<u8>>;
fn derive_dh_secret(&self, peer_public_key: &[u8]) -> anyhow::Result<Vec<u8>>;
}

// Implementation of the identity keyset where the private key is encapsulated
Expand All @@ -45,11 +45,10 @@ impl IdentityKeyHandle for IdentityKey {
Ok(self.private_key.compute_public_key().to_vec())
}

fn derive_dh_secret(&self, peer_public_key: Vec<u8>) -> anyhow::Result<Vec<u8>> {
fn derive_dh_secret(&self, peer_public_key: &[u8]) -> anyhow::Result<Vec<u8>> {
p256_scalar_mult(
&self.private_key,
peer_public_key
.as_slice()
.try_into()
.map_err(|error| anyhow!("invalid peer public key: {:?}", error))?,
)
Expand Down
29 changes: 16 additions & 13 deletions oak_crypto/src/noise_handshake/client.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ pub use crate::noise_handshake::crypto_wrapper::{
use crate::noise_handshake::{
error::Error,
noise::{HandshakeType, Noise},
Crypter, IdentityKeyHandle,
Crypter, IdentityKeyHandle, NoiseMessage,
};

pub struct HandshakeInitiator {
Expand Down Expand Up @@ -62,7 +62,7 @@ impl HandshakeInitiator {
}
}

pub fn build_initial_message(&mut self) -> Result<Vec<u8>, Error> {
pub fn build_initial_message(&mut self) -> Result<NoiseMessage, Error> {
self.noise.mix_hash(&[0; 1]);
if let Some(peer_identity_pub_key) = self.peer_identity_pub_key {
self.noise.mix_hash_point(peer_identity_pub_key.as_slice());
Expand Down Expand Up @@ -90,34 +90,37 @@ impl HandshakeInitiator {
if let Some(self_priv_key) = self.self_identity_priv_key.as_ref() {
if let Some(peer_identity_pub_key) = self.peer_identity_pub_key {
let se_ecdh_bytes = self_priv_key
.derive_dh_secret(peer_identity_pub_key.to_vec())
.derive_dh_secret(peer_identity_pub_key.as_slice())
.map_err(|_| Error::InvalidHandshake)?;
self.noise.mix_key(&se_ecdh_bytes);
} else {
return Err(Error::MissingPeerPublicKey);
}
}
let ciphertext = self.noise.encrypt_and_hash(&[]);
Ok([ephemeral_pub_key_bytes, &ciphertext].concat())
Ok(NoiseMessage { ciphertext, ephemeral_public_key: ephemeral_pub_key.to_vec() })
}

pub fn process_response(
&mut self,
handshake_response: &[u8],
handshake_response: &NoiseMessage,
) -> Result<([u8; SHA256_OUTPUT_LEN], Crypter), Error> {
let peer_public_key_bytes = &handshake_response[..P256_X962_LEN];
let ciphertext = &handshake_response[P256_X962_LEN..];

let ee_ecdh_bytes = p256_scalar_mult(
&self.ephemeral_priv_key,
peer_public_key_bytes.try_into().map_err(|_| Error::DecryptFailed)?,
handshake_response
.ephemeral_public_key
.as_slice()
.try_into()
.map_err(|_| Error::DecryptFailed)?,
)
.map_err(|_| Error::InvalidHandshake)?;
self.noise.mix_hash(peer_public_key_bytes);
self.noise.mix_key(peer_public_key_bytes);
self.noise.mix_hash(&handshake_response.ephemeral_public_key);
self.noise.mix_key(&handshake_response.ephemeral_public_key);
self.noise.mix_key(&ee_ecdh_bytes);
let plaintext =
self.noise.decrypt_and_hash(ciphertext).map_err(|_| Error::DecryptFailed)?;
let plaintext = self
.noise
.decrypt_and_hash(&handshake_response.ciphertext)
.map_err(|_| Error::DecryptFailed)?;
assert_eq!(plaintext.len(), 0);
let (write_key, read_key) = self.noise.traffic_keys();
Ok((self.noise.handshake_hash(), Crypter::new(&read_key, &write_key)))
Expand Down
135 changes: 71 additions & 64 deletions oak_crypto/src/noise_handshake/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,11 +24,10 @@ mod error;
mod noise;
#[cfg(test)]
mod tests;

use alloc::vec::Vec;

use anyhow::anyhow;
use oak_proto_rust::oak::crypto::v1::SessionKeys;
use oak_proto_rust::oak::{crypto::v1::SessionKeys, session::v1::NoiseHandshakeMessage};

use crate::noise_handshake::{
error::Error,
Expand Down Expand Up @@ -62,6 +61,20 @@ impl Nonce {
}
}

pub struct NoiseMessage {
pub ephemeral_public_key: Vec<u8>,
pub ciphertext: Vec<u8>,
}

impl From<&NoiseHandshakeMessage> for NoiseMessage {
fn from(value: &NoiseHandshakeMessage) -> Self {
NoiseMessage {
ephemeral_public_key: value.ephemeral_public_key.clone(),
ciphertext: value.ciphertext.clone(),
}
}
}

pub struct Crypter {
read_key: [u8; SYMMETRIC_KEY_LEN],
write_key: [u8; SYMMETRIC_KEY_LEN],
Expand Down Expand Up @@ -165,73 +178,37 @@ impl TryFrom<SessionKeys> for Crypter {
pub struct Response {
pub crypter: Crypter,
pub handshake_hash: [u8; SHA256_OUTPUT_LEN],
pub response: Vec<u8>,
pub response: NoiseMessage,
}

pub fn respond_nk(identity_key: &dyn IdentityKeyHandle, in_data: &[u8]) -> Result<Response, Error> {
if in_data.len() < P256_X962_LEN {
return Err(Error::InvalidHandshake);
}

pub fn respond_nk(
identity_key: &dyn IdentityKeyHandle,
in_message: &NoiseMessage,
) -> Result<Response, Error> {
let mut noise = Noise::new(HandshakeType::Nk);
noise.mix_hash(&[0; 1]); // Prologue
noise.mix_hash_point(
identity_key.get_public_key().map_err(|_| Error::InvalidPrivateKey)?.as_slice(),
);

// unwrap: we know that `in_data` is `P256_X962_LENGTH` bytes long.
let peer_pub: [u8; P256_X962_LEN] = (&in_data[..P256_X962_LEN]).try_into().unwrap();
noise.mix_hash(peer_pub.as_slice());
noise.mix_key(peer_pub.as_slice());
noise.mix_hash(in_message.ephemeral_public_key.as_slice());
noise.mix_key(in_message.ephemeral_public_key.as_slice());

let es_ecdh_bytes =
identity_key.derive_dh_secret(peer_pub.to_vec()).map_err(|_| Error::InvalidHandshake)?;
let es_ecdh_bytes = identity_key
.derive_dh_secret(in_message.ephemeral_public_key.as_slice())
.map_err(|_| Error::InvalidHandshake)?;
noise.mix_key(es_ecdh_bytes.as_slice());
finish_response(&mut noise, in_data, &peer_pub)
finish_response(&mut noise, in_message)
}

pub fn respond_nn(in_data: &[u8]) -> Result<Response, Error> {
if in_data.len() < P256_X962_LEN {
return Err(Error::InvalidHandshake);
}

pub fn respond_nn(in_message: &NoiseMessage) -> Result<Response, Error> {
let mut noise = Noise::new(HandshakeType::Nn);
noise.mix_hash(&[0; 1]); // Prologue

// unwrap: we know that `in_data` is `P256_X962_LENGTH` bytes long.
let peer_pub: [u8; P256_X962_LEN] = (&in_data[..P256_X962_LEN]).try_into().unwrap();
noise.mix_hash(peer_pub.as_slice());
noise.mix_key(peer_pub.as_slice());
finish_response(&mut noise, in_data, &peer_pub)
}

fn finish_response(
noise: &mut Noise,
in_data: &[u8],
peer_pub: &[u8; P256_X962_LEN],
) -> Result<Response, Error> {
let plaintext = noise.decrypt_and_hash(&in_data[P256_X962_LEN..])?;
if !plaintext.is_empty() {
return Err(Error::InvalidHandshake);
}

// Generate ephemeral key pair
let ephemeral_priv = P256Scalar::generate();
let ephemeral_pub_key_bytes = ephemeral_priv.compute_public_key();
noise.mix_hash(ephemeral_pub_key_bytes.as_slice());
noise.mix_key(ephemeral_pub_key_bytes.as_slice());
let ee_ecdh_bytes = crypto_wrapper::p256_scalar_mult(&ephemeral_priv, peer_pub)
.map_err(|_| Error::InvalidHandshake)?;
noise.mix_key(ee_ecdh_bytes.as_slice());

let response_ciphertext = noise.encrypt_and_hash(&[]);

let keys = noise.traffic_keys();
Ok(Response {
crypter: Crypter::new(&keys.0, &keys.1),
handshake_hash: noise.handshake_hash(),
response: [ephemeral_pub_key_bytes.as_slice(), &response_ciphertext].concat(),
})
noise.mix_hash(in_message.ephemeral_public_key.as_slice());
noise.mix_key(&in_message.ephemeral_public_key.as_slice());
finish_response(&mut noise, in_message)
}

pub fn respond_kk(
Expand All @@ -240,24 +217,19 @@ pub fn respond_kk(
// initiator s [not used for Nk]
initiator_static_pub: &[u8],
// e, es, (ss for Kk only)
in_data: &[u8],
in_message: &NoiseMessage,
) -> Result<Response, Error> {
if in_data.len() < P256_X962_LEN {
return Err(Error::InvalidHandshake);
}

let mut noise = Noise::new(HandshakeType::Kk);
noise.mix_hash(&[0; 1]); // Prologue
noise.mix_hash_point(
identity_priv.get_public_key().map_err(|_| Error::InvalidPrivateKey)?.as_slice(),
);

let initiator_pub: [u8; P256_X962_LEN] = (&in_data[..P256_X962_LEN]).try_into().unwrap();
noise.mix_hash(initiator_pub.as_slice());
noise.mix_key(initiator_pub.as_slice());
noise.mix_hash(in_message.ephemeral_public_key.as_slice());
noise.mix_key(in_message.ephemeral_public_key.as_slice());

let es_ecdh_bytes = identity_priv
.derive_dh_secret(initiator_pub.to_vec())
.derive_dh_secret(in_message.ephemeral_public_key.as_slice())
.map_err(|_| Error::InvalidHandshake)?;
noise.mix_key(es_ecdh_bytes.as_slice());

Expand All @@ -269,8 +241,43 @@ pub fn respond_kk(
noise.mix_key(&ss_ecdh_bytes);

let se_ecdh_bytes = identity_priv
.derive_dh_secret(initiator_static_pub_bytes.to_vec())
.derive_dh_secret(initiator_static_pub_bytes.as_slice())
.map_err(|_| Error::InvalidHandshake)?;
noise.mix_key(&se_ecdh_bytes);
finish_response(&mut noise, in_data, &initiator_pub)
finish_response(&mut noise, in_message)
}

fn finish_response(noise: &mut Noise, in_message: &NoiseMessage) -> Result<Response, Error> {
let plaintext = noise.decrypt_and_hash(&in_message.ciphertext)?;
if !plaintext.is_empty() {
return Err(Error::InvalidHandshake);
}

// Generate ephemeral key pair
let ephemeral_priv = P256Scalar::generate();
let ephemeral_pub_key_bytes = ephemeral_priv.compute_public_key();
noise.mix_hash(ephemeral_pub_key_bytes.as_slice());
noise.mix_key(ephemeral_pub_key_bytes.as_slice());
let ee_ecdh_bytes = crypto_wrapper::p256_scalar_mult(
&ephemeral_priv,
in_message
.ephemeral_public_key
.as_slice()
.try_into()
.map_err(|_| Error::InvalidPublicKey)?,
)
.map_err(|_| Error::InvalidHandshake)?;
noise.mix_key(ee_ecdh_bytes.as_slice());

let response_ciphertext = noise.encrypt_and_hash(&[]);

let keys = noise.traffic_keys();
Ok(Response {
crypter: Crypter::new(&keys.0, &keys.1),
handshake_hash: noise.handshake_hash(),
response: NoiseMessage {
ciphertext: response_ciphertext,
ephemeral_public_key: ephemeral_pub_key_bytes.to_vec(),
},
})
}
Loading

0 comments on commit b137c25

Please sign in to comment.