[aks-preview] az aks bastion: fix tunnel when Bastion is in a different subscription

src/aks-preview/HISTORY.rst

@@ -12,6 +12,7 @@ To release a new version, please select a new version number (usually plus 1 to
 Pending
 +++++++
 * `az aks upgrade`: Fix `--node-image-only` to skip Machines mode agent pools, which do not support node image version upgrade.
+* `az aks bastion`: Fix failure when the bastion host is in a different subscription than the cluster by using the subscription from the bastion resource ID for the internal `az network bastion tunnel` command.
 
 21.0.0b4
 +++++++

src/aks-preview/azext_aks_preview/bastion/bastion.py

@@ -30,9 +30,10 @@
 
 # pylint: disable=too-few-public-methods
 class BastionResource:
-    def __init__(self, name, resource_group):
+    def __init__(self, name, resource_group, subscription=None):
         self.name = name
         self.resource_group = resource_group
+        self.subscription = subscription
 
 
 def aks_bastion_parse_bastion_resource(
@@ -45,7 +46,11 @@ def aks_bastion_parse_bastion_resource(
         if is_valid_resource_id(bastion):
             parsed_id = parse_resource_id(bastion)
             return BastionResource(
-                name=parsed_id["name"], resource_group=parsed_id["resource_group"]
+                name=parsed_id["name"],
+                resource_group=parsed_id["resource_group"],
+                # the bastion may live in a different subscription than the cluster,
+                # so honor the subscription embedded in the resource ID when present
+                subscription=parsed_id.get("subscription") or subscription_id,
             )
         if is_valid_resource_name(bastion):
             for resource_group in resource_groups:
@@ -85,7 +90,11 @@ def aks_bastion_parse_bastion_resource(
                     bastion,
                     resource_group,
                 )
-                return BastionResource(name=bastion, resource_group=resource_group)
+                return BastionResource(
+                    name=bastion,
+                    resource_group=resource_group,
+                    subscription=subscription_id,
+                )
     logger.warning(
         "No valid bastion resource provided: '%s'. Attempting to locate one from resource groups: '%s'.",
         bastion,
@@ -127,7 +136,11 @@ def aks_bastion_parse_bastion_resource(
             bastions[0]["name"],
             resource_group,
         )
-        return BastionResource(name=bastions[0]["name"], resource_group=resource_group)
+        return BastionResource(
+            name=bastions[0]["name"],
+            resource_group=resource_group,
+            subscription=subscription_id,
+        )
     raise ResourceNotFoundError(
         "No bastion found in the provided resource groups: "
         f"{', '.join(resource_groups)}. Please provide a valid bastion name or resource ID."
@@ -484,8 +497,11 @@ async def _aks_bastion_launch_tunnel(bastion_resource, port, mc_id, subscription
             f"{az_cmd_name} network bastion tunnel --resource-group {bastion_resource.resource_group} "
             f"--name {bastion_resource.name} --port {port} --target-resource-id {mc_id} --resource-port 443"
         )
-        if subscription_id:
-            cmd += f" --subscription {subscription_id}"
+        # the bastion may live in a different subscription than the cluster; prefer the
+        # subscription resolved from the bastion resource over the cluster subscription
+        bastion_subscription_id = getattr(bastion_resource, "subscription", None) or subscription_id
+        if bastion_subscription_id:
+            cmd += f" --subscription {bastion_subscription_id}"
         logger.warning("Creating bastion tunnel with command: '%s'", cmd)
 
         # Use start_new_session on Unix to create a new process group

src/aks-preview/azext_aks_preview/tests/latest/test_aks_bastion.py

@@ -0,0 +1,111 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+import unittest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from azext_aks_preview.bastion.bastion import (
+    BastionResource,
+    _aks_bastion_launch_tunnel,
+    aks_bastion_parse_bastion_resource,
+)
+
+
+class TestAksBastionParseResource(unittest.TestCase):
+    def test_cross_subscription_resource_id_preserves_bastion_subscription(self):
+        # the bastion lives in a different (hub) subscription than the cluster
+        bastion_id = (
+            "/subscriptions/hub-sub-id/resourceGroups/bastion-rg/"
+            "providers/Microsoft.Network/bastionHosts/my-bastion"
+        )
+        resource = aks_bastion_parse_bastion_resource(
+            bastion_id, ["node-rg"], subscription_id="aks-sub-id"
+        )
+        self.assertEqual(resource.name, "my-bastion")
+        self.assertEqual(resource.resource_group, "bastion-rg")
+        # the subscription from the bastion resource ID must win over the cluster subscription
+        self.assertEqual(resource.subscription, "hub-sub-id")
+
+    def test_resource_id_without_subscription_falls_back_to_cluster_subscription(self):
+        # parse_resource_id may not yield a subscription; fall back to the cluster one
+        bastion_id = (
+            "/subscriptions/hub-sub-id/resourceGroups/bastion-rg/"
+            "providers/Microsoft.Network/bastionHosts/my-bastion"
+        )
+        with patch(
+            "azext_aks_preview.bastion.bastion.parse_resource_id",
+            return_value={"name": "my-bastion", "resource_group": "bastion-rg"},
+        ):
+            resource = aks_bastion_parse_bastion_resource(
+                bastion_id,
+                ["node-rg"],
+                subscription_id="aks-sub-id",
+            )
+        self.assertEqual(resource.subscription, "aks-sub-id")
+
+
+class TestAksBastionLaunchTunnel(unittest.IsolatedAsyncioTestCase):
+    async def test_tunnel_uses_bastion_subscription(self):
+        # regression guard: the inner `az network bastion tunnel` must be scoped to the
+        # bastion's subscription, not the cluster subscription (see azure-cli#33579)
+        bastion_resource = BastionResource(
+            name="my-bastion",
+            resource_group="bastion-rg",
+            subscription="hub-sub-id",
+        )
+        mock_process = MagicMock()
+        mock_process.wait = AsyncMock(return_value=0)
+        mock_process.returncode = 0
+
+        with patch(
+            "azext_aks_preview.bastion.bastion.asyncio.create_subprocess_exec",
+            new=AsyncMock(return_value=mock_process),
+        ) as mock_exec:
+            await _aks_bastion_launch_tunnel(
+                bastion_resource,
+                port=12345,
+                mc_id="/subscriptions/aks-sub-id/resourceGroups/aks-rg/"
+                "providers/Microsoft.ContainerService/managedClusters/cluster",
+                subscription_id="aks-sub-id",
+            )
+
+        args = mock_exec.call_args.args
+        self.assertIn("--subscription", args)
+        sub_index = args.index("--subscription")
+        self.assertEqual(args[sub_index + 1], "hub-sub-id")
+        self.assertNotIn("aks-sub-id", args[sub_index + 1])
+
+    async def test_tunnel_falls_back_to_cluster_subscription(self):
+        # when the bastion has no subscription of its own (name-based discovery in the
+        # node resource group), the cluster subscription is used
+        bastion_resource = BastionResource(
+            name="my-bastion",
+            resource_group="node-rg",
+            subscription=None,
+        )
+        mock_process = MagicMock()
+        mock_process.wait = AsyncMock(return_value=0)
+        mock_process.returncode = 0
+
+        with patch(
+            "azext_aks_preview.bastion.bastion.asyncio.create_subprocess_exec",
+            new=AsyncMock(return_value=mock_process),
+        ) as mock_exec:
+            await _aks_bastion_launch_tunnel(
+                bastion_resource,
+                port=12345,
+                mc_id="/subscriptions/aks-sub-id/resourceGroups/aks-rg/"
+                "providers/Microsoft.ContainerService/managedClusters/cluster",
+                subscription_id="aks-sub-id",
+            )
+
+        args = mock_exec.call_args.args
+        self.assertIn("--subscription", args)
+        sub_index = args.index("--subscription")
+        self.assertEqual(args[sub_index + 1], "aks-sub-id")
+
+
+if __name__ == "__main__":
+    unittest.main()