From e5531a2c90e8f206b748c76d8901459b691648c2 Mon Sep 17 00:00:00 2001 From: Christien Rioux Date: Sun, 5 Jan 2025 20:39:12 -0500 Subject: [PATCH] [skip ci] more cleanup and pick routines --- .../virtual_network/router_server/config.rs | 1 - .../machine_registry/address_pool.rs | 1 - .../machine_registry/state/blueprint_state.rs | 240 +--------------- .../state/network_locations_list.rs | 256 ++++++++++++++++++ .../machine_registry/state/network_state.rs | 1 - .../src/virtual_network/router_server/mod.rs | 1 + 6 files changed, 265 insertions(+), 235 deletions(-) diff --git a/veilid-tools/src/virtual_network/router_server/config.rs b/veilid-tools/src/virtual_network/router_server/config.rs index baa8c7c3..ab3702b2 100644 --- a/veilid-tools/src/virtual_network/router_server/config.rs +++ b/veilid-tools/src/virtual_network/router_server/config.rs @@ -1,5 +1,4 @@ use super::*; -use ipnet::*; use serde::*; use std::path::Path; diff --git a/veilid-tools/src/virtual_network/router_server/machine_registry/address_pool.rs b/veilid-tools/src/virtual_network/router_server/machine_registry/address_pool.rs index b7870385..50cc04cf 100644 --- a/veilid-tools/src/virtual_network/router_server/machine_registry/address_pool.rs +++ b/veilid-tools/src/virtual_network/router_server/machine_registry/address_pool.rs @@ -1,5 +1,4 @@ use super::*; -use ipnet::*; #[derive(Debug, Clone)] pub struct AddressPool { diff --git a/veilid-tools/src/virtual_network/router_server/machine_registry/state/blueprint_state.rs b/veilid-tools/src/virtual_network/router_server/machine_registry/state/blueprint_state.rs index 0ba8a68c..2bff723e 100644 --- a/veilid-tools/src/virtual_network/router_server/machine_registry/state/blueprint_state.rs +++ b/veilid-tools/src/virtual_network/router_server/machine_registry/state/blueprint_state.rs @@ -248,123 +248,11 @@ impl BlueprintState { return Ok(()); }; - // Get maximum prefix - let max_prefix = ipv4 - .params - .prefix - .iter() - .max() - .copied() - .expect("must have at least one element"); - // Get addresses for network - let (subnet, super_net) = match &ipv4.params.locations { - NetworkLocationsList::Allocations { allocations } => { - // Get allocations which have subnets that would fit - // our maximum requested prefix - let Some(alloc_subnets) = allocations.try_filter_map(|allocation_name| { - let allocation = machine_registry_inner - .config() - .allocations - .get(allocation_name) - .expect("must exist"); - Ok(allocation - .subnets - .subnet4 - .as_ref() - .and_then(|subnet| subnet.filter(|p| p.prefix_len() <= max_prefix))) - })? - else { - return Err(MachineRegistryError::NoAllocation); - }; - - // Pick an allocation - let subnets = machine_registry_inner - .srng() - .weighted_choice_ref(&alloc_subnets); - - // Pick a subnet - let net = *machine_registry_inner.srng().weighted_choice_ref(subnets); - - // Pick a prefix length that would fit in the subnet - let opt_subnet = ipv4 - .params - .prefix - .filter(|p| *p >= net.prefix_len()) - .as_ref() - .map(|wl| { - let subnet_prefix = *machine_registry_inner.srng().weighted_choice_ref(wl); - - // Use an address pool temporarily to pick a subnet - let mut address_pool = AddressPool::<()>::new(); - address_pool.add_scope_v4(net); - address_pool.allocate_random_v4( - machine_registry_inner.srng(), - subnet_prefix, - (), - ) - }) - .transpose()? - .flatten(); - let Some(subnet) = opt_subnet else { - return Err(MachineRegistryError::NoAllocation); - }; - (subnet, None) - } - NetworkLocationsList::Networks { networks } => { - // Get networks which have subnets that would fit - // our maximum requested prefix - let Some(available_networks) = networks.try_filter(|network_id| { - let super_network_state = machine_registry_inner - .network_states() - .get_state(*network_id) - .expect("must exist"); - - Ok(super_network_state.can_allocate_subnet_v4(None, max_prefix)) - })? - else { - return Err(MachineRegistryError::NoAllocation); - }; - - // Pick a network - let super_network_id = *machine_registry_inner - .srng() - .weighted_choice_ref(&available_networks); - let mut super_network_state = machine_registry_inner - .network_states() - .get_state(super_network_id) - .expect("must exist"); - - // Pick a prefix that fits in this network and allocate from it - let opt_subnet = ipv4 - .params - .prefix - .filter(|p| super_network_state.can_allocate_subnet_v4(None, *p)) - .as_ref() - .map(|wl| { - let subnet_prefix = *machine_registry_inner.srng().weighted_choice_ref(wl); - - // Allocate subnet from this network - super_network_state.allocate_subnet_v4( - machine_registry_inner, - OwnerTag::Network(super_network_state.id()), - None, - subnet_prefix, - ) - }) - .transpose()?; - let Some(subnet) = opt_subnet else { - return Err(MachineRegistryError::NoAllocation); - }; - - // Update network state - machine_registry_inner - .network_states_mut() - .set_state(super_network_state); - - (subnet, Some(super_network_id)) - } - }; + let NetworkLocation { subnet, super_net } = ipv4 + .params + .locations + .pick_v4(machine_registry_inner, &ipv4.params.prefix)?; let params = NetworkStateIpv4Params { allocation: subnet, @@ -442,123 +330,11 @@ impl BlueprintState { return Ok(()); }; - // Get maximum prefix - let max_prefix = ipv6 - .params - .prefix - .iter() - .max() - .copied() - .expect("must have at least one element"); - // Get addresses for network - let (subnet, super_net) = match &ipv6.params.locations { - NetworkLocationsList::Allocations { allocations } => { - // Get allocations which have subnets that would fit - // our maximum requested prefix - let Some(alloc_subnets) = allocations.try_filter_map(|allocation_name| { - let allocation = machine_registry_inner - .config() - .allocations - .get(allocation_name) - .expect("must exist"); - Ok(allocation - .subnets - .subnet6 - .as_ref() - .and_then(|subnet| subnet.filter(|p| p.prefix_len() <= max_prefix))) - })? - else { - return Err(MachineRegistryError::NoAllocation); - }; - - // Pick an allocation - let subnets = machine_registry_inner - .srng() - .weighted_choice_ref(&alloc_subnets); - - // Pick a subnet - let net = *machine_registry_inner.srng().weighted_choice_ref(subnets); - - // Pick a prefix length that would fit in the subnet - let opt_subnet = ipv6 - .params - .prefix - .filter(|p| *p >= net.prefix_len()) - .as_ref() - .map(|wl| { - let subnet_prefix = *machine_registry_inner.srng().weighted_choice_ref(wl); - - // Use an address pool temporarily to pick a subnet - let mut address_pool = AddressPool::<()>::new(); - address_pool.add_scope_v6(net); - address_pool.allocate_random_v6( - machine_registry_inner.srng(), - subnet_prefix, - (), - ) - }) - .transpose()? - .flatten(); - let Some(subnet) = opt_subnet else { - return Err(MachineRegistryError::NoAllocation); - }; - (subnet, None) - } - NetworkLocationsList::Networks { networks } => { - // Get networks which have subnets that would fit - // our maximum requested prefix - let Some(available_networks) = networks.try_filter(|network_id| { - let super_network_state = machine_registry_inner - .network_states() - .get_state(*network_id) - .expect("must exist"); - - Ok(super_network_state.can_allocate_subnet_v6(None, max_prefix)) - })? - else { - return Err(MachineRegistryError::NoAllocation); - }; - - // Pick a network - let super_network_id = *machine_registry_inner - .srng() - .weighted_choice_ref(&available_networks); - let mut super_network_state = machine_registry_inner - .network_states() - .get_state(super_network_id) - .expect("must exist"); - - // Pick a prefix that fits in this network and allocate from it - let opt_subnet = ipv6 - .params - .prefix - .filter(|p| super_network_state.can_allocate_subnet_v6(None, *p)) - .as_ref() - .map(|wl| { - let subnet_prefix = *machine_registry_inner.srng().weighted_choice_ref(wl); - - // Allocate subnet from this network - super_network_state.allocate_subnet_v6( - machine_registry_inner, - OwnerTag::Network(super_network_state.id()), - None, - subnet_prefix, - ) - }) - .transpose()?; - let Some(subnet) = opt_subnet else { - return Err(MachineRegistryError::NoAllocation); - }; - - // Update network state - machine_registry_inner - .network_states_mut() - .set_state(super_network_state); - - (subnet, Some(super_network_id)) - } - }; + let NetworkLocation { subnet, super_net } = ipv6 + .params + .locations + .pick_v6(machine_registry_inner, &ipv6.params.prefix)?; let params = NetworkStateIpv6Params { allocation: subnet, diff --git a/veilid-tools/src/virtual_network/router_server/machine_registry/state/network_locations_list.rs b/veilid-tools/src/virtual_network/router_server/machine_registry/state/network_locations_list.rs index 8699807c..a7ff021c 100644 --- a/veilid-tools/src/virtual_network/router_server/machine_registry/state/network_locations_list.rs +++ b/veilid-tools/src/virtual_network/router_server/machine_registry/state/network_locations_list.rs @@ -10,3 +10,259 @@ pub enum NetworkLocationsList { networks: WeightedList, }, } + +#[derive(Debug, Clone)] +pub struct NetworkLocation { + pub subnet: T, + pub super_net: Option, +} + +impl NetworkLocationsList { + pub fn pick_v4( + &self, + machine_registry_inner: &mut MachineRegistryInner, + prefix: &WeightedList, + ) -> MachineRegistryResult> { + // Get maximum prefix + let max_prefix = prefix + .iter() + .max() + .copied() + .expect("must have at least one element"); + + // Get addresses for network + match self { + NetworkLocationsList::Allocations { allocations } => { + // Get allocations which have subnets that would fit + // our maximum requested prefix + let Some(alloc_subnets) = allocations.try_filter_map(|allocation_name| { + let allocation = machine_registry_inner + .config() + .allocations + .get(allocation_name) + .expect("must exist"); + Ok(allocation + .subnets + .subnet4 + .as_ref() + .and_then(|subnet| subnet.filter(|p| p.prefix_len() <= max_prefix))) + })? + else { + return Err(MachineRegistryError::NoAllocation); + }; + + // Pick an allocation + let subnets = machine_registry_inner + .srng() + .weighted_choice_ref(&alloc_subnets); + + // Pick a subnet + let net = *machine_registry_inner.srng().weighted_choice_ref(subnets); + + // Pick a prefix length that would fit in the subnet + let opt_subnet = prefix + .filter(|p| *p >= net.prefix_len()) + .as_ref() + .map(|wl| { + let subnet_prefix = *machine_registry_inner.srng().weighted_choice_ref(wl); + + // Use an address pool temporarily to pick a subnet + let mut address_pool = AddressPool::<()>::new(); + address_pool.add_scope_v4(net); + address_pool.allocate_random_v4( + machine_registry_inner.srng(), + subnet_prefix, + (), + ) + }) + .transpose()? + .flatten(); + let Some(subnet) = opt_subnet else { + return Err(MachineRegistryError::NoAllocation); + }; + Ok(NetworkLocation { + subnet, + super_net: None, + }) + } + NetworkLocationsList::Networks { networks } => { + // Get networks which have subnets that would fit + // our maximum requested prefix + let Some(available_networks) = networks.try_filter(|network_id| { + let super_network_state = machine_registry_inner + .network_states() + .get_state(*network_id) + .expect("must exist"); + + Ok(super_network_state.can_allocate_subnet_v4(None, max_prefix)) + })? + else { + return Err(MachineRegistryError::NoAllocation); + }; + + // Pick a network + let super_network_id = *machine_registry_inner + .srng() + .weighted_choice_ref(&available_networks); + let mut super_network_state = machine_registry_inner + .network_states() + .get_state(super_network_id) + .expect("must exist"); + + // Pick a prefix that fits in this network and allocate from it + let opt_subnet = prefix + .filter(|p| super_network_state.can_allocate_subnet_v4(None, *p)) + .as_ref() + .map(|wl| { + let subnet_prefix = *machine_registry_inner.srng().weighted_choice_ref(wl); + + // Allocate subnet from this network + super_network_state.allocate_subnet_v4( + machine_registry_inner, + OwnerTag::Network(super_network_state.id()), + None, + subnet_prefix, + ) + }) + .transpose()?; + let Some(subnet) = opt_subnet else { + return Err(MachineRegistryError::NoAllocation); + }; + + // Update network state + machine_registry_inner + .network_states_mut() + .set_state(super_network_state); + + Ok(NetworkLocation { + subnet, + super_net: Some(super_network_id), + }) + } + } + } + + pub fn pick_v6( + &self, + machine_registry_inner: &mut MachineRegistryInner, + prefix: &WeightedList, + ) -> MachineRegistryResult> { + // Get maximum prefix + let max_prefix = prefix + .iter() + .max() + .copied() + .expect("must have at least one element"); + + // Get addresses for network + match self { + NetworkLocationsList::Allocations { allocations } => { + // Get allocations which have subnets that would fit + // our maximum requested prefix + let Some(alloc_subnets) = allocations.try_filter_map(|allocation_name| { + let allocation = machine_registry_inner + .config() + .allocations + .get(allocation_name) + .expect("must exist"); + Ok(allocation + .subnets + .subnet6 + .as_ref() + .and_then(|subnet| subnet.filter(|p| p.prefix_len() <= max_prefix))) + })? + else { + return Err(MachineRegistryError::NoAllocation); + }; + + // Pick an allocation + let subnets = machine_registry_inner + .srng() + .weighted_choice_ref(&alloc_subnets); + + // Pick a subnet + let net = *machine_registry_inner.srng().weighted_choice_ref(subnets); + + // Pick a prefix length that would fit in the subnet + let opt_subnet = prefix + .filter(|p| *p >= net.prefix_len()) + .as_ref() + .map(|wl| { + let subnet_prefix = *machine_registry_inner.srng().weighted_choice_ref(wl); + + // Use an address pool temporarily to pick a subnet + let mut address_pool = AddressPool::<()>::new(); + address_pool.add_scope_v6(net); + address_pool.allocate_random_v6( + machine_registry_inner.srng(), + subnet_prefix, + (), + ) + }) + .transpose()? + .flatten(); + let Some(subnet) = opt_subnet else { + return Err(MachineRegistryError::NoAllocation); + }; + Ok(NetworkLocation { + subnet, + super_net: None, + }) + } + NetworkLocationsList::Networks { networks } => { + // Get networks which have subnets that would fit + // our maximum requested prefix + let Some(available_networks) = networks.try_filter(|network_id| { + let super_network_state = machine_registry_inner + .network_states() + .get_state(*network_id) + .expect("must exist"); + + Ok(super_network_state.can_allocate_subnet_v6(None, max_prefix)) + })? + else { + return Err(MachineRegistryError::NoAllocation); + }; + + // Pick a network + let super_network_id = *machine_registry_inner + .srng() + .weighted_choice_ref(&available_networks); + let mut super_network_state = machine_registry_inner + .network_states() + .get_state(super_network_id) + .expect("must exist"); + + // Pick a prefix that fits in this network and allocate from it + let opt_subnet = prefix + .filter(|p| super_network_state.can_allocate_subnet_v6(None, *p)) + .as_ref() + .map(|wl| { + let subnet_prefix = *machine_registry_inner.srng().weighted_choice_ref(wl); + + // Allocate subnet from this network + super_network_state.allocate_subnet_v6( + machine_registry_inner, + OwnerTag::Network(super_network_state.id()), + None, + subnet_prefix, + ) + }) + .transpose()?; + let Some(subnet) = opt_subnet else { + return Err(MachineRegistryError::NoAllocation); + }; + + // Update network state + machine_registry_inner + .network_states_mut() + .set_state(super_network_state); + + Ok(NetworkLocation { + subnet, + super_net: Some(super_network_id), + }) + } + } + } +} diff --git a/veilid-tools/src/virtual_network/router_server/machine_registry/state/network_state.rs b/veilid-tools/src/virtual_network/router_server/machine_registry/state/network_state.rs index 9a2541db..b1673ba5 100644 --- a/veilid-tools/src/virtual_network/router_server/machine_registry/state/network_state.rs +++ b/veilid-tools/src/virtual_network/router_server/machine_registry/state/network_state.rs @@ -1,5 +1,4 @@ use super::*; -use ipnet::*; #[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)] pub enum OwnerTag { diff --git a/veilid-tools/src/virtual_network/router_server/mod.rs b/veilid-tools/src/virtual_network/router_server/mod.rs index 0926f35c..42aec569 100644 --- a/veilid-tools/src/virtual_network/router_server/mod.rs +++ b/veilid-tools/src/virtual_network/router_server/mod.rs @@ -14,6 +14,7 @@ use weighted_list::*; use async_tungstenite::accept_async; use futures_codec::{Bytes, BytesCodec, FramedRead, FramedWrite}; use futures_util::{stream::FuturesUnordered, AsyncReadExt, StreamExt, TryStreamExt}; +use ipnet::*; use postcard::{from_bytes, to_stdvec}; use std::io; use stop_token::future::FutureExt as _;