Merge pull request #4775 from opsmill/lgu-fix-search-ipv6

Search anywhere supports IPV6 extended format

Author: LucasG0

backend/infrahub/core/attribute.py

@@ -872,7 +872,7 @@ def validate_format(cls, value: Any, name: str, schema: AttributeSchema) -> None
             raise ValidationError({name: f"{value} is not a valid {schema.kind}"}) from exc
 
     def serialize_value(self) -> str:
-        """Serialize the value before storing it in the database."""
+        """Serialize the value before storing it in the database. If network is an IPv6 network, it is converted to collapsed form."""
 
         return ipaddress.ip_network(self.value).with_prefixlen
 
@@ -998,7 +998,7 @@ def validate_format(cls, value: Any, name: str, schema: AttributeSchema) -> None
             raise ValidationError({name: f"{value} is not a valid {schema.kind}"}) from exc
 
     def serialize_value(self) -> str:
-        """Serialize the value before storing it in the database."""
+        """Adds a prefix to address before storing it in the database. If address in an IPv6 address, it is converted to collapsed form."""
 
         return ipaddress.ip_interface(self.value).with_prefixlen
 

backend/infrahub/graphql/queries/search.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import ipaddress
 from typing import TYPE_CHECKING, Any, Optional
 
 from graphene import Boolean, Field, Int, List, ObjectType, String
@@ -29,6 +30,72 @@ class NodeEdges(ObjectType):
     edges = Field(List(of_type=NodeEdge, required=True), required=False)
 
 
