[SSH] Fix Arc proxy registry hostname for active cloud

src/ssh/HISTORY.md

@@ -1,5 +1,9 @@
 Release History
 ===============
+2.0.10
+------
+* [Bug Fix] Resolve the Arc SSH proxy registry hostname from the active cloud endpoint.
+
 2.0.9
 -----
 * Migrate code from Azure SDK to AAZ based commands for compute operations (VM)

src/ssh/azext_ssh/connectivity_utils.py

@@ -219,7 +219,7 @@ def _handle_relay_connection_delay(cmd, message):
 
 
 # Downloads client side proxy to connect to Arc Connectivity Platform
-def install_client_side_proxy(arc_proxy_folder):
+def install_client_side_proxy(cmd, arc_proxy_folder):
 
     client_operating_system = _get_client_operating_system()
     client_architecture = _get_client_architeture()
@@ -241,14 +241,40 @@ def install_client_side_proxy(arc_proxy_folder):
             for f in older_version_files:
                 file_utils.delete_file(f, f"failed to delete older version file {f}", warning=True)
 
-        _download_proxy_from_MCR(install_dir, proxy_name, client_operating_system, client_architecture)
+        _download_proxy_from_MCR(cmd, install_dir, proxy_name, client_operating_system, client_architecture)
         _check_proxy_installation(install_dir, proxy_name)
 
     return install_location
 
 
-def _download_proxy_from_MCR(dest_dir, proxy_name, operating_system, architecture):
-    mar_target = f"{consts.CLIENT_PROXY_MCR_TARGET}/{operating_system.lower()}/{architecture}/ssh-proxy"
+def get_mcr_path(active_directory_endpoint: str) -> str:
+    active_directory_array = active_directory_endpoint.split(".")
+
+    if active_directory_endpoint.endswith((".us", ".cn")):
+        return "mcr.microsoft.com"
+
+    mcr_postfix = "com"
+    if len(active_directory_array) == 4 and active_directory_array[2] == "microsoft":
+        mcr_postfix = active_directory_array[3].strip("/")
+    elif len(active_directory_array) == 5:
+        mcr_postfix = (
+            active_directory_array[2]
+            + "."
+            + active_directory_array[3]
+            + "."
+            + active_directory_array[4].strip("/")
+        )
+
+    mcr_url = f"mcr.microsoft.{mcr_postfix}"
+    return mcr_url
+
+
+def _download_proxy_from_MCR(cmd, dest_dir, proxy_name, operating_system, architecture):
+    hostname = get_mcr_path(cmd.cli_ctx.cloud.endpoints.active_directory)
+    mar_target = (
+        f"{hostname}/{consts.CLIENT_PROXY_MCR_REPOSITORY}/"
+        f"{operating_system.lower()}/{architecture}/ssh-proxy"
+    )
     logger.debug("Downloading Arc Connectivity Proxy from %s in Microsoft Artifact Regristy.", mar_target)
 
     client = oras.client.OrasClient()

src/ssh/azext_ssh/constants.py

@@ -9,7 +9,7 @@
 AGENT_MINIMUM_VERSION_MAJOR = 1
 AGENT_MINIMUM_VERSION_MINOR = 31
 CLIENT_PROXY_VERSION = "1.3.033291"
-CLIENT_PROXY_MCR_TARGET = "mcr.microsoft.com/azureconnectivity/proxy"
+CLIENT_PROXY_MCR_REPOSITORY = "azureconnectivity/proxy"
 CLEANUP_TOTAL_TIME_LIMIT_IN_SECONDS = 120
 CLEANUP_TIME_INTERVAL_IN_SECONDS = 10
 CLEANUP_AWAIT_TERMINATION_IN_SECONDS = 30

src/ssh/azext_ssh/custom.py

@@ -187,7 +187,7 @@ def _do_ssh_op(cmd, op_info, op_call):
 
     try:
         if op_info.is_arc():
-            op_info.proxy_path = connectivity_utils.install_client_side_proxy(op_info.ssh_proxy_folder)
+            op_info.proxy_path = connectivity_utils.install_client_side_proxy(cmd, op_info.ssh_proxy_folder)
             (op_info.relay_info, op_info.new_service_config) = connectivity_utils.get_relay_information(
                 cmd, op_info.resource_group_name, op_info.vm_name, op_info.resource_type,
                 cert_lifetime, op_info.port, op_info.yes_without_prompt)

src/ssh/azext_ssh/tests/latest/test_connectivity_utils.py

@@ -11,6 +11,48 @@
 
 class SshConnectivityUtilsCommandTest(unittest.TestCase):
 
