A Go Middleware Thought

Mon 26 February 2018 | -- (permalink)

At work we just implemented a few Go HTTP middlewares for handling different kinds of authentication. They pull values from headers in the request (e.g. JWTs or api keys), check their validity, then fetch the associated account and stick it into the request context for use by the handlers they wrap. Here's a simple but complete example of this pattern:

package main

import (
  "context"
  "fmt"
  "net/http"
)

func main() {
  http.Handle("/", userAuthMiddleware(handler))
  http.ListenAndServe(":8080", nil)
}

type user struct {
  Name string
}

// handler says hello to the user provided in the request context.  This requires that the handler
// be wrapped in the userAuthMiddleware.
var handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  // get user from the request context.
  u, ok := r.Context().Value("user").(*user)

  // handle the user not being provided.  This should never happen!
  if !ok {
    http.Error(w, "Internal Server Error", http.StatusInternalServerError)
  }

  fmt.Fprintf(w, "Hello %q\n", u.Name)
})

// userAuthMiddleware takes the user name from the Authorization header, creates a user{} with it,
// and sticks it into the request context.
func userAuthMiddleware(next http.Handler) http.Handler {
  return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    authHeader := r.Header.Get("Authorization")
    // In a real auth middleware, you'd check the validity of some token or key here.
    if authHeader == "" {
      http.Error(w, "invalid Authorization header", http.StatusUnauthorized)
      return
    }
    u := user{Name: authHeader}
    // using the bare "user" string as our key is a little dirty.  In production code you should use
    // your own string type so it's impossible for some other package to clobber your "user" value.
    ctx := context.WithValue(r.Context(), "user", &u)
    next.ServeHTTP(w, r.WithContext(ctx))
  })
}

It pulls the "Authorization" header from the request and says hello to the user specified there:

$ curl http://localhost:8080 -H Authorization:joe
Hello "joe"

If you omit the header, the middleware returns a 401:

$ curl -i http://localhost:8080
HTTP/1.1 401 Unauthorized
Content-Type: text/plain; charset=utf-8
X-Content-Type-Options: nosniff
Date: Mon, 26 Feb 2018 17:12:39 GMT
Content-Length: 29

invalid Authorization header

There are a couple nice things about this pattern:

  1. It's re-usable. If I have a bunch of endpoints that require a user, I can just wrap them all in this middleware and I don't have to repeat the code for dealing with the Authorization header in each of them.
  2. It's composable. I could have middlewares for other things too (like logging), and they'll all happily wrap each other.

But there are also a couple things I'm uncomfortable with:

  1. It's possible that I could forget to wrap my handler in the middleware. The compiler won't catch this; only at runtime will we see the error.
  2. Every single handler that's wrapped by this middleware needs code for dealing with the case where no user was put into the request context.

What if my handler, instead of just taking a response writer and a request, also took a user as a function argument? Then I'd get a compile time guarantee that a user was passed to the handler (which I greatly prefer to a run-time failure). And my handlers would no longer need to include code for the case where the user is missing from the request context.

All we need to make that happen is to define our own handler type:

package main

import (
  "fmt"
  "net/http"
)

func main() {
  http.Handle("/", userAuthMiddleware(handler))
  http.ListenAndServe(":8080", nil)
}

type user struct {
  Name string
}

type userHandler interface {
  ServeHTTP(http.ResponseWriter, *http.Request, *user)
}

// The userHandlerFunc type is an adapter to allow the use of
// ordinary functions as HTTP handlers. If f is a function
// with the appropriate signature, userHandlerFunc(f) is a
// Handler that calls f.
type userHandlerFunc func(http.ResponseWriter, *http.Request, *user)

// ServeHTTP calls f(w, r, u).
func (f userHandlerFunc) ServeHTTP(w http.ResponseWriter, r *http.Request, u *user) {
  f(w, r, u)
}

// handler says hello to the user provided.  This requires that the handler be wrapped in the
// userAuthMiddleware.
var handler = userHandlerFunc(func(w http.ResponseWriter, r *http.Request, u *user) {
  fmt.Fprintf(w, "Hello %q\n", u.Name)
})

// userAuthMiddleware takes a special handler that needs a user to be passed in, and wraps it with
// the logic to check the Authorization header and pass the appropriate user.
func userAuthMiddleware(next userHandler) http.Handler {
  return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    authHeader := r.Header.Get("Authorization")
    // In a real auth middleware, you'd check the validity of some token or key here.
    if authHeader == "" {
      http.Error(w, "invalid Authorization header", http.StatusUnauthorized)
      return
    }
    u := user{Name: authHeader}
    next.ServeHTTP(w, r, &u)
  })
}

Instead of being a http.Handler, my handler is now of my own userHandler type. The compiler guarantees that it's passed a user when called; if I tried routing my handler without wrapping it in userAuthMiddleware, it wouldn't compile. The handler no longer needs to deal with the case of the user not being provided.

This approach means writing less code. In a simple example like this, the context handling code amounts to about the same number of lines as my custom handler type. But if I had more handlers then the custom type would be a clear winner, since each of the old-style handlers would need lines for pulling the user from the context and dealing with cases where it wasn't there.

There is a slight loss of composability with the custom handler type. Instead of being able to wrap any handler or middleware, my userAuthMiddleware can now only wrap a userHandler. In other words, it now must be the innermost middleware. It can still be wrapped by generic HTTP middlewares, however.

Given that we will end up writing less code, and we trade runtime errors for compile time errors, I think the custom handler type's tradeoffs are worth making.