Merge ~mpontillo/maas:workaround-netaddr-slash128-bug--2.4 into maas:2.4

[Mike Pontillo]

Backport 80dcada - LP: #1789721 - Workaround for StopIteration seen via get_source_address().

This works around a bug in netaddr that causes StopIteration to be raised when a /128 IPv6 network is passed in.

[Mike Pontillo (mpontillo)]

Self-approve backport.

review: Approve

diff --git a/src/provisioningserver/utils/network.py b/src/provisioningserver/utils/network.py
index 6093871..06ed74e 100644
--- a/src/provisioningserver/utils/network.py
+++ b/src/provisioningserver/utils/network.py
@@ -1421,20 +1421,41 @@ def coerce_to_valid_hostname(hostname):
     return hostname
 
 
-def get_source_address(destination_ip: IPAddressOrNetwork):
-    """Returns the local source address for the specified destination IP.
+def get_source_address(destination: IPAddressOrNetwork):
+    """Returns the local source address for the specified destination.
 
-    :param destination_ip: Can be an IP address in string format, an IPNetwork,
+    :param destination: Can be an IP address in string format, an IPNetwork,
         or an IPAddress object.
     :return: the string representation of the local IP address that would be
         used for communication with the specified destination.
     """
-    if isinstance(destination_ip, IPNetwork):
-        destination_ip = IPAddress(next(destination_ip.iter_hosts()))
+    if isinstance(destination, IPNetwork):
+        if destination.prefixlen == 128:
+            # LP: #1789721 - netaddr raises an exception when using
+            # iter_hosts() with a /128 address.
+            destination = destination.ip
+        else:
+            # Make sure to return a host (not a network) if possible.
+            # Using iter_hosts() here is a general way to accomplish that.
+            destination = IPAddress(next(destination.iter_hosts()))
     else:
-        destination_ip = make_ipaddress(destination_ip)
-    if destination_ip.is_ipv4_mapped():
-        destination_ip = destination_ip.ipv4()
+        destination = make_ipaddress(destination)
+    if destination.is_ipv4_mapped():
+        destination = destination.ipv4()
+    return get_source_address_for_ipaddress(destination)
+
+
+def get_source_address_for_ipaddress(destination_ip: IPAddress):
+    """Returns the local source address for the specified IPAddress.
+
+    Callers should generally use the `get_source_address()` utility instead;
+    it is more flexible about the type of the input parameter.
+
+    :param destination_ip: Can be an IP address in string format, an IPNetwork,
+        or an IPAddress object.
+    :return: the string representation of the local IP address that would be
+        used for communication with the specified destination.
+    """
     af = AF_INET if destination_ip.version == 4 else AF_INET6
     with socket.socket(af, socket.SOCK_DGRAM) as sock:
         peername = str(destination_ip)
diff --git a/src/provisioningserver/utils/tests/test_network.py b/src/provisioningserver/utils/tests/test_network.py
index f23f31b..b7472ef 100644
--- a/src/provisioningserver/utils/tests/test_network.py
+++ b/src/provisioningserver/utils/tests/test_network.py
@@ -2294,6 +2294,33 @@ class TestGetSourceAddress(MAASTestCase):
         self.assertThat(
             get_source_address(IPNetwork("127.0.0.1/8")), Equals("127.0.0.1"))
 
+    def test__ipnetwork_works_for_ipv4_slash32(self):
+        mock = self.patch(network_module, 'get_source_address_for_ipaddress')
+        mock.return_value = "10.0.0.1"
+        self.assertThat(
+            get_source_address(IPNetwork("10.0.0.1/32")), Equals("10.0.0.1"))
+
+    def test__ipnetwork_works_for_ipv6_slash128(self):
+        mock = self.patch(network_module, 'get_source_address_for_ipaddress')
+        mock.return_value = "2001:67c:1560::1"
+        self.assertThat(
+            get_source_address(
+                IPNetwork("2001:67c:1560::/128")), Equals("2001:67c:1560::1"))
+
+    def test__ipnetwork_works_for_ipv6_slash48(self):
+        mock = self.patch(network_module, 'get_source_address_for_ipaddress')
+        mock.return_value = "2001:67c:1560::1"
+        self.assertThat(
+            get_source_address(
+                IPNetwork("2001:67c:1560::/48")), Equals("2001:67c:1560::1"))
+
+    def test__ipnetwork_works_for_ipv6_slash64(self):
+        mock = self.patch(network_module, 'get_source_address_for_ipaddress')
+        mock.return_value = "2001:67c:1560::1"
+        self.assertThat(
+            get_source_address(
+                IPNetwork("2001:67c:1560::1/64")), Equals("2001:67c:1560::1"))
+
     def test__accepts_ipaddress(self):
         self.assertThat(
             get_source_address(IPAddress("127.0.0.1")), Equals("127.0.0.1"))