Adding Cookie Based Auth To Servant

One of the applications I’ve been working on exclusively lately is a surprisingly large Job Board application. Not the most exciting concept, I know, but it’s quickly becoming the culmination of some of the best development (and consequently “engineering”) work that I’ve ever done. The codebase is growing huge – 11,000+ lines counting both front and backend code, and while I haven’t gone on any large refactoring sprees, most of the code does not make me feel bad. I use such a vague description for code quality because it’s such a relative target from person to person, but generally when you write code that doesn’t fit even your own standards, I don’t know about other people but I kind of start to feel bad/uneasy. Let’s get into the breakdown of my thought process while I was feeling my way around Servant and writing this authentication combinator in.

Considerations

As I often do with most tickets, I make a quick to-do list trying to capture all the concerns that I can think of off the top of my head when I write the Gitlab issue. My to-do list looked something like this:

  1. Creating a roles table
  2. Creating a permissions table
  3. Add HTTP cookie based auth as supported by Servant – which would enable cookies to be saved on the requests
  4. Add combinator (monad transformer?) that makes user available to handlers (possibly related to above) – something that would enable … :> AuthCheck Administrator :> Post ….
  5. Add a /me endpoint that just gives back session info for current user
  6. Augment session info to always include user permissions & role (+ a JOIN somewhere)
  7. Create combinator to restrict endpoints by user role (just for admin interface)

Of course, I didn’t end up needing/doing all of these points, but I think seeing this listing shows just how much still needed to be done starting on this adventure. Again, this i sjust a list of things I could think of, while I still didn’t know exactly what the implementation would look like.

Getting started writing a nice little combinator

I started the journey by looking at servant’s documentation for “experimental” general purpose authentication combinators. It made me a tad uneasy that the support was “experimental”, but luckily despite being experimental, it served as an OK guide for learning how to extend Servant in the way I wanted to. By using this guide, I was able to get 99% of the way towards the combinator that I really wanted – I came up against the problem of actually storing the cookies somewhere.

Resting on what I already knew from other languages, my first thought was to start writing som ecode to store the session in the database (or in memory, whatever), sign & encrypt my data, and serve it in an HttpOnly header. Of course, the only bits that were cloudy were everything:

  • How easy is signing/encryption in Haskell (already worked with Crypto.Bcrypt), and it was super fun to work with.
  • How to make sure the cookie gets tacked onto the request after interacting with storage (however that works out)

During looking at this I discovered Web.Cookie which handles the cookie saving and parsing stuff that I needed to do.

Handling Sessions in a SessionStore

