void/ZeroTierOne
2
0
Fork 0
mirror of https://github.com/zerotier/ZeroTierOne.git synced 2025-04-19 13:36:54 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

saved

Browse source
This commit is contained in:
monica 2023-03-22 11:51:10 -04:00
parent da1bd74bdc
commit 98470d6373
No known key found for this signature in database
GPG key ID: ADCCDBBE0E3D3B3B
1 changed files with 25 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

55
zssp/src/fragged.rs
View file

@ -38,17 +38,13 @@ 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.
debug_assert!(MAX_FRAGMENTS <= 64);
debug_assert_eq!(size_of::<MaybeUninit<Fragment>>(), size_of::<Fragment>());
debug_assert_eq!(
assert!(MAX_FRAGMENTS <= 64);
assert_eq!(size_of::<MaybeUninit<Fragment>>(), size_of::<Fragment>());
assert_eq!(
size_of::<[MaybeUninit<Fragment>; MAX_FRAGMENTS]>(),
size_of::<[Fragment; MAX_FRAGMENTS]>()
);
Self {
have: AtomicU64::new(0),
counter_want: RwLock::new((0, 0)),
frags: unsafe { zeroed() },
}
unsafe { zeroed() }
}
/// Add a fragment and return an assembled packet container if all fragments have been received.
@ -56,21 +52,22 @@ impl<Fragment, const MAX_FRAGMENTS: usize> Fragged<Fragment, MAX_FRAGMENTS> {
/// When a fully assembled packet is returned the internal state is reset and this object can
/// be reused to assemble another packet.
#[inline(always)]
pub fn assemble(&self, counter: u64, fragment: Fragment, fragment_no: u8, fragment_count: u8) -> Option<Assembled<Fragment, MAX_FRAGMENTS>> {
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 r = self.counter_want.read().unwrap();
let (cur_counter, want) = *r;
let mut r_guard = Some(r);
let (mut r_counter, mut r_want) = *r;
let mut _r_guard = Some(r);
let mut _w_guard = None;
// If the counter has changed, reset the structure to receive a new packet.
if counter != cur_counter {
r_guard.take();
if counter != r_counter {
_r_guard.take();
let mut w = self.counter_want.write().unwrap();
let (cur_counter, _) = *w;
if counter != cur_counter {
let (w_counter, w_want) = *w;
if counter != w_counter {
if needs_drop::<Fragment>() {
let mut have = self.have.load(Ordering::Relaxed);
if have != want {
if have != w_want {
let mut i = 0;
while have != 0 {
if (have & 1) != 0 {
@ -82,28 +79,26 @@ impl<Fragment, const MAX_FRAGMENTS: usize> Fragged<Fragment, MAX_FRAGMENTS> {
}
}
}
r_counter = counter;
r_want = 1u64.wrapping_shl(fragment_count as u32) - 1;
*w = (r_counter, r_want);
self.have.store(0, Ordering::Relaxed);
let want = 1u64.wrapping_shl(fragment_count as u32) - 1;
*w = (counter, want);
}
_w_guard = Some(w);
}
// Validate inputs, mainly that the fragment_count matches what we expect
// and that fragment_no is actually a wanted fragment
let got = 1u64.wrapping_shl(fragment_no as u32);
if want & got > 0 && 1u64.wrapping_shl(fragment_count as u32) == want + 1 {
// Now we check whether fragment_no is missing
if got & r_want > 0 && 1u64.wrapping_shl(fragment_count as u32) - 1 == r_want {
let have = self.have.fetch_or(got, Ordering::Relaxed);
if have & got == 0 {
unsafe {
(*self.frags.get()).get_unchecked_mut(fragment_no as usize).write(fragment);
}
if (have | got) == want {
// The fragments are effectively moved into the Assembled<> container and returned.
// That container will drop them when it is dropped.
if have | got == r_want {
// 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));
}
}
@ -117,10 +112,10 @@ impl<Fragment, const MAX_FRAGMENTS: usize> Drop for Fragged<Fragment, MAX_FRAGME
#[inline(always)]
fn drop(&mut self) {
if needs_drop::<Fragment>() {
let r = self.counter_want.read().unwrap();
let (_, want) = *r;
let mut have = self.have.load(Ordering::Relaxed);
if have != want {
let r = self.counter_want.read().unwrap();
let (_, r_want) = *r;
if have != r_want {
let mut i = 0;
while have != 0 {
if (have & 1) != 0 {