begun changes

This commit is contained in:
mamoniot 2023-03-22 18:29:14 -04:00
parent d471138eb5
commit 1e32e0ad2c
No known key found for this signature in database
GPG key ID: ADCCDBBE0E3D3B3B
3 changed files with 29 additions and 38 deletions

View file

@ -3,6 +3,7 @@ use std::ptr::slice_from_raw_parts;
/// Fast packet defragmenter
pub struct Fragged<Fragment, const MAX_FRAGMENTS: usize> {
count: u32,
have: u64,
counter: u64,
frags: [MaybeUninit<Fragment>; MAX_FRAGMENTS],
@ -35,9 +36,9 @@ impl<Fragment, const MAX_FRAGMENTS: usize> Fragged<Fragment, MAX_FRAGMENTS> {
// that the array of MaybeUninit<Fragment> can be freely cast into an array of
// Fragment. They also check that the maximum number of fragments is not too large
// for the fact that we use bits in a u64 to track which fragments are received.
assert!(MAX_FRAGMENTS <= 64);
assert_eq!(size_of::<MaybeUninit<Fragment>>(), size_of::<Fragment>());
assert_eq!(
debug_assert!(MAX_FRAGMENTS <= 64);
debug_assert_eq!(size_of::<MaybeUninit<Fragment>>(), size_of::<Fragment>());
debug_assert_eq!(
size_of::<[MaybeUninit<Fragment>; MAX_FRAGMENTS]>(),
size_of::<[Fragment; MAX_FRAGMENTS]>()
);
@ -51,12 +52,11 @@ impl<Fragment, const MAX_FRAGMENTS: usize> Fragged<Fragment, MAX_FRAGMENTS> {
#[inline(always)]
pub fn assemble(&mut self, counter: u64, fragment: Fragment, fragment_no: u8, fragment_count: u8) -> Option<Assembled<Fragment, MAX_FRAGMENTS>> {
if fragment_no < fragment_count && (fragment_count as usize) <= MAX_FRAGMENTS {
let mut have = self.have;
// If the counter has changed, reset the structure to receive a new packet.
if counter != self.counter {
self.counter = counter;
if needs_drop::<Fragment>() {
let mut have = self.have;
let mut i = 0;
while have != 0 {
if (have & 1) != 0 {
@ -66,26 +66,29 @@ impl<Fragment, const MAX_FRAGMENTS: usize> Fragged<Fragment, MAX_FRAGMENTS> {
have = have.wrapping_shr(1);
i += 1;
}
} else {
have = 0;
}
self.have = 0;
self.count = fragment_count as u32;
self.counter = counter;
}
let got = 1u64.wrapping_shl(fragment_no as u32);
if got & self.have == 0 && self.count as u8 == fragment_count {
self.have |= got;
unsafe {
self.frags.get_unchecked_mut(fragment_no as usize).write(fragment);
}
if self.have == 1u64.wrapping_shl(self.count) - 1 {
self.have = 0;
self.count = 0;
self.counter = 0;
// Setting 'have' to 0 resets the state of this object, and the fragments
// are effectively moved into the Assembled<> container and returned. That
// container will drop them when it is dropped.
return Some(Assembled(unsafe { std::mem::transmute_copy(&self.frags) }, fragment_count as usize));
}
}
unsafe {
self.frags.get_unchecked_mut(fragment_no as usize).write(fragment);
}
let want = 0xffffffffffffffffu64.wrapping_shr((64 - fragment_count) as u32);
have |= 1u64.wrapping_shl(fragment_no as u32);
if (have & want) == want {
self.have = 0;
// Setting 'have' to 0 resets the state of this object, and the fragments
// are effectively moved into the Assembled<> container and returned. That
// container will drop them when it is dropped.
return Some(Assembled(unsafe { std::mem::transmute_copy(&self.frags) }, fragment_count as usize));
} else {
self.have = have;
}
}
return None;
}

View file

@ -63,6 +63,7 @@ pub(crate) const KBKDF_KEY_USAGE_LABEL_AES_GCM_ALICE_TO_BOB: u8 = b'A'; // AES-G
pub(crate) const KBKDF_KEY_USAGE_LABEL_AES_GCM_BOB_TO_ALICE: u8 = b'B'; // AES-GCM in B->A direction
pub(crate) const KBKDF_KEY_USAGE_LABEL_RATCHET: u8 = b'R'; // Key used in derivatin of next session key
pub(crate) const MAX_INCOMPLETE_SESSION_QUEUE_SIZE: usize = 32;
pub(crate) const MAX_FRAGMENTS: usize = 48; // hard protocol max: 63
pub(crate) const MAX_NOISE_HANDSHAKE_FRAGMENTS: usize = 16; // enough room for p384 + ZT identity + kyber1024 + tag/hmac/etc.
pub(crate) const MAX_NOISE_HANDSHAKE_SIZE: usize = MAX_NOISE_HANDSHAKE_FRAGMENTS * MIN_TRANSPORT_MTU;

View file

@ -36,17 +36,8 @@ const GCM_CIPHER_POOL_SIZE: usize = 4;
/// Each application using ZSSP must create an instance of this to own sessions and
/// defragment incoming packets that are not yet associated with a session.
pub struct Context<Application: ApplicationLayer> {
max_incomplete_session_queue_size: usize,
default_physical_mtu: AtomicUsize,
defrag: Mutex<
HashMap<
(Application::PhysicalPath, u64),
Arc<(
Mutex<Fragged<Application::IncomingPacketBuffer, MAX_NOISE_HANDSHAKE_FRAGMENTS>>,
i64, // creation timestamp
)>,
>,
>,
defrag: Mutex<[Fragged<Application::IncomingPacketBuffer, MAX_NOISE_HANDSHAKE_FRAGMENTS>; MAX_INCOMPLETE_SESSION_QUEUE_SIZE]>,
sessions: RwLock<SessionsById<Application>>,
}
@ -154,11 +145,10 @@ impl<Application: ApplicationLayer> Context<Application> {
/// Create a new session context.
///
/// * `max_incomplete_session_queue_size` - Maximum number of incomplete sessions in negotiation phase
pub fn new(max_incomplete_session_queue_size: usize, default_physical_mtu: usize) -> Self {
pub fn new(default_physical_mtu: usize) -> Self {
Self {
max_incomplete_session_queue_size,
default_physical_mtu: AtomicUsize::new(default_physical_mtu),
defrag: Mutex::new(HashMap::new()),
defrag: Mutex::new(std::array::from_fn(|_| Fragged::new())),
sessions: RwLock::new(SessionsById {
active: HashMap::with_capacity(64),
incoming: HashMap::with_capacity(64),
@ -262,9 +252,6 @@ impl<Application: ApplicationLayer> Context<Application> {
}
}
// Delete any expired defragmentation queue items not associated with a session.
self.defrag.lock().unwrap().retain(|_, fragged| fragged.1 > negotiation_timeout_cutoff);
Application::INCOMING_SESSION_NEGOTIATION_TIMEOUT_MS.min(Application::RETRY_INTERVAL)
}