diff --git a/backend/infrahub/core/node/resource_manager/ip_prefix_pool.py b/backend/infrahub/core/node/resource_manager/ip_prefix_pool.py index 17eb587ccd..fbdb917145 100644 --- a/backend/infrahub/core/node/resource_manager/ip_prefix_pool.py +++ b/backend/infrahub/core/node/resource_manager/ip_prefix_pool.py @@ -3,6 +3,8 @@ import ipaddress from typing import TYPE_CHECKING, Any, Optional +from netaddr import IPSet + from infrahub.core import registry from infrahub.core.ipam.reconciler import IpamReconciler from infrahub.core.query.ipam import get_subnets @@ -10,7 +12,7 @@ PrefixPoolGetReserved, PrefixPoolSetReserved, ) -from infrahub.pools.prefix import PrefixPool +from infrahub.pools.prefix import get_next_available_prefix from .. import Node @@ -95,14 +97,15 @@ async def get_next(self, db: InfrahubDatabase, prefixlen: int) -> IPNetworkType: branch_agnostic=True, ) - pool = PrefixPool(resource.prefix.value) # type: ignore[attr-defined] + pool = IPSet([resource.prefix.value]) for subnet in subnets: - pool.reserve(subnet=str(subnet.prefix)) + pool.remove(addr=str(subnet.prefix)) try: - next_available = pool.get(prefixlen=prefixlen) + prefix_ver = ipaddress.ip_network(resource.prefix.value).version + next_available = get_next_available_prefix(pool=pool, prefix_length=prefixlen, prefix_ver=prefix_ver) return next_available - except IndexError: + except ValueError: continue raise IndexError("No more resources available") diff --git a/backend/infrahub/graphql/queries/ipam.py b/backend/infrahub/graphql/queries/ipam.py index df27b364f2..a68808f90f 100644 --- a/backend/infrahub/graphql/queries/ipam.py +++ b/backend/infrahub/graphql/queries/ipam.py @@ -4,13 +4,14 @@ from typing import TYPE_CHECKING, Optional from graphene import Field, Int, ObjectType, String +from netaddr import IPSet from infrahub.core.constants import InfrahubKind from infrahub.core.manager import NodeManager from infrahub.core.query.ipam import get_ip_addresses, get_subnets from infrahub.exceptions import NodeNotFoundError, ValidationError from infrahub.pools.address import get_available -from infrahub.pools.prefix import PrefixPool +from infrahub.pools.prefix import get_next_available_prefix if TYPE_CHECKING: from graphql import GraphQLResolveInfo @@ -92,11 +93,13 @@ async def resolve( branch=context.branch, ) - pool = PrefixPool(prefix.prefix.value) # type: ignore[attr-defined] + pool = IPSet([prefix.prefix.value]) for subnet in subnets: - pool.reserve(subnet=str(subnet.prefix)) + pool.remove(addr=str(subnet.prefix)) + + prefix_ver = ipaddress.ip_network(prefix.prefix.value).version + next_available = get_next_available_prefix(pool=pool, prefix_length=prefix_length, prefix_ver=prefix_ver) - next_available = pool.get(prefixlen=prefix_length) return {"prefix": str(next_available)} diff --git a/backend/infrahub/pools/prefix.py b/backend/infrahub/pools/prefix.py index f28a984f71..28ff6338c1 100644 --- a/backend/infrahub/pools/prefix.py +++ b/backend/infrahub/pools/prefix.py @@ -1,174 +1,38 @@ from __future__ import annotations import ipaddress -from collections import OrderedDict, defaultdict -from ipaddress import IPv4Network, IPv6Network -from typing import Optional, Union +from typing import TYPE_CHECKING, Literal +from netaddr import IPSet -class PrefixPool: - """ - Class to automatically manage Prefixes and help to carve out sub-prefixes - """ - - def __init__(self, network: str) -> None: - self.network = ipaddress.ip_network(network) - - # Define biggest and smallest possible masks - self.mask_biggest = self.network.prefixlen + 1 - if self.network.version == 4: - self.mask_smallest = 32 - else: - self.mask_smallest = 128 - - self.available_subnets = defaultdict(list) - self.sub_by_key: dict[str, Optional[str]] = OrderedDict() - self.sub_by_id: dict[str, str] = OrderedDict() - - # Save the top level available subnet - for subnet in list(self.network.subnets(new_prefix=self.mask_biggest)): - self.available_subnets[self.mask_biggest].append(str(subnet)) - - def reserve(self, subnet: str, identifier: Optional[str] = None) -> bool: - """ - Indicate that a specific subnet is already reserved/used - """ - - # TODO Add check to make sure the subnet provided has the right size - sub = ipaddress.ip_network(subnet) - - if int(sub.prefixlen) <= int(self.network.prefixlen): - raise ValueError(f"{subnet} do not have the right size ({sub.prefixlen},{self.network.prefixlen})") - - if sub.supernet(new_prefix=self.network.prefixlen) != self.network: - raise ValueError(f"{subnet} is not part of this network") - - # Check first if this ID as already done a reservation - if identifier and identifier in self.sub_by_id.keys(): - if self.sub_by_id[identifier] == str(sub): - return True - raise ValueError( - f"this identifier ({identifier}) is already used but for a different resource ({self.sub_by_id[identifier]})" - ) - - if identifier and str(sub) in self.sub_by_key.keys(): - raise ValueError(f"this subnet is already reserved but not with this identifier ({identifier})") - - if str(sub) in self.sub_by_key.keys(): - self.remove_subnet_from_available_list(sub) - return True - - # Check if the subnet itself is available - # if available reserve and return - if subnet in self.available_subnets[sub.prefixlen]: - if identifier: - self.sub_by_id[identifier] = subnet - self.sub_by_key[subnet] = identifier - else: - self.sub_by_key[subnet] = None - - self.remove_subnet_from_available_list(sub) - return True - - # If not reserved already, check if the subnet is available - # start at sublen and check all available subnet - # increase 1 by 1 until we find the closer supernet available - # break it down and keep track of the other available subnets +if TYPE_CHECKING: + from infrahub.core.ipam.constants import IPNetworkType - for sublen in range(sub.prefixlen - 1, self.network.prefixlen, -1): - supernet = sub.supernet(new_prefix=sublen) - if str(supernet) in self.available_subnets[sublen]: - self.split_supernet(supernet=supernet, subnet=sub) - return self.reserve(subnet=subnet, identifier=identifier) - return False +def get_next_available_prefix(pool: IPSet, prefix_length: int, prefix_ver: Literal[4, 6] = 4) -> IPNetworkType: + """Get the next available prefix of a given prefix length from an IPSet. - def get(self, prefixlen: int, identifier: Optional[str] = None) -> Union[IPv4Network, IPv6Network]: - """Return the next available Subnet.""" + Args: + pool: netaddr IPSet object with available subnets + prefix_length: length of the desired prefix + prefix_ver: IPSet can contain a mix of IPv4 and IPv6 subnets. This parameter specifies the IP version of prefix to acquire. - clean_prefixlen = int(prefixlen) - - if identifier and identifier in self.sub_by_id.keys(): - net = ipaddress.ip_network(self.sub_by_id[identifier]) - if net.prefixlen == clean_prefixlen: - return net - raise ValueError() - - if len(self.available_subnets[clean_prefixlen]) != 0: - sub = self.available_subnets[clean_prefixlen][0] - self.reserve(subnet=sub, identifier=identifier) - return ipaddress.ip_network(sub) - - # if a subnet of this size is not available - # we need to find the closest subnet available and split it - for i in range(clean_prefixlen - 1, self.mask_biggest - 1, -1): - if len(self.available_subnets[i]) != 0: - supernet = ipaddress.ip_network(self.available_subnets[i][0]) - # supernet available, will split it - subs = supernet.subnets(new_prefix=clean_prefixlen) - next_sub: Union[IPv4Network, IPv6Network] = next(subs) # type: ignore[assignment] - self.split_supernet(supernet=supernet, subnet=next_sub) - self.reserve(subnet=str(next_sub), identifier=identifier) - return next_sub - - raise IndexError("No More subnet available") - - def get_nbr_available_subnets(self) -> dict[int, int]: - tmp = {} - for i in range(self.mask_biggest, self.mask_smallest + 1): - tmp[i] = len(self.available_subnets[i]) - - return tmp - - def check_if_already_allocated(self, identifier: str) -> bool: - """ - Check if a subnet has already been allocated based on an identifier - - Need to add the same capability based on Network address - If both identifier and subnet are provided, identifier take precedence - """ - if identifier in self.sub_by_id.keys(): - return True - return False - - def split_supernet( - self, supernet: Union[IPv4Network, IPv6Network], subnet: Union[IPv4Network, IPv6Network] - ) -> None: - """Split a supernet into smaller networks""" - - # TODO ensure subnet is small than supernet - # TODO ensure that subnet is part of supernet - parent_net = supernet - for i in range(supernet.prefixlen + 1, subnet.prefixlen + 1): - tmp_net: list[Union[IPv4Network, IPv6Network]] = list(parent_net.subnets(new_prefix=i)) - - if i == subnet.prefixlen: - for net in tmp_net: - self.available_subnets[i].append(str(net)) - else: - if subnet.subnet_of(other=tmp_net[0]): # type: ignore[arg-type] - parent = 0 - other = 1 - else: - parent = 1 - other = 0 - - parent_net = tmp_net[parent] - self.available_subnets[i].append(str(tmp_net[other])) - - self.remove_subnet_from_available_list(supernet) - - def remove_subnet_from_available_list(self, subnet: Union[IPv4Network, IPv6Network]) -> None: - """Remove a subnet from the list of available Subnet.""" - try: - idx = self.available_subnets[subnet.prefixlen].index(str(subnet)) - del self.available_subnets[subnet.prefixlen][idx] - except ValueError: - # Already removed - pass - - # if idx: - # return True - # except: - # log.warn("Unable to remove %s from list of available subnets" % str(subnet)) - # return False + Raises: + ValueError: If there are no available subnets in the pool + """ + prefix_ver_map = { + 4: ipaddress.IPv4Network, + 6: ipaddress.IPv6Network, + } + + filtered_pool = IPSet([]) + for subnet in pool.iter_cidrs(): + if isinstance(ipaddress.ip_network(str(subnet)), prefix_ver_map[prefix_ver]): + filtered_pool.add(subnet) + + for cidr in filtered_pool.iter_cidrs(): + if cidr.prefixlen <= prefix_length: + next_available = ipaddress.ip_network(f"{cidr.network}/{prefix_length}") + return next_available + + raise ValueError("No available subnets in pool") diff --git a/backend/tests/unit/graphql/mutations/test_resource_manager.py b/backend/tests/unit/graphql/mutations/test_resource_manager.py index 89bebf72a0..c906493c14 100644 --- a/backend/tests/unit/graphql/mutations/test_resource_manager.py +++ b/backend/tests/unit/graphql/mutations/test_resource_manager.py @@ -430,7 +430,7 @@ async def test_prefix_pool_get_resource_with_prefix_length( assert not result.errors assert result.data assert result.data["IPPrefixPoolGetResource"]["ok"] - assert result.data["IPPrefixPoolGetResource"]["node"] == {"display_label": "10.10.3.32/31", "kind": "IpamIPPrefix"} + assert result.data["IPPrefixPoolGetResource"]["node"] == {"display_label": "10.10.0.0/31", "kind": "IpamIPPrefix"} async def test_address_pool_get_resource( diff --git a/backend/tests/unit/graphql/queries/test_ipam.py b/backend/tests/unit/graphql/queries/test_ipam.py index e54bb4c000..3df6cc2a39 100644 --- a/backend/tests/unit/graphql/queries/test_ipam.py +++ b/backend/tests/unit/graphql/queries/test_ipam.py @@ -107,8 +107,8 @@ async def ip_dataset_02( @pytest.mark.parametrize( "prefix,prefix_length,response", [ - ("net146", 16, "10.11.0.0/16"), - ("net146", 24, "10.11.0.0/24"), + ("net146", 16, "10.0.0.0/16"), + ("net146", 24, "10.0.0.0/24"), ("net142", 26, "10.10.1.64/26"), ("net142", 27, "10.10.1.32/27"), ], diff --git a/backend/tests/unit/pools/test_prefix.py b/backend/tests/unit/pools/test_prefix.py index a8f10d98e1..12c15b26e9 100644 --- a/backend/tests/unit/pools/test_prefix.py +++ b/backend/tests/unit/pools/test_prefix.py @@ -1,148 +1,47 @@ import ipaddress import pytest +from netaddr import IPSet -from infrahub.pools.prefix import PrefixPool +from infrahub.pools.prefix import get_next_available_prefix -def test_init_v4(): - sub = PrefixPool("192.168.0.0/28") - avail_subs = sub.get_nbr_available_subnets() +def test_get_next_available_prefix_v4(): + """Tests getting the next available IPv4 prefix.""" + pool = IPSet(["192.0.2.0/24"]) + pool.remove("192.0.2.0/30") + next_prefix = get_next_available_prefix(pool, 30, 4) + assert next_prefix == ipaddress.IPv4Network("192.0.2.4/30") - assert avail_subs == {29: 2, 30: 0, 31: 0, 32: 0} - assert sub.available_subnets[29] == ["192.168.0.0/29", "192.168.0.8/29"] +def test_get_next_available_prefix_v6(): + """Tests getting the next available IPv6 prefix.""" + pool = IPSet(["2001:db8::/32"]) + pool.remove("2001:db8::/64") + next_prefix = get_next_available_prefix(pool, 64, 6) + assert next_prefix == ipaddress.IPv6Network("2001:db8:0:1::/64") -def test_init_v6(): - sub = PrefixPool("2001:db8::0/48") - avail_subs = sub.get_nbr_available_subnets() - assert avail_subs[49] == 2 - assert avail_subs[50] == 0 +def test_get_next_available_prefix_mixed(): + """Tests getting the next available prefix from a mixed pool.""" + pool = IPSet(["2001:db8::/32", "192.0.2.0/24"]) + next_prefix = get_next_available_prefix(pool, 30, 4) + assert next_prefix == ipaddress.IPv4Network("192.0.2.0/30") + pool = IPSet(["192.0.2.0/24", "2001:db8::/32"]) + next_prefix = get_next_available_prefix(pool, 64, 6) + assert next_prefix == ipaddress.IPv6Network("2001:db8::/64") -def test_split_supernet_first(): - sub = PrefixPool("192.168.0.0/24") - - sub.split_supernet( - supernet=ipaddress.ip_network("192.168.0.128/25"), - subnet=ipaddress.ip_network("192.168.0.128/27"), - ) - avail_subs = sub.get_nbr_available_subnets() - - assert avail_subs[25] == 1 - assert avail_subs[26] == 1 - assert avail_subs[27] == 2 - assert avail_subs[28] == 0 - - assert sub.available_subnets[25] == ["192.168.0.0/25"] - assert sub.available_subnets[26] == ["192.168.0.192/26"] - assert sub.available_subnets[27] == ["192.168.0.128/27", "192.168.0.160/27"] - - -def test_split_supernet_middle(): - sub = PrefixPool("192.168.0.0/24") - - sub.split_supernet( - supernet=ipaddress.ip_network("192.168.0.128/25"), - subnet=ipaddress.ip_network("192.168.0.192/27"), - ) - avail_subs = sub.get_nbr_available_subnets() - - assert avail_subs[25] == 1 - assert avail_subs[26] == 1 - assert avail_subs[27] == 2 - assert avail_subs[28] == 0 - - assert sub.available_subnets[25] == ["192.168.0.0/25"] - assert sub.available_subnets[26] == ["192.168.0.128/26"] - assert sub.available_subnets[27] == ["192.168.0.192/27", "192.168.0.224/27"] - - -def test_split_supernet_end(): - sub = PrefixPool("192.168.0.0/16") - - sub.split_supernet( - supernet=ipaddress.ip_network("192.168.128.0/17"), - subnet=ipaddress.ip_network("192.168.255.192/27"), - ) - avail_subs = sub.get_nbr_available_subnets() - - assert avail_subs[17] == 1 - assert avail_subs[25] == 1 - assert avail_subs[26] == 1 - assert avail_subs[27] == 2 - assert avail_subs[28] == 0 - - assert sub.available_subnets[17] == ["192.168.0.0/17"] - assert sub.available_subnets[26] == ["192.168.255.128/26"] - assert sub.available_subnets[27] == ["192.168.255.192/27", "192.168.255.224/27"] - - -def test_get_subnet_v4_no_owner(): - sub = PrefixPool("192.168.0.0/16") - - assert str(sub.get(prefixlen=24)) == "192.168.0.0/24" - assert str(sub.get(prefixlen=25)) == "192.168.1.0/25" - assert str(sub.get(prefixlen=17)) == "192.168.128.0/17" - assert str(sub.get(prefixlen=24)) == "192.168.2.0/24" - assert str(sub.get(prefixlen=25)) == "192.168.1.128/25" - - -def test_get_subnet_v4_with_owner(): - sub = PrefixPool("192.168.0.0/16") - - assert str(sub.get(prefixlen=24, identifier="first")) == "192.168.0.0/24" - assert str(sub.get(prefixlen=25, identifier="second")) == "192.168.1.0/25" - assert str(sub.get(prefixlen=17, identifier="third")) == "192.168.128.0/17" - assert str(sub.get(prefixlen=25, identifier="second")) == "192.168.1.0/25" - assert str(sub.get(prefixlen=17, identifier="third")) == "192.168.128.0/17" - - -def test_get_subnet_no_more_subnet(): - sub = PrefixPool("192.0.0.0/22") - - assert str(sub.get(prefixlen=24)) == "192.0.0.0/24" - assert str(sub.get(prefixlen=24)) == "192.0.1.0/24" - assert str(sub.get(prefixlen=24)) == "192.0.2.0/24" - assert str(sub.get(prefixlen=24)) == "192.0.3.0/24" - with pytest.raises(IndexError): - assert sub.get(prefixlen=24) is False - - -def test_get_subnet_v6_no_owner(): - sub = PrefixPool("2620:135:6000:fffe::/64") - assert str(sub.get(prefixlen=127)), "2620:135:6000:fffe::/127" - - -def test_already_allocated_v4_no_owner(): - sub = PrefixPool("192.168.0.0/16") - - assert str(sub.get(prefixlen=24, identifier="first")) == "192.168.0.0/24" - assert str(sub.get(prefixlen=24, identifier="second")) == "192.168.1.0/24" - - assert sub.check_if_already_allocated(identifier="second") is True - assert sub.check_if_already_allocated(identifier="third") is False - - -def test_reserve_no_owner(): - sub = PrefixPool("192.168.0.0/16") - - assert sub.reserve("192.168.0.0/24") is True - assert str(sub.get(prefixlen=24)) == "192.168.1.0/24" - - -def test_reserve_wrong_input(): - sub = PrefixPool("192.168.0.0/16") +def test_get_next_available_prefix_exhausted_v4_pool(): + """Tests getting the next available prefix from an exhausted IPv4 pool.""" + pool = IPSet(["192.0.2.0/24"]) with pytest.raises(ValueError): - assert sub.reserve("192.192.1.0/24", identifier="first") - + get_next_available_prefix(pool, 23, 4) -def test_reserve_with_owner(): - sub = PrefixPool("192.192.0.0/16") - assert sub.reserve("192.192.0.0/24", identifier="first") is True - assert sub.reserve("192.192.1.0/24", identifier="second") is True - - assert str(sub.get(prefixlen=24)) == "192.192.2.0/24" +def test_get_next_available_prefix_exhausted_v6_pool(): + """Tests getting the next available prefix from an exhausted IPv6 pool.""" + pool = IPSet(["2001:db8::/32"]) + with pytest.raises(ValueError): + get_next_available_prefix(pool, 30, 6) diff --git a/changelog/3547.changed.md b/changelog/3547.changed.md new file mode 100644 index 0000000000..82f051b480 --- /dev/null +++ b/changelog/3547.changed.md @@ -0,0 +1 @@ +Replace PrefixPool with netaddr.IPSet