{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE OverloadedStrings #-}

{- | Adds cookie-based session management to simple 'Controller's. To add to an
  application, declare the Controller setting\'s type an instance of
  'HasSession', and wrap routes with 'withSession'. For example:

  > data AppSettings = ...
  >
  > instance HasSession AppSettings where
  >   ...

  > controllerApp settings $ withSessions $ do
  >   routeName \"posts\" $ ...

 -}
module Web.Simple.Session
  ( Session
  -- * Class and Middleware
  , HasSession(..), withSession
  -- * Accessors
  , sessionLookup, sessionInsert, sessionDelete, sessionClear

  -- * Utilities
  , session, parseSession, dumpSession, addCookie
  ) where

import Control.Monad
import Control.Monad.IO.Class
import Blaze.ByteString.Builder
import Crypto.Hash
import Data.Byteable
import Data.ByteString.Base64
import qualified Data.ByteString as S
import qualified Data.ByteString.Char8 as S8
import Data.Map (Map)
import qualified Data.Map as M
import Network.HTTP.Types.Header
import Network.HTTP.Types.URI
import Network.Wai.Internal (Response(..))
import Network.Wai.Request (appearsSecure)
import Web.Cookie
import Web.Simple.Controller
import System.Environment

-- | Plaintext mapping of the session map. Both keys and values are
-- 'S.ByteString's.
type Session = Map S.ByteString S.ByteString

-- | Instances of this class can be used as states by a 'Controller' to manage
-- cookie-based user sessions. Instances must minimally implement 'getSession'
-- and 'setSession'.
--
-- By default, the secret session key is taken from the
-- environment variable \"SESSION_KEY\", or a default dummy key is used if the
-- environment variable \"ENV\" is set to \"development\". You can override
-- this behaviour by implementing the 'sessionKey' method.
--
-- The default generated cookie always uses the `httponly` option, and the
-- `secure` option if the request is over HTTPS. You can override this behavior,
-- as well as other cookie options (e.g. the path, expiration and domain) by
-- implementing the `sessionBaseCookie` method.
--
-- If the controller
-- state contains a dedicated field of type 'Maybe Session', a reasonable
-- implementation would be:
--
-- > data MyAppSettings = MyAppSettings { myAppSess :: Maybe Session, ...}
-- >
-- > instance HasSession MyAppSettings where
-- >    getSession = myAppSess <$> controllerState
-- >    setSession sess = do
-- >      cs <- controllerState
-- >      putState $ cs { myAppSess = sess }
--
class HasSession hs where
  -- | Returns the secret session key. The default implementation uses the
  -- \"SESSION_KEY\" environment variable. If it is not present, and the
  -- \"ENV\" environment variable is set to \"development\", a dummy, hardcoded
  -- key is used.
  sessionKey :: Controller hs S.ByteString
  sessionKey = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    [(String, String)]
env <- IO [(String, String)]
getEnvironment
    case forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup String
"SESSION_KEY" [(String, String)]
env of
      Just String
key -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ String -> ByteString
S8.pack String
key
      Maybe String
Nothing ->
        case forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup String
"ENV" [(String, String)]
env of
          Just String
e | String
e forall a. Eq a => a -> a -> Bool
== String
"development" -> forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
"test-session-key"
          Maybe String
_ -> (forall a. HasCallStack => String -> a
error String
"SESSION_KEY environment variable not set")

  -- | Returns the cached session for the current request, or nothing if the
  -- session has not been set yet for this request.
  getSession :: hs -> Maybe Session

  -- | Stores a parsed or changed session for the remainder of the request.This
  -- is used both for cached a parsed session cookie as well as for serializing
  -- to the \"Set-Cookie\" header when responding.
  setSession :: Session -> Controller hs ()

  sessionBaseCookie :: Controller hs SetCookie
  sessionBaseCookie = forall hs. Controller hs SetCookie
defaultSessionBaseCookie

defaultSessionBaseCookie :: Controller hs SetCookie
defaultSessionBaseCookie :: forall hs. Controller hs SetCookie
defaultSessionBaseCookie = do
  Request
req <- forall s. Controller s Request
request
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Default a => a
def { setCookieSecure :: Bool
setCookieSecure = Request -> Bool
appearsSecure Request
req, setCookieHttpOnly :: Bool
setCookieHttpOnly = Bool
True }

-- | A trivial implementation if the 'Controller' settings is just a Session
-- store.
instance HasSession (Maybe Session) where
  getSession :: Maybe Session -> Maybe Session
getSession = forall a. a -> a
id
  setSession :: Session -> Controller (Maybe Session) ()
setSession = forall s. s -> Controller s ()
putState forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> Maybe a
Just

-- | A middleware wrapper around a 'Controller' that sets the \"Set-Cookie\"
-- header in the HTTP response if the Session is present, i.e. if it was
-- accessed/modified by the 'Controller'.
withSession :: HasSession hs
            => Controller hs a -> Controller hs a
withSession :: forall hs a. HasSession hs => Controller hs a -> Controller hs a
withSession (ControllerT hs -> Request -> IO (Either Response a, hs)
act) = do
  ByteString
sk <- forall hs. HasSession hs => Controller hs ByteString
sessionKey
  SetCookie
baseCookie <- forall hs. HasSession hs => Controller hs SetCookie
sessionBaseCookie
  forall s (m :: * -> *) a.
(s -> Request -> m (Either Response a, s)) -> ControllerT s m a
ControllerT forall a b. (a -> b) -> a -> b
$ \hs
st0 Request
req -> do
    (Either Response a
eres, hs
st) <- hs -> Request -> IO (Either Response a, hs)
act hs
st0 Request
req
    case Either Response a
eres of
      Left Response
resp0 -> do
        let resp :: Response
resp = case forall hs. HasSession hs => hs -> Maybe Session
getSession hs
st of
                     Just Session
sess -> (ByteString, ByteString) -> SetCookie -> Response -> Response
addCookie
                                   (ByteString
"session", ByteString -> Session -> ByteString
dumpSession ByteString
sk Session
sess)
                                   SetCookie
baseCookie
                                   Response
resp0
                     Maybe Session
Nothing -> Response
resp0
        forall (m :: * -> *) a. Monad m => a -> m a
return (forall a b. a -> Either a b
Left Response
resp, hs
st)
      Right a
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return (Either Response a
eres, hs
st)

-- | Adds a \"Set-Cookie\" with the given key-value tuple to the 'Response'.
-- The path set on the cookie is \"/\", meaning it applies to all routes on the
-- domain, and no expiration is set.
addCookie :: (S.ByteString, S.ByteString) -> SetCookie -> Response -> Response

addCookie :: (ByteString, ByteString) -> SetCookie -> Response -> Response
addCookie (ByteString
key, ByteString
value) SetCookie
baseCookie (ResponseFile Status
stat ResponseHeaders
hdrs String
fl Maybe FilePart
mfp) =
  Status -> ResponseHeaders -> String -> Maybe FilePart -> Response
ResponseFile Status
stat ((HeaderName
"Set-Cookie", SetCookie -> ByteString -> ByteString -> ByteString
cookie SetCookie
baseCookie ByteString
key ByteString
value)forall a. a -> [a] -> [a]
:ResponseHeaders
hdrs) String
fl Maybe FilePart
mfp

addCookie (ByteString
key, ByteString
value) SetCookie
baseCookie (ResponseBuilder Status
stat ResponseHeaders
hdrs Builder
bldr) =
  Status -> ResponseHeaders -> Builder -> Response
ResponseBuilder Status
stat ((HeaderName
"Set-Cookie", SetCookie -> ByteString -> ByteString -> ByteString
cookie SetCookie
baseCookie ByteString
key ByteString
value)forall a. a -> [a] -> [a]
:ResponseHeaders
hdrs) Builder
bldr

addCookie (ByteString
key, ByteString
value) SetCookie
baseCookie (ResponseStream Status
stat ResponseHeaders
hdrs StreamingBody
src) =
  Status -> ResponseHeaders -> StreamingBody -> Response
ResponseStream Status
stat ((HeaderName
"Set-Cookie", SetCookie -> ByteString -> ByteString -> ByteString
cookie SetCookie
baseCookie ByteString
key ByteString
value)forall a. a -> [a] -> [a]
:ResponseHeaders
hdrs) StreamingBody
src

addCookie (ByteString, ByteString)
_ SetCookie
_ Response
resp = Response
resp -- Can't do anything for ResponseRaw

cookie :: SetCookie -> S.ByteString -> S.ByteString -> S.ByteString
cookie :: SetCookie -> ByteString -> ByteString -> ByteString
cookie SetCookie
baseCookie ByteString
key ByteString
value = Builder -> ByteString
toByteString forall b c a. (b -> c) -> (a -> b) -> a -> c
. SetCookie -> Builder
renderSetCookie forall a b. (a -> b) -> a -> b
$
    SetCookie
baseCookie { setCookieName :: ByteString
setCookieName = ByteString
key
        , setCookieValue :: ByteString
setCookieValue = ByteString
value
        , setCookiePath :: Maybe ByteString
setCookiePath = forall a. a -> Maybe a
Just ByteString
"/" }

-- | Returns the current 'Session', either from the 'getSession' cache or by
-- parsing the cookie from the 'Request' using 'sessionFromCookie'.
session :: HasSession hs => Controller hs Session
session :: forall hs. HasSession hs => Controller hs Session
session = do
  hs
cs <- forall s. Controller s s
controllerState
  case forall hs. HasSession hs => hs -> Maybe Session
getSession hs
cs of
    Just Session
sess -> forall (m :: * -> *) a. Monad m => a -> m a
return Session
sess
    Maybe Session
Nothing -> do
      Session
sess <- forall hs. HasSession hs => Controller hs Session
sessionFromCookie
      forall hs. HasSession hs => Session -> Controller hs ()
setSession Session
sess
      forall (m :: * -> *) a. Monad m => a -> m a
return Session
sess

-- | Get and parse a 'Session' from the current 'Request'.
sessionFromCookie :: HasSession hs => Controller hs Session
sessionFromCookie :: forall hs. HasSession hs => Controller hs Session
sessionFromCookie = do
  [(ByteString, ByteString)]
cookies <- (forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] ByteString -> [(ByteString, ByteString)]
parseCookies) forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` forall s. HeaderName -> Controller s (Maybe ByteString)
requestHeader HeaderName
hCookie
  Session
sess <- case forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup ByteString
"session" [(ByteString, ByteString)]
cookies of
            Just ByteString
sessionCookie -> do
              ByteString
sk <- forall hs. HasSession hs => Controller hs ByteString
sessionKey
              forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString -> Session
parseSession ByteString
sk ByteString
sessionCookie
            Maybe ByteString
Nothing -> forall (m :: * -> *) a. Monad m => a -> m a
return forall k a. Map k a
M.empty
  forall (m :: * -> *) a. Monad m => a -> m a
return Session
sess

-- | Parses and validates a serialized 'Session' given the secret. If the
-- 'Session' is invalid (i.e. the hmac does not match), an empty 'Session' is
-- returned.
parseSession :: S.ByteString -> S.ByteString -> Session
parseSession :: ByteString -> ByteString -> Session
parseSession ByteString
secret ByteString
sessionCookie =
  let (ByteString
b64, ByteString
serial) = Int -> ByteString -> (ByteString, ByteString)
S.splitAt Int
44 ByteString
sessionCookie
      mdigest :: Maybe (Digest SHA256)
mdigest = forall a. HashAlgorithm a => ByteString -> Maybe (Digest a)
digestFromByteString forall a b. (a -> b) -> a -> b
$ forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall a b. a -> b -> a
const ByteString
S.empty) forall a. a -> a
id forall a b. (a -> b) -> a -> b
$ ByteString -> Either String ByteString
decode ByteString
b64
  in case Maybe (Digest SHA256)
mdigest of
       Maybe (Digest SHA256)
Nothing -> forall k a. Map k a
M.empty
       Just Digest SHA256
digest ->
         if forall a. HMAC a -> Digest a
hmacGetDigest (forall a.
HashAlgorithm a =>
a -> ByteString -> ByteString -> HMAC a
hmacAlg SHA256
SHA256 ByteString
secret ByteString
serial) forall a. Eq a => a -> a -> Bool
== Digest SHA256
digest then
           forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall a b. (a -> b) -> a -> b
$ ByteString -> [(ByteString, ByteString)]
parseSimpleQuery ByteString
serial
           else forall k a. Map k a
M.empty

-- | Serializes a 'Session' by applying a sha256 hmac with the given secret
-- key to the serialized 'Session' (using 'renderSimpleQuery'), base64 encoding
-- the result, and prepending it to the serialized 'Session'.
dumpSession :: S.ByteString -> Session -> S.ByteString
dumpSession :: ByteString -> Session -> ByteString
dumpSession ByteString
secret Session
sess =
  let serial :: ByteString
serial = Bool -> [(ByteString, ByteString)] -> ByteString
renderSimpleQuery Bool
False forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
M.toList Session
sess
      digest :: Digest SHA256
digest = forall a. HMAC a -> Digest a
hmacGetDigest forall a b. (a -> b) -> a -> b
$ forall a.
HashAlgorithm a =>
a -> ByteString -> ByteString -> HMAC a
hmacAlg SHA256
SHA256 ByteString
secret ByteString
serial
      b64 :: ByteString
b64 = ByteString -> ByteString
encode forall a b. (a -> b) -> a -> b
$ forall a. Byteable a => a -> ByteString
toBytes Digest SHA256
digest
  in ByteString
b64 ByteString -> ByteString -> ByteString
`S.append` ByteString
serial

-- | Lookup a key from the current 'Request's session.
sessionLookup :: HasSession hs
              => S.ByteString -> Controller hs (Maybe S.ByteString)
sessionLookup :: forall hs.
HasSession hs =>
ByteString -> Controller hs (Maybe ByteString)
sessionLookup ByteString
key = forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup ByteString
key forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
`liftM` forall hs. HasSession hs => Controller hs Session
session

-- | Insert or replace a key in the current 'Request's session.
sessionInsert :: HasSession hs
              => S.ByteString -> S.ByteString -> Controller  hs ()
sessionInsert :: forall hs.
HasSession hs =>
ByteString -> ByteString -> Controller hs ()
sessionInsert ByteString
key ByteString
value = do
  Session
sess <- forall hs. HasSession hs => Controller hs Session
session
  forall hs. HasSession hs => Session -> Controller hs ()
setSession (forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert ByteString
key ByteString
value Session
sess)

-- | Remove a key from the current 'Request's session.
sessionDelete :: HasSession hs
              => S.ByteString -> Controller hs ()
sessionDelete :: forall hs. HasSession hs => ByteString -> Controller hs ()
sessionDelete ByteString
key = do
  Session
sess <- forall hs. HasSession hs => Controller hs Session
session
  forall hs. HasSession hs => Session -> Controller hs ()
setSession forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Map k a
M.delete ByteString
key Session
sess

-- | Clear the entire 'Session'.
sessionClear :: HasSession hs => Controller hs ()
sessionClear :: forall hs. HasSession hs => Controller hs ()
sessionClear = forall hs. HasSession hs => Session -> Controller hs ()
setSession forall k a. Map k a
M.empty