feat(auth): add OAuth2 client credentials support (alpha)

aws_lambda_powertools/utilities/auth_alpha/__init__.py

@@ -6,15 +6,21 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    from aws_lambda_powertools.utilities.auth_alpha.exceptions import AuthFailureReason as AuthFailureReason
     from aws_lambda_powertools.utilities.auth_alpha.jwt import AuthErrorContext as AuthErrorContext
-    from aws_lambda_powertools.utilities.auth_alpha.jwt import AuthFailureReason as AuthFailureReason
     from aws_lambda_powertools.utilities.auth_alpha.jwt import JWTVerifier as JWTVerifier
+    from aws_lambda_powertools.utilities.auth_alpha.oauth2 import OAuth2Client as OAuth2Client
 
-__all__ = ["AuthErrorContext", "AuthFailureReason", "JWTVerifier"]
+__all__ = ["AuthErrorContext", "AuthFailureReason", "JWTVerifier", "OAuth2Client"]
 
 
 def __getattr__(name: str) -> object:
-    modules = {"AuthErrorContext": "jwt", "AuthFailureReason": "jwt", "JWTVerifier": "jwt"}
+    modules = {
+        "AuthErrorContext": "jwt",
+        "AuthFailureReason": "exceptions",
+        "JWTVerifier": "jwt",
+        "OAuth2Client": "oauth2",
+    }
     if name in modules:
         value = getattr(importlib.import_module(f"{__name__}.{modules[name]}"), name)
         globals()[name] = value

aws_lambda_powertools/utilities/auth_alpha/_internal/errors.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from functools import wraps
+from typing import TYPE_CHECKING, ParamSpec, TypeVar
+
+from aws_lambda_powertools.utilities.auth_alpha.exceptions import AuthError
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+
+
+def sanitize_errors(operation: Callable[_P, _T]) -> Callable[_P, _T]:
+    """Detach provider exceptions before an Auth error leaves a public operation."""
+
+    @wraps(operation)
+    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _T:
+        try:
+            return operation(*args, **kwargs)
+        except AuthError as error:
+            # `raise ... from None` only suppresses display of the context.
+            # Clear both references and use a bare re-raise so Python does not
+            # attach the active exception again.
+            error.__context__ = None
+            error.__cause__ = None
+            raise
+
+    return wrapper

aws_lambda_powertools/utilities/auth_alpha/_internal/http.py

@@ -1,16 +1,22 @@
 from __future__ import annotations
 
 import json
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 import urllib3
-from urllib3.connection import HTTPConnection
+from urllib3.poolmanager import pool_classes_by_scheme
 
 from aws_lambda_powertools.utilities.auth_alpha._internal.deadline import Deadline, RequestError
+from aws_lambda_powertools.utilities.auth_alpha._internal.transport import (
+    DeadlineHTTPSConnectionPool,
+    response_deadline,
+)
 
 if TYPE_CHECKING:
     from collections.abc import Mapping
 
+    from urllib3.connectionpool import HTTPConnectionPool
+
 _MAX_JSON_BYTES = 1024 * 1024
 
 
@@ -19,6 +25,27 @@ class HTTPClient:
 
     def __init__(self) -> None:
         self.pool = urllib3.PoolManager(cert_reqs="CERT_REQUIRED")
