commit 959815fb1eae4712fc4173bd7fa9d6dc1a8bff11 Author: Adam Ierymenko Date: Wed Aug 2 10:09:53 2023 -0400 Initial commit of utils broken out of next-gen monorepo. diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..7dd2bd4 --- /dev/null +++ b/.gitignore @@ -0,0 +1,15 @@ +/target +/**/target +/**/Cargo.lock + +.DS_* +.Icon* +._* +*.o +*.so +*.dylib +*.dSYM +*.a +/.idea +/.nova +*.secret diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..7a45489 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,22 @@ +[package] +authors = ["ZeroTier, Inc. "] +edition = "2021" +license = "MPL-2.0" +name = "zerotier-common-utils" +version = "0.1.0" + +[features] + +[dependencies] +base64 = "^0" +serde = { version = "^1", features = ["derive"], default-features = false } + +[dev-dependencies] +rand = "*" + +[target."cfg(windows)".dependencies] +winapi = { version = "^0", features = ["handleapi", "ws2ipdef", "ws2tcpip"] } + +[target."cfg(not(windows))".dependencies] +libc = "^0" +signal-hook = "^0" diff --git a/rustfmt.toml b/rustfmt.toml new file mode 100644 index 0000000..bc80098 --- /dev/null +++ b/rustfmt.toml @@ -0,0 +1,16 @@ +#unstable_features = true +max_width = 150 +#use_small_heuristics = "Max" +edition = "2021" +#empty_item_single_line = true +newline_style = "Unix" +struct_lit_width = 60 +tab_spaces = 4 +use_small_heuristics = "Default" +#fn_single_line = true +#hex_literal_case = "Lower" +#merge_imports = false +group_imports = "StdExternalCrate" +single_line_if_else_max_width = 0 +use_try_shorthand = true +imports_granularity = "Module" diff --git a/src/arrayvec.rs b/src/arrayvec.rs new file mode 100644 index 0000000..e304c83 --- /dev/null +++ b/src/arrayvec.rs @@ -0,0 +1,434 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +use std::fmt::Debug; +use std::io::Write; +use std::mem::{needs_drop, size_of, MaybeUninit}; +use std::ptr::{slice_from_raw_parts, slice_from_raw_parts_mut}; + +use serde::ser::SerializeSeq; +use serde::{Deserialize, Deserializer, Serialize, Serializer}; + +#[derive(Clone, Copy, Debug)] +pub struct OutOfCapacityError(pub T); + +impl std::fmt::Display for OutOfCapacityError { + fn fmt(&self, stream: &mut std::fmt::Formatter) -> std::fmt::Result { + std::fmt::Display::fmt("ArrayVec out of space", stream) + } +} + +impl ::std::error::Error for OutOfCapacityError { + fn description(&self) -> &str { + "ArrayVec out of space" + } +} + +/// A simple vector backed by a static sized array with no memory allocations and no overhead construction. +pub struct ArrayVec { + pub(crate) s: usize, + pub(crate) a: [MaybeUninit; C], +} + +impl Default for ArrayVec { + #[inline(always)] + fn default() -> Self { + Self::new() + } +} + +impl Clone for ArrayVec { + #[inline] + fn clone(&self) -> Self { + debug_assert!(self.s <= C); + Self { + s: self.s, + a: unsafe { + let mut tmp: [MaybeUninit; C] = MaybeUninit::uninit().assume_init(); + for i in 0..self.s { + tmp.get_unchecked_mut(i).write(self.a[i].assume_init_ref().clone()); + } + tmp + }, + } + } +} + +impl From<[T; S]> for ArrayVec { + #[inline] + fn from(v: [T; S]) -> Self { + if S <= C { + let mut tmp = Self::new(); + for i in 0..S { + tmp.push(v[i].clone()); + } + tmp + } else { + panic!(); + } + } +} + +impl ToString for ArrayVec { + #[inline] + fn to_string(&self) -> String { + crate::hex::to_string(self.as_bytes()) + } +} + +impl Write for ArrayVec { + #[inline] + fn write(&mut self, buf: &[u8]) -> std::io::Result { + for i in buf.iter() { + if self.try_push(*i).is_err() { + return Err(std::io::Error::new(std::io::ErrorKind::Other, "ArrayVec out of space")); + } + } + Ok(buf.len()) + } + + #[inline(always)] + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +impl TryFrom> for ArrayVec { + type Error = OutOfCapacityError; + + #[inline(always)] + fn try_from(mut value: Vec) -> Result { + let mut tmp = Self::new(); + for x in value.drain(..) { + tmp.try_push(x)?; + } + Ok(tmp) + } +} + +impl TryFrom<&Vec> for ArrayVec { + type Error = OutOfCapacityError; + + #[inline(always)] + fn try_from(value: &Vec) -> Result { + let mut tmp = Self::new(); + for x in value.iter() { + tmp.try_push(x.clone())?; + } + Ok(tmp) + } +} + +impl TryFrom<&[T]> for ArrayVec { + type Error = OutOfCapacityError; + + #[inline(always)] + fn try_from(value: &[T]) -> Result { + let mut tmp = Self::new(); + for x in value.iter() { + tmp.try_push(x.clone())?; + } + Ok(tmp) + } +} + +impl ArrayVec { + #[inline(always)] + pub fn new() -> Self { + assert_eq!(size_of::<[T; C]>(), size_of::<[MaybeUninit; C]>()); + Self { s: 0, a: unsafe { MaybeUninit::uninit().assume_init() } } + } + + #[inline] + pub fn push(&mut self, v: T) { + let i = self.s; + if i < C { + unsafe { self.a.get_unchecked_mut(i).write(v) }; + self.s = i + 1; + } else { + panic!(); + } + } + + #[inline] + pub fn try_push(&mut self, v: T) -> Result<(), OutOfCapacityError> { + if self.s < C { + let i = self.s; + unsafe { self.a.get_unchecked_mut(i).write(v) }; + self.s = i + 1; + Ok(()) + } else { + Err(OutOfCapacityError(v)) + } + } + + /// Get a raw byte slice view of the contents of this vector. + /// This is only available for Copy types and will panic if the type needs_drop(). + #[inline(always)] + pub fn as_bytes(&self) -> &[T] + where + T: Copy, + { + assert!(!std::mem::needs_drop::()); + unsafe { &*slice_from_raw_parts(self.a.as_ptr().cast(), self.s) } + } + + #[inline(always)] + pub fn is_empty(&self) -> bool { + self.s == 0 + } + + #[inline(always)] + pub fn len(&self) -> usize { + self.s + } + + #[inline(always)] + pub const fn capacity(&self) -> usize { + C + } + + #[inline(always)] + pub fn capacity_remaining(&self) -> usize { + C - self.s + } + + #[inline(always)] + pub fn iter(&self) -> impl DoubleEndedIterator { + self.as_ref().iter() + } + + #[inline(always)] + pub fn iter_mut(&mut self) -> impl DoubleEndedIterator { + self.as_mut().iter_mut() + } + + #[inline(always)] + pub fn first(&self) -> Option<&T> { + if self.s != 0 { + Some(unsafe { self.a.get_unchecked(0).assume_init_ref() }) + } else { + None + } + } + + #[inline(always)] + pub fn last(&self) -> Option<&T> { + if self.s != 0 { + Some(unsafe { self.a.get_unchecked(self.s - 1).assume_init_ref() }) + } else { + None + } + } + + #[inline] + pub fn pop(&mut self) -> Option { + if self.s > 0 { + let i = self.s - 1; + debug_assert!(i < C); + self.s = i; + Some(unsafe { self.a.get_unchecked(i).assume_init_read() }) + } else { + None + } + } + + #[inline] + pub fn clear(&mut self) { + if needs_drop::() { + for i in 0..self.s { + unsafe { self.a.get_unchecked_mut(i).assume_init_drop() }; + } + } + self.s = 0; + } + + #[inline] + pub fn sort(&mut self) + where + T: Ord, + { + self.as_mut().sort(); + } + + #[inline] + pub fn sort_unstable(&mut self) + where + T: Ord, + { + self.as_mut().sort_unstable(); + } +} + +impl ArrayVec +where + T: Copy, +{ + /// Push a slice of copyable objects, panic if capacity exceeded. + #[inline] + pub fn push_slice(&mut self, v: &[T]) { + let start = self.s; + let end = self.s + v.len(); + if end <= C { + for i in start..end { + unsafe { self.a.get_unchecked_mut(i).write(*v.get_unchecked(i - start)) }; + } + self.s = end; + } else { + panic!(); + } + } +} + +impl Drop for ArrayVec { + #[inline(always)] + fn drop(&mut self) { + if needs_drop::() { + for i in 0..self.s { + unsafe { self.a.get_unchecked_mut(i).assume_init_drop() }; + } + } + } +} + +impl AsRef<[T]> for ArrayVec { + #[inline(always)] + fn as_ref(&self) -> &[T] { + unsafe { &*slice_from_raw_parts(self.a.as_ptr().cast(), self.s) } + } +} + +impl AsMut<[T]> for ArrayVec { + #[inline(always)] + fn as_mut(&mut self) -> &mut [T] { + unsafe { &mut *slice_from_raw_parts_mut(self.a.as_mut_ptr().cast(), self.s) } + } +} + +impl PartialEq for ArrayVec +where + T: PartialEq, +{ + #[inline(always)] + fn eq(&self, other: &Self) -> bool { + let tmp: &[T] = self.as_ref(); + tmp.eq(other.as_ref()) + } +} + +impl Eq for ArrayVec where T: Eq {} + +impl PartialOrd for ArrayVec +where + T: PartialOrd, +{ + fn partial_cmp(&self, other: &Self) -> Option { + self.iter().partial_cmp(other.iter()) + } +} + +impl Ord for ArrayVec +where + T: Ord, +{ + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.iter().cmp(other.iter()) + } +} + +impl Debug for ArrayVec +where + T: Debug, +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("[")?; + for x in self.iter() { + x.fmt(f)?; + } + f.write_str("]") + } +} + +impl Serialize for ArrayVec +where + T: Serialize, +{ + #[inline] + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + let mut seq = serializer.serialize_seq(Some(self.len()))?; + let sl: &[T] = self.as_ref(); + for i in 0..self.s { + seq.serialize_element(&sl[i])?; + } + seq.end() + } +} + +struct ArrayVecVisitor<'de, T: Deserialize<'de>, const L: usize>(std::marker::PhantomData<&'de T>); + +impl<'de, T, const L: usize> serde::de::Visitor<'de> for ArrayVecVisitor<'de, T, L> +where + T: Deserialize<'de>, +{ + type Value = ArrayVec; + + #[inline] + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str(format!("up to {} elements", L).as_str()) + } + + #[inline] + fn visit_seq(self, mut seq: A) -> Result + where + A: serde::de::SeqAccess<'de>, + { + let mut a = ArrayVec::::new(); + while let Some(x) = seq.next_element()? { + if !a.try_push(x).is_ok() { + return Err(serde::de::Error::custom("capacity exceeded")); + } + } + return Ok(a); + } +} + +impl<'de, T: Deserialize<'de> + 'de, const L: usize> Deserialize<'de> for ArrayVec +where + T: Deserialize<'de>, +{ + #[inline] + fn deserialize(deserializer: D) -> Result, D::Error> + where + D: Deserializer<'de>, + { + deserializer.deserialize_seq(ArrayVecVisitor(std::marker::PhantomData::default())) + } +} + +#[cfg(test)] +mod tests { + use super::ArrayVec; + + #[test] + fn array_vec() { + let mut v = ArrayVec::::new(); + for i in 0..128 { + v.push(i); + } + assert_eq!(v.len(), 128); + assert!(v.try_push(1000).is_err()); + assert_eq!(v.len(), 128); + for _ in 0..128 { + assert!(v.pop().is_some()); + } + assert!(v.pop().is_none()); + } +} diff --git a/src/base24.rs b/src/base24.rs new file mode 100644 index 0000000..99ae389 --- /dev/null +++ b/src/base24.rs @@ -0,0 +1,53 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +use crate::error::InvalidParameterError; + +/// All unambiguous letters, thus easy to type on the alphabetic keyboards on phones without extra shift taps. +/// The letters 'l' and 'u' are skipped. +const BASE24_ALPHABET: [u8; 24] = [ + b'a', b'b', b'c', b'd', b'e', b'f', b'g', b'h', b'i', b'j', b'k', b'm', b'n', b'o', b'p', b'q', b'r', b's', b't', b'v', b'w', b'x', b'y', b'z', +]; +/// Reverse table for BASE24 alphabet, indexed relative to 'a' or 'A'. +const BASE24_ALPHABET_INV: [u8; 26] = [ + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 255, 11, 12, 13, 14, 15, 16, 17, 18, 255, 19, 20, 21, 22, 23, +]; + +/// Encode 4 binary bytes into 7 base24 characters. +pub fn encode_4to7(b: &[u8], s: &mut String) { + let mut n = u32::from_be_bytes(b[..4].try_into().unwrap()); + for _ in 0..6 { + s.push(BASE24_ALPHABET[(n % 24) as usize] as char); + n /= 24; + } + s.push(BASE24_ALPHABET[n as usize] as char); +} + +pub fn decode_7to4(mut s: &[u8]) -> Result<[u8; 4], InvalidParameterError> { + let mut n = 0u32; + if s.len() > 7 { + s = &s[..7]; + } + for c in s.iter().rev() { + let mut c = *c; + if c >= 97 && c <= 122 { + c -= 97; + } else if c >= 65 && c <= 90 { + c -= 65; + } else { + return Err(InvalidParameterError("invalid base24 character")); + } + let i = BASE24_ALPHABET_INV[c as usize]; + if i == 255 { + return Err(InvalidParameterError("invalid base24 character")); + } + n *= 24; + n = n.wrapping_add(i as u32); + } + return Ok(n.to_be_bytes()); +} diff --git a/src/base64.rs b/src/base64.rs new file mode 100644 index 0000000..e0a504f --- /dev/null +++ b/src/base64.rs @@ -0,0 +1,21 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +use base64::{engine::general_purpose, Engine as _}; + +/// Encode a byte slice as Base64 using the URL-safe alphabet without padding +#[inline(always)] +pub fn to_string(b: &[u8]) -> String { + general_purpose::URL_SAFE_NO_PAD.encode(b) +} + +/// Decode a byte slcie using the URL-safe alphabet without padding +#[inline(always)] +pub fn from_string(s: &[u8]) -> Option> { + general_purpose::URL_SAFE_NO_PAD.decode(s).ok() +} diff --git a/src/blob.rs b/src/blob.rs new file mode 100644 index 0000000..47a905b --- /dev/null +++ b/src/blob.rs @@ -0,0 +1,160 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +use std::array::TryFromSliceError; +use std::fmt::Debug; +use std::hash::Hash; + +use serde::{Deserialize, Deserializer, Serialize, Serializer}; + +use crate::base64; +use crate::hex; + +/// Fixed size Serde serializable byte array. +/// This makes it easier to deal with blobs larger than 32 bytes (due to serde array limitations) +#[repr(transparent)] +#[derive(Clone, Eq, PartialEq, PartialOrd, Ord, Hash)] +pub struct Blob([u8; L]); + +impl Blob { + #[inline(always)] + pub fn as_bytes(&self) -> &[u8; L] { + &self.0 + } + + #[inline(always)] + pub const fn len(&self) -> usize { + L + } +} + +impl From> for [u8; L] { + #[inline(always)] + fn from(value: Blob) -> Self { + value.0 + } +} + +impl From<[u8; L]> for Blob { + #[inline(always)] + fn from(a: [u8; L]) -> Self { + Self(a) + } +} + +impl From<&[u8; L]> for &Blob { + #[inline(always)] + fn from(a: &[u8; L]) -> Self { + // Blob is a transparent wrapper around an array, so this is just a type cast. + unsafe { std::mem::transmute(a) } + } +} + +impl TryFrom<&[u8]> for Blob { + type Error = TryFromSliceError; + + #[inline(always)] + fn try_from(value: &[u8]) -> Result { + value.try_into().map(|b| Self(b)) + } +} + +impl Default for Blob { + #[inline(always)] + fn default() -> Self { + unsafe { std::mem::zeroed() } + } +} + +impl AsRef<[u8; L]> for Blob { + #[inline(always)] + fn as_ref(&self) -> &[u8; L] { + &self.0 + } +} + +impl AsMut<[u8; L]> for Blob { + #[inline(always)] + fn as_mut(&mut self) -> &mut [u8; L] { + &mut self.0 + } +} + +impl ToString for Blob { + #[inline(always)] + fn to_string(&self) -> String { + hex::to_string(&self.0) + } +} + +impl Debug for Blob { + #[inline] + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.to_string().as_str()) + } +} + +impl Serialize for Blob { + #[inline] + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + if serializer.is_human_readable() { + base64::to_string(&self.0).serialize(serializer) + } else { + serializer.serialize_bytes(&self.0) + } + } +} + +struct BlobVisitor; + +impl<'de, const L: usize> serde::de::Visitor<'de> for BlobVisitor { + type Value = Blob; + + #[inline] + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str(format!("{} bytes", L).as_str()) + } + + #[inline] + fn visit_str(self, v: &str) -> Result + where + E: serde::de::Error, + { + let b = base64::from_string(v.trim().as_bytes()).ok_or(serde::de::Error::custom("invalid base64"))?; + b.as_slice() + .try_into() + .map(|b| Blob::(b)) + .map_err(|_| serde::de::Error::invalid_length(b.len(), &self)) + } + + fn visit_bytes(self, v: &[u8]) -> Result + where + E: serde::de::Error, + { + v.try_into() + .map(|b| Blob::(b)) + .map_err(|_| serde::de::Error::invalid_length(v.len(), &self)) + } +} + +impl<'de, const L: usize> Deserialize<'de> for Blob { + #[inline] + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + if deserializer.is_human_readable() { + deserializer.deserialize_str(BlobVisitor::) + } else { + deserializer.deserialize_bytes(BlobVisitor::) + } + } +} diff --git a/src/buffer.rs b/src/buffer.rs new file mode 100644 index 0000000..939d9ae --- /dev/null +++ b/src/buffer.rs @@ -0,0 +1,703 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +use std::error::Error; +use std::fmt::{Debug, Display}; +use std::io::{Read, Write}; +use std::mem::{size_of, MaybeUninit}; + +use crate::pool::PoolFactory; +use crate::unlikely_branch; + +const OUT_OF_BOUNDS_MSG: &str = "Buffer access out of bounds"; + +pub struct OutOfBoundsError; + +impl Display for OutOfBoundsError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(OUT_OF_BOUNDS_MSG) + } +} + +impl Debug for OutOfBoundsError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + Display::fmt(self, f) + } +} + +impl Error for OutOfBoundsError {} + +impl From for std::io::Error { + fn from(_: OutOfBoundsError) -> Self { + std::io::Error::new(std::io::ErrorKind::Other, OUT_OF_BOUNDS_MSG) + } +} + +/// An I/O buffer with extensions for efficiently reading and writing various objects. +/// +/// WARNING: Structures can only be handled through raw read/write here if they are +/// tagged a Copy, meaning they are safe to just copy as raw memory. Care must also +/// be taken to ensure that access to them is safe on architectures that do not support +/// unaligned access. In vl1/protocol.rs this is accomplished by only using byte arrays +/// (including for integers) and accessing via things like u64::from_be_bytes() etc. +/// +/// Needless to say anything with non-Copy internal members or that depends on Drop to +/// not leak resources or other higher level semantics won't work here, but Rust should +/// not let you tag that as Copy in safe code. +#[derive(Clone, Copy, PartialEq, Eq)] +pub struct Buffer(usize, [u8; L]); + +impl Default for Buffer { + #[inline(always)] + fn default() -> Self { + unsafe { std::mem::zeroed() } + } +} + +impl Buffer { + pub const CAPACITY: usize = L; + + /// Create an empty zeroed buffer. + #[inline(always)] + pub fn new() -> Self { + unsafe { std::mem::zeroed() } + } + + /// Create an empty zeroed buffer on the heap without intermediate stack allocation. + /// This can be used to allocate buffers too large for the stack. + #[inline(always)] + pub fn new_boxed() -> Box { + unsafe { Box::from_raw(std::alloc::alloc_zeroed(std::alloc::Layout::new::()).cast()) } + } + + /// Create an empty buffer without internally zeroing its memory. + /// + /// This is unsafe because unwritten memory in the buffer will have undefined contents. + /// This means that some of the append_X_get_mut() functions may return mutable references to + /// undefined memory contents rather than zeroed memory. + #[inline(always)] + pub unsafe fn new_without_memzero() -> Self { + Self(0, MaybeUninit::uninit().assume_init()) + } + + pub const fn capacity(&self) -> usize { + Self::CAPACITY + } + + #[inline] + pub fn from_bytes(b: &[u8]) -> Result { + let l = b.len(); + if l <= L { + let mut tmp = Self::new(); + tmp.0 = l; + tmp.1[0..l].copy_from_slice(b); + Ok(tmp) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn as_bytes(&self) -> &[u8] { + &self.1[..self.0] + } + + #[inline(always)] + pub fn as_bytes_mut(&mut self) -> &mut [u8] { + &mut self.1[..self.0] + } + + #[inline(always)] + pub fn as_ptr(&self) -> *const u8 { + self.1.as_ptr() + } + + #[inline(always)] + pub fn as_mut_ptr(&mut self) -> *mut u8 { + self.1.as_mut_ptr() + } + + #[inline(always)] + pub fn as_bytes_after(&self, start: usize) -> Result<&[u8], OutOfBoundsError> { + if start <= self.0 { + Ok(&self.1[start..self.0]) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn as_bytes_after_mut(&mut self, start: usize) -> Result<&mut [u8], OutOfBoundsError> { + if start <= self.0 { + Ok(&mut self.1[start..self.0]) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn as_byte_range(&self, start: usize, end: usize) -> Result<&[u8], OutOfBoundsError> { + if end <= self.0 { + Ok(&self.1[start..end]) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn clear(&mut self) { + self.1[0..self.0].fill(0); + self.0 = 0; + } + + /// Load array into buffer. + /// This will panic if the array is larger than L. + #[inline(always)] + pub fn set_to(&mut self, b: &[u8]) { + let len = b.len(); + self.0 = len; + self.1[0..len].copy_from_slice(b); + } + + #[inline(always)] + pub fn len(&self) -> usize { + self.0 + } + + #[inline(always)] + pub fn is_empty(&self) -> bool { + self.0 == 0 + } + + /// Set the size of this buffer's data. + /// + /// This will panic if the specified size is larger than L. If the size is larger + /// than the current size uninitialized space will be zeroed. + #[inline] + pub fn set_size(&mut self, s: usize) { + let prev_len = self.0; + self.0 = s; + if s > prev_len { + self.1[prev_len..s].fill(0); + } + } + + /// Get a mutable reference to the entire buffer regardless of the current 'size'. + #[inline(always)] + pub unsafe fn entire_buffer_mut(&mut self) -> &mut [u8; L] { + &mut self.1 + } + + /// Set the size of the data in this buffer without checking bounds or zeroing new space. + #[inline(always)] + pub unsafe fn set_size_unchecked(&mut self, s: usize) { + self.0 = s; + } + + /// Get a byte from this buffer without checking bounds. + #[inline(always)] + pub unsafe fn get_unchecked(&self, i: usize) -> u8 { + *self.1.get_unchecked(i) + } + + /// Append a structure and return a mutable reference to its memory. + #[inline(always)] + pub fn append_struct_get_mut(&mut self) -> Result<&mut T, OutOfBoundsError> { + let ptr = self.0; + let end = ptr + size_of::(); + if end <= L { + self.0 = end; + Ok(unsafe { &mut *self.1.as_mut_ptr().add(ptr).cast() }) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + /// Append a fixed size array and return a mutable reference to its memory. + #[inline(always)] + pub fn append_bytes_fixed_get_mut(&mut self) -> Result<&mut [u8; S], OutOfBoundsError> { + let ptr = self.0; + let end = ptr + S; + if end <= L { + self.0 = end; + Ok(unsafe { &mut *self.1.as_mut_ptr().add(ptr).cast() }) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + /// Append a runtime sized array and return a mutable reference to its memory. + #[inline(always)] + pub fn append_bytes_get_mut(&mut self, s: usize) -> Result<&mut [u8], OutOfBoundsError> { + let ptr = self.0; + let end = ptr + s; + if end <= L { + self.0 = end; + Ok(&mut self.1[ptr..end]) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn append_padding(&mut self, b: u8, count: usize) -> Result<(), OutOfBoundsError> { + let ptr = self.0; + let end = ptr + count; + if end <= L { + self.0 = end; + self.1[ptr..end].fill(b); + Ok(()) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn append_bytes(&mut self, buf: &[u8]) -> Result<(), OutOfBoundsError> { + let ptr = self.0; + let end = ptr + buf.len(); + if end <= L { + self.0 = end; + self.1[ptr..end].copy_from_slice(buf); + Ok(()) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn append_bytes_fixed(&mut self, buf: &[u8; S]) -> Result<(), OutOfBoundsError> { + let ptr = self.0; + let end = ptr + S; + if end <= L { + self.0 = end; + self.1[ptr..end].copy_from_slice(buf); + Ok(()) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn append_u8(&mut self, i: u8) -> Result<(), OutOfBoundsError> { + let ptr = self.0; + if ptr < L { + self.0 = ptr + 1; + self.1[ptr] = i; + Ok(()) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn append_u16(&mut self, i: u16) -> Result<(), OutOfBoundsError> { + let i = i.to_be_bytes(); + let ptr = self.0; + let end = ptr + 2; + if end <= L { + self.0 = end; + self.1[ptr..end].copy_from_slice(&i); + Ok(()) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn append_u32(&mut self, i: u32) -> Result<(), OutOfBoundsError> { + let i = i.to_be_bytes(); + let ptr = self.0; + let end = ptr + 4; + if end <= L { + self.0 = end; + self.1[ptr..end].copy_from_slice(&i); + Ok(()) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn append_u64(&mut self, i: u64) -> Result<(), OutOfBoundsError> { + let i = i.to_be_bytes(); + let ptr = self.0; + let end = ptr + 8; + if end <= L { + self.0 = end; + self.1[ptr..end].copy_from_slice(&i); + Ok(()) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn bytes_fixed_at(&self, ptr: usize) -> Result<&[u8; S], OutOfBoundsError> { + if (ptr + S) <= self.0 { + unsafe { Ok(&*self.1.as_ptr().cast::().add(ptr).cast::<[u8; S]>()) } + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn bytes_fixed_at_mut(&mut self, ptr: usize) -> Result<&mut [u8; S], OutOfBoundsError> { + if (ptr + S) <= self.0 { + unsafe { Ok(&mut *self.1.as_mut_ptr().cast::().add(ptr).cast::<[u8; S]>()) } + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn struct_at(&self, ptr: usize) -> Result<&T, OutOfBoundsError> { + if (ptr + size_of::()) <= self.0 { + unsafe { Ok(&*self.1.as_ptr().cast::().add(ptr).cast::()) } + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn struct_mut_at(&mut self, ptr: usize) -> Result<&mut T, OutOfBoundsError> { + if (ptr + size_of::()) <= self.0 { + unsafe { Ok(&mut *self.1.as_mut_ptr().cast::().add(ptr).cast::()) } + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn u8_at(&self, ptr: usize) -> Result { + if ptr < self.0 { + Ok(self.1[ptr]) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn u16_at(&self, ptr: usize) -> Result { + let end = ptr + 2; + debug_assert!(end <= L); + if end <= self.0 { + Ok(u16::from_be_bytes(unsafe { *self.1.as_ptr().add(ptr).cast::<[u8; 2]>() })) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn u32_at(&self, ptr: usize) -> Result { + let end = ptr + 4; + debug_assert!(end <= L); + if end <= self.0 { + Ok(u32::from_be_bytes(unsafe { *self.1.as_ptr().add(ptr).cast::<[u8; 4]>() })) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn u64_at(&self, ptr: usize) -> Result { + let end = ptr + 8; + debug_assert!(end <= L); + if end <= self.0 { + Ok(u64::from_be_bytes(unsafe { *self.1.as_ptr().add(ptr).cast::<[u8; 8]>() })) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn read_struct(&self, cursor: &mut usize) -> Result<&T, OutOfBoundsError> { + let ptr = *cursor; + let end = ptr + size_of::(); + debug_assert!(end <= L); + if end <= self.0 { + *cursor = end; + unsafe { Ok(&*self.1.as_ptr().cast::().add(ptr).cast::()) } + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn read_bytes_fixed(&self, cursor: &mut usize) -> Result<&[u8; S], OutOfBoundsError> { + let ptr = *cursor; + let end = ptr + S; + debug_assert!(end <= L); + if end <= self.0 { + *cursor = end; + unsafe { Ok(&*self.1.as_ptr().cast::().add(ptr).cast::<[u8; S]>()) } + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn read_bytes(&self, l: usize, cursor: &mut usize) -> Result<&[u8], OutOfBoundsError> { + let ptr = *cursor; + let end = ptr + l; + debug_assert!(end <= L); + if end <= self.0 { + *cursor = end; + Ok(&self.1[ptr..end]) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn read_u8(&self, cursor: &mut usize) -> Result { + let ptr = *cursor; + debug_assert!(ptr < L); + if ptr < self.0 { + *cursor = ptr + 1; + Ok(self.1[ptr]) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn read_u16(&self, cursor: &mut usize) -> Result { + let ptr = *cursor; + let end = ptr + 2; + debug_assert!(end <= L); + if end <= self.0 { + *cursor = end; + Ok(u16::from_be_bytes(unsafe { *self.1.as_ptr().add(ptr).cast::<[u8; 2]>() })) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn read_u32(&self, cursor: &mut usize) -> Result { + let ptr = *cursor; + let end = ptr + 4; + debug_assert!(end <= L); + if end <= self.0 { + *cursor = end; + Ok(u32::from_be_bytes(unsafe { *self.1.as_ptr().add(ptr).cast::<[u8; 4]>() })) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } + + #[inline(always)] + pub fn read_u64(&self, cursor: &mut usize) -> Result { + let ptr = *cursor; + let end = ptr + 8; + debug_assert!(end <= L); + if end <= self.0 { + *cursor = end; + Ok(u64::from_be_bytes(unsafe { *self.1.as_ptr().add(ptr).cast::<[u8; 8]>() })) + } else { + unlikely_branch(); + Err(OutOfBoundsError) + } + } +} + +impl Write for Buffer { + #[inline(always)] + fn write(&mut self, buf: &[u8]) -> std::io::Result { + let ptr = self.0; + let end = ptr + buf.len(); + if end <= L { + self.0 = end; + self.1[ptr..end].copy_from_slice(buf); + Ok(buf.len()) + } else { + unlikely_branch(); + Err(std::io::Error::new(std::io::ErrorKind::Other, OUT_OF_BOUNDS_MSG)) + } + } + + #[inline(always)] + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +impl AsRef<[u8]> for Buffer { + #[inline(always)] + fn as_ref(&self) -> &[u8] { + self.as_bytes() + } +} + +impl AsMut<[u8]> for Buffer { + #[inline(always)] + fn as_mut(&mut self) -> &mut [u8] { + self.as_bytes_mut() + } +} + +impl From<[u8; L]> for Buffer { + #[inline(always)] + fn from(a: [u8; L]) -> Self { + Self(L, a) + } +} + +impl From<&[u8; L]> for Buffer { + #[inline(always)] + fn from(a: &[u8; L]) -> Self { + Self(L, *a) + } +} + +/// Implements std::io::Read for a buffer and a cursor. +pub struct BufferReader<'a, 'b, const L: usize>(&'a Buffer, &'b mut usize); + +impl<'a, 'b, const L: usize> BufferReader<'a, 'b, L> { + #[inline(always)] + pub fn new(b: &'a Buffer, cursor: &'b mut usize) -> Self { + Self(b, cursor) + } +} + +impl<'a, 'b, const L: usize> Read for BufferReader<'a, 'b, L> { + #[inline(always)] + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + buf.copy_from_slice(self.0.read_bytes(buf.len(), self.1)?); + Ok(buf.len()) + } +} + +pub struct PooledBufferFactory; + +impl PooledBufferFactory { + #[inline(always)] + pub fn new() -> Self { + Self {} + } +} + +impl PoolFactory> for PooledBufferFactory { + #[inline(always)] + fn create(&self) -> Buffer { + Buffer::new() + } + + #[inline(always)] + fn reset(&self, obj: &mut Buffer) { + obj.clear(); + } +} + +#[cfg(test)] +mod tests { + use super::Buffer; + + #[test] + fn buffer_basic_u64() { + let mut b = Buffer::<8>::new(); + assert_eq!(b.len(), 0); + assert!(b.is_empty()); + assert!(b.append_u64(1234).is_ok()); + assert_eq!(b.len(), 8); + assert!(!b.is_empty()); + assert_eq!(b.read_u64(&mut 0).unwrap(), 1234); + b.clear(); + assert_eq!(b.len(), 0); + assert!(b.is_empty()); + } + + #[test] + fn buffer_basic_u32() { + let mut b = Buffer::<4>::new(); + assert_eq!(b.len(), 0); + assert!(b.is_empty()); + assert!(b.append_u32(1234).is_ok()); + assert_eq!(b.len(), 4); + assert!(!b.is_empty()); + assert_eq!(b.read_u32(&mut 0).unwrap(), 1234); + b.clear(); + assert_eq!(b.len(), 0); + assert!(b.is_empty()); + } + + #[test] + fn buffer_basic_u16() { + let mut b = Buffer::<2>::new(); + assert_eq!(b.len(), 0); + assert!(b.is_empty()); + assert!(b.append_u16(1234).is_ok()); + assert_eq!(b.len(), 2); + assert!(!b.is_empty()); + assert_eq!(b.read_u16(&mut 0).unwrap(), 1234); + b.clear(); + assert_eq!(b.len(), 0); + assert!(b.is_empty()); + } + + #[test] + fn buffer_basic_u8() { + let mut b = Buffer::<1>::new(); + assert_eq!(b.len(), 0); + assert!(b.is_empty()); + assert!(b.append_u8(128).is_ok()); + assert_eq!(b.len(), 1); + assert!(!b.is_empty()); + assert_eq!(b.read_u8(&mut 0).unwrap(), 128); + b.clear(); + assert_eq!(b.len(), 0); + assert!(b.is_empty()); + } + + #[test] + fn buffer_sizing() { + const SIZE: usize = 100; + + for _ in 0..1000 { + let v = [0u8; SIZE]; + let mut b = Buffer::::new(); + assert!(b.append_bytes(&v).is_ok()); + assert_eq!(b.len(), SIZE); + b.set_size(10); + assert_eq!(b.len(), 10); + unsafe { + b.set_size_unchecked(8675309); + } + assert_eq!(b.len(), 8675309); + } + } +} diff --git a/src/cast.rs b/src/cast.rs new file mode 100644 index 0000000..38406aa --- /dev/null +++ b/src/cast.rs @@ -0,0 +1,36 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +use std::any::TypeId; +use std::mem::size_of; + +/// Returns true if two types are in fact the same type. +#[inline(always)] +pub fn same_type() -> bool { + TypeId::of::() == TypeId::of::() && size_of::() == size_of::() +} + +/// Cast a reference if the types are equal, such as from a specific type to a generic that it implements. +#[inline(always)] +pub fn cast_ref(u: &U) -> Option<&V> { + if same_type::() { + Some(unsafe { std::mem::transmute::<&U, &V>(u) }) + } else { + None + } +} + +/// Cast a reference if the types are equal, such as from a specific type to a generic that it implements. +#[inline(always)] +pub fn cast_mut(u: &mut U) -> Option<&mut V> { + if same_type::() { + Some(unsafe { std::mem::transmute::<&mut U, &mut V>(u) }) + } else { + None + } +} diff --git a/src/dictionary.rs b/src/dictionary.rs new file mode 100644 index 0000000..393b1c8 --- /dev/null +++ b/src/dictionary.rs @@ -0,0 +1,209 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +use std::collections::BTreeMap; +use std::io::Write; + +use crate::hex; + +const BOOL_TRUTH: &str = "1tTyY"; + +/// Dictionary is an extremely simple key=value serialization format. +/// +/// It's designed for extreme parsing simplicity and is human readable if keys and values are strings. +/// It also supports binary keys and values which will be minimally escaped but render the result not +/// entirely human readable. Keys are serialized in natural sort order so the result can be consistently +/// checksummed or hashed. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Dictionary(pub(crate) BTreeMap>); + +fn write_escaped(mut b: &[u8], w: &mut W) -> std::io::Result<()> { + while !b.is_empty() { + match b[0] { + 0 => { + w.write_all(&[b'\\', b'0'])?; + } + b'\n' => { + w.write_all(&[b'\\', b'n'])?; + } + b'\r' => { + w.write_all(&[b'\\', b'r'])?; + } + b'=' => { + w.write_all(&[b'\\', b'e'])?; + } + b'\\' => { + w.write_all(&[b'\\', b'\\'])?; + } + _ => { + w.write_all(&b[..1])?; + } + } + b = &b[1..]; + } + Ok(()) +} + +fn append_printable(s: &mut String, b: &[u8]) { + for c in b { + let c = *c as char; + if c.is_alphanumeric() || c.is_whitespace() { + s.push(c); + } else { + s.push('\\'); + s.push('x'); + s.push(hex::HEX_CHARS[((c as u8) >> 4) as usize] as char); + s.push(hex::HEX_CHARS[((c as u8) & 0xf) as usize] as char); + } + } +} + +impl Dictionary { + pub fn new() -> Self { + Self(BTreeMap::new()) + } + + pub fn clear(&mut self) { + self.0.clear() + } + + #[inline(always)] + pub fn len(&self) -> usize { + self.0.len() + } + + #[inline(always)] + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } + + pub fn get_str(&self, k: &str) -> Option<&str> { + self.0.get(k).and_then(|v| std::str::from_utf8(v.as_slice()).ok()) + } + + pub fn get_bytes(&self, k: &str) -> Option<&[u8]> { + self.0.get(k).map(|v| v.as_slice()) + } + + pub fn get_u64(&self, k: &str) -> Option { + self.get_str(k).and_then(|s| u64::from_str_radix(s, 16).ok()) + } + + pub fn get_i64(&self, k: &str) -> Option { + self.get_str(k).and_then(|s| i64::from_str_radix(s, 16).ok()) + } + + pub fn get_bool(&self, k: &str) -> Option { + self.0 + .get(k) + .and_then(|v| v.first().map_or(Some(false), |c| Some(BOOL_TRUTH.contains(*c as char)))) + } + + pub fn set_str(&mut self, k: &str, v: &str) { + let _ = self.0.insert(String::from(k), v.as_bytes().to_vec()); + } + + pub fn set_u64(&mut self, k: &str, v: u64) { + let _ = self.0.insert(String::from(k), hex::to_vec_u64(v, true)); + } + + pub fn set_bytes(&mut self, k: &str, v: Vec) { + let _ = self.0.insert(String::from(k), v); + } + + pub fn set_bool(&mut self, k: &str, v: bool) { + let _ = self.0.insert( + String::from(k), + vec![if v { + b'1' + } else { + b'0' + }], + ); + } + + pub fn write_to(&self, w: &mut W) -> std::io::Result<()> { + for kv in self.0.iter() { + write_escaped(kv.0.as_bytes(), w)?; + w.write_all(&[b'='])?; + write_escaped(kv.1.as_slice(), w)?; + w.write_all(&[b'\n'])?; + } + Ok(()) + } + + pub fn to_bytes(&self) -> Vec { + let mut b: Vec = Vec::with_capacity(32 * self.0.len()); + let _ = self.write_to(&mut b); + b + } + + pub fn from_bytes(b: &[u8]) -> Option { + let mut d = Dictionary::new(); + let mut kv: [Vec; 2] = [Vec::new(), Vec::new()]; + let mut state = 0; + let mut escape = false; + for c in b { + let c = *c; + if escape { + escape = false; + kv[state].push(match c { + b'0' => 0, + b'n' => b'\n', + b'r' => b'\r', + b'e' => b'=', + _ => c, // =, \, and escapes before other characters are unnecessary but not errors + }); + } else if c == b'\\' { + escape = true; + } else if c == b'=' { + if state != 0 { + return None; + } + state = 1; + } else if c == b'\n' { + if state != 1 { + return None; + } + state = 0; + if !kv[0].is_empty() + && String::from_utf8(kv[0].clone()).map_or(true, |key| { + d.0.insert(key, kv[1].clone()); + false + }) + { + return None; + } + kv[0].clear(); + kv[1].clear(); + } else if c != b'\r' { + kv[state].push(c); + } + } + Some(d) + } + + pub fn iter(&self) -> impl Iterator)> { + self.0.iter() + } +} + +impl ToString for Dictionary { + /// Get the dictionary in an always readable format with non-printable characters replaced by '\xXX'. + /// This is not a serializable output that can be re-imported. Use write_to() for that. + fn to_string(&self) -> String { + let mut s = String::new(); + for kv in self.0.iter() { + append_printable(&mut s, kv.0.as_bytes()); + s.push('='); + append_printable(&mut s, kv.1.as_slice()); + s.push('\n'); + } + s + } +} diff --git a/src/error.rs b/src/error.rs new file mode 100644 index 0000000..723099f --- /dev/null +++ b/src/error.rs @@ -0,0 +1,44 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +use std::error::Error; +use std::fmt::{Debug, Display}; + +pub struct InvalidFormatError; + +impl Display for InvalidFormatError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("InvalidFormatError") + } +} + +impl Debug for InvalidFormatError { + #[inline(always)] + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + ::fmt(self, f) + } +} + +impl Error for InvalidFormatError {} + +pub struct InvalidParameterError(pub &'static str); + +impl Display for InvalidParameterError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "InvalidParameterError: {}", self.0) + } +} + +impl Debug for InvalidParameterError { + #[inline(always)] + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + ::fmt(self, f) + } +} + +impl Error for InvalidParameterError {} diff --git a/src/exitcode.rs b/src/exitcode.rs new file mode 100644 index 0000000..7cb96c3 --- /dev/null +++ b/src/exitcode.rs @@ -0,0 +1,22 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +// These were taken from BSD sysexits.h to provide some standard for process exit codes. + +pub const OK: i32 = 0; + +pub const ERR_USAGE: i32 = 64; +pub const ERR_DATA_FORMAT: i32 = 65; +pub const ERR_NO_INPUT: i32 = 66; +pub const ERR_SERVICE_UNAVAILABLE: i32 = 69; +pub const ERR_INTERNAL: i32 = 70; +pub const ERR_OSERR: i32 = 71; +pub const ERR_OSFILE: i32 = 72; +pub const ERR_IOERR: i32 = 74; +pub const ERR_NOPERM: i32 = 77; +pub const ERR_CONFIG: i32 = 78; diff --git a/src/gate.rs b/src/gate.rs new file mode 100644 index 0000000..e3a8aa6 --- /dev/null +++ b/src/gate.rs @@ -0,0 +1,35 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +/// Boolean rate limiter with normal (non-atomic) semantics. +#[repr(transparent)] +pub struct IntervalGate(i64); + +impl Default for IntervalGate { + #[inline(always)] + fn default() -> Self { + Self(crate::NEVER_HAPPENED_TICKS) + } +} + +impl IntervalGate { + #[inline(always)] + pub fn new(initial_ts: i64) -> Self { + Self(initial_ts) + } + + #[inline(always)] + pub fn gate(&mut self, time: i64) -> bool { + if (time - self.0) >= FREQ { + self.0 = time; + true + } else { + false + } + } +} diff --git a/src/hex.rs b/src/hex.rs new file mode 100644 index 0000000..0fbad1d --- /dev/null +++ b/src/hex.rs @@ -0,0 +1,124 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +pub const HEX_CHARS: [u8; 16] = [ + b'0', b'1', b'2', b'3', b'4', b'5', b'6', b'7', b'8', b'9', b'a', b'b', b'c', b'd', b'e', b'f', +]; + +/// Encode a byte slice to a hexadecimal string. +pub fn to_string(b: &[u8]) -> String { + let mut s = String::with_capacity(b.len() * 2); + s.reserve(b.len() * 2); + for c in b { + let x = *c as usize; + s.push(HEX_CHARS[x >> 4] as char); + s.push(HEX_CHARS[x & 0xf] as char); + } + s +} + +/// Encode an unsigned 64-bit value as a hexadecimal string. +pub fn to_string_u64(mut i: u64, skip_leading_zeroes: bool) -> String { + let mut s = String::with_capacity(16); + for _ in 0..16 { + let ii = i >> 60; + if ii != 0 || !s.is_empty() || !skip_leading_zeroes { + s.push(HEX_CHARS[ii as usize] as char); + } + i = i.wrapping_shl(4); + } + s +} + +/// Encode an unsigned 64-bit value as a hexadecimal ASCII string. +pub fn to_vec_u64(mut i: u64, skip_leading_zeroes: bool) -> Vec { + let mut s = Vec::with_capacity(16); + for _ in 0..16 { + let ii = i >> 60; + if ii != 0 || !s.is_empty() || !skip_leading_zeroes { + s.push(HEX_CHARS[ii as usize]); + } + i = i.wrapping_shl(4); + } + s +} + +/// Decode a hex string, ignoring all non-hexadecimal characters. +pub fn from_string(s: &str) -> Vec { + let mut b: Vec = Vec::with_capacity((s.len() / 2) + 1); + let mut byte = 0_u8; + let mut have_8: bool = false; + for cc in s.as_bytes() { + let c = *cc; + if (48..=57).contains(&c) { + byte = (byte.wrapping_shl(4)) | (c - 48); + if have_8 { + b.push(byte); + } + have_8 = !have_8; + } else if (65..=70).contains(&c) { + byte = (byte.wrapping_shl(4)) | (c - 55); + if have_8 { + b.push(byte); + } + have_8 = !have_8; + } else if (97..=102).contains(&c) { + byte = (byte.wrapping_shl(4)) | (c - 87); + if have_8 { + b.push(byte); + } + have_8 = !have_8; + } + } + b +} + +pub fn from_string_u64(s: &str) -> u64 { + let mut n = 0u64; + let mut byte = 0_u8; + let mut have_8: bool = false; + for cc in s.as_bytes() { + let c = *cc; + if (48..=57).contains(&c) { + byte = (byte.wrapping_shl(4)) | (c - 48); + if have_8 { + n = n.wrapping_shl(8); + n |= byte as u64; + } + have_8 = !have_8; + } else if (65..=70).contains(&c) { + byte = (byte.wrapping_shl(4)) | (c - 55); + if have_8 { + n = n.wrapping_shl(8); + n |= byte as u64; + } + have_8 = !have_8; + } else if (97..=102).contains(&c) { + byte = (byte.wrapping_shl(4)) | (c - 87); + if have_8 { + n = n.wrapping_shl(8); + n |= byte as u64; + } + have_8 = !have_8; + } + } + n +} + +/// Encode bytes from 'b' into hex characters in 'dest' and return the number of hex characters written. +/// This will panic if the destination slice is smaller than twice the length of the source. +pub fn to_hex_bytes(b: &[u8], dest: &mut [u8]) -> usize { + let mut j = 0; + for c in b { + let x = *c as usize; + dest[j] = HEX_CHARS[x >> 4]; + dest[j + 1] = HEX_CHARS[x & 0xf]; + j += 2; + } + j +} diff --git a/src/inetaddress.rs b/src/inetaddress.rs new file mode 100644 index 0000000..d1861b2 --- /dev/null +++ b/src/inetaddress.rs @@ -0,0 +1,1191 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +use std::cmp::Ordering; +use std::hash::{Hash, Hasher}; +use std::mem::{size_of, transmute_copy, zeroed, MaybeUninit}; +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6, ToSocketAddrs}; +use std::ptr::{copy_nonoverlapping, null, slice_from_raw_parts, write_bytes}; +use std::str::FromStr; + +use serde::{Deserialize, Deserializer, Serialize, Serializer}; + +use crate::error::InvalidParameterError; +use crate::tofrombytes::ToFromBytes; + +#[allow(non_camel_case_types)] +#[cfg(windows)] +type sockaddr = winapi::shared::ws2def::SOCKADDR; + +#[allow(non_camel_case_types)] +#[cfg(windows)] +type sockaddr_in = winapi::shared::ws2def::SOCKADDR_IN; + +#[allow(non_camel_case_types)] +#[cfg(windows)] +type sockaddr_in6 = winapi::shared::ws2ipdef::SOCKADDR_IN6; + +#[allow(non_camel_case_types)] +#[cfg(windows)] +type sockaddr_storage = winapi::shared::ws2def::SOCKADDR_STORAGE; + +#[allow(non_camel_case_types)] +#[cfg(windows)] +type in6_addr = winapi::shared::in6addr::in6_addr; + +#[allow(non_camel_case_types)] +#[cfg(not(windows))] +type sockaddr = libc::sockaddr; + +#[allow(non_camel_case_types)] +#[cfg(not(windows))] +type sockaddr_in = libc::sockaddr_in; + +#[allow(non_camel_case_types)] +#[cfg(not(windows))] +type sockaddr_in6 = libc::sockaddr_in6; + +#[allow(non_camel_case_types)] +#[cfg(not(windows))] +type sockaddr_storage = libc::sockaddr_storage; + +#[allow(non_camel_case_types)] +#[cfg(not(windows))] +type in6_addr = libc::in6_addr; + +#[cfg(all(not(target_os = "windows"), not(target_os = "linux")))] +pub type AddressFamilyType = u8; + +#[cfg(target_os = "linux")] +pub type AddressFamilyType = u16; + +#[cfg(windows)] +pub type AddressFamilyType = winapi::ctypes::c_int; + +#[cfg(windows)] +pub const AF_INET: AddressFamilyType = winapi::shared::ws2def::AF_INET as AddressFamilyType; + +#[cfg(windows)] +pub const AF_INET6: AddressFamilyType = winapi::shared::ws2def::AF_INET6 as AddressFamilyType; + +#[cfg(not(windows))] +pub const AF_INET: AddressFamilyType = libc::AF_INET as AddressFamilyType; + +#[cfg(not(windows))] +pub const AF_INET6: AddressFamilyType = libc::AF_INET6 as AddressFamilyType; + +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +#[repr(u8)] +pub enum IpScope { + None = 0, + Multicast = 1, + Loopback = 2, + PseudoPrivate = 3, + Global = 4, + LinkLocal = 5, + Shared = 6, + Private = 7, +} + +#[cfg(windows)] +#[inline(always)] +fn get_s_addr(sa: &sockaddr_in) -> u32 { + unsafe { *sa.sin_addr.S_un.S_addr() } +} + +#[cfg(not(windows))] +#[inline(always)] +fn get_s_addr(sa: &sockaddr_in) -> u32 { + sa.sin_addr.s_addr +} + +#[cfg(windows)] +#[inline(always)] +fn get_s6_addr(&sa: &sockaddr_in6) -> &[u8; 16] { + unsafe { &*(sa.sin6_addr.u.Byte() as *const [u8; 16]) } +} + +#[cfg(not(windows))] +#[inline(always)] +fn get_s6_addr(&sa: &sockaddr_in6) -> &[u8; 16] { + unsafe { &*(&sa.sin6_addr.s6_addr as *const [u8; 16]) } +} + +/// An IPv4 or IPv6 socket address that directly encapsulates C sockaddr types. +/// +/// The ZeroTier core uses this in preference to std::net stuff so this can be +/// directly used via the C API or with C socket I/O functions. +/// +/// This supports into() and from() the std::net types and also has a custom +/// serde serializer that serializes it in ZT protocol binary form for binary +/// formats and canonical ZT string form for string formats. +/// +/// Unfortunately this is full of unsafe because it's a union, but the code is +/// not complex and doesn't allocate anything. +#[repr(C)] +pub union InetAddress { + sa: sockaddr, + sin: sockaddr_in, + sin6: sockaddr_in6, + ss: sockaddr_storage, // some external code may expect the struct to be this full length +} + +impl ToSocketAddrs for InetAddress { + type Iter = std::iter::Once; + + #[inline(always)] + fn to_socket_addrs(&self) -> std::io::Result { + self.try_into().map_or_else( + |_| Err(std::io::Error::new(std::io::ErrorKind::Other, "not an IP address")), + |sa| Ok(std::iter::once(sa)), + ) + } +} + +impl TryInto for InetAddress { + type Error = InvalidParameterError; + + #[inline(always)] + fn try_into(self) -> Result { + (&self).try_into() + } +} + +impl TryInto for &InetAddress { + type Error = InvalidParameterError; + + #[inline(always)] + fn try_into(self) -> Result { + unsafe { + match self.sa.sa_family as AddressFamilyType { + AF_INET => Ok(IpAddr::V4(Ipv4Addr::from(get_s_addr(&self.sin).to_ne_bytes()))), + AF_INET6 => Ok(IpAddr::V6(Ipv6Addr::from(*get_s6_addr(&self.sin6)))), + _ => Err(InvalidParameterError("not an IP address")), + } + } + } +} + +impl TryInto for InetAddress { + type Error = InvalidParameterError; + + #[inline(always)] + fn try_into(self) -> Result { + (&self).try_into() + } +} + +impl TryInto for &InetAddress { + type Error = InvalidParameterError; + + #[inline(always)] + fn try_into(self) -> Result { + match unsafe { self.sa.sa_family } as AddressFamilyType { + AF_INET => Ok(Ipv4Addr::from(unsafe { get_s_addr(&self.sin).to_ne_bytes() })), + _ => Err(InvalidParameterError("not an IPv4 address")), + } + } +} + +impl TryInto for InetAddress { + type Error = InvalidParameterError; + + #[inline(always)] + fn try_into(self) -> Result { + (&self).try_into() + } +} + +impl TryInto for &InetAddress { + type Error = InvalidParameterError; + + #[inline(always)] + fn try_into(self) -> Result { + match unsafe { self.sa.sa_family } as AddressFamilyType { + AF_INET6 => Ok(Ipv6Addr::from(*get_s6_addr(unsafe { &self.sin6 }))), + _ => Err(InvalidParameterError("not an IPv6 address")), + } + } +} + +impl TryInto for InetAddress { + type Error = InvalidParameterError; + + #[inline(always)] + fn try_into(self) -> Result { + (&self).try_into() + } +} + +impl TryInto for &InetAddress { + type Error = InvalidParameterError; + + #[inline(always)] + fn try_into(self) -> Result { + unsafe { + match self.sa.sa_family as AddressFamilyType { + AF_INET => Ok(SocketAddr::V4(SocketAddrV4::new( + Ipv4Addr::from(get_s_addr(&self.sin).to_ne_bytes()), + u16::from_be(self.sin.sin_port), + ))), + AF_INET6 => Ok(SocketAddr::V6(SocketAddrV6::new( + Ipv6Addr::from(*get_s6_addr(&self.sin6)), + u16::from_be(self.sin6.sin6_port), + 0, + 0, + ))), + _ => Err(InvalidParameterError("not an IP address")), + } + } + } +} + +impl TryInto for InetAddress { + type Error = InvalidParameterError; + + #[inline(always)] + fn try_into(self) -> Result { + (&self).try_into() + } +} + +impl TryInto for &InetAddress { + type Error = InvalidParameterError; + + #[inline(always)] + fn try_into(self) -> Result { + unsafe { + match self.sa.sa_family as AddressFamilyType { + AF_INET => Ok(SocketAddrV4::new( + Ipv4Addr::from(get_s_addr(&self.sin).to_ne_bytes()), + u16::from_be(self.sin.sin_port), + )), + _ => Err(InvalidParameterError("not an IPv4 address")), + } + } + } +} + +impl TryInto for InetAddress { + type Error = InvalidParameterError; + + #[inline(always)] + fn try_into(self) -> Result { + (&self).try_into() + } +} + +impl TryInto for &InetAddress { + type Error = InvalidParameterError; + + #[inline] + fn try_into(self) -> Result { + unsafe { + match self.sa.sa_family as AddressFamilyType { + AF_INET6 => Ok(SocketAddrV6::new( + Ipv6Addr::from(*get_s6_addr(&self.sin6)), + u16::from_be(self.sin6.sin6_port), + 0, + 0, + )), + _ => Err(InvalidParameterError("not an IPv6 address")), + } + } + } +} + +impl From<&IpAddr> for InetAddress { + #[inline] + fn from(ip: &IpAddr) -> Self { + match ip { + IpAddr::V4(ip4) => Self::from(ip4), + IpAddr::V6(ip6) => Self::from(ip6), + } + } +} + +impl From for InetAddress { + #[inline(always)] + fn from(ip: IpAddr) -> Self { + Self::from(&ip) + } +} + +impl From<&Ipv4Addr> for InetAddress { + #[inline(always)] + fn from(ip4: &Ipv4Addr) -> Self { + Self::from_ip_port(&ip4.octets(), 0) + } +} + +impl From for InetAddress { + #[inline(always)] + fn from(ip4: Ipv4Addr) -> Self { + Self::from_ip_port(&ip4.octets(), 0) + } +} + +impl From<&Ipv6Addr> for InetAddress { + #[inline(always)] + fn from(ip6: &Ipv6Addr) -> Self { + Self::from_ip_port(&ip6.octets(), 0) + } +} + +impl From for InetAddress { + #[inline(always)] + fn from(ip6: Ipv6Addr) -> Self { + Self::from_ip_port(&ip6.octets(), 0) + } +} + +impl From<&SocketAddr> for InetAddress { + #[inline] + fn from(sa: &SocketAddr) -> Self { + match sa { + SocketAddr::V4(sa4) => Self::from(sa4), + SocketAddr::V6(sa6) => Self::from(sa6), + } + } +} + +impl From for InetAddress { + #[inline(always)] + fn from(sa: SocketAddr) -> Self { + Self::from(&sa) + } +} + +impl From<&SocketAddrV4> for InetAddress { + #[inline(always)] + fn from(sa: &SocketAddrV4) -> Self { + Self::from_ip_port(&sa.ip().octets(), sa.port()) + } +} + +impl From for InetAddress { + #[inline(always)] + fn from(sa: SocketAddrV4) -> Self { + Self::from_ip_port(&sa.ip().octets(), sa.port()) + } +} + +impl From<&SocketAddrV6> for InetAddress { + #[inline(always)] + fn from(sa: &SocketAddrV6) -> Self { + Self::from_ip_port(&sa.ip().octets(), sa.port()) + } +} + +impl From for InetAddress { + #[inline(always)] + fn from(sa: SocketAddrV6) -> Self { + Self::from_ip_port(&sa.ip().octets(), sa.port()) + } +} + +impl Clone for InetAddress { + #[inline(always)] + fn clone(&self) -> Self { + unsafe { transmute_copy(self) } + } +} + +impl Default for InetAddress { + #[inline(always)] + fn default() -> Self { + Self::new() + } +} + +impl std::fmt::Debug for InetAddress { + #[inline(always)] + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.to_string().as_str()) + } +} + +impl Serialize for InetAddress { + #[inline] + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + if serializer.is_human_readable() { + serializer.serialize_str(self.to_string().as_str()) + } else { + serializer.serialize_bytes(self.to_bytes_on_stack::<32>().as_bytes()) + } + } +} + +struct InetAddressVisitor; + +impl<'de> serde::de::Visitor<'de> for InetAddressVisitor { + type Value = InetAddress; + + #[inline] + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("InetAddress") + } + + #[inline] + fn visit_bytes(self, mut v: &[u8]) -> Result + where + E: serde::de::Error, + { + InetAddress::read_bytes(&mut v).map_err(|e| serde::de::Error::custom(e.to_string())) + } + + #[inline] + fn visit_str(self, v: &str) -> Result + where + E: serde::de::Error, + { + InetAddress::from_str(v).map_err(|e| E::custom(e.to_string())) + } +} + +impl<'de> Deserialize<'de> for InetAddress { + #[inline] + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + if deserializer.is_human_readable() { + deserializer.deserialize_str(InetAddressVisitor) + } else { + deserializer.deserialize_bytes(InetAddressVisitor) + } + } +} + +impl InetAddress { + /// Get a new zero/nil InetAddress. + #[inline(always)] + pub fn new() -> InetAddress { + unsafe { zeroed() } + } + + /// Construct from IP and port. + /// If the IP is not either 4 or 16 bytes in length, a nil/0 InetAddress is returned. + #[inline] + pub fn from_ip_port(ip: &[u8], port: u16) -> InetAddress { + unsafe { + let mut c = MaybeUninit::::uninit().assume_init(); // gets zeroed in set() + c.set(ip, port); + c + } + } + + /// Zero the contents of this InetAddress. + #[inline(always)] + pub fn zero(&mut self) { + unsafe { write_bytes((self as *mut Self).cast::(), 0, size_of::()) }; + } + + /// Get an instance of 127.0.0.1/port + pub fn ipv4_loopback(port: u16) -> InetAddress { + let mut addr = Self::new(); + addr.sin.sin_family = AF_INET.into(); + addr.sin.sin_port = port.to_be().into(); + #[cfg(not(windows))] + { + addr.sin.sin_addr.s_addr = 0x7f000001_u32.to_be(); + } + #[cfg(windows)] + unsafe { + *addr.sin.sin_addr.S_un.S_addr_mut() = (0x7f000001 as u32).to_be(); + } + addr + } + + /// Get an instance of 0.0.0.0/0 + pub fn ipv4_any() -> InetAddress { + let mut addr = Self::new(); + addr.sin.sin_family = AF_INET.into(); + addr + } + + /// Get an instance of ::1/port + pub fn ipv6_loopback(port: u16) -> InetAddress { + let mut addr = Self::new(); + addr.sin6.sin6_family = AF_INET6.into(); + addr.sin6.sin6_port = port.to_be().into(); + unsafe { + *((&mut (addr.sin6.sin6_addr) as *mut in6_addr).cast::().offset(15)) = 1; + } + addr + } + + /// Get an instance of ::0/0 + pub fn ipv6_any() -> InetAddress { + let mut addr = Self::new(); + addr.sin6.sin6_family = AF_INET6.into(); + addr + } + + /// Returns true if this InetAddress is the nil value (zero). + #[inline(always)] + pub fn is_nil(&self) -> bool { + unsafe { self.sa.sa_family == 0 } + } + + /// Check if this is an IPv4 address. + #[inline(always)] + pub fn is_ipv4(&self) -> bool { + unsafe { self.sa.sa_family as AddressFamilyType == AF_INET } + } + + /// Check if this is an IPv6 address. + #[inline(always)] + pub fn is_ipv6(&self) -> bool { + unsafe { self.sa.sa_family as AddressFamilyType == AF_INET6 } + } + + /// Check if this is either an IPv4 or an IPv6 address. + #[inline(always)] + pub fn is_ip(&self) -> bool { + let family = unsafe { self.sa.sa_family } as AddressFamilyType; + family == AF_INET || family == AF_INET6 + } + + /// Get the address family of this InetAddress: AF_INET, AF_INET6, or 0 if uninitialized. + #[inline(always)] + pub fn family(&self) -> AddressFamilyType { + unsafe { self.sa.sa_family as AddressFamilyType } + } + + /// Get a pointer to the C "sockaddr" structure and the size of the returned structure in bytes. + /// This is useful for interacting with C-level socket APIs. This returns a null pointer if + /// the address is not initialized. + #[inline(always)] + pub fn c_sockaddr(&self) -> (*const (), usize) { + unsafe { + match self.sa.sa_family as AddressFamilyType { + AF_INET => ((&self.sin as *const sockaddr_in).cast(), size_of::()), + AF_INET6 => ((&self.sin6 as *const sockaddr_in6).cast(), size_of::()), + _ => (null(), 0), + } + } + } + + /// Set the IP and port of this InetAddress. + /// Whether this is IPv4 or IPv6 is inferred from the size of ip[], which must be + /// either 4 or 16 bytes. The family (AF_INET or AF_INET6) is returned, or zero on + /// failure. + pub fn set>(&mut self, ip: T, port: u16) -> AddressFamilyType { + self.zero(); + let port = port.to_be(); + let ip2 = ip.as_ref(); + unsafe { + if ip2.len() == 4 { + self.sin.sin_family = AF_INET.into(); + self.sin.sin_port = port.into(); + #[cfg(windows)] + { + self.sin.sin_addr.S_un.S_un_b_mut().s_b1 = ip2[0]; + self.sin.sin_addr.S_un.S_un_b_mut().s_b2 = ip2[1]; + self.sin.sin_addr.S_un.S_un_b_mut().s_b3 = ip2[2]; + self.sin.sin_addr.S_un.S_un_b_mut().s_b4 = ip2[3]; + } + #[cfg(not(windows))] + { + copy_nonoverlapping(ip2.as_ptr(), (&mut self.sin.sin_addr.s_addr as *mut u32).cast::(), 4); + } + AF_INET + } else if ip2.len() == 16 { + self.sin6.sin6_family = AF_INET6.into(); + self.sin6.sin6_port = port.into(); + copy_nonoverlapping(ip2.as_ptr(), (&mut self.sin6.sin6_addr as *mut in6_addr).cast::(), 16); + AF_INET6 + } else { + 0 + } + } + } + + /// Get raw IP bytes, with length dependent on address family (4 or 16). + pub fn ip_bytes(&self) -> &[u8] { + unsafe { + match self.sa.sa_family as AddressFamilyType { + AF_INET => { + #[cfg(windows)] + { + &*(self.sin.sin_addr.S_un.S_addr() as *const u32).cast::<[u8; 4]>() + } + #[cfg(not(windows))] + { + &*(&self.sin.sin_addr.s_addr as *const u32).cast::<[u8; 4]>() + } + } + AF_INET6 => &*(&self.sin6.sin6_addr as *const in6_addr).cast::<[u8; 16]>(), + _ => &[], + } + } + } + + /// Get an efficient local lookup key from this IP address. + pub fn key(&self) -> (u128, u16) { + unsafe { + match self.sa.sa_family as AddressFamilyType { + AF_INET => { + #[cfg(windows)] + { + (*(self.sin.sin_addr.S_un.S_addr() as *const u32) as u128, self.sin.sin_port as u16) + } + #[cfg(not(windows))] + { + (*(&self.sin.sin_addr.s_addr as *const u32) as u128, self.sin.sin_port as u16) + } + } + AF_INET6 => ( + u128::from_ne_bytes(*(&self.sin6.sin6_addr as *const in6_addr).cast::<[u8; 16]>()), + self.sin6.sin6_port as u16, + ), + _ => (0, 0), + } + } + } + + /// Get a Rust stdlib SocketAddr structure from this InetAddress. + pub fn to_socketaddr(&self) -> Option { + unsafe { + match self.sa.sa_family as AddressFamilyType { + AF_INET => Some(SocketAddr::V4(SocketAddrV4::new( + Ipv4Addr::from(get_s_addr(&self.sin).to_ne_bytes()), + u16::from_be(self.sin.sin_port), + ))), + AF_INET6 => Some(SocketAddr::V6(SocketAddrV6::new( + Ipv6Addr::from(*get_s6_addr(&self.sin6)), + u16::from_be(self.sin6.sin6_port), + 0, + 0, + ))), + _ => None, + } + } + } + + /// Get the IP port for this InetAddress. + #[inline] + pub fn port(&self) -> u16 { + unsafe { + u16::from_be(match self.sa.sa_family as AddressFamilyType { + AF_INET => self.sin.sin_port, + AF_INET6 => self.sin6.sin6_port, + _ => 0, + }) + } + } + + /// Set the IP port. + /// + /// This does nothing on uninitialized InetAddress objects. An address must first + /// be initialized with an IP to select the correct address type. + pub fn set_port(&mut self, port: u16) { + let port = port.to_be(); + unsafe { + match self.sa.sa_family as AddressFamilyType { + AF_INET => self.sin.sin_port = port.into(), + AF_INET6 => self.sin6.sin6_port = port.into(), + _ => {} + } + } + } + + /// Check whether this IP address is within a CIDR range + /// + /// The argument is a CIDR range in which the port is interpreted as the number of bits, e.g. 10.0.0.0/24. + /// Any bits beyond the bit count are ignored, e.g. 10.0.0.1/24 is the same as 10.0.0.0/24. + pub fn is_within(&self, cidr: &InetAddress) -> bool { + unsafe { + if self.sa.sa_family == cidr.sa.sa_family { + let mut cidr_bits = cidr.port() as u32; + match self.sa.sa_family as AddressFamilyType { + AF_INET => { + if cidr_bits <= 32 { + let discard_bits = 32 - cidr_bits; + if u32::from_be(get_s_addr(&self.sin)).wrapping_shr(discard_bits) + == u32::from_be(get_s_addr(&cidr.sin)).wrapping_shr(discard_bits) + { + return true; + } + } + } + AF_INET6 => { + if cidr_bits <= 128 { + let a = get_s6_addr(&self.sin6); + let b = get_s6_addr(&cidr.sin6); + let mut p = 0; + while cidr_bits >= 8 { + cidr_bits -= 8; + if a[p] != b[p] { + return false; + } + p += 1; + } + let discard_bits = 8 - cidr_bits; + if a[p].wrapping_shr(discard_bits) == b[p].wrapping_shr(discard_bits) { + return true; + } + } + } + _ => {} + } + } + } + false + } + + /// Get this IP address's scope as per RFC documents and what is advertised via BGP. + pub fn scope(&self) -> IpScope { + unsafe { + match self.sa.sa_family as AddressFamilyType { + AF_INET => { + let ip = get_s_addr(&self.sin); + let class_a = (ip >> 24) as u8; + match class_a { + 0x00 | 0xff => IpScope::None, // 0.0.0.0/8 and 255.0.0.0/8 are not usable + 0x0a => IpScope::Private, // 10.0.0.0/8 + 0x7f => IpScope::Loopback, // 127.0.0.0/8 + 0x64 => { + if (ip & 0xffc00000) == 0x64400000 { + // 100.64.0.0/10 + IpScope::Private + } else { + IpScope::Global + } + } + 0xa9 => { + if (ip & 0xffff0000) == 0xa9fe0000 { + // 169.254.0.0/16 + IpScope::LinkLocal + } else { + IpScope::Global + } + } + 0xac => { + if (ip & 0xfff00000) == 0xac100000 { + // 172.16.0.0/12 + IpScope::Private + } else { + IpScope::Global + } + } + 0xc0 => { + if (ip & 0xffff0000) == 0xc0a80000 || (ip & 0xffffff00) == 0xc0000200 { + // 192.168.0.0/16 and 192.0.2.0/24 + IpScope::Private + } else { + IpScope::Global + } + } + 0xc6 => { + if (ip & 0xfffe0000) == 0xc6120000 || (ip & 0xffffff00) == 0xc6336400 { + // 198.18.0.0/15 and 198.51.100.0/24 + IpScope::Private + } else { + IpScope::Global + } + } + 0xcb => { + if (ip & 0xffffff00) == 0xcb007100 { + // 203.0.113.0/24 + IpScope::Private + } else { + IpScope::Global + } + } + _ => { + if [ + 0x06_u8, // 6.0.0.0/8 (US Army) + 0x15_u8, // 21.0.0.0/8 (US DDN-RVN) + 0x16_u8, // 22.0.0.0/8 (US DISA) + 0x19_u8, // 25.0.0.0/8 (UK Ministry of Defense) + 0x1a_u8, // 26.0.0.0/8 (US DISA) + 0x1c_u8, // 28.0.0.0/8 (US DSI-North) + 0x1d_u8, // 29.0.0.0/8 (US DISA) + 0x1e_u8, // 30.0.0.0/8 (US DISA) + 0x33_u8, // 51.0.0.0/8 (UK Department of Social Security) + 0x37_u8, // 55.0.0.0/8 (US DoD) + 0x38_u8, // 56.0.0.0/8 (US Postal Service) + ] + .contains(&class_a) + { + IpScope::PseudoPrivate + } else { + match ip >> 28 { + 0xe => IpScope::Multicast, // 224.0.0.0/4 + 0xf => IpScope::Private, // 240.0.0.0/4 ("reserved," usually unusable) + _ => IpScope::Global, + } + } + } + } + } + AF_INET6 => { + let ip = &*(&(self.sin6.sin6_addr) as *const in6_addr).cast::<[u8; 16]>(); + if (ip[0] & 0xf0) == 0xf0 { + if ip[0] == 0xff { + return IpScope::Multicast; // ff00::/8 + } + if ip[0] == 0xfe && (ip[1] & 0xc0) == 0x80 { + let mut k: usize = 2; + while ip[k] == 0 && k < 15 { + k += 1; + } + return if k == 15 && ip[15] == 0x01 { + IpScope::Loopback // fe80::1/128 + } else { + IpScope::LinkLocal // fe80::/10 + }; + } + if (ip[0] & 0xfe) == 0xfc { + return IpScope::Private; // fc00::/7 + } + } + let mut k: usize = 0; + while ip[k] == 0 && k < 15 { + k += 1; + } + if k == 15 { + if ip[15] == 0x01 { + return IpScope::Loopback; // ::1/128 + } else if ip[15] == 0x00 { + return IpScope::None; // ::/128 + } + } + IpScope::Global + } + _ => IpScope::None, + } + } + } + + /// Get only the IP portion of this address as a string. + pub fn to_ip_string(&self) -> String { + unsafe { + match self.sa.sa_family as AddressFamilyType { + AF_INET => { + #[cfg(not(windows))] + { + let ip = &*(&self.sin.sin_addr.s_addr as *const u32).cast::<[u8; 4]>(); + format!("{}.{}.{}.{}", ip[0], ip[1], ip[2], ip[3]) + } + #[cfg(windows)] + { + let ip = self.sin.sin_addr.S_un.S_addr().to_ne_bytes(); + format!("{}.{}.{}.{}", ip[0], ip[1], ip[2], ip[3]) + } + } + AF_INET6 => Ipv6Addr::from(*(&(self.sin6.sin6_addr) as *const in6_addr).cast::<[u8; 16]>()).to_string(), + _ => String::from("(null)"), + } + } + } +} + +impl ToFromBytes for InetAddress { + fn read_bytes(r: &mut R) -> std::io::Result { + let mut b = [0u8; 18]; + r.read_exact(&mut b[..1])?; + if b[0] == 4 { + r.read_exact(&mut b[..6])?; + Ok(InetAddress::from_ip_port(&b[0..4], u16::from_be_bytes(b[4..6].try_into().unwrap()))) + } else if b[0] == 6 { + r.read_exact(&mut b[..18])?; + Ok(InetAddress::from_ip_port(&b[0..16], u16::from_be_bytes(b[16..18].try_into().unwrap()))) + } else { + return Ok(InetAddress::new()); + } + } + + fn write_bytes(&self, w: &mut W) -> std::io::Result<()> { + unsafe { + match self.sa.sa_family as AddressFamilyType { + AF_INET => { + let mut b = [0u8; 7]; + b[0] = 4; + b[1..5].copy_from_slice(&get_s_addr(&self.sin).to_ne_bytes()); + b[5] = *(&self.sin.sin_port as *const u16).cast::(); + b[6] = *(&self.sin.sin_port as *const u16).cast::().offset(1); + return w.write_all(&b); + } + AF_INET6 => { + let mut b = [0u8; 19]; + b[0] = 6; + copy_nonoverlapping((&(self.sin6.sin6_addr) as *const in6_addr).cast::(), b.as_mut_ptr().offset(1), 16); + b[17] = *(&self.sin6.sin6_port as *const u16).cast::(); + b[18] = *(&self.sin6.sin6_port as *const u16).cast::().offset(1); + return w.write_all(&b); + } + _ => return w.write_all(&[0]), + } + } + } +} + +impl ToString for InetAddress { + fn to_string(&self) -> String { + unsafe { + let mut s = self.to_ip_string(); + match self.sa.sa_family as AddressFamilyType { + AF_INET => { + s.push('/'); + s.push_str(u16::from_be(self.sin.sin_port).to_string().as_str()) + } + AF_INET6 => { + s.push('/'); + s.push_str(u16::from_be(self.sin6.sin6_port).to_string().as_str()) + } + _ => {} + } + s + } + } +} + +impl FromStr for InetAddress { + type Err = InvalidParameterError; + + fn from_str(ip_string: &str) -> Result { + let mut addr = InetAddress::new(); + let s = ip_string.trim(); + if !s.is_empty() { + let (ip_str, port) = s.find('/').map_or_else( + || (s, 0), + |pos| { + let ss = s.split_at(pos); + let mut port_str = ss.1; + if port_str.starts_with('/') { + port_str = &port_str[1..]; + } + (ss.0, u16::from_str_radix(port_str, 10).unwrap_or(0).to_be()) + }, + ); + IpAddr::from_str(ip_str).map_or_else( + |_| Err(InvalidParameterError("invalid IP/CIDR")), + |ip| { + unsafe { + match ip { + IpAddr::V4(v4) => { + addr.sin.sin_family = AF_INET.into(); + addr.sin.sin_port = port.into(); + #[cfg(windows)] + { + let oct = v4.octets(); + addr.sin.sin_addr.S_un.S_un_b_mut().s_b1 = oct[0]; + addr.sin.sin_addr.S_un.S_un_b_mut().s_b2 = oct[1]; + addr.sin.sin_addr.S_un.S_un_b_mut().s_b3 = oct[2]; + addr.sin.sin_addr.S_un.S_un_b_mut().s_b4 = oct[3]; + } + #[cfg(not(windows))] + { + copy_nonoverlapping(v4.octets().as_ptr(), (&mut (addr.sin.sin_addr.s_addr) as *mut u32).cast(), 4); + } + } + IpAddr::V6(v6) => { + addr.sin6.sin6_family = AF_INET6.into(); + addr.sin6.sin6_port = port.into(); + copy_nonoverlapping(v6.octets().as_ptr(), (&mut (addr.sin6.sin6_addr) as *mut in6_addr).cast(), 16); + } + } + } + Ok(addr) + }, + ) + } else { + Ok(addr) + } + } +} + +impl PartialEq for InetAddress { + fn eq(&self, other: &Self) -> bool { + unsafe { + if self.sa.sa_family == other.sa.sa_family { + match self.sa.sa_family as AddressFamilyType { + AF_INET => self.sin.sin_port == other.sin.sin_port && get_s_addr(&self.sin) == get_s_addr(&other.sin), + AF_INET6 => { + if self.sin6.sin6_port == other.sin6.sin6_port { + (*(&(self.sin6.sin6_addr) as *const in6_addr).cast::<[u8; 16]>()) + .eq(&*(&(other.sin6.sin6_addr) as *const in6_addr).cast::<[u8; 16]>()) + } else { + false + } + } + _ => true, + } + } else { + false + } + } + } +} + +impl Eq for InetAddress {} + +impl PartialOrd for InetAddress { + #[inline(always)] + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +// Manually implement Ord to ensure consistent sort order across platforms, since we don't know exactly +// how sockaddr structs will be laid out. +impl Ord for InetAddress { + fn cmp(&self, other: &Self) -> Ordering { + unsafe { + if self.sa.sa_family == other.sa.sa_family { + match self.sa.sa_family as AddressFamilyType { + 0 => Ordering::Equal, + AF_INET => { + let ip_ordering = u32::from_be(get_s_addr(&self.sin)).cmp(&u32::from_be(get_s_addr(&other.sin))); + if ip_ordering == Ordering::Equal { + u16::from_be(self.sin.sin_port).cmp(&u16::from_be(other.sin.sin_port)) + } else { + ip_ordering + } + } + AF_INET6 => { + let a = &*(&(self.sin6.sin6_addr) as *const in6_addr).cast::<[u8; 16]>(); + let b = &*(&(other.sin6.sin6_addr) as *const in6_addr).cast::<[u8; 16]>(); + let ip_ordering = a.cmp(b); + if ip_ordering == Ordering::Equal { + u16::from_be(self.sin6.sin6_port).cmp(&u16::from_be(other.sin6.sin6_port)) + } else { + ip_ordering + } + } + _ => { + // This shouldn't be possible, but handle it for correctness. + (*slice_from_raw_parts((self as *const Self).cast::(), size_of::())) + .cmp(&*slice_from_raw_parts((other as *const Self).cast::(), size_of::())) + } + } + } else { + match self.sa.sa_family as AddressFamilyType { + 0 => Ordering::Less, + AF_INET => { + if other.sa.sa_family as AddressFamilyType == AF_INET6 { + Ordering::Less + } else { + self.sa.sa_family.cmp(&other.sa.sa_family) + } + } + AF_INET6 => { + if other.sa.sa_family as AddressFamilyType == AF_INET { + Ordering::Greater + } else { + self.sa.sa_family.cmp(&other.sa.sa_family) + } + } + _ => { + // This likewise should not be possible. + self.sa.sa_family.cmp(&other.sa.sa_family) + } + } + } + } + } +} + +impl Hash for InetAddress { + fn hash(&self, state: &mut H) { + unsafe { + state.write_u8(self.sa.sa_family as u8); + + match self.sa.sa_family as AddressFamilyType { + AF_INET => { + state.write_u16(self.sin.sin_port); + state.write_u32(get_s_addr(&self.sin)); + } + AF_INET6 => { + state.write_u16(self.sin6.sin6_port); + state.write(&*(&(self.sin6.sin6_addr) as *const in6_addr).cast::<[u8; 16]>()); + } + _ => {} + } + } + } +} + +unsafe impl Send for InetAddress {} + +#[cfg(test)] +mod tests { + use std::mem::size_of; + use std::str::FromStr; + + use super::*; + + #[test] + fn values() { + assert_ne!(AF_INET, 0); + assert_ne!(AF_INET6, 0); + assert_ne!(AF_INET, AF_INET6); + } + + #[test] + fn set_get() { + // ipv4 + let mut v = [0_u8; 4]; + for port in 0..=65535 { + v.fill_with(rand::random); + let mut addr = InetAddress::new(); + assert_ne!(addr.set(v, port), 0); + assert_eq!(addr.ip_bytes(), &v); + assert_eq!(addr.port(), port); + } + + // ipv6 + let mut v = [0_u8; 16]; + for port in 0..=65535 { + v.fill_with(rand::random); + let mut addr = InetAddress::new(); + assert_ne!(addr.set(v, port), 0); + assert_eq!(addr.ip_bytes(), &v); + assert_eq!(addr.port(), port); + } + } + + #[test] + fn layout() { + unsafe { + assert_eq!(size_of::(), size_of::()); + + let mut tmp = InetAddress::new(); + tmp.sa.sa_family = 0xab; + if tmp.sin.sin_family != 0xab { + panic!("sin_family misaligned in union"); + } + if tmp.sin6.sin6_family != 0xab { + panic!("sin6_family misaligned in union"); + } + if tmp.ss.ss_family != 0xab { + panic!("ss_family misaligned in union"); + } + } + } + + #[test] + fn ipv6_string() { + let ip = InetAddress::from_str("2603:6010:6e00:1118:d92a:ab88:4dfb:670a/1234").unwrap(); + assert_eq!("2603:6010:6e00:1118:d92a:ab88:4dfb:670a/1234", ip.to_string()); + let ip = InetAddress::from_str("fd80::1/1234").unwrap(); + assert_eq!("fd80::1/1234", ip.to_string()); + } + + #[test] + fn ipv4_string() { + let ip = InetAddress::from_str("1.2.3.4/1234").unwrap(); + assert_eq!("1.2.3.4/1234", ip.to_string()); + } +} diff --git a/src/io.rs b/src/io.rs new file mode 100644 index 0000000..6262d35 --- /dev/null +++ b/src/io.rs @@ -0,0 +1,56 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +use std::fs::File; +use std::io::{Read, Write}; +use std::path::Path; + +/// Default sanity limit parameter for read_limit() used throughout the service. +pub const DEFAULT_FILE_IO_READ_LIMIT: usize = 262144; + +/// Convenience function to read up to limit bytes from a file. +/// +/// If the file is larger than limit, the excess is not read. +pub fn read_limit>(path: P, limit: usize) -> std::io::Result> { + let mut f = File::open(path)?; + let bytes = f.metadata()?.len().min(limit as u64) as usize; + let mut v: Vec = Vec::with_capacity(bytes); + v.resize(bytes, 0); + f.read_exact(v.as_mut_slice())?; + Ok(v) +} + +/// Set permissions on a file or directory to be most restrictive (visible only to the service's user). +#[cfg(unix)] +pub fn fs_restrict_permissions>(path: P) -> bool { + unsafe { + let c_path = std::ffi::CString::new(path.as_ref().to_str().unwrap()).unwrap(); + libc::chmod( + c_path.as_ptr(), + if path.as_ref().is_dir() { + 0o700 + } else { + 0o600 + }, + ) == 0 + } +} + +#[inline] +pub fn write_all_multi(w: &mut W, s: &[&[u8]]) -> std::io::Result<()> { + for ss in s { + w.write_all(*ss)?; + } + Ok(()) +} + +/// Set permissions on a file or directory to be most restrictive (visible only to the service's user). +#[cfg(windows)] +pub fn fs_restrict_permissions>(path: P) -> bool { + todo!() +} diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..2cdd961 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,119 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +pub mod arrayvec; +pub mod base24; +pub mod base64; +pub mod blob; +pub mod buffer; +pub mod cast; +pub mod dictionary; +pub mod error; +pub mod exitcode; +pub mod gate; +pub mod hex; +pub mod inetaddress; +pub mod io; +pub mod memory; +pub mod pool; +pub mod ringbuffer; +pub mod str; +pub mod sync; +pub mod tofrombytes; + +/// Initial value that should be used for monotonic tick time variables. +pub const NEVER_HAPPENED_TICKS: i64 = i64::MIN / 2; + +/// Get milliseconds since unix epoch. +#[inline] +pub fn ms_since_epoch() -> i64 { + std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_millis() as i64 +} + +/// Get an estimate of the number of CPU cores in the system. +/// This defaults to 1 if information is not available. +pub fn parallelism() -> usize { + static mut PARALLELISM: usize = 0; + let mut p = unsafe { PARALLELISM }; + // It's perfectly fine if this runs more than once due to concurrent calls as it should always yield the same value. + if p == 0 { + p = std::thread::available_parallelism().map(|p| p.get()).unwrap_or(1); + unsafe { + PARALLELISM = p; + } + } + p +} + +/// Get milliseconds since an arbitrary time in the past, guaranteed to monotonically increase within a given process. +#[inline] +pub fn ms_monotonic() -> i64 { + static STARTUP_INSTANT: std::sync::RwLock> = std::sync::RwLock::new(None); + let si = *STARTUP_INSTANT.read().unwrap(); + if let Some(si) = si { + si.elapsed().as_millis() as i64 + } else { + STARTUP_INSTANT + .write() + .unwrap() + .get_or_insert(std::time::Instant::now()) + .elapsed() + .as_millis() as i64 + } +} + +/// Wait for a kill signal (e.g. SIGINT or OS-equivalent) sent to this process and return when received. +#[cfg(unix)] +pub fn wait_for_process_abort() { + if let Ok(mut signals) = signal_hook::iterator::Signals::new([libc::SIGINT, libc::SIGTERM, libc::SIGQUIT]) { + 'wait_for_exit: loop { + for signal in signals.wait() { + match signal as libc::c_int { + libc::SIGINT | libc::SIGTERM | libc::SIGQUIT => { + break 'wait_for_exit; + } + _ => {} + } + } + std::thread::sleep(std::time::Duration::from_millis(100)); + } + } else { + panic!("unable to listen for OS signals"); + } +} + +/// Helper for use in serde directives +#[inline(always)] +pub fn slice_is_empty(s: &[T]) -> bool { + s.is_empty() +} + +#[cold] +#[inline(never)] +pub extern "C" fn unlikely_branch() {} + +#[cfg(test)] +mod tests { + use super::ms_monotonic; + use std::time::Duration; + + #[test] + fn monotonic_clock_sanity_check() { + let start = ms_monotonic(); + assert!(start >= 0); + std::thread::sleep(Duration::from_millis(500)); + let end = ms_monotonic(); + // per docs: + // + // The thread may sleep longer than the duration specified due to scheduling specifics or + // platform-dependent functionality. It will never sleep less. + // + assert!((end - start).abs() >= 500); + assert!((end - start).abs() < 750); + } +} diff --git a/src/memory.rs b/src/memory.rs new file mode 100644 index 0000000..4344b90 --- /dev/null +++ b/src/memory.rs @@ -0,0 +1,80 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +// This is a collection of functions that use "unsafe" to do things with memory that should in fact +// be safe. Some of these may eventually get stable standard library replacements. + +#[allow(unused_imports)] +use std::mem::{needs_drop, size_of, MaybeUninit}; + +#[allow(unused_imports)] +use std::ptr::copy_nonoverlapping; + +/// Implement this trait to mark a struct as safe to cast (in place) from a byte array. +/// To be safe it must contain alignment-neutral objects which basically means bytes. For +/// integers they should be represented as byte arrays e.g. [u8; 4] for u32. +pub unsafe trait FlatBuffer: Sized {} + +/// Our version of the not-yet-stable array_chunks method in slice. +#[inline(always)] +pub fn array_chunks_exact(a: &[T]) -> impl Iterator { + let mut i = 0; + let l = a.len(); + std::iter::from_fn(move || { + let j = i + S; + if j <= l { + let next = unsafe { &*a.as_ptr().add(i).cast() }; + i = j; + Some(next) + } else { + None + } + }) +} + +/// Obtain a view into an array cast as another array. +/// This will panic if the template parameters would result in out of bounds access. +#[inline(always)] +pub fn array_range(a: &[T; S]) -> &[T; LEN] { + assert!((START + LEN) <= S); + unsafe { &*a.as_ptr().add(START).cast::<[T; LEN]>() } +} + +/// Get a reference to a raw object as a byte array. +/// The template parameter S must be less than or equal to the size of the object in bytes or this will panic. +#[inline(always)] +pub fn as_byte_array(o: &T) -> &[u8; S] { + assert!(S <= size_of::()); + unsafe { &*(o as *const T).cast() } +} + +/// Get a reference to a raw object as a byte array. +/// The template parameter S must be less than or equal to the size of the object in bytes or this will panic. +#[inline(always)] +pub fn as_byte_array_mut(o: &mut T) -> &mut [u8; S] { + assert!(S <= size_of::()); + unsafe { &mut *(o as *mut T).cast() } +} + +/// Transmute an object to a byte array. +/// The template parameter S must equal the size of the object in bytes or this will panic. +#[inline(always)] +pub fn to_byte_array(o: T) -> [u8; S] { + assert_eq!(S, size_of::()); + assert!(!std::mem::needs_drop::()); + unsafe { *(&o as *const T).cast() } +} + +/// Cast a byte slice into a flat struct. +/// This will panic if the slice is too small or the struct requires drop. +#[inline(always)] +pub fn cast_to_struct(b: &[u8]) -> &T { + assert!(b.len() >= size_of::()); + assert!(!std::mem::needs_drop::()); + unsafe { &*b.as_ptr().cast() } +} diff --git a/src/pool.rs b/src/pool.rs new file mode 100644 index 0000000..d8a1a5a --- /dev/null +++ b/src/pool.rs @@ -0,0 +1,251 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +use std::ops::{Deref, DerefMut}; +use std::ptr::NonNull; +use std::sync::{Arc, Mutex, Weak}; + +/// Each pool requires a factory that creates and resets (for re-use) pooled objects. +pub trait PoolFactory { + fn create(&self) -> O; + fn reset(&self, obj: &mut O); +} + +/// Container for pooled objects that have been checked out of the pool. +/// +/// Objects are automagically returned to the pool when Pooled<> is dropped if the pool still exists. +/// If the pool itself is gone objects are freed. Two methods for conversion to/from raw pointers are +/// available for interoperation with foreign APIs. +#[repr(transparent)] +pub struct Pooled>(NonNull>); + +#[repr(C)] +struct PoolEntry> { + obj: O, // must be first + return_pool: Weak>, +} + +impl> Pooled { + /// Create a pooled object wrapper around an object but with no pool to return it to. + /// The object will be freed when this pooled container is dropped. + #[inline] + pub fn naked(o: O) -> Self { + unsafe { + Self(NonNull::new_unchecked(Box::into_raw(Box::new(PoolEntry:: { + obj: o, + return_pool: Weak::new(), + })))) + } + } + + /// Get a raw pointer to the object wrapped by this pooled object container. + /// + /// The returned pointer MUST be returned to the pooling system with from_raw() or memory + /// will leak. + #[inline] + pub unsafe fn into_raw(self) -> *mut O { + // Verify that the structure is not padded before 'obj'. + assert_eq!( + (&self.0.as_ref().obj as *const O).cast::(), + (self.0.as_ref() as *const PoolEntry).cast::() + ); + + let ptr = self.0.as_ptr().cast::(); + std::mem::forget(self); + ptr + } + + /// Restore a raw pointer from into_raw() into a Pooled object. + /// + /// The supplied pointer MUST have been obtained from a Pooled object. None is returned + /// if the pointer is null. + #[inline] + pub unsafe fn from_raw(raw: *mut O) -> Option { + if !raw.is_null() { + Some(Self(NonNull::new_unchecked(raw.cast()))) + } else { + None + } + } +} + +impl> Clone for Pooled +where + O: Clone, +{ + #[inline] + fn clone(&self) -> Self { + let internal = unsafe { &mut *self.0.as_ptr() }; + if let Some(p) = internal.return_pool.upgrade() { + if let Some(o) = p.pool.lock().unwrap().pop() { + let mut o = Self(o); + *o.as_mut() = self.as_ref().clone(); + o + } else { + Pooled::(unsafe { + NonNull::new_unchecked(Box::into_raw(Box::new(PoolEntry:: { + obj: self.as_ref().clone(), + return_pool: Arc::downgrade(&p), + }))) + }) + } + } else { + Self::naked(self.as_ref().clone()) + } + } +} + +unsafe impl> Send for Pooled where O: Send {} +unsafe impl> Sync for Pooled where O: Sync {} + +impl> Deref for Pooled { + type Target = O; + + #[inline(always)] + fn deref(&self) -> &Self::Target { + unsafe { &self.0.as_ref().obj } + } +} + +impl> DerefMut for Pooled { + #[inline(always)] + fn deref_mut(&mut self) -> &mut Self::Target { + unsafe { &mut self.0.as_mut().obj } + } +} + +impl> AsRef for Pooled { + #[inline(always)] + fn as_ref(&self) -> &O { + unsafe { &self.0.as_ref().obj } + } +} + +impl> AsMut for Pooled { + #[inline(always)] + fn as_mut(&mut self) -> &mut O { + unsafe { &mut self.0.as_mut().obj } + } +} + +impl> Drop for Pooled { + #[inline] + fn drop(&mut self) { + let internal = unsafe { &mut *self.0.as_ptr() }; + if let Some(p) = internal.return_pool.upgrade() { + p.factory.reset(&mut internal.obj); + p.pool.lock().unwrap().push(self.0); + } else { + drop(unsafe { Box::from_raw(self.0.as_ptr()) }); + } + } +} + +/// An object pool for Reusable objects. +/// Checked out objects are held by a guard object that returns them when dropped if +/// the pool still exists or drops them if the pool has itself been dropped. +pub struct Pool>(Arc>); + +struct PoolInner> { + factory: F, + pool: Mutex>>>, +} + +impl> Pool { + #[inline] + pub fn new(initial_stack_capacity: usize, factory: F) -> Self { + Self(Arc::new(PoolInner:: { + factory, + pool: Mutex::new(Vec::with_capacity(initial_stack_capacity)), + })) + } + + /// Get a pooled object, or allocate one if the pool is empty. + #[inline] + pub fn get(&self) -> Pooled { + if let Some(o) = self.0.pool.lock().unwrap().pop() { + return Pooled::(o); + } + Pooled::(unsafe { + NonNull::new_unchecked(Box::into_raw(Box::new(PoolEntry:: { + obj: self.0.factory.create(), + return_pool: Arc::downgrade(&self.0), + }))) + }) + } + + /// Dispose of all pooled objects, freeing any memory they use. + /// + /// If get() is called after this new objects will be allocated, and any outstanding + /// objects will still be returned on drop unless the pool itself is dropped. This can + /// be done to free some memory if there has been a spike in memory use. + #[inline] + pub fn purge(&self) { + for o in self.0.pool.lock().unwrap().drain(..) { + drop(unsafe { Box::from_raw(o.as_ptr()) }) + } + } +} + +impl> Drop for Pool { + #[inline(always)] + fn drop(&mut self) { + self.purge(); + } +} + +unsafe impl> Send for Pool {} +unsafe impl> Sync for Pool {} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + use std::time::Duration; + + use super::*; + + struct TestPoolFactory; + + impl PoolFactory for TestPoolFactory { + fn create(&self) -> String { + String::new() + } + + fn reset(&self, obj: &mut String) { + obj.clear(); + } + } + + #[test] + fn threaded_pool_use() { + let p: Arc> = Arc::new(Pool::new(2, TestPoolFactory {})); + let ctr = Arc::new(AtomicUsize::new(0)); + for _ in 0..64 { + let p2 = p.clone(); + let ctr2 = ctr.clone(); + let _ = std::thread::spawn(move || { + for _ in 0..16384 { + let mut o1 = p2.get(); + o1.push('a'); + let o2 = p2.get(); + drop(o1); + let mut o2 = unsafe { Pooled::::from_raw(o2.into_raw()).unwrap() }; + o2.push('b'); + ctr2.fetch_add(1, Ordering::Relaxed); + } + }); + } + loop { + std::thread::sleep(Duration::from_millis(100)); + if ctr.load(Ordering::Relaxed) >= 16384 * 64 { + break; + } + } + } +} diff --git a/src/ringbuffer.rs b/src/ringbuffer.rs new file mode 100644 index 0000000..225fbf5 --- /dev/null +++ b/src/ringbuffer.rs @@ -0,0 +1,120 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +use std::mem::MaybeUninit; + +/// A FIFO ring buffer. +pub struct RingBuffer { + a: [MaybeUninit; C], + p: usize, +} + +impl RingBuffer { + #[inline] + pub fn new() -> Self { + #[allow(invalid_value)] + let mut tmp: Self = unsafe { MaybeUninit::uninit().assume_init() }; + tmp.p = 0; + tmp + } + + /// Add an element to the buffer, replacing old elements if full. + #[inline] + pub fn add(&mut self, o: T) { + let p = self.p; + if p < C { + unsafe { self.a.get_unchecked_mut(p).write(o) }; + } else { + unsafe { *self.a.get_unchecked_mut(p % C).assume_init_mut() = o }; + } + self.p = p.wrapping_add(1); + } + + /// Clear the buffer and drop all elements. + #[inline] + pub fn clear(&mut self) { + for i in 0..C.min(self.p) { + unsafe { self.a.get_unchecked_mut(i).assume_init_drop() }; + } + self.p = 0; + } + + /// Gets an iterator that dumps the contents of the buffer in FIFO order. + #[inline] + pub fn iter(&self) -> RingBufferIterator<'_, T, C> { + let s = C.min(self.p); + RingBufferIterator { b: self, s, i: self.p.wrapping_sub(s) } + } +} + +impl Default for RingBuffer { + #[inline(always)] + fn default() -> Self { + Self::new() + } +} + +impl Drop for RingBuffer { + #[inline(always)] + fn drop(&mut self) { + self.clear(); + } +} + +pub struct RingBufferIterator<'a, T, const C: usize> { + b: &'a RingBuffer, + s: usize, + i: usize, +} + +impl<'a, T, const C: usize> Iterator for RingBufferIterator<'a, T, C> { + type Item = &'a T; + + #[inline] + fn next(&mut self) -> Option { + let s = self.s; + if s > 0 { + let i = self.i; + self.s = s.wrapping_sub(1); + self.i = i.wrapping_add(1); + Some(unsafe { self.b.a.get_unchecked(i % C).assume_init_ref() }) + } else { + None + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn fifo() { + let mut tmp: RingBuffer = RingBuffer::new(); + let mut tmp2 = Vec::new(); + for i in 0..4 { + tmp.add(i); + tmp2.push(i); + } + for (i, j) in tmp.iter().zip(tmp2.iter()) { + assert_eq!(*i, *j); + } + tmp.clear(); + tmp2.clear(); + for i in 0..23 { + tmp.add(i); + tmp2.push(i); + } + while tmp2.len() > 8 { + tmp2.remove(0); + } + for (i, j) in tmp.iter().zip(tmp2.iter()) { + assert_eq!(*i, *j); + } + } +} diff --git a/src/str.rs b/src/str.rs new file mode 100644 index 0000000..adcc3f8 --- /dev/null +++ b/src/str.rs @@ -0,0 +1,56 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +use crate::hex::HEX_CHARS; + +/// Escape non-ASCII-printable characters in a string. +/// This also escapes quotes and other sensitive characters that cause issues on terminals. +pub fn escape(b: &[u8]) -> String { + let mut s = String::with_capacity(b.len() * 2); + for b in b.iter() { + let b = *b; + if (43..=126).contains(&b) && b != 92 && b != 96 { + s.push(b as char); + } else { + s.push('\\'); + s.push(HEX_CHARS[(b.wrapping_shr(4) & 0xf) as usize] as char); + s.push(HEX_CHARS[(b & 0xf) as usize] as char); + } + } + s +} + +/// Unescape a string with \XX hexadecimal escapes. +pub fn unescape(s: &str) -> Vec { + let mut b = Vec::with_capacity(s.len()); + let mut s = s.as_bytes(); + while let Some(c) = s.first() { + let c = *c; + if c == b'\\' { + if s.len() < 3 { + break; + } + let mut cc = 0u8; + for c in [s[1], s[2]] { + if (48..=57).contains(&c) { + cc = cc.wrapping_shl(4) | (c - 48); + } else if (65..=70).contains(&c) { + cc = cc.wrapping_shl(4) | (c - 55); + } else if (97..=102).contains(&c) { + cc = cc.wrapping_shl(4) | (c - 87); + } + } + b.push(cc); + s = &s[3..]; + } else { + b.push(c); + s = &s[1..]; + } + } + b +} diff --git a/src/sync.rs b/src/sync.rs new file mode 100644 index 0000000..ea065b3 --- /dev/null +++ b/src/sync.rs @@ -0,0 +1,50 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +use std::{mem::MaybeUninit, ops::Deref, sync::RwLockReadGuard}; + +#[allow(unused)] +pub struct MappedReadGuard<'a, S, T: ?Sized> { + r: *const T, + rg: RwLockReadGuard<'a, S>, +} + +impl<'a, S, T: ?Sized> MappedReadGuard<'a, S, T> { + /// Get a read guard that points to a field in an object another read guard locks. + #[inline(always)] + pub fn map &T>(rg: RwLockReadGuard<'a, S>, f: F) -> Self { + #[allow(invalid_value)] + let mut m = MappedReadGuard { r: unsafe { MaybeUninit::uninit().assume_init() }, rg }; + m.r = f(&m.rg); + m + } + + /// Get a read guard that points to a field in an object another read guard locks. + #[inline(always)] + pub fn maybe_map Option<&T>>(rg: RwLockReadGuard<'a, S>, f: F) -> Option { + #[allow(invalid_value)] + let mut m = MappedReadGuard { r: unsafe { MaybeUninit::uninit().assume_init() }, rg }; + if let Some(ptr) = f(&m.rg) { + m.r = ptr; + return Some(m); + } + return None; + } +} + +impl<'a, S, T: ?Sized> Deref for MappedReadGuard<'a, S, T> { + type Target = T; + + #[inline(always)] + fn deref(&self) -> &Self::Target { + unsafe { &*self.r } + } +} + +unsafe impl<'a, S, T: ?Sized> Send for MappedReadGuard<'a, S, T> where T: Send {} +unsafe impl<'a, S, T: ?Sized> Sync for MappedReadGuard<'a, S, T> where T: Sync {} diff --git a/src/tofrombytes.rs b/src/tofrombytes.rs new file mode 100644 index 0000000..df37fc3 --- /dev/null +++ b/src/tofrombytes.rs @@ -0,0 +1,44 @@ +/* This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this + * file, You can obtain one at https://mozilla.org/MPL/2.0/. + * + * (c) ZeroTier, Inc. + * https://www.zerotier.com/ + */ + +use std::io::{Read, Write}; + +use super::arrayvec::ArrayVec; +use super::buffer::{Buffer, BufferReader}; + +/// Trait for types that implement easy conversion to/from bytes either as slices or I/O streams. +/// The associated implement_serialize and implement_deserialize macros implement serde serialization +/// for types implementing both this and ToString / FromStr. +pub trait ToFromBytes: Sized { + fn read_bytes(r: &mut R) -> std::io::Result; + fn write_bytes(&self, w: &mut W) -> std::io::Result<()>; + + fn read_from_buffer(&self, b: &Buffer, cursor: &mut usize) -> std::io::Result { + Self::read_bytes(&mut BufferReader::new(b, cursor)) + } + + fn write_to_buffer(&self, b: &mut Buffer) -> std::io::Result<()> { + self.write_bytes(b) + } + + fn from_bytes(mut b: &[u8]) -> std::io::Result { + Self::read_bytes(&mut b) + } + + fn to_bytes(&self) -> Vec { + let mut v = Vec::new(); + self.write_bytes(&mut v).expect("write_bytes() failed"); + v + } + + fn to_bytes_on_stack(&self) -> ArrayVec { + let mut v = ArrayVec::new(); + self.write_bytes(&mut v).expect("write_bytes() failed"); + v + } +}