+def _collapse_ipv6(s: str) -> str:
+    """Collapse an ipv6 address, ipv6 network, or a partial ipv6 address in extended format, into its collapsed form.
+    Raises an error if input does not resemble an IPv6 address in extended format. It means this function also raises
+    an error if input string is the start of an IPv6 address in collapsed format.
+    """
+
+    try:
+        return str(ipaddress.IPv6Address(s))
+    except ipaddress.AddressValueError:
+        pass
+
+    try:
+        return ipaddress.IPv6Network(s).with_prefixlen
+    except ipaddress.AddressValueError:
+        pass
+
+    # Input string might be an incomplete address in IPv6 format,
+    # in which case we would like the collapsed form equivalent of this incomplete address for matching purposes.
+    # To get it, we first try to pad the incomplete address with zeros, then we retrieve the collapsed form
+    # of the full address, and we remove extra "::" or ":0" at the end of it.
+
+    error_message = "Input string does not match IPv6 extended format"
+
+    # Input string cannot be an IPv6 in extended format if it contains ":"
+    if "::" in s:
+        raise ValueError(error_message)
+
+    # Add padding to complete the address if needed
+    segments = s.split(":")
+
+    if len(segments) == 0:
+        raise ValueError(error_message)
+
+    # If any of the non-last segments has less than 4 characters it means we deal with
+    # a IPv6 collapsed form or an invalid address
+    for segment in segments[:-1]:
+        if len(segment) != 4:
+            raise ValueError(error_message)
+
+    # Add 0 padding to last segment
+    if len(segments[-1]) > 4:
+        raise ValueError(error_message)
+
+    segments[-1] += "0" * (4 - len(segments[-1]))
+
+    # Complete the address to have 8 segments by padding with zeros
+    while len(segments) < 8:
+        segments.append("0000")
+
+    # Create a full IPv6 address from the partial input
+    full_address = ":".join(segments)
+
+    # Create an IPv6Address object for validation and to build IPv6 collapsed form.
+    ipv6_address = ipaddress.IPv6Address(full_address)
+
+    compressed_address = ipv6_address.compressed
+
+    # We padded with zeros so address might endswith "::" or ":0".
+    if compressed_address.endswith(("::", ":0")):
+        return compressed_address[:-2]
+
+    # Otherwise, it means 8th segment of ipv6 address was not full and not composed of 0 only
+    # e.g. 2001:0db8:0000:0000:0000:0000:03
+    return compressed_address
+
+
 async def search_resolver(
     root: dict,  # pylint: disable=unused-argument
     info: GraphQLResolveInfo,
@@ -49,6 +116,12 @@ async def search_resolver(
         if matching:
             result.append(matching)
     else:
+        try:
+            # Convert any IPv6 address, network or partial address to collapsed format as it might be stored in db.
+            q = _collapse_ipv6(q)
+        except (ValueError, ipaddress.AddressValueError):
+            pass
+
         result.extend(
             await NodeManager.query(
                 db=context.db,

backend/tests/unit/graphql/queries/test_search.py

@@ -1,9 +1,11 @@
+import pytest
 from graphql import graphql
 
 from infrahub.core.branch import Branch
 from infrahub.core.node import Node
 from infrahub.database import InfrahubDatabase
 from infrahub.graphql.initialization import prepare_graphql_params
+from infrahub.graphql.queries.search import _collapse_ipv6
 
 SEARCH_QUERY = """
 query ($search: String!) {
@@ -93,3 +95,190 @@ async def test_search_anywhere_by_string(
 
     assert sorted(node_ids) == sorted([person_john_main.id, person_jane_main.id])
     assert sorted(node_kinds) == sorted([person_john_main.get_kind(), person_jane_main.get_kind()])
+
+
+async def test_search_ipv6_address_extended_format(
+    db: InfrahubDatabase,
+    ip_dataset_01,
+    branch: Branch,
+):
+    gql_params = prepare_graphql_params(db=db, include_subscription=False, branch=branch)
+
+    res_collapsed = await graphql(
+        schema=gql_params.schema,
+        source=SEARCH_QUERY,
+        context_value=gql_params.context,
+        root_value=None,
+        variable_values={"search": "2001:db8::"},
+    )
+
+    res_extended = await graphql(
+        schema=gql_params.schema,
+        source=SEARCH_QUERY,
+        context_value=gql_params.context,
+        root_value=None,
+        variable_values={"search": "2001:0db8:0000:0000:0000:0000:0000:0000"},
+    )
+
+    assert (
+        res_extended.data["InfrahubSearchAnywhere"]["count"]
+        == res_collapsed.data["InfrahubSearchAnywhere"]["count"]
+        == 2
+    )
+
+    assert (
+        res_extended.data["InfrahubSearchAnywhere"]["edges"][0]["node"]["id"]
+        == res_collapsed.data["InfrahubSearchAnywhere"]["edges"][0]["node"]["id"]
+    )
+
+    assert (
+        res_extended.data["InfrahubSearchAnywhere"]["edges"][1]["node"]["id"]
+        == res_collapsed.data["InfrahubSearchAnywhere"]["edges"][1]["node"]["id"]
+    )
+
+
+async def test_search_ipv6_network_extended_format(
+    db: InfrahubDatabase,
+    ip_dataset_01,
+    branch: Branch,
+):
+    gql_params = prepare_graphql_params(db=db, include_subscription=False, branch=branch)
+
+    res_collapsed = await graphql(
+        schema=gql_params.schema,
+        source=SEARCH_QUERY,
+        context_value=gql_params.context,
+        root_value=None,
+        variable_values={"search": "2001:db8::/48"},
+    )
+
+    res_extended = await graphql(
+        schema=gql_params.schema,
+        source=SEARCH_QUERY,
+        context_value=gql_params.context,
+        root_value=None,
+        variable_values={"search": "2001:0db8:0000:0000:0000:0000:0000:0000/48"},
+    )
+
+    assert (
+        res_extended.data["InfrahubSearchAnywhere"]["count"]
+        == res_collapsed.data["InfrahubSearchAnywhere"]["count"]
+        == 1
+    )
+
+    assert (
+        res_extended.data["InfrahubSearchAnywhere"]["edges"][0]["node"]["id"]
+        == res_collapsed.data["InfrahubSearchAnywhere"]["edges"][0]["node"]["id"]
+    )
+
+
+async def test_search_ipv6_partial_address(
+    db: InfrahubDatabase,
+    ip_dataset_01,
+    branch: Branch,
+):
+    gql_params = prepare_graphql_params(db=db, include_subscription=False, branch=branch)
+
+    res_two_segments = await graphql(
+        schema=gql_params.schema,
+        source=SEARCH_QUERY,
+        context_value=gql_params.context,
+        root_value=None,
+        variable_values={"search": "2001:0db8"},
+    )
+
+    res_partial_segment_1 = await graphql(
+        schema=gql_params.schema,
+        source=SEARCH_QUERY,
+        context_value=gql_params.context,
+        root_value=None,
+        variable_values={"search": "2001:0db8:0"},
+    )
+
+    res_partial_segment_2 = await graphql(
+        schema=gql_params.schema,
+        source=SEARCH_QUERY,
+        context_value=gql_params.context,
+        root_value=None,
+        variable_values={"search": "2001:0db8:0000:0"},
+    )
+
+    assert (
+        res_two_segments.data["InfrahubSearchAnywhere"]["count"]
+        == res_partial_segment_1.data["InfrahubSearchAnywhere"]["count"]
+        == res_partial_segment_2.data["InfrahubSearchAnywhere"]["count"]
+        == 2
+    )
+
+    assert (
+        res_two_segments.data["InfrahubSearchAnywhere"]["edges"][0]["node"]["id"]
+        == res_partial_segment_1.data["InfrahubSearchAnywhere"]["edges"][0]["node"]["id"]
+        == res_partial_segment_2.data["InfrahubSearchAnywhere"]["edges"][0]["node"]["id"]
+    )
+
+
+async def test_search_ipv4(
+    db: InfrahubDatabase,
+    ip_dataset_01,
+    branch: Branch,
+):
+    """
+    This only tests that ipv6 search specific behavior does not break ipv4 search.
+    """
+
+    gql_params = prepare_graphql_params(db=db, include_subscription=False, branch=branch)
+
+    result_address = await graphql(
+        schema=gql_params.schema,
+        source=SEARCH_QUERY,
+        context_value=gql_params.context,
+        root_value=None,
+        variable_values={"search": "10.0.0.0"},
+    )
+
+    result_network = await graphql(
+        schema=gql_params.schema,
+        source=SEARCH_QUERY,
+        context_value=gql_params.context,
+        root_value=None,
+        variable_values={"search": "10.0.0.0/8"},
+    )
+
+    assert (
+        result_address.data["InfrahubSearchAnywhere"]["count"]
+        == result_network.data["InfrahubSearchAnywhere"]["count"]
+        == 1
+    )
+
+    assert (
+        result_address.data["InfrahubSearchAnywhere"]["edges"][0]["node"]["id"]
+        == result_network.data["InfrahubSearchAnywhere"]["edges"][0]["node"]["id"]
+    )
+
+
+@pytest.mark.parametrize(
+    "query,expected",
+    [
+        ("2001:0db8:0000:0000:0000:0000:0000:0000/48", "2001:db8::/48"),
+        ("2001:0db8:0000:0000:0000:0000:0000:0000", "2001:db8::"),
+        ("2001:0db8", "2001:db8"),
+        ("2001:0db8:0", "2001:db8"),
+        ("2001:0db8:0000", "2001:db8"),
+        ("2001:0db8:0000:0", "2001:db8"),
+        ("2001:0db8:0000:0000:00", "2001:db8"),
+        ("2001:0db8:0000:0001:00", "2001:db8:0:1"),
+        ("2001:0db8:0001:0002:00", "2001:db8:1:2"),
+        ("2001:0db8:0001:0000:0002:0000:0003", "2001:db8:1:0:2:0:3"),
+    ],
+)
+def test_collapse_ipv6_address_or_network(query, expected):
+    assert _collapse_ipv6(query) == expected
+
+
+@pytest.mark.parametrize(
+    "query",
+    ["invalid", "invalid:case", "2001:invalid", "2001:0db81:0000", "10.0.0.0", "2001:db8:1"],
+)
+def test_collapse_ipv6_address_or_network_invalid_cases(query):
+    with pytest.raises(ValueError):
+        _collapse_ipv6(query)

changelog/4613.fixed.md

@@ -0,0 +1 @@
+Search anywhere now supports IPv6 extended format