+        self.pool.pool_classes_by_scheme = pool_classes_by_scheme.copy()
+        pool_classes = cast("dict[str, type[HTTPConnectionPool]]", self.pool.pool_classes_by_scheme)
+        pool_classes["https"] = DeadlineHTTPSConnectionPool
+
+    def _open_response(
+        self,
+        method: str,
+        url: str,
+        deadline: Deadline,
+        **options: Any,
+    ) -> urllib3.response.BaseHTTPResponse:
+        with response_deadline(deadline):
+            return self.pool.request(
+                method,
+                url,
+                timeout=urllib3.Timeout(total=deadline.remaining()),
+                retries=False,
+                redirect=False,
+                preload_content=False,
+                **options,
+            )
 
     def json_request(
         self,
@@ -31,15 +58,12 @@ def json_request(
     ) -> tuple[int, dict[str, Any]]:
         response = None
         try:
-            response = self.pool.request(
+            response = self._open_response(
                 method,
                 url,
+                deadline,
                 body=body,
                 headers=headers,
-                timeout=urllib3.Timeout(total=deadline.remaining()),
-                retries=False,
-                redirect=False,
-                preload_content=False,
             )
             if response.status != 200:
                 deadline.remaining()
@@ -55,23 +79,67 @@ def json_request(
 
     @staticmethod
     def _read_json(response: urllib3.response.BaseHTTPResponse, deadline: Deadline) -> dict[str, Any]:
+        content = HTTPClient._read_body(response, deadline, limit=_MAX_JSON_BYTES, decode_content=False)
+        try:
+            data = json.loads(content)
+        except (ValueError, UnicodeError, RecursionError):
+            raise RequestError() from None
+        if not isinstance(data, dict):
+            raise RequestError()
+        return data
+
+    def request(
+        self,
+        method: str,
+        url: str,
+        deadline: Deadline,
+        *,
+        headers: Mapping[str, str],
+        **options: Any,
+    ) -> urllib3.response.HTTPResponse:
+        """Buffer an authenticated response within one network time budget."""
+        response = None
+        try:
+            response = self._open_response(
+                method,
+                url,
+                deadline,
+                headers=headers,
+                **options,
+            )
+            content = self._read_body(response, deadline)
+            return urllib3.HTTPResponse(
+                body=content,
+                status=response.status,
+                headers=response.headers,
+                reason=response.reason,
+                version=response.version,
+                request_method=method,
+                request_url=url,
+                decode_content=False,
+            )
+        finally:
+            if response is not None:
+                response.close()
+                response.release_conn()
+
+    @staticmethod
+    def _read_body(
+        response: urllib3.response.BaseHTTPResponse,
+        deadline: Deadline,
+        *,
+        limit: int | None = None,
+        decode_content: bool = True,
+    ) -> bytes:
         chunks = bytearray()
         while True:
-            remaining = deadline.remaining()
-            connection = response.connection
-            if isinstance(connection, HTTPConnection) and connection.sock is not None:
-                connection.sock.settimeout(remaining)
-            chunk = response.read1(min(65536, _MAX_JSON_BYTES + 1 - len(chunks)), decode_content=False)
+            deadline.remaining()
+            size = 65536 if limit is None else min(65536, limit + 1 - len(chunks))
+            chunk = response.read1(size, decode_content=decode_content)
             deadline.remaining()
             if not chunk:
                 break
             chunks.extend(chunk)
-            if len(chunks) > _MAX_JSON_BYTES:
+            if limit is not None and len(chunks) > limit:
                 raise RequestError()
-        try:
-            data = json.loads(chunks)
-        except (ValueError, UnicodeError, RecursionError):
-            raise RequestError() from None
-        if not isinstance(data, dict):
-            raise RequestError()
-        return data
+        return bytes(chunks)

aws_lambda_powertools/utilities/auth_alpha/_internal/scopes.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from aws_lambda_powertools.utilities.auth_alpha._internal.validation import string_list
+
+
+def valid_scope(value: str) -> bool:
+    return bool(value) and all(33 <= ord(character) <= 126 and character not in {'"', "\\"} for character in value)
+
+
+def required_scopes(scopes: list[str] | None) -> tuple[str, ...]:
+    values = string_list(scopes if scopes is not None else [])
+    if not all(valid_scope(value) for value in values):
+        raise ValueError("Scopes must be valid OAuth scope tokens")
+    return values

aws_lambda_powertools/utilities/auth_alpha/_internal/transport.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from contextvars import ContextVar
+from http.client import HTTPResponse
+from io import BufferedReader, RawIOBase
+from typing import TYPE_CHECKING
+
+from urllib3.connection import HTTPSConnection
+from urllib3.connectionpool import HTTPSConnectionPool
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+    from socket import socket
+    from typing import Protocol
+
+    from typing_extensions import Buffer
+
+    from aws_lambda_powertools.utilities.auth_alpha._internal.deadline import Deadline
+
+    class _ResponseStream(Protocol):
+        def readinto1(self, buffer: Buffer, /) -> int: ...
+        def close(self) -> None: ...
+
+
+_current_deadline: ContextVar[Deadline] = ContextVar("auth_response_deadline")
+
+
+@contextmanager
+def response_deadline(deadline: Deadline) -> Iterator[None]:
+    """Pass the operation's deadline to responses created by this synchronous call."""
+    token = _current_deadline.set(deadline)
+    try:
+        yield
+    finally:
+        _current_deadline.reset(token)
+
+
+class _DeadlineReader(RawIOBase):
+    """Check the original budget on every refill, including HTTP framing reads."""
+
+    def __init__(self, stream: _ResponseStream, sock: socket, deadline: Deadline) -> None:
+        super().__init__()
+        self._stream = stream
+        self._socket = sock
+        self._deadline = deadline
+
+    def readable(self) -> bool:
+        return True
+
+    def readinto(self, buffer: Buffer, /) -> int:
+        self._socket.settimeout(self._deadline.remaining())
+        # Unlike readinto(), readinto1() performs at most one underlying read.
+        count = self._stream.readinto1(buffer)
+        self._deadline.remaining()
+        return count
+
+    def close(self) -> None:
+        try:
+            self._stream.close()
+        finally:
+            super().close()
+
+
+class _DeadlineResponse(HTTPResponse):
+    def __init__(
+        self,
+        sock: socket,
+        debuglevel: int = 0,
+        method: str | None = None,
+        url: str | None = None,
+    ) -> None:
+        deadline = _current_deadline.get()
+        super().__init__(sock, debuglevel=debuglevel, method=method, url=url)
+        # Retain the wrapper through body consumption: read1() can also parse
+        # chunk-size lines, delimiters and trailers before returning to our loop.
+        # The stream owns the socket reference even for Connection: close.
+        self.fp = BufferedReader(_DeadlineReader(self.fp, sock, deadline), buffer_size=8192)
+
+
+class _DeadlineHTTPSConnection(HTTPSConnection):
+    response_class = _DeadlineResponse
+
+
+class DeadlineHTTPSConnectionPool(HTTPSConnectionPool):
+    """Use the stdlib response hook without overriding urllib3's request machinery."""
+
+    ConnectionCls = _DeadlineHTTPSConnection

aws_lambda_powertools/utilities/auth_alpha/exceptions.py

@@ -0,0 +1,29 @@
+"""Credential-free errors raised by the Auth utility."""
+
+from enum import Enum
+
+
+class AuthFailureReason(str, Enum):
+    """Stable, credential-free reasons suitable for application logs and metrics."""
+
+    MISSING_TOKEN = "missing_token"  # nosec B105
+    INVALID_TOKEN = "invalid_token"  # nosec B105
+    INVALID_CLAIMS = "invalid_claims"
+    TOKEN_EXPIRED = "token_expired"  # nosec B105
+    INVALID_SIGNATURE = "invalid_signature"
+    INSUFFICIENT_SCOPE = "insufficient_scope"
+    FORBIDDEN = "forbidden"
+    JWKS_UNAVAILABLE = "jwks_unavailable"
+    TOKEN_EXCHANGE_FAILED = "token_exchange_failed"  # nosec B105
+    DOWNSTREAM_REQUEST_FAILED = "downstream_request_failed"
+
+
+class AuthError(Exception):
+    """Base error with a fixed message that never includes credential material."""
+
+    message = "Authentication failed"
+    reason = AuthFailureReason.INVALID_TOKEN
+    retryable = False
+
+    def __init__(self) -> None:
+        super().__init__(self.message)

aws_lambda_powertools/utilities/auth_alpha/jwt/_internal/authorization.py

@@ -3,14 +3,16 @@
 from collections.abc import Mapping
 from typing import Any
 
-from aws_lambda_powertools.utilities.auth_alpha._internal.validation import string_list
+from aws_lambda_powertools.utilities.auth_alpha._internal.scopes import required_scopes, valid_scope
 from aws_lambda_powertools.utilities.auth_alpha.jwt.exceptions import (
     AuthError,
     AuthFailureReason,
     InvalidClaimsError,
     InvalidTokenError,
 )
 
+__all__ = ["required_scopes"]
+
 
 class MissingTokenError(InvalidTokenError):
     """No authorization header was supplied."""
@@ -65,17 +67,6 @@ def _authorization_values(headers: Any) -> list[Any]:
     return values
 
 
-def valid_scope(value: str) -> bool:
-    return bool(value) and all(33 <= ord(character) <= 126 and character not in {'"', "\\"} for character in value)
-
-
-def required_scopes(scopes: list[str] | None) -> tuple[str, ...]:
-    values = string_list(scopes if scopes is not None else [])
-    if not all(valid_scope(value) for value in values):
-        raise ValueError("Scopes must be valid OAuth scope tokens")
-    return values
-
-
 def enforce_scopes(claims: dict[str, Any], expected: tuple[str, ...]) -> None:
     value: Any = next((claims[name] for name in ("scope", "scp", "scopes") if name in claims), [])
     if isinstance(value, str):