From cd6d8d36b08e4c01f14ec0ecd1d525473d888ab6 Mon Sep 17 00:00:00 2001 From: Adam Ierymenko Date: Wed, 8 Mar 2023 15:03:27 -0500 Subject: [PATCH] Simplify some packet building code. --- zssp/src/error.rs | 8 -------- zssp/src/zssp.rs | 35 +++++++++++++++++++---------------- 2 files changed, 19 insertions(+), 24 deletions(-) diff --git a/zssp/src/error.rs b/zssp/src/error.rs index a8448bc35..85f24f50c 100644 --- a/zssp/src/error.rs +++ b/zssp/src/error.rs @@ -47,14 +47,6 @@ pub enum Error { UnexpectedBufferOverrun, } -// An I/O error in the parser means an invalid packet. -impl From for Error { - #[inline(always)] - fn from(_: std::io::Error) -> Self { - Self::UnexpectedBufferOverrun - } -} - impl std::fmt::Display for Error { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.write_str(match self { diff --git a/zssp/src/zssp.rs b/zssp/src/zssp.rs index 89bdf490e..87f28563f 100644 --- a/zssp/src/zssp.rs +++ b/zssp/src/zssp.rs @@ -10,7 +10,6 @@ // FIPS compliant Noise_XK with Jedi powers (Kyber1024) and built-in attack-resistant large payload (fragmentation) support. use std::collections::{HashMap, HashSet}; -use std::io::Write; use std::num::NonZeroU64; use std::sync::atomic::{AtomicI64, AtomicU64, Ordering}; use std::sync::{Arc, Mutex, RwLock, Weak}; @@ -923,15 +922,14 @@ impl Context { // up forward secrecy. Also return Bob's opaque note. let mut reply_buffer = [0u8; MAX_NOISE_HANDSHAKE_SIZE]; reply_buffer[HEADER_SIZE] = SESSION_PROTOCOL_VERSION; - let mut rw = &mut reply_buffer[HEADER_SIZE + 1..]; + let mut reply_len = HEADER_SIZE + 1; let alice_s_public_blob = app.get_local_s_public_blob(); assert!(alice_s_public_blob.len() <= (u16::MAX as usize)); - rw.write_all(&(alice_s_public_blob.len() as u16).to_le_bytes())?; - let mut enc_start = MAX_NOISE_HANDSHAKE_SIZE - rw.len(); - rw.write_all(alice_s_public_blob)?; + reply_len = append_to_slice(&mut reply_buffer, reply_len, &(alice_s_public_blob.len() as u16).to_le_bytes())?; + let mut enc_start = reply_len; + reply_len = append_to_slice(&mut reply_buffer, reply_len, alice_s_public_blob)?; - let mut reply_len = MAX_NOISE_HANDSHAKE_SIZE - rw.len(); let mut gcm = AesGcm::new( kbkdf::(&hmac_sha512( noise_es_ee.as_bytes(), @@ -943,20 +941,17 @@ impl Context { gcm.reset_init_gcm(&reply_message_nonce); gcm.aad(&noise_h_next); gcm.crypt_in_place(&mut reply_buffer[enc_start..reply_len]); - let mut rw = &mut reply_buffer[reply_len..]; - rw.write_all(&gcm.finish_encrypt())?; + reply_len = append_to_slice(&mut reply_buffer, reply_len, &gcm.finish_encrypt())?; let metadata = outgoing_offer.metadata.as_ref().map_or(&[][..0], |md| md.as_slice()); assert!(metadata.len() <= (u16::MAX as usize)); - rw.write_all(&(metadata.len() as u16).to_le_bytes())?; + reply_len = append_to_slice(&mut reply_buffer, reply_len, &(metadata.len() as u16).to_le_bytes())?; - reply_len = MAX_NOISE_HANDSHAKE_SIZE - rw.len(); let noise_h_next = mix_hash(&noise_h_next, &reply_buffer[HEADER_SIZE..reply_len]); - let mut rw = &mut reply_buffer[reply_len..]; enc_start = reply_len; - rw.write_all(metadata.as_ref())?; + reply_len = append_to_slice(&mut reply_buffer, reply_len, metadata)?; let mut gcm = AesGcm::new( kbkdf::(noise_es_ee_se_hk_psk.as_bytes()).as_bytes(), @@ -964,11 +959,8 @@ impl Context { ); gcm.reset_init_gcm(&reply_message_nonce); gcm.aad(&noise_h_next); - reply_len = MAX_NOISE_HANDSHAKE_SIZE - rw.len(); gcm.crypt_in_place(&mut reply_buffer[enc_start..reply_len]); - let mut rw = &mut reply_buffer[reply_len..]; - rw.write_all(&gcm.finish_encrypt())?; - reply_len = MAX_NOISE_HANDSHAKE_SIZE - rw.len(); + reply_len = append_to_slice(&mut reply_buffer, reply_len, &gcm.finish_encrypt())?; drop(state); { @@ -1665,6 +1657,17 @@ impl<'a> PktReader<'a> { } } +/// Helper function to append to a slice when we still want to be able to look back at it. +fn append_to_slice(s: &mut [u8], p: usize, d: &[u8]) -> Result { + let tmp = p + d.len(); + if tmp <= s.len() { + s[p..tmp].copy_from_slice(d); + Ok(tmp) + } else { + Err(Error::UnexpectedBufferOverrun) + } +} + /// MixHash to update 'h' during negotiation. fn mix_hash(h: &[u8; SHA384_HASH_SIZE], m: &[u8]) -> [u8; SHA384_HASH_SIZE] { let mut hasher = SHA384::new();