+    def test_get_mcr_path(self):
+        input_active_directory = "login.microsoftonline.com"
+        expected_output = "mcr.microsoft.com"
+        self.assertEqual(connectivity_utils.get_mcr_path(input_active_directory), expected_output)
+
+        input_active_directory = "login.microsoftonline.us"
+        expected_output = "mcr.microsoft.com"
+        self.assertEqual(connectivity_utils.get_mcr_path(input_active_directory), expected_output)
+
+        input_active_directory = "login.chinacloudapi.cn"
+        expected_output = "mcr.microsoft.com"
+        self.assertEqual(connectivity_utils.get_mcr_path(input_active_directory), expected_output)
+
+        input_active_directory = "https://login.microsoftonline.microsoft.foo"
+        expected_output = "mcr.microsoft.foo"
+        self.assertEqual(connectivity_utils.get_mcr_path(input_active_directory), expected_output)
+
+        input_active_directory = "https://login.microsoftonline.some.cloud.bar"
+        expected_output = "mcr.microsoft.some.cloud.bar"
+        self.assertEqual(connectivity_utils.get_mcr_path(input_active_directory), expected_output)
+
+    @mock.patch('azext_ssh.connectivity_utils.file_utils.delete_file')
+    @mock.patch('azext_ssh.connectivity_utils._extract_proxy_tar_files')
+    @mock.patch('azext_ssh.connectivity_utils._get_proxy_package_path_from_oras_response')
+    @mock.patch('azext_ssh.connectivity_utils.telemetry.add_extension_event')
+    @mock.patch('azext_ssh.connectivity_utils.oras.client.OrasClient')
+    def test_download_proxy_cloud_target(self, mock_client, mock_telemetry, mock_package, mock_extract, mock_delete):
+        cmd = mock.Mock()
+        cmd.cli_ctx.cloud.endpoints.active_directory = 'https://login.microsoftonline.microsoft.foo/'
+        mock_package.return_value = 'proxy.tar.gz'
+
+        connectivity_utils._download_proxy_from_MCR(cmd, '/dir', 'proxy', 'Linux', 'amd64')
+
+        mock_client.return_value.pull.assert_called_once_with(
+            target='mcr.microsoft.foo/azureconnectivity/proxy/linux/amd64/ssh-proxy:'
+                   + connectivity_utils.consts.CLIENT_PROXY_VERSION,
+            outdir='/dir')
+        mock_package.assert_called_once_with(mock_client.return_value.pull.return_value)
+        mock_extract.assert_called_once_with('proxy.tar.gz', '/dir', 'proxy')
+        mock_delete.assert_called_once()
+        mock_telemetry.assert_called_once()
+
     @mock.patch('platform.machine')
     def test_get_client_architecture_arm64(self, mock_plat):
         mock_plat.return_value = 'arm64'
@@ -133,8 +175,9 @@ def test_install_proxy_create_dir(self, mock_check, mock_download, mock_dir, moc
         mock_get_proxy_dir.return_value = "/dir/proxy"
         mock_isfile.return_value = False
 
-        connectivity_utils.install_client_side_proxy(None)
+        cmd = mock.Mock()
+        connectivity_utils.install_client_side_proxy(cmd, None)
 
         mock_dir.assert_called_once_with("/dir/proxy", "Failed to create client proxy directory \'/dir/proxy\'.")
-        mock_download.assert_called_once_with("/dir/proxy", "sshProxy_linux_arm64_1_3_033291", "linux", "arm64")
+        mock_download.assert_called_once_with(cmd, "/dir/proxy", "sshProxy_linux_arm64_1_3_033291", "linux", "arm64")
         mock_check.assert_called_once_with("/dir/proxy", "sshProxy_linux_arm64_1_3_033291")

src/ssh/azext_ssh/tests/latest/test_custom.py

@@ -432,7 +432,7 @@ def test_do_ssh_op_arc_local_user(self, mock_get_cert, mock_check_keys, mock_sta
 
         custom._do_ssh_op(cmd, op_info, mock_op)
 
-        mock_get_proxy.assert_called_once_with('proxy')
+        mock_get_proxy.assert_called_once_with(cmd, 'proxy')
         mock_get_relay_info.assert_called_once_with(cmd, 'rg', 'vm', 'Microsoft.HybridCompute/machines', None, "port", False)
         mock_op.assert_called_once_with(op_info, False, False)
         mock_get_cert.assert_not_called()
@@ -483,7 +483,7 @@ def test_do_ssh_arc_op_aad_user(self, mock_cert_exp, mock_start_ssh, mock_write_
         mock_check_files.assert_called_once_with("publicfile", "privatefile", None, "client")
         mock_get_mod_exp.assert_called_once_with("public")
         mock_write_cert.assert_called_once_with("certificate", "public-aadcert.pub")
-        mock_get_proxy.assert_called_once_with('proxy')
+        mock_get_proxy.assert_called_once_with(cmd, 'proxy')
         mock_get_relay_info.assert_called_once_with(cmd, 'rg', 'vm', 'Microsoft.HybridCompute/machines', 3600, 'port', False)
         mock_op.assert_called_once_with(op_info, False, True)
 

src/ssh/setup.py

@@ -7,7 +7,7 @@
 
 from setuptools import setup, find_packages
 
-VERSION = "2.0.9"
+VERSION = "2.0.10"
 
 CLASSIFIERS = [
     'Development Status :: 4 - Beta',