void/ZeroTierOne
2
0
Fork 0
mirror of https://github.com/zerotier/ZeroTierOne.git synced 2025-06-06 20:43:44 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

A bunch of ZSSP cleanup and optimization. Runs a bit faster now.

Browse source
This commit is contained in:
Adam Ierymenko 2023-03-10 16:58:38 -05:00
parent 7072338037
commit f83bf41427
2 changed files with 292 additions and 282 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
zssp/src/sessionid.rs
View file

@ -10,7 +10,6 @@ use std::fmt::Display;
use std::num::NonZeroU64; use std::num::NonZeroU64;
use zerotier_crypto::random; use zerotier_crypto::random;
use zerotier_utils::memory::{array_range, as_byte_array};
/// 48-bit session ID (most significant 16 bits of u64 are unused) /// 48-bit session ID (most significant 16 bits of u64 are unused)
#[derive(Copy, Clone, PartialEq, Eq, Hash)] #[derive(Copy, Clone, PartialEq, Eq, Hash)]
@ -25,6 +24,7 @@ impl SessionId {
pub const MAX: u64 = 0xffffffffffff; pub const MAX: u64 = 0xffffffffffff;
/// Create a new session ID, panicing if 'i' is zero or exceeds MAX. /// Create a new session ID, panicing if 'i' is zero or exceeds MAX.
#[inline(always)]
pub fn new(i: u64) -> SessionId { pub fn new(i: u64) -> SessionId {
assert!(i <= Self::MAX); assert!(i <= Self::MAX);
Self(NonZeroU64::new(i.to_le()).unwrap()) Self(NonZeroU64::new(i.to_le()).unwrap())
@ -35,22 +35,23 @@ impl SessionId {
Self(NonZeroU64::new(((random::xorshift64_random() % (Self::MAX - 1)) + 1).to_le()).unwrap()) Self(NonZeroU64::new(((random::xorshift64_random() % (Self::MAX - 1)) + 1).to_le()).unwrap())
} }
pub(crate) fn new_from_bytes(b: &[u8; Self::SIZE]) -> Option<SessionId> { #[inline(always)]
let mut tmp = [0u8; 8]; pub fn to_bytes(&self) -> [u8; Self::SIZE] {
self.0.get().to_ne_bytes()[..Self::SIZE].try_into().unwrap()
}
#[inline(always)]
pub fn new_from_bytes(b: &[u8]) -> Option<SessionId> {
let mut tmp = 0u64.to_ne_bytes();
tmp[..SESSION_ID_SIZE_BYTES].copy_from_slice(b); tmp[..SESSION_ID_SIZE_BYTES].copy_from_slice(b);
Self::new_from_u64_le(u64::from_ne_bytes(tmp)) NonZeroU64::new(u64::from_ne_bytes(tmp)).map(|i| Self(i))
} }
/// Create from a u64 that is already in little-endian byte order.
#[inline(always)] #[inline(always)]
pub(crate) fn new_from_u64_le(i: u64) -> Option<SessionId> { pub fn new_from_array(b: &[u8; Self::SIZE]) -> Option<SessionId> {
NonZeroU64::new(i & Self::MAX.to_le()).map(|i| Self(i)) let mut tmp = 0u64.to_ne_bytes();
} tmp[..SESSION_ID_SIZE_BYTES].copy_from_slice(b);
NonZeroU64::new(u64::from_ne_bytes(tmp)).map(|i| Self(i))
/// Get this session ID as a little-endian byte array.
#[inline(always)]
pub(crate) fn as_bytes(&self) -> &[u8; Self::SIZE] {
array_range::<u8, 8, 0, SESSION_ID_SIZE_BYTES>(as_byte_array(&self.0))
} }
} }

233
zssp/src/zssp.rs
View file

