From d182b3bd692811894bb4502221998460be64fd02 Mon Sep 17 00:00:00 2001
From: Olivier Gayot <olivier.gayot@canonical.com>
Date: Tue, 23 Aug 2022 17:49:48 +0200
Subject: [PATCH] integration: speed up check for listening interfaces in
 system_setup

Signed-off-by: Olivier Gayot <olivier.gayot@canonical.com>
---
 scripts/runtests.sh                        | 36 +--------
 scripts/test-system-setup-loopback-only.py | 93 ++++++++++++++++++++++
 2 files changed, 97 insertions(+), 32 deletions(-)
 create mode 100755 scripts/test-system-setup-loopback-only.py

diff --git a/scripts/runtests.sh b/scripts/runtests.sh
index fd4337f0..cb53a39c 100755
--- a/scripts/runtests.sh
+++ b/scripts/runtests.sh
@@ -248,38 +248,10 @@ if [ "${RELEASE%.*}" -ge 20 ]; then
         kill $subiquity_pid || true
         exit 1
     fi
-    loopback_failed=0
-    unallowed_failed=0
-    # Assert that only loopback interface is accepted.
-    interfaces=($(ip --json link show up | jq -r '.[]["ifname"] | select ( . != null )'))
-    for if in ${interfaces[@]}; do
-        for ipv in 4 6; do
-            curl_ec=0
-            timeout 10s \
-                curl -$ipv "http://localhost:$port/meta/status" --interface $if \
-                || curl_ec=$?
-            # Loopback should exit 0 on IPv4
-            if [ $if = "lo" ]; then
-                if [ $curl_ec -ne 0 -a $ipv -eq 4 ]; then
-                    loopback_failed=1
-                fi
-            # Everything else should not.
-            else
-                if [ $curl_ec -eq 0 ]; then
-                    unallowed_failed=1
-                fi
-            fi
-        done
-    done
-    kill $subiquity_pid || true
-    if [ $loopback_failed -ne 0 ]; then
-        echo "Loopback was expected to connect"
-        exit 1
-    fi
-    if [ $unallowed_failed -ne 0 ]; then
-        echo "Only the loopback interface should be allowed."
-        exit 1
-    fi
+
+    scripts/test-system-setup-loopback-only.py --port "$port" --debug
+
+    kill -- "$subiquity_pid" || true
 
     # Test system_setup autoinstall.
     for mode in "" "-full" "-no-shutdown"; do
diff --git a/scripts/test-system-setup-loopback-only.py b/scripts/test-system-setup-loopback-only.py
new file mode 100755
index 00000000..7ffa8ec5
--- /dev/null
+++ b/scripts/test-system-setup-loopback-only.py
@@ -0,0 +1,93 @@
+#!/usr/bin/env python3
+
+""" Makes sure system_setup is only listening to the loopback interface. """
+
+import argparse
+import asyncio
+from dataclasses import dataclass
+import json
+import logging
+import subprocess
+from typing import List
+
+
+class FailedTestCase(Exception):
+    pass
+
+
+@dataclass
+class Test:
+    interface: str
+    url: str
+    family: int
+    expect_success: bool
+
+
+def read_network_interfaces() -> List[str]:
+    """ Return a list of network interfaces that are up. """
+    cmd = ["ip", "--json", "link", "show", "up"]
+    output = subprocess.check_output(cmd, text=True)
+    data = json.loads(output)
+    return [iface["ifname"] for iface in data if iface.get("ifname")]
+
+
+async def test_connect(cmd: List[str]) -> bool:
+    """ Return true if the command specified exits with status 0 within 10
+    seconds. """
+    proc = await asyncio.create_subprocess_exec(
+            *cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    try:
+        await asyncio.wait_for(proc.wait(), 10)
+    except asyncio.TimeoutError:
+        return False
+    return proc.returncode == 0
+
+
+async def run_test(test: Test) -> None:
+    """ Execute a test and raise a FailedTestCase if it fails. """
+    logging.debug("Test: %s", test)
+    cmd = ["curl", f"-{test.family}", test.url, "--interface", test.interface]
+    status = await test_connect(cmd)
+    if status != test.expect_success:
+        logging.error("cmd %s exited %s but we expected %s", cmd,
+                      "successfully" if status else "unsuccessfully",
+                      "success" if test.expect_success else "failure")
+        raise FailedTestCase
+
+
+async def main() -> None:
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--debug", action="store_true")
+    parser.add_argument("--port", type=int, default=50321)
+
+    args = vars(parser.parse_args())
+
+    if args["debug"]:
+        logging.getLogger().level = logging.DEBUG
+
+    interfaces = read_network_interfaces()
+    logging.debug("interfaces = %s", interfaces)
+
+    coroutines = []
+
+    url = f"http://localhost:{args['port']}/meta/status"
+    for iface in interfaces:
+        for family in 4, 6:
+            if family == 4 and iface == "lo":
+                # Loopback should succeed on IPv4
+                expect_success=True
+            else:
+                # Everything else should not
+                expect_success = False
+            coroutines.append(run_test(Test(
+                interface=iface, url=url, family=family,
+                expect_success=expect_success)))
+
+    results = await asyncio.gather(*coroutines, return_exceptions=True)
+    if any(map(lambda x: isinstance(x, FailedTestCase), results)):
+        raise FailedTestCase
+
+
+if __name__ == "__main__":
+    asyncio.run(main())