HTTP in Swift, Part 15: OAuth

Part 15 in a series on building a Swift HTTP framework:

The last post covered the basics of an OAuth flow: how we check to see if tokens, how we ask the user to log in, how we refresh the tokens, and so on. In this post we’re going to take that state machine and integrate it into an HTTPLoader subclass.

The Loader

We’ve defined the state machine for the authorization flow, but we need another simple state machine to describe how the loader will interact with it. Let’s consider the various “states” our loader can be in when it is asked to load a task:

idle (nothing has happened yet, or the state machine failed) → we should spin up the state machine and start running through the authorization flow
authorizing (state machine is running) → we should make the task wait for the state machine to finish
authorized (we have valid credentials) → load the task
authorized (we have expired credentials) → we need to refresh the tokens

If we take a closer look at this, we’ll see that the “idle” state is really the same thing as the “authorized + expired tokens” state: in either case, we’ll need to spin up the state machine so that we can get new tokens (Recall that the state machine has logic already to refresh expired tokens). With this in mind, let’s stub out our loader:

public class OAuth: HTTPLoader {

    private var stateMachine: OAuthStateMachine?
    private var credentials: OAuthCredentials?
    private var pendingTasks = Array<HTTPTask>()

    public override func load(task: HTTPTask) {
        // TODO: make everything threadsafe

        if stateMachine != nil {
            // "AUTHORIZING" state
            // we are running the state machine; load this task later
            self.enqueueTask(task)

        } else if let tokens = credentials {
            // we are not running the state machine
            // we have tokens, but they might be expired
            if tokens.expired == true {
                // "AUTHORIZED+EXPIRED" state
                // we need new tokens
                self.enqueueTask(task)
                self.runStateMachine()
            } else {
                // "AUTHORIZED+VALID" state
                // we have valid tokens!
                self.authorizeTask(task, with: tokens)
                super.load(task: task)
            }

        } else {
            // "IDLE" state
            // we are not running the state machine, but we also do not have tokens
            self.enqueueTask(task)
            self.runStateMachine()
        }
    }

}

We can see the four possible states encoded in the if statement. We’re missing some pieces, so let’s take a look at those:

public class OAuth: HTTPLoader {
    ... // the stuff above

    private func authorizeTask(_ task: HTTPTask, with credentials: OAuthCredentials) {
        // TODO: create the "Authorization" header value
        // TODO: set the header value on the task
    }

    private func enqueueTask(_ task: HTTPTask) {
        self.pendingTasks.append(task)
        // TODO: how should we react if the task is cancelled while it's pending?
    }

    private func runStateMachine() {
        self.stateMachine = OAuthStateMachine(...)
        self.stateMachine?.delegate = self
        self.stateMachine?.run()
    }
}

extension OAuth: OAuthStateMachineDelegate {

    // TODO: the OAuth loader itself needs a delegate for some of these to work

    func stateMachine(_ machine: OAuthStateMachine, wantsPersistedCredentials: @escaping (OAuthCredentials?) -> Void) {
        // The state machine is asking if we have any credentials
        // TODO: if self.credentials != nil, use those
        // TODO: if self.credentials == nil, ask a delegate
    }

    func stateMachine(_ machine: OAuthStateMachine, persistCredentials: OAuthCredentials?) {
        // The state machine has found new tokens for us to save (nil = delete tokens)
        // TODO: save them to self.credentials
        // TODO: also pass them on to our delegate
    }

    func stateMachine(_ machine: OAuthStateMachine, displayLoginURL: URL, completion: @escaping (URL?) -> Void) {
        // The state machine needs us to display a login UI
        // TODO: pass this on to our delegate
    }

    func stateMachine(_ machine: OAuthStateMachine, displayLogoutURL: URL, completion: @escaping () -> Void) {
        // The state machine needs us to display a logout UI
        // This happens when the loader is reset. Some OAuth flows need to display a webpage to clear cookies from the browser session
        // However, this is not always necessary. For example, an ephemeral ASWebAuthenticationSession does not need this
        // TODO: pass this on to our delegate
    }

