Fix IBLT read/write and remove some warnings.

This commit is contained in:
Adam Ierymenko 2022-01-29 15:04:26 -08:00
parent 1a3d9a28a0
commit 06c0d452f4
No known key found for this signature in database
GPG key ID: C8877CF2D7A5D7F3
2 changed files with 41 additions and 18 deletions

View file

@ -7,7 +7,8 @@
*/ */
use std::io::{Read, Write}; use std::io::{Read, Write};
use std::mem::{transmute, zeroed}; use std::mem::{size_of, transmute, zeroed};
use std::ptr::write_bytes;
use zerotier_core_crypto::varint; use zerotier_core_crypto::varint;
@ -61,20 +62,25 @@ impl<const HS: usize, const B: usize> IBLT<HS, B> {
pub fn new() -> Self { unsafe { zeroed() } } pub fn new() -> Self { unsafe { zeroed() } }
#[inline(always)]
pub fn clear(&mut self) {
unsafe { write_bytes((self as *mut Self).cast::<u8>(), 0, size_of::<Self>()) };
}
pub fn read<R: Read>(&mut self, r: &mut R) -> std::io::Result<()> { pub fn read<R: Read>(&mut self, r: &mut R) -> std::io::Result<()> {
let mut tmp = [0_u8; 8]; let mut tmp = [0_u8; 8];
let mut prev_c = 0_i64; let mut prev_c = 0_i64;
for i in 0..B { for b in self.map.iter_mut() {
let b = &mut self.map[i];
r.read_exact(&mut b.key_sum)?; r.read_exact(&mut b.key_sum)?;
r.read_exact(&mut tmp)?; r.read_exact(&mut tmp)?;
b.check_hash_sum = u64::from_le_bytes(tmp); b.check_hash_sum = u64::from_le_bytes(tmp);
let c = ((varint::read(r)?.0 as i64) + prev_c) as u64; let mut c = varint::read(r)?.0 as i64;
if (c & 1) == 0 { if (c & 1) == 0 {
b.count = c.wrapping_shr(1) as i64; c = c.wrapping_shr(1);
} else { } else {
b.count = -(c.wrapping_shr(1) as i64); c = -c.wrapping_shr(1);
} }
b.count = c + prev_c;
prev_c = b.count; prev_c = b.count;
} }
Ok(()) Ok(())
@ -82,17 +88,15 @@ impl<const HS: usize, const B: usize> IBLT<HS, B> {
pub fn write<W: Write>(&self, w: &mut W) -> std::io::Result<()> { pub fn write<W: Write>(&self, w: &mut W) -> std::io::Result<()> {
let mut prev_c = 0_i64; let mut prev_c = 0_i64;
for i in 0..B { for b in self.map.iter() {
let b = &self.map[i];
w.write_all(&b.key_sum)?; w.write_all(&b.key_sum)?;
w.write_all(&b.check_hash_sum.to_le_bytes())?; w.write_all(&b.check_hash_sum.to_le_bytes())?;
let c = b.count - prev_c; let mut c = (b.count - prev_c).wrapping_shl(1);
prev_c = b.count; prev_c = b.count;
if c >= 0 { if c < 0 {
varint::write(w, (c as u64).wrapping_shl(1))?; c = -c | 1;
} else {
varint::write(w, (-c as u64).wrapping_shl(1) | 1)?;
} }
varint::write(w, c as u64)?;
} }
Ok(()) Ok(())
} }
@ -103,7 +107,7 @@ impl<const HS: usize, const B: usize> IBLT<HS, B> {
let iteration_indices: [u64; 8] = unsafe { transmute(zerotier_core_crypto::hash::SHA512::hash(key)) }; let iteration_indices: [u64; 8] = unsafe { transmute(zerotier_core_crypto::hash::SHA512::hash(key)) };
let check_hash = calc_check_hash::<HS>(&key); let check_hash = calc_check_hash::<HS>(&key);
for i in 0..KEY_MAPPING_ITERATIONS { for i in 0..KEY_MAPPING_ITERATIONS {
let b = &mut self.map[(u64::from_le(iteration_indices[i]) as usize) % B]; let b = unsafe { self.map.get_unchecked_mut((u64::from_le(iteration_indices[i]) as usize) % B) };
for x in 0..HS { for x in 0..HS {
b.key_sum[x] ^= key[x]; b.key_sum[x] ^= key[x];
} }
@ -200,12 +204,31 @@ mod tests {
t.insert(&k); t.insert(&k);
e.insert(k); e.insert(k);
} }
let t_backup = t.clone();
assert!(t == t_backup);
let mut cnt = 0; let mut cnt = 0;
t.list(|k| { t.list(|k| {
assert!(e.contains(k)); assert!(e.contains(k));
cnt += 1; cnt += 1;
}); });
assert_eq!(cnt, expected_cnt); assert_eq!(cnt, expected_cnt);
let mut test_buf: Vec<u8> = Vec::new();
assert!(t_backup.write(&mut test_buf).is_ok());
let mut t_restore: IBLT<48, 1152> = IBLT::new();
let mut test_read = test_buf.as_slice();
assert!(t_restore.read(&mut test_read).is_ok());
assert!(t_restore == t_backup);
cnt = 0;
t_restore.list(|k| {
assert!(e.contains(k));
cnt += 1;
});
assert_eq!(cnt, expected_cnt);
} }
} }
} }

View file

@ -362,18 +362,18 @@ impl Identity {
*/ */
buf.append_u8(0x03)?; buf.append_u8(0x03)?;
let p521_has_private = if include_private && secret.map_or(false, |s| s.p521.is_some()) { let p521_has_private = if include_private && secret.map_or(false, |s| s.p521.is_some()) {
buf.append_u16(Self::P521_PUBLIC_AND_PRIVATE_BUNDLE_SIZE + 1 + 2); buf.append_u16(Self::P521_PUBLIC_AND_PRIVATE_BUNDLE_SIZE + 1 + 2)?;
true true
} else { } else {
buf.append_u16(Self::P521_PUBLIC_ONLY_BUNDLE_SIZE + 1 + 2); buf.append_u16(Self::P521_PUBLIC_ONLY_BUNDLE_SIZE + 1 + 2)?;
false false
}; };
buf.append_u8(IDENTITY_ALGORITHM_EC_NIST_P521)?; buf.append_u8(IDENTITY_ALGORITHM_EC_NIST_P521)?;
if p521_has_private { if p521_has_private {
buf.append_u16(Self::P521_PUBLIC_AND_PRIVATE_BUNDLE_SIZE); buf.append_u16(Self::P521_PUBLIC_AND_PRIVATE_BUNDLE_SIZE)?;
} else { } else {
buf.append_u16(Self::P521_PUBLIC_ONLY_BUNDLE_SIZE); buf.append_u16(Self::P521_PUBLIC_ONLY_BUNDLE_SIZE)?;
} }
buf.append_bytes_fixed(&p521.ecdh)?; buf.append_bytes_fixed(&p521.ecdh)?;
buf.append_bytes_fixed(&p521.ecdsa)?; buf.append_bytes_fixed(&p521.ecdsa)?;