[skip ci] more cleanup and pick routines

veilid-tools/src/virtual_network/router_server/config.rs

@@ -1,5 +1,4 @@
 use super::*;
-use ipnet::*;
 use serde::*;
 use std::path::Path;
 

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<T: fmt::Debug + Clone> {

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 {
+        let NetworkLocation { subnet, super_net } = ipv4
-            NetworkLocationsList::Allocations { allocations } => {
+            .params
-                // Get allocations which have subnets that would fit
+            .locations
-                // our maximum requested prefix
+            .pick_v4(machine_registry_inner, &ipv4.params.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 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 {
+        let NetworkLocation { subnet, super_net } = ipv6
-            NetworkLocationsList::Allocations { allocations } => {
+            .params
-                // Get allocations which have subnets that would fit
+            .locations
-                // our maximum requested prefix
+            .pick_v6(machine_registry_inner, &ipv6.params.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 params = NetworkStateIpv6Params {
             allocation: subnet,

veilid-tools/src/virtual_network/router_server/machine_registry/state/network_locations_list.rs

@@ -10,3 +10,259 @@ pub enum NetworkLocationsList {
         networks: WeightedList<NetworkStateId>,
     },
 }
+
+#[derive(Debug, Clone)]
+pub struct NetworkLocation<T> {
+    pub subnet: T,
+    pub super_net: Option<NetworkStateId>,
+}
+
+impl NetworkLocationsList {
+    pub fn pick_v4(
+        &self,
+        machine_registry_inner: &mut MachineRegistryInner,
+        prefix: &WeightedList<u8>,
+    ) -> MachineRegistryResult<NetworkLocation<Ipv4Net>> {
+        // 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<u8>,
+    ) -> MachineRegistryResult<NetworkLocation<Ipv6Net>> {
+        // 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),
+                })
+            }
+        }
+    }
+}

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 {

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 _;