AddSetCookie.hs

{-# LANGUAGE PolyKinds            #-}
{-# LANGUAGE UndecidableInstances #-}
module Servant.Auth.Server.Internal.AddSetCookie where

import           Blaze.ByteString.Builder (toByteString)
import qualified Data.ByteString          as BS
import qualified Data.ByteString.Base64   as BS64
import           Data.Monoid
import           Servant
import           System.Entropy           (getEntropy)
import           Web.Cookie

-- What are we doing here? Well, the idea is to add headers to the response,
-- but the headers come from the authentication check. In order to do that, we
-- tweak a little the general theme of recursing down the API tree; this time,
-- we recurse down a variation of it that adds headers to all the endpoints.
-- This involves the usual type-level checks.
--
-- TODO: If the endpoints already have headers, this will not work as is.


type family AddSetCookieApi a where
  AddSetCookieApi (a :> b) = a :> AddSetCookieApi b
  AddSetCookieApi (a :<|> b) = AddSetCookieApi a :<|> AddSetCookieApi b
  AddSetCookieApi (Verb method stat ctyps a)
     = Verb method stat ctyps (Headers '[Header "Set-Cookie" BS.ByteString] a)
  AddSetCookieApi (Headers ls a) = Headers ((Header "Set-Cookie" BS.ByteString) ': ls) a


class AddSetCookie orig new where
  addSetCookie :: [SetCookie] -> orig -> new

instance {-# OVERLAPS #-} AddSetCookie oldb newb
  => AddSetCookie (a -> oldb) (a -> newb) where
  addSetCookie cookie oldfn = \val -> addSetCookie cookie $ oldfn val

instance {-# OVERLAPPABLE #-}
  ( Functor m
  ) => AddSetCookie (m a) (m (Headers '[Header "Set-Cookie" BS.ByteString] a))  where
  addSetCookie cookie v = addSetCookie cookie <$> v

instance {-# OVERLAPS #-}
  (AddSetCookie a a', AddSetCookie b b')
  => AddSetCookie (a :<|> b) (a' :<|> b') where
  addSetCookie cookie (a :<|> b) = addSetCookie cookie a :<|> addSetCookie cookie b

instance {-# OVERLAPPABLE #-}
  (AddHeader "Set-Cookie" BS.ByteString old new)
  => AddSetCookie old new where
  addSetCookie cookie val
    -- What is happening here is sheer awfulness. Look the other way.
    = addHeader (foldr1 go $ toByteString . renderSetCookie <$> cookie) val
    where
      go new old = old <> "\r\nSet-Cookie: " <> new


csrfCookie :: IO BS.ByteString
csrfCookie = BS64.encode <$> getEntropy 32