{-# OPTIONS_GHC -fno-warn-orphans #-}

module Web.ServerSession.Frontend.Yesod.Jwt
  ( backend
  , siteApproot
  , ServerSessionJwtConfig(..)
  , ServerSessionJwtException(..)
  , forceInvalidate
  ) where

-- Module heavily inspired by:
-- serversession-frontend-yesod-1.0@sha256:8ddb112a1ef6ee863f5ea13978dd08e1c39444c1a252f775a780013430bcc884,1230


import Import.NoModel hiding (State, state, Header, deleteCookie)
import Yesod.Core.Types
import Model.Types.Common

import Model.Tokens.Session

import Jose.Jwt (JwtEncoding(..))
import qualified Jose.Jwt as Jose
import qualified Jose.Jwk as Jose

import qualified Network.Wai as Wai

import Web.Cookie (parseCookies, SetCookie(..))

import Web.ServerSession.Core hiding (SessionMap, setCookieName)

import qualified Data.Map as Map

import qualified Data.Aeson as JSON


instance Universe ForceInvalidate
instance Finite ForceInvalidate
finitePathPiece ''ForceInvalidate 
  [ "current", "all", "none" ]


data ServerSessionJwtConfig = ServerSessionJwtConfig
  { sJwtJwkSet     :: JwkSet
  , sJwtStart
  , sJwtExpiration :: Maybe NominalDiffTime
  , sJwtEncoding   :: JwtEncoding
  , sJwtIssueBy    :: InstanceId
  , sJwtIssueFor   :: ClusterId
  }


data ServerSessionJwtException
  = SessionTokenJwtError Jose.JwtError
  | SessionTokenUnsecured
  | SessionTokenInvalidFormat String
  | SessionTokenExpired | SessionTokenNotStarted
  deriving (Eq, Show, Generic, Typeable)
instance Exception ServerSessionJwtException


backend :: ( Applicative m
           , Storage sto
           , SessionData sto ~ Map Text ByteString
           )
        => ServerSessionJwtConfig
        -> (Wai.Request -> Maybe Text)
        -> State sto
        -> m (Maybe SessionBackend)
backend jwtCfg getApprootText' state = pure $ Just SessionBackend{..}
  where
    sbLoadSession :: Wai.Request -> IO (SessionMap, SaveSession)
    sbLoadSession req = do
      session <- runMaybeT . catchMPlus (Proxy @ServerSessionJwtException) $
        decodeSession jwtCfg =<< hoistMaybe (findSession state req)
      (sessionData, saveSessionToken) <- loadSession state $ encodeUtf8 . toPathPiece . sessionId <$> session

      let save :: SessionMap -> IO [Header]
          save sessMap = pure <$> do
            saveRes <- saveSession state saveSessionToken sessMap
            case saveRes of
              Nothing   ->
                return $ deleteCookie state approot'
              Just sess ->
                fmap (createCookie state approot' sess) . encodeSession jwtCfg =<< mkSessionToken jwtCfg sess

          approot' = getApprootText' req

      return (sessionData, save)
      
findSession :: State sto
            -> Wai.Request
            -> Maybe Jwt
findSession state req = do
  [raw] <- return $ do
    ("Cookie", header) <- Wai.requestHeaders req
    (k, v) <- parseCookies header
    guard $ k == encodeUtf8 (getCookieName state)
    return v
  return $ Jwt raw


mkSessionToken :: MonadIO m
               => ServerSessionJwtConfig
               -> Session sess
               -> m (SessionToken sess)
mkSessionToken ServerSessionJwtConfig{..} Session{..} = liftIO $
    mkSessionToken' <$> getCurrentTime <*> getRandom
  where
    mkSessionToken' now sessionIdentifier
      = let sessionId = sessionKey
            sessionIssuedAt = now
            sessionIssuedBy = sJwtIssueBy
            sessionIssuedFor = sJwtIssueFor
            sessionExpiresAt = flip addUTCTime now <$> sJwtExpiration
            sessionStartsAt = flip addUTCTime now <$> sJwtStart
         in SessionToken{..}


deleteCookie :: State sto -> Maybe Text -> Header
deleteCookie state approot' = DeleteCookie (encodeUtf8 $ getCookieName state) $ cookiePath approot'

createCookie :: State sto -> Maybe Text -> Session sess -> Jwt -> Header
createCookie state approot' session (Jwt payload) = AddCookie def
  { setCookieName     = encodeUtf8 $ getCookieName state
  , setCookieValue    = payload
  , setCookiePath     = Just $ cookiePath approot'
  , setCookieExpires  = cookieExpires state session
  , setCookieDomain   = Nothing -- Setting anything here would have browsers include subdomains, which might be wrong
  , setCookieHttpOnly = getHttpOnlyCookies state
  , setCookieSecure   = getSecureCookies state
  }

  
decodeSession :: ( MonadThrow m
                 , MonadIO m
                 )
              => ServerSessionJwtConfig
              -> Jwt
              -> m (SessionToken sess)
decodeSession ServerSessionJwtConfig{..} (Jwt bs) = do
  content <- either (throwM . SessionTokenJwtError) return =<< liftIO (Jose.decode (Jose.keys sJwtJwkSet) Nothing bs)
  content' <- case content of
    Jose.Unsecured _ -> throwM SessionTokenUnsecured
    Jose.Jws (_header, payload) -> return payload
    Jose.Jwe (_header, payload) -> return payload
  session@SessionToken{..} <- either (throwM . SessionTokenInvalidFormat) return $ JSON.eitherDecodeStrict content'

  now <- liftIO getCurrentTime
  unless (NTop sessionExpiresAt > NTop (Just now)) $
    throwM SessionTokenExpired
  unless (sessionStartsAt <= Just now) $
    throwM SessionTokenNotStarted

  return session

encodeSession :: MonadIO m
              => ServerSessionJwtConfig
              -> SessionToken sess
              -> m Jwt
encodeSession ServerSessionJwtConfig{..} token = liftIO $ 
    either throwM return =<< Jose.encode (Jose.keys sJwtJwkSet) sJwtEncoding payload
  where payload = Jose.Claims . toStrict $ JSON.encode token


-- | Invalidate the current session ID (and possibly more, check
-- 'ForceInvalidate').  This is useful to avoid session fixation
-- attacks (cf. <http://www.acrossecurity.com/papers/session_fixation.pdf>).
--
-- Note that the invalidate /does not/ occur when the call to
-- this action is made!  The sessions will be invalidated on the
-- end of the handler processing.  This means that later calls to
-- 'forceInvalidate' on the same handler will override earlier
-- calls.
--
-- This function works by setting a session variable that is
-- checked when saving the session.  The session variable set by
-- this function is then discarded and is not persisted across
-- requests.
forceInvalidate :: MonadHandler m => ForceInvalidate -> m ()
forceInvalidate = setSessionBS forceInvalidateKey . encodeUtf8 . toPathPiece


instance IsSessionData (Map Text ByteString) where
  type Decomposed (Map Text ByteString) = Map Text ByteString
  emptySession = mempty
  decomposeSession authKey session
    = let dsAuthId = Map.lookup authKey session
          dsForceInvalidate = fromMaybe DoNotForceInvalidate
            $ fromPathPiece . decodeUtf8 =<< Map.lookup forceInvalidateKey session
          dsDecomposed = session
            & Map.delete authKey
            & Map.delete forceInvalidateKey
       in DecomposedSession{..}
  recomposeSession authKey mAuthId
    = maybe id (Map.insert authKey) mAuthId
  isDecomposedEmpty _ = Map.null
  isSameDecomposed _ = (==)