Auth handler per connection

[jborean93]

I'm currently trying out httpx as part of a WinRM based library which has a unique requirement for encryption over HTTP. The authentication is typically Negotiate authentication (Kerberos or NTLM) and while I've able to create a handler class that implements httpx.Auth that will successfully authenticate against the host I need to preserve the security context so I can encrypt the data sent to the client. This is a very basic mock up of what I currently have in place

import httpx

class NegotiateAuth(httpx.Auth):

    def __init__(self, username, password):
        self.username = username
        self.password = None
        self.context = None

    def sync_auth_flow(self, request):
        response = yield request
        if response.status_code != 401:
            return

         self.context = MySecContext(self.username, self.password)
        out_token = self.context.step()
        while not self.context.complete:
            request.headers['Authorization'] = out_token

            # send the request with the auth token and get the response
            response = yield request

            # Feed the token back into the context until it's set up

    def wrap(self, data);
        return self.context.wrap(data)

    def unwrap(self, data):
        return self.context.unwrap(data)


class WinRMClient:

    def __init__(self):
        auth = NegotiateAuth('username', 'password')
        client = httpx.AsyncClient(auth=auth)

    async def send(self, data):
        # set up initial auth with blank request
        if not self.auth.context:
            await self.client.post('endpoint', content=b'')

        enc_data = self.client.auth.wrap(data)
        resp = await self.client.post('endpoint', content=enc_data)
        return self.client.auth.unwrap(resp.content)

This works fine if I've only ever got 1 connection being used in the pool as the context in the auth object always relates to the connection that was opened. Unfortunately as soon as I start sending multiple requests concurrently httpx will create a new connection from the pool starting the authentication process again for the new connection. Ultimately this will overwrite the context for the initial connection meaning I'm no longer able to wrap/unwrap data.

What I'm hoping to find out is a way to:

• Create a context for each connection
  • I've looked into subclassing httpcore.AsyncConnectionPool as a transport and overriding _create_connection() but it's not a public API so this is not ideal
  • Even though it's not public I think I could set up the auth context here and just use setattr() on the connection to store it but I really don't want to be messing around with these internals
• Intercepting something that prepares the requests and processes the response so that I can wrap/unwrap the data
  • I would need to be able to link the socket/connection that the request will be sent over to the security context for that connection
  • A lot of the arequest() functions I can see in some of the internal classes aren't really tied to an actual
  • I think I could override the arequest() method on the connection created by _create_connection() which transforms the data but like the above I prefer not to mess around with this stuff if I can avoid it

Do you have any ideas how I could potentially implement these requirements?

[florimondmanca]

Hey @davidbgk

This is a very interesting scenario.

I believe the transport route is the most promising. The "wrap / unwrap" thingy you do there for example screams for "wrap the request / response stream" at the transport level. :)

But I wouldn't do subclassing. Instead I think you might be able to get something working by composing with an existing transport, and then wrapping things however is necessary to accommodate your auth scheme. You might not even need to modify the pool, and instead do the tracking of origins inside your own transport.

I can draft some code snippets to help you out with this intuition later if you'd like. :)

[jborean93]

So after thinking about this a bit more I think the simplest solution is to not use a connection pool and to instantiate my httpx.Client for every task I know will run in parallel with the main connection. This ensures that the auth context is purely for that connection as there will only be one per client. It's more work on the users of the httpx (Async)Client but it's less brittle than drilling down into private interfaces and overriding some of the stuff I mentioned above.

Appreciate everything you've shared so far and I'm happy to answer any questions you have for me.

[jborean93]

This brings to my next problem that somewhat ties into this issue. When a connection has gone past the Keep-Alive expiry the connection pool will drop the connection and create a new one which invalidates the security context that was for the original connection. Unfortunately there does not seem to be any indicator that this has happened and the connection is dropped after I've already prepared my request. I've got a hack which aligns with one of the options above to hijack into the aclose() method on the new connection but it relies on _create_connection().

The code I'm currently using is:

class _AsyncWinRMTransport(httpcore.AsyncConnectionPool):

    def _create_connection(self, *args, **kwargs):
        # We need to hook into the connection close() method so we can be alerted when we've started a new connection
        # due to a keep-alive timeout.
        connection = super()._create_connection(*args, **kwargs)
        orig_close = connection.aclose

        async def new_close():
            await orig_close()
            raise _NoAuthenticationContext()
        connection.aclose = new_close
        return connection

When the connection pool closes and connections that were timed out this will still close the connection but it raises a custom Exception that lets me know this occured and I need to re set up my authentication context and rewrap my data for the next request. Can you see any other way this might be achievable?

[jborean93]

Ok hopefully final update/question I have. I've think I've found a solution to the keep-alive problem, I can check if the response is a 401 which would indicate a new connection was opened and the authentication has not happened. It's not perfect but it's workable and I don't need to touch any of the private methods to achieve it so I'm going for that.

One problem I have is that I cannot seem to indicate the connection is now in an IDLE state after receiving a response so subsequent requests will either use a new connection or hang because the pool is full of "active" connections. I tried doing

# self._transport is `httpcore.AsyncConnectionPool`
response = await self._transport.arequest(method, url, headers, stream, ext)
await response[2].aclose()

Which removes the active connection from the pool but instead of making it IDLE so the subsequent request uses that connection it's closed breaking the authentication process. How do I mark a response is read and the connection is free to service a new request?

[florimondmanca]

@jborean93 To move the connection to an IDLE state rather than close it, I think you'd need to read through the stream, rather than .aclose() it:

response = await self._transport.arequest(method, url, headers, stream, ext)

# Read response body to terminate request-response cycle and allow reusing the connection.
_, _, stream, _ = response
async for _ in stream:
    pass

[jborean93]

Turns out I had 2 issues

• The initial 401 I get back when sending no auth has the Connection: close header which closed the connection so it always created a new one
• You need to both read and close the stream to turn it into IDLE

Thanks for all your help.