Raise CredentialUnavailableError when CLI subprocess times out (Azure#18509)

Author: chlowell

sdk/identity/azure-identity/azure/identity/_credentials/azure_cli.py

@@ -7,11 +7,12 @@
 import os
 import platform
 import re
+import subprocess
 import sys
 import time
 from typing import TYPE_CHECKING
 
-import subprocess
+import six
 
 from azure.core.credentials import AccessToken
 from azure.core.exceptions import ClientAuthenticationError
@@ -53,9 +54,7 @@ def get_token(self, *scopes, **kwargs):  # pylint:disable=no-self-use,unused-arg
         """
 
         resource = _scopes_to_resource(*scopes)
-        output, error = _run_command(COMMAND_LINE.format(resource))
-        if error:
-            raise error
+        output = _run_command(COMMAND_LINE.format(resource))
 
         token = parse_token(output)
         if not token:
@@ -120,25 +119,25 @@ def _run_command(command):
         if platform.python_version() >= "3.3":
             kwargs["timeout"] = 10
 
-        output = subprocess.check_output(args, **kwargs)
-        return output, None
+        return subprocess.check_output(args, **kwargs)
     except subprocess.CalledProcessError as ex:
         # non-zero return from shell
         if ex.returncode == 127 or ex.output.startswith("'az' is not recognized"):
-            error = CredentialUnavailableError(message=CLI_NOT_FOUND)
-        elif "az login" in ex.output or "az account set" in ex.output:
-            error = CredentialUnavailableError(message=NOT_LOGGED_IN)
+            raise CredentialUnavailableError(message=CLI_NOT_FOUND)
+        if "az login" in ex.output or "az account set" in ex.output:
+            raise CredentialUnavailableError(message=NOT_LOGGED_IN)
+
+        # return code is from the CLI -> propagate its output
+        if ex.output:
+            message = sanitize_output(ex.output)
         else:
-            # return code is from the CLI -> propagate its output
-            if ex.output:
-                message = sanitize_output(ex.output)
-            else:
-                message = "Failed to invoke Azure CLI"
-            error = ClientAuthenticationError(message=message)
+            message = "Failed to invoke Azure CLI"
+        raise ClientAuthenticationError(message=message)
     except OSError as ex:
         # failed to execute 'cmd' or '/bin/sh'; CLI may or may not be installed
         error = CredentialUnavailableError(message="Failed to execute '{}'".format(args[0]))
+        six.raise_from(error, ex)
     except Exception as ex:  # pylint:disable=broad-except
-        error = ex
-
-    return None, error
+        # could be a timeout, for example
+        error = CredentialUnavailableError(message="Failed to invoke the Azure CLI")
+        six.raise_from(error, ex)

sdk/identity/azure-identity/azure/identity/aio/_credentials/azure_cli.py

@@ -81,13 +81,15 @@ async def _run_command(command: str) -> str:
             cwd=working_directory,
             env=dict(os.environ, AZURE_CORE_NO_COLOR="true")
         )
+        stdout, _ = await asyncio.wait_for(proc.communicate(), 10)
+        output = stdout.decode()
     except OSError as ex:
         # failed to execute 'cmd' or '/bin/sh'; CLI may or may not be installed
         error = CredentialUnavailableError(message="Failed to execute '{}'".format(args[0]))
         raise error from ex
-
-    stdout, _ = await asyncio.wait_for(proc.communicate(), 10)
-    output = stdout.decode()
+    except asyncio.TimeoutError as ex:
+        proc.kill()
+        raise CredentialUnavailableError(message="Timed out waiting for Azure CLI") from ex
 
     if proc.returncode == 0:
         return output

sdk/identity/azure-identity/tests/test_cli_credential.py

@@ -4,6 +4,7 @@
 # ------------------------------------
 from datetime import datetime
 import json
+import sys
 
 from azure.identity import AzureCliCredential, CredentialUnavailableError
 from azure.identity._credentials.azure_cli import CLI_NOT_FOUND, NOT_LOGGED_IN
@@ -136,3 +137,14 @@ def test_subprocess_error_does_not_expose_token(output):
 
     assert "secret value" not in str(ex.value)
     assert "secret value" not in repr(ex.value)
+
+
+@pytest.mark.skipif(sys.version_info < (3, 3), reason="Python 3.3 added timeout support")
+def test_timeout():
+    """The credential should raise CredentialUnavailableError when the subprocess times out"""
+
+    from subprocess import TimeoutExpired
+
+    with mock.patch(CHECK_OUTPUT, mock.Mock(side_effect=TimeoutExpired("", 42))):
+        with pytest.raises(CredentialUnavailableError):
+            AzureCliCredential().get_token("scope")

sdk/identity/azure-identity/tests/test_cli_credential_async.py

@@ -2,6 +2,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
+import asyncio
 from datetime import datetime
 import json
 import sys
@@ -168,3 +169,15 @@ async def test_subprocess_error_does_not_expose_token(output):
 
     assert "secret value" not in str(ex.value)
     assert "secret value" not in repr(ex.value)
+
+
+async def test_timeout():
+    """The credential should kill the subprocess after a timeout"""
+
+    proc = mock.Mock(communicate=mock.Mock(side_effect=asyncio.TimeoutError), returncode=None)
+    with mock.patch(SUBPROCESS_EXEC, mock.Mock(return_value=get_completed_future(proc))):
+        with pytest.raises(CredentialUnavailableError):
+            await AzureCliCredential().get_token("scope")
+
+    assert proc.communicate.call_count == 1
+    assert proc.kill.call_count == 1