Refactor authentication flow

Solves the concurrency issue brought up in pythongssapi/requests-gssapi#8 by passing the context around through the flow instead of storing it in a global dictionary.

Inlined authenticate_user() because it didn't do anything more than set the auth header and handle_401() because it became unnecessarily complicated.

.github/workflows/tests.yml

@@ -26,7 +26,7 @@ jobs:
           python -m pip install -U pip
           python -m pip install flake8
       - name: Check
-        run: python -m flake8
+        run: python -m flake8 httpx_gssapi
   test:
     name: Test
     needs: flake8

httpx_gssapi/exceptions.py

@@ -19,7 +19,3 @@ def __repr__(self):
 
 class SPNEGOExchangeError(HTTPError):
     """SPNEGO Exchange Failed Error"""
-
-
-""" Deprecated compatability shim """
-KerberosExchangeError = SPNEGOExchangeError

httpx_gssapi/gssapi_.py

@@ -6,6 +6,7 @@
 from base64 import b64encode, b64decode
 
 import gssapi
+from gssapi import SecurityContext as SecCtx
 from gssapi.exceptions import GSSError
 
 import httpx
@@ -34,7 +35,7 @@
 _find_auth = re.compile(r'Negotiate\s*([^,]*)', re.I).search
 
 
-def _negotiate_value(response: Response) -> Optional[str]:
+def _negotiate_value(response: Response) -> Optional[bytes]:
     """Extracts the gssapi authentication token from the appropriate header"""
     authreq = response.headers.get('www-authenticate', None)
     if authreq:
@@ -71,23 +72,27 @@ def _handle_gsserror(*, gss_stage: str, result: Any):
         The result to return if a GSSError is raised. If it's an Exception
         type, then it will be raised with the logged message.
     """
+
     def _decor(func):
         @wraps(func)
         def _wrapper(*args, **kwargs):
             try:
                 return func(*args, **kwargs)
-            except gssapi.exceptions.GSSError as error:
+            except GSSError as error:
                 msg = f"{gss_stage} context failed: {error.gen_message()}"
                 log.exception(f"{func.__name__}(): {msg}")
                 if isinstance(result, type) and issubclass(result, Exception):
                     raise result(msg)
                 return result
+
         return _wrapper
+
     return _decor
 
 
 class HTTPSPNEGOAuth(Auth):
-    """Attaches HTTP GSSAPI Authentication to the given Request object.
+    """
+    Attaches HTTP GSSAPI Authentication to the given Request object.
 
     `mutual_authentication` controls whether GSSAPI should attempt mutual
     authentication.  It may be `REQUIRED`, `OPTIONAL`, or `DISABLED`
@@ -112,6 +117,7 @@ class HTTPSPNEGOAuth(Auth):
     server responses.  See the `SanitizedResponse` class.
 
     """
