From 2025f955bd06b805ac8229cee173482d2eedf90d Mon Sep 17 00:00:00 2001 From: shima Date: Tue, 14 Jan 2025 10:31:33 +0900 Subject: [PATCH] fix: handle URL-safe base64 decoding for JWT (#38991) * fix: handle URL-safe base64 decoding for JWT - Updated the JWT decoding logic to use URL-safe base64 decoding. - Added padding to the base64 encoded string to ensure proper decoding. - This fixes the issue where UTF-8 decoding errors occurred due to missing padding in the base64 string. Changes: - Replaced `base64.decodebytes` with `base64.urlsafe_b64decode`. - Added logic to calculate and append necessary padding to the base64 string. * More concise way as requested * Extend changes to aio decorators.py as requested * format by black * Update sdk/identity/azure-identity/azure/identity/_internal/decorators.py Co-authored-by: Paul Van Eck * Update sdk/identity/azure-identity/azure/identity/aio/_internal/decorators.py Co-authored-by: Paul Van Eck * Formatted code using Black as specified in ../../../eng/tox/tox.ini with the designated version --------- Co-authored-by: Paul Van Eck --- .../azure/identity/_internal/decorators.py | 11 ++++++++--- .../azure/identity/aio/_internal/decorators.py | 11 ++++++++--- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/sdk/identity/azure-identity/azure/identity/_internal/decorators.py b/sdk/identity/azure-identity/azure/identity/_internal/decorators.py index d56edd9efede..67cffd0b7518 100644 --- a/sdk/identity/azure-identity/azure/identity/_internal/decorators.py +++ b/sdk/identity/azure-identity/azure/identity/_internal/decorators.py @@ -22,12 +22,17 @@ def wrapper(*args, **kwargs): try: token = fn(*args, **kwargs) _LOGGER.log( - logging.DEBUG if within_credential_chain.get() else logging.INFO, "%s succeeded", fn.__qualname__ + logging.DEBUG if within_credential_chain.get() else logging.INFO, + "%s succeeded", + fn.__qualname__, ) if _LOGGER.isEnabledFor(logging.DEBUG): try: - base64_meta_data = token.token.split(".")[1].encode("utf-8") + b"==" - json_bytes = base64.decodebytes(base64_meta_data) + base64_meta_data = token.token.split(".")[1] + padding_needed = -len(base64_meta_data) % 4 + if padding_needed: + base64_meta_data += "=" * padding_needed + json_bytes = base64.urlsafe_b64decode(base64_meta_data) json_string = json_bytes.decode("utf-8") json_dict = json.loads(json_string) upn = json_dict.get("upn", "unavailableUpn") diff --git a/sdk/identity/azure-identity/azure/identity/aio/_internal/decorators.py b/sdk/identity/azure-identity/azure/identity/aio/_internal/decorators.py index 8d92881b36cb..d63f293ca479 100644 --- a/sdk/identity/azure-identity/azure/identity/aio/_internal/decorators.py +++ b/sdk/identity/azure-identity/azure/identity/aio/_internal/decorators.py @@ -21,12 +21,17 @@ async def wrapper(*args, **kwargs): try: token = await fn(*args, **kwargs) _LOGGER.log( - logging.DEBUG if within_credential_chain.get() else logging.INFO, "%s succeeded", fn.__qualname__ + logging.DEBUG if within_credential_chain.get() else logging.INFO, + "%s succeeded", + fn.__qualname__, ) if _LOGGER.isEnabledFor(logging.DEBUG): try: - base64_meta_data = token.token.split(".")[1].encode("utf-8") + b"==" - json_bytes = base64.decodebytes(base64_meta_data) + base64_meta_data = token.token.split(".")[1] + padding_needed = -len(base64_meta_data) % 4 + if padding_needed: + base64_meta_data += "=" * padding_needed + json_bytes = base64.urlsafe_b64decode(base64_meta_data) json_string = json_bytes.decode("utf-8") json_dict = json.loads(json_string) upn = json_dict.get("upn", "unavailableUpn")