@ -56,7 +56,7 @@ struct SessionsById<Application: ApplicationLayer> {
active: HashMap<SessionId, Weak<Session<Application>>>, active: HashMap<SessionId, Weak<Session<Application>>>,
// Incomplete sessions in the middle of three-phase Noise_XK negotiation, expired after timeout. // Incomplete sessions in the middle of three-phase Noise_XK negotiation, expired after timeout.
incoming: HashMap<SessionId, Arc<IncomingIncompleteSession>>, incoming: HashMap<SessionId, Arc<IncomingIncompleteSession<Application>>>,
} }
/// Result generated by the context packet receive function, with possible payloads. /// Result generated by the context packet receive function, with possible payloads.
@ -97,10 +97,10 @@ struct State {
remote_session_id: Option<SessionId>, remote_session_id: Option<SessionId>,
keys: [Option<SessionKey>; 2], keys: [Option<SessionKey>; 2],
current_key: usize, current_key: usize,
current_offer: Offer, outgoing_offer: Offer,
} }
struct IncomingIncompleteSession { struct IncomingIncompleteSession<Application: ApplicationLayer> {
timestamp: i64, timestamp: i64,
alice_session_id: SessionId, alice_session_id: SessionId,
bob_session_id: SessionId, bob_session_id: SessionId,
@ -109,6 +109,7 @@ struct IncomingIncompleteSession {
hk: Secret<KYBER_SSBYTES>, hk: Secret<KYBER_SSBYTES>,
header_protection_key: Secret<AES_HEADER_PROTECTION_KEY_SIZE>, header_protection_key: Secret<AES_HEADER_PROTECTION_KEY_SIZE>,
bob_noise_e_secret: P384KeyPair, bob_noise_e_secret: P384KeyPair,
defrag: [Mutex<Fragged<Application::IncomingPacketBuffer, MAX_FRAGMENTS>>; MAX_NOISE_HANDSHAKE_FRAGMENTS],
} }
struct OutgoingSessionOffer { struct OutgoingSessionOffer {
@ -184,7 +185,7 @@ impl<Application: ApplicationLayer> Context<Application> {
for (id, s) in sessions.active.iter() { for (id, s) in sessions.active.iter() {
if let Some(session) = s.upgrade() { if let Some(session) = s.upgrade() {
let state = session.state.read().unwrap(); let state = session.state.read().unwrap();
if match &state.current_offer { if match &state.outgoing_offer {
Offer::None => true, Offer::None => true,
Offer::NoiseXKInit(offer) => { Offer::NoiseXKInit(offer) => {
// If there's an outstanding attempt to open a session, retransmit this periodically // If there's an outstanding attempt to open a session, retransmit this periodically
@ -324,7 +325,7 @@ impl<Application: ApplicationLayer> Context<Application> {
remote_session_id: None, remote_session_id: None,
keys: [None, None], keys: [None, None],
current_key: 0, current_key: 0,
current_offer: Offer::NoiseXKInit(Box::new(OutgoingSessionOffer { outgoing_offer: Offer::NoiseXKInit(Box::new(OutgoingSessionOffer {
last_retry_time: AtomicI64::new(current_time), last_retry_time: AtomicI64::new(current_time),
psk, psk,
noise_h: mix_hash(&mix_hash(&INITIAL_H, remote_s_public_blob), &alice_noise_e), noise_h: mix_hash(&mix_hash(&INITIAL_H, remote_s_public_blob), &alice_noise_e),
@ -345,7 +346,7 @@ impl<Application: ApplicationLayer> Context<Application> {
{ {
let mut state = session.state.write().unwrap(); let mut state = session.state.write().unwrap();
let offer = if let Offer::NoiseXKInit(offer) = &mut state.current_offer { let offer = if let Offer::NoiseXKInit(offer) = &mut state.outgoing_offer {
offer offer
} else { } else {
panic!(); // should be impossible as this is what we initialized with panic!(); // should be impossible as this is what we initialized with
@ -357,7 +358,7 @@ impl<Application: ApplicationLayer> Context<Application> {
let init: &mut AliceNoiseXKInit = byte_array_as_proto_buffer_mut(init_packet).unwrap(); let init: &mut AliceNoiseXKInit = byte_array_as_proto_buffer_mut(init_packet).unwrap();
init.session_protocol_version = SESSION_PROTOCOL_VERSION; init.session_protocol_version = SESSION_PROTOCOL_VERSION;
init.alice_noise_e = alice_noise_e; init.alice_noise_e = alice_noise_e;
init.alice_session_id = *local_session_id.as_bytes(); init.alice_session_id = local_session_id.to_bytes();
init.alice_hk_public = alice_hk_secret.public; init.alice_hk_public = alice_hk_secret.public;
init.header_protection_key = header_protection_key.0; init.header_protection_key = header_protection_key.0;
} }
@ -417,6 +418,7 @@ impl<Application: ApplicationLayer> Context<Application> {
/// * `data_buf` - Buffer to receive decrypted and authenticated object data (an error is returned if too small) /// * `data_buf` - Buffer to receive decrypted and authenticated object data (an error is returned if too small)
/// * `incoming_packet_buf` - Buffer containing incoming wire packet (receive() takes ownership) /// * `incoming_packet_buf` - Buffer containing incoming wire packet (receive() takes ownership)
/// * `current_time` - Current monotonic time in milliseconds /// * `current_time` - Current monotonic time in milliseconds
#[inline]
pub fn receive< pub fn receive<
'b, 'b,
SendFunction: FnMut(Option<&Arc<Session<Application>>>, &mut [u8]), SendFunction: FnMut(Option<&Arc<Session<Application>>>, &mut [u8]),
@ -430,48 +432,42 @@ impl<Application: ApplicationLayer> Context<Application> {
mut send: SendFunction, mut send: SendFunction,
source: &Application::PhysicalPath, source: &Application::PhysicalPath,
data_buf: &'b mut [u8], data_buf: &'b mut [u8],
mut incoming_packet_buf: Application::IncomingPacketBuffer, mut incoming_physical_packet_buf: Application::IncomingPacketBuffer,
current_time: i64, current_time: i64,
) -> Result<ReceiveResult<'b, Application>, Error> { ) -> Result<ReceiveResult<'b, Application>, Error> {
let incoming_packet: &mut [u8] = incoming_packet_buf.as_mut(); let incoming_physical_packet: &mut [u8] = incoming_physical_packet_buf.as_mut();
if incoming_packet.len() < MIN_PACKET_SIZE { if incoming_physical_packet.len() < MIN_PACKET_SIZE {
return Err(Error::InvalidPacket); return Err(Error::InvalidPacket);
} }
let mut incoming = None; if let Some(local_session_id) = SessionId::new_from_bytes(&incoming_physical_packet[0..SessionId::SIZE]) {
if let Some(local_session_id) = SessionId::new_from_u64_le(u64::from_le_bytes(incoming_packet[0..8].try_into().unwrap())) { let sessions = self.sessions.read().unwrap();
if let Some(session) = self.sessions.read().unwrap().active.get(&local_session_id).and_then(|s| s.upgrade()) { if let Some(session) = sessions.active.get(&local_session_id).and_then(|s| s.upgrade()) {
drop(sessions);
debug_assert!(!self.sessions.read().unwrap().incoming.contains_key(&local_session_id)); debug_assert!(!self.sessions.read().unwrap().incoming.contains_key(&local_session_id));
session session
.header_protection_cipher .header_protection_cipher
.decrypt_block_in_place(&mut incoming_packet[HEADER_PROTECT_ENCRYPT_START..HEADER_PROTECT_ENCRYPT_END]); .decrypt_block_in_place(&mut incoming_physical_packet[HEADER_PROTECT_ENCRYPT_START..HEADER_PROTECT_ENCRYPT_END]);
let (key_index, packet_type, fragment_count, fragment_no, incoming_counter) = parse_packet_header(&incoming_packet); let (key_index, packet_type, fragment_count, fragment_no, incoming_counter) = parse_packet_header(&incoming_physical_packet);
if session.check_receive_window(incoming_counter) { if session.check_receive_window(incoming_counter) {
if fragment_count > 1 { let (assembled_packet, incoming_packet_buf_arr);
let mut fragged = session.defrag[(incoming_counter as usize) % COUNTER_WINDOW_MAX_OOO].lock().unwrap(); let incoming_packet = if fragment_count > 1 {
if let Some(assembled_packet) = fragged.assemble(incoming_counter, incoming_packet_buf, fragment_no, fragment_count) { assembled_packet = session.defrag[(incoming_counter as usize) % COUNTER_WINDOW_MAX_OOO]
drop(fragged); .lock()
return self.process_complete_incoming_packet( .unwrap()
app, .assemble(incoming_counter, incoming_physical_packet_buf, fragment_no, fragment_count);
&mut send, if let Some(assembled_packet) = assembled_packet.as_ref() {
&mut check_allow_incoming_session, assembled_packet.as_ref()
&mut check_accept_session,
data_buf,
incoming_counter,
assembled_packet.as_ref(),
packet_type,
Some(session),
None,
key_index,
current_time,
);
} else { } else {
drop(fragged);
return Ok(ReceiveResult::Ok(Some(session))); return Ok(ReceiveResult::Ok(Some(session)));
} }
} else { } else {
incoming_packet_buf_arr = [incoming_physical_packet_buf];
&incoming_packet_buf_arr
};
return self.process_complete_incoming_packet( return self.process_complete_incoming_packet(
app, app,
&mut send, &mut send,
@ -479,34 +475,63 @@ impl<Application: ApplicationLayer> Context<Application> {
&mut check_accept_session, &mut check_accept_session,
data_buf, data_buf,
incoming_counter, incoming_counter,
&[incoming_packet_buf], incoming_packet,
packet_type, packet_type,
Some(session), Some(session),
None, None,
key_index, key_index,
current_time, current_time,
); );
}
} else { } else {
return Err(Error::OutOfSequence); return Err(Error::OutOfSequence);
} }
} else if let Some(incoming) = sessions.incoming.get(&local_session_id).cloned() {
drop(sessions);
debug_assert!(!self.sessions.read().unwrap().active.contains_key(&local_session_id));
Aes::new(&incoming.header_protection_key)
.decrypt_block_in_place(&mut incoming_physical_packet[HEADER_PROTECT_ENCRYPT_START..HEADER_PROTECT_ENCRYPT_END]);
let (key_index, packet_type, fragment_count, fragment_no, incoming_counter) = parse_packet_header(&incoming_physical_packet);
let (assembled_packet, incoming_packet_buf_arr);
let incoming_packet = if fragment_count > 1 {
assembled_packet = incoming.defrag[(incoming_counter as usize) % COUNTER_WINDOW_MAX_OOO]
.lock()
.unwrap()
.assemble(incoming_counter, incoming_physical_packet_buf, fragment_no, fragment_count);
if let Some(assembled_packet) = assembled_packet.as_ref() {
assembled_packet.as_ref()
} else { } else {
if let Some(i) = self.sessions.read().unwrap().incoming.get(&local_session_id).cloned() { return Ok(ReceiveResult::Ok(None));
Aes::new(&i.header_protection_key) }
.decrypt_block_in_place(&mut incoming_packet[HEADER_PROTECT_ENCRYPT_START..HEADER_PROTECT_ENCRYPT_END]); } else {
incoming = Some(i); incoming_packet_buf_arr = [incoming_physical_packet_buf];
&incoming_packet_buf_arr
};
return self.process_complete_incoming_packet(
app,
&mut send,
&mut check_allow_incoming_session,
&mut check_accept_session,
data_buf,
incoming_counter,
incoming_packet,
packet_type,
None,
Some(incoming),
key_index,
current_time,
);
} else { } else {
return Err(Error::UnknownLocalSessionId); return Err(Error::UnknownLocalSessionId);
} }
} } else {
} let (key_index, packet_type, fragment_count, fragment_no, incoming_counter) = parse_packet_header(&incoming_physical_packet);
// If we make it here the packet is not associated with a session or is associated with an let (assembled_packet, incoming_packet_buf_arr);
// incoming session (Noise_XK mid-negotiation). let incoming_packet = if fragment_count > 1 {
assembled_packet = {
let (key_index, packet_type, fragment_count, fragment_no, incoming_counter) = parse_packet_header(&incoming_packet);
if fragment_count > 1 {
let f = {
let mut defrag = self.defrag.lock().unwrap(); let mut defrag = self.defrag.lock().unwrap();
let f = defrag let f = defrag
.entry((source.clone(), incoming_counter)) .entry((source.clone(), incoming_counter))
@ -518,7 +543,8 @@ impl<Application: ApplicationLayer> Context<Application> {
// First, drop all entries that are timed out or whose physical source duplicates another entry. // First, drop all entries that are timed out or whose physical source duplicates another entry.
let mut sources = HashSet::with_capacity(defrag.len()); let mut sources = HashSet::with_capacity(defrag.len());
let negotiation_timeout_cutoff = current_time - Application::INCOMING_SESSION_NEGOTIATION_TIMEOUT_MS; let negotiation_timeout_cutoff = current_time - Application::INCOMING_SESSION_NEGOTIATION_TIMEOUT_MS;
defrag.retain(|k, fragged| (fragged.1 > negotiation_timeout_cutoff && sources.insert(k.0.clone())) || Arc::ptr_eq(fragged, &f)); defrag
.retain(|k, fragged| (fragged.1 > negotiation_timeout_cutoff && sources.insert(k.0.clone())) || Arc::ptr_eq(fragged, &f));
// Then, if we are still at or over the limit, drop 10% of remaining entries at random. // Then, if we are still at or over the limit, drop 10% of remaining entries at random.
if defrag.len() >= self.max_incomplete_session_queue_size { if defrag.len() >= self.max_incomplete_session_queue_size {
@ -531,27 +557,22 @@ impl<Application: ApplicationLayer> Context<Application> {
} }
f f
}; }
let mut fragged = f.0.lock().unwrap(); .0
.lock()
if let Some(assembled_packet) = fragged.assemble(incoming_counter, incoming_packet_buf, fragment_no, fragment_count) { .unwrap()
.assemble(incoming_counter, incoming_physical_packet_buf, fragment_no, fragment_count);
if let Some(assembled_packet) = assembled_packet.as_ref() {
self.defrag.lock().unwrap().remove(&(source.clone(), incoming_counter)); self.defrag.lock().unwrap().remove(&(source.clone(), incoming_counter));
return self.process_complete_incoming_packet( assembled_packet.as_ref()
app, } else {
&mut send, return Ok(ReceiveResult::Ok(None));
&mut check_allow_incoming_session,
&mut check_accept_session,
data_buf,
incoming_counter,
assembled_packet.as_ref(),
packet_type,
None,
incoming,
key_index,
current_time,
);
} }
} else { } else {
incoming_packet_buf_arr = [incoming_physical_packet_buf];
&incoming_packet_buf_arr
};
return self.process_complete_incoming_packet( return self.process_complete_incoming_packet(
app, app,
&mut send, &mut send,
@ -559,16 +580,14 @@ impl<Application: ApplicationLayer> Context<Application> {
&mut check_accept_session, &mut check_accept_session,
data_buf, data_buf,
incoming_counter, incoming_counter,
&[incoming_packet_buf], incoming_packet,
packet_type, packet_type,
None, None,
incoming, None,
key_index, key_index,
current_time, current_time,
); );
} }
return Ok(ReceiveResult::Ok(None));
} }
fn process_complete_incoming_packet< fn process_complete_incoming_packet<
@ -587,7 +606,7 @@ impl<Application: ApplicationLayer> Context<Application> {
fragments: &[Application::IncomingPacketBuffer], fragments: &[Application::IncomingPacketBuffer],
packet_type: u8, packet_type: u8,
session: Option<Arc<Session<Application>>>, session: Option<Arc<Session<Application>>>,
incoming: Option<Arc<IncomingIncompleteSession>>, incoming: Option<Arc<IncomingIncompleteSession<Application>>>,
key_index: usize, key_index: usize,
current_time: i64, current_time: i64,
) -> Result<ReceiveResult<'b, Application>, Error> { ) -> Result<ReceiveResult<'b, Application>, Error> {
@ -651,9 +670,9 @@ impl<Application: ApplicationLayer> Context<Application> {
// If we got a valid data packet from Bob, this means we can cancel any offers // If we got a valid data packet from Bob, this means we can cancel any offers
// that are still oustanding for initialization. // that are still oustanding for initialization.
match &state.current_offer { match &state.outgoing_offer {
Offer::NoiseXKInit(_) | Offer::NoiseXKAck(_) => { Offer::NoiseXKInit(_) | Offer::NoiseXKAck(_) => {
state.current_offer = Offer::None; state.outgoing_offer = Offer::None;
} }
_ => {} _ => {}
} }
@ -730,7 +749,7 @@ impl<Application: ApplicationLayer> Context<Application> {
} }
let pkt: &AliceNoiseXKInit = byte_array_as_proto_buffer(pkt_assembled)?; let pkt: &AliceNoiseXKInit = byte_array_as_proto_buffer(pkt_assembled)?;
let alice_session_id = SessionId::new_from_bytes(&pkt.alice_session_id).ok_or(Error::InvalidPacket)?; let alice_session_id = SessionId::new_from_array(&pkt.alice_session_id).ok_or(Error::InvalidPacket)?;
let header_protection_key = Secret(pkt.header_protection_key); let header_protection_key = Secret(pkt.header_protection_key);
// Create Bob's ephemeral keys and derive noise_es_ee by agreeing with Alice's. Also create // Create Bob's ephemeral keys and derive noise_es_ee by agreeing with Alice's. Also create
@ -761,7 +780,7 @@ impl<Application: ApplicationLayer> Context<Application> {
let ack: &mut BobNoiseXKAck = byte_array_as_proto_buffer_mut(&mut ack_packet)?; let ack: &mut BobNoiseXKAck = byte_array_as_proto_buffer_mut(&mut ack_packet)?;
ack.session_protocol_version = SESSION_PROTOCOL_VERSION; ack.session_protocol_version = SESSION_PROTOCOL_VERSION;
ack.bob_noise_e = bob_noise_e; ack.bob_noise_e = bob_noise_e;
ack.bob_session_id = *bob_session_id.as_bytes(); ack.bob_session_id = bob_session_id.to_bytes();
ack.bob_hk_ciphertext = bob_hk_ciphertext; ack.bob_hk_ciphertext = bob_hk_ciphertext;
// Encrypt main section of reply and attach tag. // Encrypt main section of reply and attach tag.
@ -802,6 +821,7 @@ impl<Application: ApplicationLayer> Context<Application> {
hk, hk,
bob_noise_e_secret, bob_noise_e_secret,
header_protection_key: Secret(pkt.header_protection_key), header_protection_key: Secret(pkt.header_protection_key),
defrag: std::array::from_fn(|_| Mutex::new(Fragged::new())),
}), }),
); );
debug_assert!(!sessions.active.contains_key(&bob_session_id)); debug_assert!(!sessions.active.contains_key(&bob_session_id));
@ -847,7 +867,7 @@ impl<Application: ApplicationLayer> Context<Application> {
return Err(Error::OutOfSequence); return Err(Error::OutOfSequence);
} }
if let Offer::NoiseXKInit(outgoing_offer) = &state.current_offer { if let Offer::NoiseXKInit(outgoing_offer) = &state.outgoing_offer {
let pkt: &BobNoiseXKAck = byte_array_as_proto_buffer(pkt_assembled)?; let pkt: &BobNoiseXKAck = byte_array_as_proto_buffer(pkt_assembled)?;
// Derive noise_es_ee from Bob's ephemeral public key. // Derive noise_es_ee from Bob's ephemeral public key.
@ -875,7 +895,7 @@ impl<Application: ApplicationLayer> Context<Application> {
let pkt: &BobNoiseXKAck = byte_array_as_proto_buffer(pkt_assembled)?; let pkt: &BobNoiseXKAck = byte_array_as_proto_buffer(pkt_assembled)?;
if let Some(bob_session_id) = SessionId::new_from_bytes(&pkt.bob_session_id) { if let Some(bob_session_id) = SessionId::new_from_array(&pkt.bob_session_id) {
// Complete Noise_XKpsk3 by mixing in noise_se followed by the PSK. The PSK as far as // Complete Noise_XKpsk3 by mixing in noise_se followed by the PSK. The PSK as far as
// the Noise pattern is concerned is the result of mixing the externally supplied PSK // the Noise pattern is concerned is the result of mixing the externally supplied PSK
// with the Kyber1024 shared secret (hk). Kyber is treated as part of the PSK because // with the Kyber1024 shared secret (hk). Kyber is treated as part of the PSK because
@ -948,7 +968,7 @@ impl<Application: ApplicationLayer> Context<Application> {
)); ));
debug_assert!(state.keys[1].is_none()); debug_assert!(state.keys[1].is_none());
state.current_key = 0; state.current_key = 0;
state.current_offer = Offer::NoiseXKAck(Box::new(OutgoingSessionAck { state.outgoing_offer = Offer::NoiseXKAck(Box::new(OutgoingSessionAck {
last_retry_time: AtomicI64::new(current_time), last_retry_time: AtomicI64::new(current_time),
ack, ack,
ack_len, ack_len,
@ -1071,7 +1091,7 @@ impl<Application: ApplicationLayer> Context<Application> {
None, None,
], ],
current_key: 0, current_key: 0,
current_offer: Offer::None, outgoing_offer: Offer::None,
}), }),
defrag: std::array::from_fn(|_| Mutex::new(Fragged::new())), defrag: std::array::from_fn(|_| Mutex::new(Fragged::new())),
}); });
@ -1108,18 +1128,13 @@ impl<Application: ApplicationLayer> Context<Application> {
if let Some(session) = session { if let Some(session) = session {
let state = session.state.read().unwrap(); let state = session.state.read().unwrap();
if let Some(remote_session_id) = state.remote_session_id { if let (Some(remote_session_id), Some(key)) = (state.remote_session_id, state.keys[key_index].as_ref()) {
if let Some(key) = state.keys[key_index].as_ref() { if !key.my_turn_to_rekey && {
// Only the current "Alice" accepts rekeys initiated by the current "Bob." These roles
// flip with each rekey event.
if !key.my_turn_to_rekey {
let mut c = key.get_receive_cipher(incoming_counter); let mut c = key.get_receive_cipher(incoming_counter);
c.reset_init_gcm(&incoming_message_nonce); c.reset_init_gcm(&incoming_message_nonce);
c.crypt_in_place(&mut pkt_assembled[RekeyInit::ENC_START..RekeyInit::AUTH_START]); c.crypt_in_place(&mut pkt_assembled[RekeyInit::ENC_START..RekeyInit::AUTH_START]);
let aead_authentication_ok = c.finish_decrypt(&pkt_assembled[RekeyInit::AUTH_START..]); c.finish_decrypt(&pkt_assembled[RekeyInit::AUTH_START..])
drop(c); } {
if aead_authentication_ok {
let pkt: &RekeyInit = byte_array_as_proto_buffer(&pkt_assembled).unwrap(); let pkt: &RekeyInit = byte_array_as_proto_buffer(&pkt_assembled).unwrap();
if let Some(alice_e) = P384PublicKey::from_bytes(&pkt.alice_e) { if let Some(alice_e) = P384PublicKey::from_bytes(&pkt.alice_e) {
let bob_e_secret = P384KeyPair::generate(); let bob_e_secret = P384KeyPair::generate();
@ -1180,11 +1195,9 @@ impl<Application: ApplicationLayer> Context<Application> {
return Err(Error::OutOfSequence); return Err(Error::OutOfSequence);
} }
} }
}
return Err(Error::FailedAuthentication); return Err(Error::FailedAuthentication);
} }
} }
}
return Err(Error::OutOfSequence); return Err(Error::OutOfSequence);
} else { } else {
return Err(Error::UnknownLocalSessionId); return Err(Error::UnknownLocalSessionId);
@ -1201,19 +1214,13 @@ impl<Application: ApplicationLayer> Context<Application> {
if let Some(session) = session { if let Some(session) = session {
let state = session.state.read().unwrap(); let state = session.state.read().unwrap();
if let Offer::RekeyInit(alice_e_secret, _) = &state.current_offer { if let (Offer::RekeyInit(alice_e_secret, _), Some(key)) = (&state.outgoing_offer, state.keys[key_index].as_ref()) {
if let Some(key) = state.keys[key_index].as_ref() { if key.my_turn_to_rekey && {
// Only the current "Bob" initiates rekeys and expects this ACK.
if key.my_turn_to_rekey {
let mut c = key.get_receive_cipher(incoming_counter); let mut c = key.get_receive_cipher(incoming_counter);
c.reset_init_gcm(&incoming_message_nonce); c.reset_init_gcm(&incoming_message_nonce);
c.crypt_in_place(&mut pkt_assembled[RekeyAck::ENC_START..RekeyAck::AUTH_START]); c.crypt_in_place(&mut pkt_assembled[RekeyAck::ENC_START..RekeyAck::AUTH_START]);
let aead_authentication_ok = c.finish_decrypt(&pkt_assembled[RekeyAck::AUTH_START..]); c.finish_decrypt(&pkt_assembled[RekeyAck::AUTH_START..])
drop(c); } {
if aead_authentication_ok {
// Packet fully authenticated
let pkt: &RekeyAck = byte_array_as_proto_buffer(&pkt_assembled).unwrap(); let pkt: &RekeyAck = byte_array_as_proto_buffer(&pkt_assembled).unwrap();
if let Some(bob_e) = P384PublicKey::from_bytes(&pkt.bob_e) { if let Some(bob_e) = P384PublicKey::from_bytes(&pkt.bob_e) {
let next_session_key = hmac_sha512_secret( let next_session_key = hmac_sha512_secret(
@ -1239,7 +1246,7 @@ impl<Application: ApplicationLayer> Context<Application> {
true, true,
)); ));
state.current_key = next_key_index; // this is an ACK so it's confirmed state.current_key = next_key_index; // this is an ACK so it's confirmed
state.current_offer = Offer::None; state.outgoing_offer = Offer::None;
drop(state); drop(state);
return Ok(ReceiveResult::Ok(Some(session))); return Ok(ReceiveResult::Ok(Some(session)));
@ -1250,10 +1257,9 @@ impl<Application: ApplicationLayer> Context<Application> {
} }
} }
return Err(Error::FailedAuthentication); return Err(Error::FailedAuthentication);
} } else {
}
}
return Err(Error::OutOfSequence); return Err(Error::OutOfSequence);
}
} else { } else {
return Err(Error::UnknownLocalSessionId); return Err(Error::UnknownLocalSessionId);
} }
@ -1277,8 +1283,7 @@ impl<Application: ApplicationLayer> Session<Application> {
pub fn send<SendFunction: FnMut(&mut [u8])>(&self, mut send: SendFunction, mtu_sized_buffer: &mut [u8], mut data: &[u8]) -> Result<(), Error> { pub fn send<SendFunction: FnMut(&mut [u8])>(&self, mut send: SendFunction, mtu_sized_buffer: &mut [u8], mut data: &[u8]) -> Result<(), Error> {
debug_assert!(mtu_sized_buffer.len() >= MIN_TRANSPORT_MTU); debug_assert!(mtu_sized_buffer.len() >= MIN_TRANSPORT_MTU);
let state = self.state.read().unwrap(); let state = self.state.read().unwrap();
if let Some(remote_session_id) = state.remote_session_id { if let (Some(remote_session_id), Some(session_key)) = (state.remote_session_id, state.keys[state.current_key].as_ref()) {
if let Some(session_key) = state.keys[state.current_key].as_ref() {
let counter = self.get_next_outgoing_counter().ok_or(Error::MaxKeyLifetimeExceeded)?.get(); let counter = self.get_next_outgoing_counter().ok_or(Error::MaxKeyLifetimeExceeded)?.get();
let mut c = session_key.get_send_cipher(counter)?; let mut c = session_key.get_send_cipher(counter)?;
@ -1322,15 +1327,13 @@ impl<Application: ApplicationLayer> Session<Application> {
return Ok(()); return Ok(());
} }
}
return Err(Error::SessionNotEstablished); return Err(Error::SessionNotEstablished);
} }
/// Send a NOP to the other side (e.g. for keep alive). /// Send a NOP to the other side (e.g. for keep alive).
pub fn send_nop<SendFunction: FnMut(&mut [u8])>(&self, mut send: SendFunction) -> Result<(), Error> { pub fn send_nop<SendFunction: FnMut(&mut [u8])>(&self, mut send: SendFunction) -> Result<(), Error> {
let state = self.state.read().unwrap(); let state = self.state.read().unwrap();
if let Some(remote_session_id) = state.remote_session_id { if let (Some(remote_session_id), Some(session_key)) = (state.remote_session_id, state.keys[state.current_key].as_ref()) {
if let Some(session_key) = state.keys[state.current_key].as_ref() {
let counter = self.get_next_outgoing_counter().ok_or(Error::MaxKeyLifetimeExceeded)?.get(); let counter = self.get_next_outgoing_counter().ok_or(Error::MaxKeyLifetimeExceeded)?.get();
let mut nop = [0u8; HEADER_SIZE + AES_GCM_TAG_SIZE]; let mut nop = [0u8; HEADER_SIZE + AES_GCM_TAG_SIZE];
let mut c = session_key.get_send_cipher(counter)?; let mut c = session_key.get_send_cipher(counter)?;
@ -1342,10 +1345,14 @@ impl<Application: ApplicationLayer> Session<Application> {
.encrypt_block_in_place(&mut nop[HEADER_PROTECT_ENCRYPT_START..HEADER_PROTECT_ENCRYPT_END]); .encrypt_block_in_place(&mut nop[HEADER_PROTECT_ENCRYPT_START..HEADER_PROTECT_ENCRYPT_END]);
send(&mut nop); send(&mut nop);
} }
}
return Err(Error::SessionNotEstablished); return Err(Error::SessionNotEstablished);
} }
/// Set the current physical MTU that this session should use to send packets.
pub fn set_physical_mtu(&self, mtu: usize) {
self.state.write().unwrap().physical_mtu = mtu;
}
/// Check whether this session is established. /// Check whether this session is established.
pub fn established(&self) -> bool { pub fn established(&self) -> bool {
let state = self.state.read().unwrap(); let state = self.state.read().unwrap();
@ -1403,7 +1410,7 @@ impl<Application: ApplicationLayer> Session<Application> {
send(&mut rekey_buf); send(&mut rekey_buf);
drop(state); drop(state);
self.state.write().unwrap().current_offer = Offer::RekeyInit(rekey_e, current_time); self.state.write().unwrap().outgoing_offer = Offer::RekeyInit(rekey_e, current_time);
} }
} }
} }
@ -1578,6 +1585,7 @@ impl SessionKey {
} }
} }
#[inline(always)]
fn get_send_cipher<'a>(&'a self, counter: u64) -> Result<MutexGuard<'a, AesGcm<true>>, Error> { fn get_send_cipher<'a>(&'a self, counter: u64) -> Result<MutexGuard<'a, AesGcm<true>>, Error> {
if counter < self.expire_at_counter { if counter < self.expire_at_counter {
Ok(self.send_cipher_pool[(counter as usize) % GCM_CIPHER_POOL_SIZE].lock().unwrap()) Ok(self.send_cipher_pool[(counter as usize) % GCM_CIPHER_POOL_SIZE].lock().unwrap())
@ -1586,6 +1594,7 @@ impl SessionKey {
} }
} }
#[inline(always)]
fn get_receive_cipher<'a>(&'a self, counter: u64) -> MutexGuard<'a, AesGcm<false>> { fn get_receive_cipher<'a>(&'a self, counter: u64) -> MutexGuard<'a, AesGcm<false>> {
self.receive_cipher_pool[(counter as usize) % GCM_CIPHER_POOL_SIZE].lock().unwrap() self.receive_cipher_pool[(counter as usize) % GCM_CIPHER_POOL_SIZE].lock().unwrap()
} }