From 2121d6bb89c44cba700db72bd2710925781e3a89 Mon Sep 17 00:00:00 2001 From: Charles Lowell Date: Tue, 4 May 2021 15:40:53 -0700 Subject: [PATCH 1/2] raise CredentialUnavailableError when subprocess times out --- .../azure/identity/_credentials/azure_cli.py | 35 +++++++++---------- .../identity/aio/_credentials/azure_cli.py | 8 +++-- 2 files changed, 22 insertions(+), 21 deletions(-) diff --git a/sdk/identity/azure-identity/azure/identity/_credentials/azure_cli.py b/sdk/identity/azure-identity/azure/identity/_credentials/azure_cli.py index dc01f4518a44..c8f5cc8bd18a 100644 --- a/sdk/identity/azure-identity/azure/identity/_credentials/azure_cli.py +++ b/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) diff --git a/sdk/identity/azure-identity/azure/identity/aio/_credentials/azure_cli.py b/sdk/identity/azure-identity/azure/identity/aio/_credentials/azure_cli.py index bf875b69290c..2db249d9f750 100644 --- a/sdk/identity/azure-identity/azure/identity/aio/_credentials/azure_cli.py +++ b/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 From 691313d7015932837638fa919d3933c62d782884 Mon Sep 17 00:00:00 2001 From: Charles Lowell Date: Tue, 4 May 2021 15:41:42 -0700 Subject: [PATCH 2/2] test timeouts --- .../azure-identity/tests/test_cli_credential.py | 12 ++++++++++++ .../tests/test_cli_credential_async.py | 13 +++++++++++++ 2 files changed, 25 insertions(+) diff --git a/sdk/identity/azure-identity/tests/test_cli_credential.py b/sdk/identity/azure-identity/tests/test_cli_credential.py index fc94ce97f931..eb0bb106125a 100644 --- a/sdk/identity/azure-identity/tests/test_cli_credential.py +++ b/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") diff --git a/sdk/identity/azure-identity/tests/test_cli_credential_async.py b/sdk/identity/azure-identity/tests/test_cli_credential_async.py index 947d5370d4a6..69b1f8c1d41f 100644 --- a/sdk/identity/azure-identity/tests/test_cli_credential_async.py +++ b/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