+
     def __init__(self,
                  mutual_authentication: int = DISABLED,
                  target_name: Optional[str] = "HTTP",
@@ -120,7 +126,6 @@ def __init__(self,
                  creds: gssapi.Credentials = None,
                  mech: bytes = None,
                  sanitize_mutual_error_response: bool = True):
-        self.context = {}
         self.mutual_authentication = mutual_authentication
         self.target_name = target_name
         self.delegate = delegate
@@ -132,42 +137,41 @@ def __init__(self,
     def auth_flow(self, request: Request) -> FlowGen:
         if self.opportunistic_auth:
             # add Authorization header before we receive a 401
-            auth_header = self.generate_request_header(request.url.host)
-
-            log.debug(f"Preemptive Authorization header: {auth_header}")
-            request.headers['Authorization'] = auth_header
+            ctx = self.set_auth_header(request)
+        else:
+            ctx = None
 
         response = yield request
-        yield from self.handle_response(response)
+        yield from self.handle_response(response, ctx)
 
-    def handle_response(self, response: Response) -> FlowGen:
+    def handle_response(self,
+                        response: Response,
+                        ctx: SecCtx = None) -> FlowGen:
         num_401s = 0
         while response.status_code == 401 and num_401s < 2:
             num_401s += 1
             log.debug(f"Handling 401 response, total seen: {num_401s}")
-            try:
-                response = yield self.handle_401(response)
-            except httpx.ProtocolError:  # GSSAPI isn't supported
+
+            if _negotiate_value(response) is None:
+                log.debug("GSSAPI is not supported")
                 break
 
-        if response.status_code == 401:
-            log.debug(f"Failed to authenticate, returning 401 response")
-            return
+            log.debug("Generating user authentication header")
+            try:
+                ctx = self.set_auth_header(response.request, response)
+            except SPNEGOExchangeError:
+                log.debug("Failed to generate authentication header")
 
-        self.handle_mutual_auth(response)
+            # Try request again, hopefully with a new auth header
+            response = yield response.request
 
-    def handle_401(self, response: Response) -> Request:
-        """Handles 401's, attempts to use GSSAPI authentication"""
-        log.debug("handle_401(): Handling 401")
-        if _negotiate_value(response) is None:
-            log.debug("handle_401(): GSSAPI is not supported")
-            raise httpx.ProtocolError("GSSAPI is not supported")
+        if response.status_code == 401 or ctx is None:
+            log.debug("Failed to authenticate, returning 401 response")
+            return
 
-        request = self.authenticate_user(response)
-        log.debug(f"handle_401(): returning {request}")
-        return request
+        self.handle_mutual_auth(response, ctx)
 
-    def handle_mutual_auth(self, response: Response):
+    def handle_mutual_auth(self, response: Response, ctx: SecCtx):
         """
         Handles all responses with the exception of 401s.
 
@@ -176,14 +180,14 @@ def handle_mutual_auth(self, response: Response):
         log.debug(f"handle_mutual_auth(): Handling {response.status_code}")
 
         if self.mutual_authentication == DISABLED:
-            log.debug(f"handle_mutual_auth(): Mutual auth disabled, ignoring")
+            log.debug("handle_mutual_auth(): Mutual auth disabled, ignoring")
             return
 
         is_http_error = response.status_code >= 400
 
         if _negotiate_value(response) is not None:
             log.debug("handle_mutual_auth(): Authenticating the server")
-            if not self.authenticate_server(response):
+            if not self.authenticate_server(response, ctx):
                 # Mutual authentication failure when mutual auth is wanted,
                 # raise an exception so the user doesn't use an untrusted
                 # response.
@@ -209,52 +213,49 @@ def handle_mutual_auth(self, response: Response):
             raise MutualAuthenticationError(response=response)
 
     @_handle_gsserror(gss_stage='stepping', result=SPNEGOExchangeError)
-    def generate_request_header(self,
-                                host: str,
-                                response: Response = None) -> str:
+    def set_auth_header(self,
+                        request: Request,
+                        response: Response = None) -> SecCtx:
         """
-        Generates the GSSAPI authentication token
+        Create a new security context, generate the GSSAPI authentication
+        token, and insert it into the request header. The new security context
+        will be returned.
 
-        If any GSSAPI step fails, raise SPNEGOExchangeError
-        with failure detail.
+        If any GSSAPI step fails, raise SPNEGOExchangeError with failure detail.
         """
-        self.context[host] = self._make_context(host)
-
-        token = _negotiate_value(response) if response else None
-        gss_resp = self.context[host].step(token)
-        return f"Negotiate {b64encode(gss_resp).decode()}"
-
-    def authenticate_user(self, response: Response) -> Request:
-        """Handles user authentication with GSSAPI"""
-        host = response.url.host
-        try:
-            auth_header = self.generate_request_header(host, response)
-        except SPNEGOExchangeError:  # GSS Failure, return existing response
-            log.debug(f"authenticate_user(): Failed to generate auth header")
-        else:
-            log.debug(f"authenticate_user(): Auth header: {auth_header}")
-            response.request.headers['Authorization'] = auth_header
+        ctx = self._make_context(request.url.host)
+
+        token = _negotiate_value(response) if response is not None else None
+        gss_resp = ctx.step(token or None)
+        auth_header = f"Negotiate {b64encode(gss_resp).decode()}"
+        log.debug(
+            f"add_request_header(): "
+            f"{'Preemptive ' if token is None else ''}"
+            f"Authorization header: {auth_header}"
+        )
 
-        return response.request
+        request.headers['Authorization'] = auth_header
+        return ctx
 
     @_handle_gsserror(gss_stage="stepping", result=False)
-    def authenticate_server(self, response: Response) -> bool:
+    def authenticate_server(self, response: Response, ctx: SecCtx) -> bool:
         """
-        Uses GSSAPI to authenticate the server.
+        Uses GSSAPI to authenticate the server by extracting the negotiate
+        value from the response and stepping the security context.
 
         Returns True on success, False on failure.
         """
         auth_header = _negotiate_value(response)
         log.debug(f"authenticate_server(): Authenticate header: {auth_header}")
 
         # If the handshake isn't complete here, nothing we can do
-        self.context[response.url.host].step(auth_header)
+        ctx.step(auth_header)
 
         log.debug("authenticate_server(): authentication successful")
         return True
 
     @_handle_gsserror(gss_stage="initializing", result=SPNEGOExchangeError)
-    def _make_context(self, host: str) -> gssapi.SecurityContext:
+    def _make_context(self, host: str) -> SecCtx:
         """
         Create a GSSAPI security context for handling the authentication.
 
@@ -268,7 +269,7 @@ def _make_context(self, host: str) -> gssapi.SecurityContext:
                 name += f"@{host}"
             name = gssapi.Name(name, gssapi.NameType.hostbased_service)
 
-        return gssapi.SecurityContext(
+        return SecCtx(
             usage="initiate",
             flags=self._gssflags,
             name=name,