diff --git a/rust/connlib/snownet/src/index.rs b/rust/connlib/snownet/src/index.rs index 3b222aa60..ea93000a1 100644 --- a/rust/connlib/snownet/src/index.rs +++ b/rust/connlib/snownet/src/index.rs @@ -1,3 +1,5 @@ +use rand::Rng; + // A basic linear-feedback shift register implemented as xorshift, used to // distribute peer indexes across the 24-bit address space reserved for peer // identification. @@ -12,11 +14,20 @@ pub(crate) struct IndexLfsr { } impl IndexLfsr { + pub(crate) fn new(rng: &mut impl Rng) -> Self { + let seed = Self::random_index(rng); + IndexLfsr { + initial: seed, + lfsr: seed, + mask: Self::random_index(rng), + } + } + /// Generate a random 24-bit nonzero integer - fn random_index() -> u32 { + fn random_index(rng: &mut impl Rng) -> u32 { const LFSR_MAX: u32 = 0xffffff; // 24-bit seed loop { - let i = rand::random::() & LFSR_MAX; + let i = rng.next_u32() & LFSR_MAX; if i > 0 { // LFSR seed must be non-zero break i; @@ -36,14 +47,3 @@ impl IndexLfsr { value ^ self.mask } } - -impl Default for IndexLfsr { - fn default() -> Self { - let seed = Self::random_index(); - IndexLfsr { - initial: seed, - lfsr: seed, - mask: Self::random_index(), - } - } -} diff --git a/rust/connlib/snownet/src/node.rs b/rust/connlib/snownet/src/node.rs index 974271e49..5d41cce3d 100644 --- a/rust/connlib/snownet/src/node.rs +++ b/rust/connlib/snownet/src/node.rs @@ -165,6 +165,7 @@ where let mut rng = StdRng::from_seed(seed); let private_key = StaticSecret::random_from_rng(&mut rng); let public_key = &(&private_key).into(); + let index = IndexLfsr::new(&mut rng); Self { rng, @@ -172,7 +173,7 @@ where private_key, public_key: *public_key, mode: T::new(), - index: IndexLfsr::default(), + index, rate_limiter: Arc::new(RateLimiter::new(public_key, HANDSHAKE_RATE_LIMIT)), shared_candidates: Default::default(), buffered_transmits: VecDeque::default(),