    func stateMachine(_ machine: OAuthStateMachine, didFinishWithResult result: Result<OAuthCredentials, Error>) {
        // The state machine has finished its authorization flow

        // TODO: if the result is a success
        //       - save the credentials to self.credentials (we should already have gotten the "persistCredentials" callback)
        //       - apply these credentials to everything in self.pendingTasks
        // 
        // TODO: if the result is a failure
        //       - fail all the pending tasks as "cannot authenticate" and use the error as the "underlyingError"

        self.stateMachine = nil
    }

}

Most of the reactions to the state machine involve forwarding on the information to yet-another delegate. This is because our loader (correctly!) doesn’t know how to display a login/logout UI, nor does our loader have any idea how or where credentials are persisted. This is as it should be. Displaying UI and persisting information are unrelated to the our loader’s task of “authenticating a request”.

Resetting

Other than the TODO: items scattered around our code, the last major piece of the puzzle we’re missing is the “reset” logic. At first glance, we might think it to be this:

public func reset(with group: DispatchGroup) {
    self.stateMachine?.reset(with: group)
    super.reset(with: group)
}

As discussed in the previous post, every state in the state machine can be interrupted by a reset() call, and this is how that would happen. Thus, if our machine is currently running, this is how we can interrupt it.

… But what if it’s not running? What if we have already authenticated and have valid tokens, and then we get the call to reset()? (This would actually be the common scenario, since “reset” is largely analogous to “log out”, which typically only happens if authentication has succeeded)

In this case, we need to modify our state machine. Recall that last time we described this OAuth flow:

A basic OAuth flow

There’s nothing in this flow that handles the “log out” scenario. We need to modify this slightly so that we also have a way to invalidate tokens. This LogOut state was listed in the “caveats” section before. With it included, the state flow diagram now looks approximately like this:

OAuth with reset

Two things to note about this are:

The dashed lines from all the previous states to the new “Log Out” state represent the “interruption” of that state by calling reset() while the state machine is running
The new “Log Out” state is a possible entry point to the state machine. In other words, we can start the machine in this state.

I’ll leave the implementation of the “Log Out” state to you, but it needs to do a handful of things:

It needs to construct the URL to display the “log out” page to show to the user (the one mentioned before to clear cookies from the browser session)
It needs to contact the server and tell them that the credentials have been revoked
It needs to notify its delegate to clear any persisted credentials

With this in place, our OAuth loader should be completely functional:

public func reset(with group: DispatchGroup) {
    if let currentMachine = self.stateMachine {
       // we are currently authorizing; interrupt the flow
       currentMachine.reset(with: group)
    } else {
        // TODO: you'll want to pass the "group" into the machine here
        self.stateMachine = OAuthStateMachine(...)
        self.stateMachine?.delegate = self

        // "running" the state machine after we gave it the DispatchGroup should start it in the LogOut state
        self.stateMachine?.run()
    }
    super.reset(with: group)
}

Food for thought: There’s technically another possible situation here. What if we reset the loader before we’ve loaded any tasks, and thus do not know if we’re currently authorized? Do we still need to run the state machine? Why or why not? How would the state flow diagram change, if at all?

Conclusion

I hope that these two posts illustrate that OAuth doesn’t have to be this big scary thing. We’ve got our state machine to authorize (or un-authorize) the user, and it has six possible states. That’s not very many, and we can keep that in our heads. Similarly, the loader itself only has a handful of possible states, depending on what’s going on with the state machine. By encapsulating the respective logic in different layers of abstraction, we’re able to keep overall complexity fairly low. Each individual State subclass of our machine is straight-forward; our StateMachine class has almost no code in it; and even our OAuth loader is barely a couple dozen lines.

But from it, we’ve ended up with a full-featured OAuth flow:

we guarantee we only run a single OAuth authorization UI at a time
we allow clients to display the OAuth UI that they want
we allow clients to persist tokens how they want
we allow for interruption of authorization
we allow for un-authorizing by resetting

That’s pretty awesome!

In the next post, we’ll be combining the BasicAuth and OAuth loaders into a single composite Authentication loader.

The Loader

Resetting

Conclusion

Related️️ Posts️