Luckily (or unluckily) I came across Network.Wai.Session – and set about to grappling with trying to understand it

  • PRO: The interface is small
  • PRO: There’s a thread-safe reference implementation (yay STM reference implementation for in-memory cookie storage (which is usually the basic approach, sessions die when your app restarts)

My first serious look is at withSession, as it returns the thing I want, a Middleware. Seeing the type signature gives me lots of hope (especially because I’m using Haskell), and it’s the venerable Application -> Application, giving off the warm architectural fuzzies (in some other languages, it’s hard to know that some piece of code will do NOTHING except middleware like activity, but in Haskell you just know :).

This left me with one question: How do I get the things that withSession expects? Here’s how I figured that out (the order is a little jumbled):

  • ByteString - Just a name to use for the cookie, so I used “-auth” (X-something-something was deprecated, IIRC)

  • SetCookie - Comes from Web.cookie, all you have to do is follow the instructions to make one, not much complexity there.

  • Key (Session m k v) - Unfortunately I’m still quite “wet behind the ears” and definitely get scared when I see a parametrized type with this many parameters. Luckily, since the time that I tried to shoehorn a database connection into Vault, I immediately recognized that Key was a Vault.Key and felt quite comfortable with it. Vault.Keys are actually really easy to generate, and they can be for anything, and they always come out unique… They remind me of Symbols in ES6/ES2015 Javascript, and I imagine the implementation is somewhat similar (surely with some MVar somewhere).

  • SessionStore m k v - Again, I always get a little scared/uneasy when I see a parametrized type with this many parameters… I was stumped for a while trying to figure out how to make one of these – I could see the type definitions, but still couldn’t quite wrap my head out of how to implement one. Looking back at the docs, I found that there was an included reference implementation, which is awesome, and let me be productive while I tried to think about the types and what they really meant.

At this point, the code looked like:

   sessionKey <- V.newKey -- <---- Haskell couldn't figure out what kind of Key this is, and in turn, what kind of cookie store it's trying to open, and what kind of cookie store it has at all.
   let cookieSettings = WC.def { WC.setCookieHttpOnly=True }
   cookieStore  <- mapStore_
   let cookieMiddleware = withSession cookieStore "<appname>-auth" cookieSettings sessionKey

Figuring out the type of the key actually took me a minute, and I need to slow down and really think of what the types involved were describing. I shouldn’t have stumbled here, but I did for a little bit, but am glad I took the time to really think through and understand it.

   sessionKey <- V.newKey :: IO (V.Key (Session IO String User))
   let cookieSettings = WC.def { WC.setCookieHttpOnly=True }
   cookieStore  <- mapStore_
   let cookieMiddleware = withSession cookieStore "<appname>-auth" cookieSettings sessionKey

Yay, everything compiles! all that was left was to use that middleware on the app that I had already generated and been running, and I had my middleware all ready to use! According to the documentation (as far as I knew at this point), alll requests should have a Vault object, and that object would contain the session (if there is one), under the appropriate key

Figuring out how to get Sessions back out of the SessionStore

One small issue I realized I had - How do I get & use the Vault key I made for the session to get the session out again? I need to access the vault key in my combinators, because all I have is the request (which contains the vault, but I need the original key to pull the value that I want out). My first thought was to use the State monad (clearly I have some piece of state that I need to access later (I talk abou tthis in the previous servant-related post), to store “Application information”, and pass along the vault key. Rather than doing that however, I decided to go with using Servant’s concept of Contexts. That meant my app starting code, which looked like this:

app :: SqliteBackend -> Middleware -> Application
app backend cookieMiddleware = cookieMiddleware $ serve api appWithDB
    where
      addDB = evalStateTLNat backend
      appWithDB = enter addDB server

Had to turn into this:

makeApp :: SqliteBackend -> Middleware -> V.Key (Session IO ByteString SessionInfo) -> Application
makeApp b addCookies cookieKey = addCookies $ serveWithContext api context appWithGlobals
    where
      addGlobals = evalStateTLNat $ ApplicationGlobals b
      appWithGlobals = enter addGlobals server
      context = Auth.genAuthServerContext cookieKey

I end up kind of doing both – You can see that I change the object that I use with enter into a kind of “application globals” object, rather than just the backend. I also, however, start using Servant’s serveWithContext. The problem of where to store the vault key is most certainly (for me): use Contexts and serveWithContext. I first encountered Contexts when trying to figure out the proper way to store a DB connection in Servant, and I’m glad that learning experience was able to be of use to me so soon aferwards.

Well Enough blathering, here’s what the final code looks like:

The pudding

The code below is mostly relegated to a file called Auth.hs so I could at least feel like I am properly organizing my code. (It became clear pretty quickly that I should move this code as it was getting hairy to have right next to endpoint-specific code)

Auth.hs

type instance AuthServerData (AuthProtect "cookie-auth") = (Session IO ByteString SessionInfo)

Servant code that creates the TypeFamily that allows Servant to know where to find your custom auth implementation (this code was just copied from servant, and customized)

-- Generate a cookie auth handler, given a session store cookie
genCookieAuthHandler :: V.Key (Session IO ByteString SessionInfo) -> AuthHandler Request (Session IO ByteString SessionInfo)
genCookieAuthHandler k = mkAuthHandler handler
    where
      handler req = do
        let sessionStore = V.lookup k $ vault req
        case sessionStore of
          Nothing -> throwError Err.invalidAuthHeaders
          Just s -> return s

This code is actually a generator for the auth request handler. The Servant docs aren’t exactly clear on what’s happening, but looking at the types, what happens is that , the handler this method generates actually gets stored on the Context of a started server. This is how AuthProtect later knows what to do to authenticate. You make a handler, put it in the context, then tell Servant (using the TypeFamily) how to find it.

Note that that this function has to take the Vault.Key that pulls out the session from the request’s Vault, as that is made basically on-the-fly during server startup. Of course there are other options, you could hav used a global, but I preferred not to.

-- Given a vault key to use to retrieve the session store (provided by wai-session),
-- generate a context containing a cookie auth handler
genAuthServerContext :: V.Key (Session IO ByteString SessionInfo) -> Context (AuthHandler Request (Session IO ByteString SessionInfo) ': '[])
genAuthServerContext k = authHandler :. EmptyContext
    where
      authHandler = genCookieAuthHandler k

This function’s job is to produce the context you need to pass to serveWithContext. As you can see, it builds on the auth handler that is generated in the previous chunk of code, and just builds a simple context with it.

Lib.hs (where the servant routes and handlers are)

... lines and lines of app setup code ...
-- Create a vault key for use storing the cookie-based sessions
sessionKey <- V.newKey :: IO (V.Key (Session IO ByteString SessionInfo))

-- TODO Add debug mode as a AppConfig setting, need to add secure, expires, max age, settings differently for debug vs prod

-- Create cookie store && middleware to use with the app
let cookieSettings = WC.def { WC.setCookieHttpOnly=True }
cookieStore  <- mapStore_
let cookieMiddleware = withSession cookieStore "<appname>-auth" cookieSettings sessionKey

-- Start the app
run port $ makeApp backend cookieMiddleware sessionKey

This code is the app set up code (part of which you’ve seen before)

makeApp :: SqliteBackend -> Middleware -> V.Key (Session IO ByteString SessionInfo) -> Application
makeApp b addCookies cookieKey = applicationWithCookies
    where
      -- Inject ApplicationGlobals into Server by Monad (StateT) Transforming
      addGlobals = evalStateTLNat $ ApplicationGlobals b
      serverWithGlobals = enter addGlobals server

      context = Auth.genAuthServerContext cookieKey

      application = serveWithContext api context serverWithGlobals
      applicationWithCookies = addCookies application

To finish it up, here’s the code that creates the app object that will be used with the Servant run command (this was refactored a little bit, and while still kind of ugly, it should be very clear and easy to read).

Wrap up

Getting this done was a good learning experience, and a good way to figure out how concepts I know pretty well in other languages/frameworks (Node/Express, Python/Flask, Ruby/Sinatra, Go/http, Clojure/ring), get accomplished in Servant.

While this post was definitely more stream-of-consciousness than polished technical overview, I hope it was worth following and it contains some insight that will help someone else understand just a little bit more about how things are working.