{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE NoFieldSelectors #-}

-- | This module exports tools for safely storing encrypted data on client-side
-- cookies through "Network.Wai". Consider using it in conjunction with "Wai.CSRF".
module Wai.CryptoCookie
   ( -- * Config
    defaultConfig
   , Config (..)
   , Env
   , newEnv

    -- * Request and responses
   , middleware
   , msgFromRequestCookie
   , setCookie
   , expireCookie

    -- * Encryption
   , Encryption (..)
   , autoKeyFileBase16
   , readKeyFileBase16
   , readKeyFile
   , writeKeyFile
   ) where

import Control.Monad.IO.Class
import Data.Aeson qualified as Ae
import Data.ByteArray.Encoding qualified as BA
import Data.ByteArray.Sized qualified as BAS
import Data.ByteString qualified as B
import Data.ByteString.Lazy qualified as BL
import Data.IORef
import Data.Kind (Type)
import Data.Time.Clock.POSIX qualified as Time
import Network.Wai qualified as Wai
import Wai.CSRF qualified
import Wai.CryptoCookie.Encryption
import Wai.CryptoCookie.Encryption.AEAD_AES_128_GCM_SIV ()
import Wai.CryptoCookie.Encryption.AEAD_AES_256_GCM_SIV ()
import Web.Cookie qualified as WC

-- | Default 'Config':
--
-- * Cookie name is @SESSION@.
--
-- * Encoding and decoding of @msg@ is done through 'Ae.ToJSON' and
--   'Ae.FromJSON'.
--
-- * The 'Encryption' scheme is the nonce-misuse resistant @AEAD_AES_256_GCM_SIV@
--   as defined in in <https://tools.ietf.org/html/rfc8452 RFC 8452>, using
--   a "Wai.CSRF".'Wai.CSRF.Token' as AEAD associated data.
--
--     As an AEAD encryption scheme, you can be confident that a successfully
--     decrypted cookie could only have been encrypted by the same
--     'Key' known only to your server, and associated with a specific
--     "Wai.CSRF".'Wai.CSRF.Token', expected to have been sent with the
--     incoming request.
--
--     In principle, this makes this encryption scheme suitable for storing
--     server-generated user session data in the @msg@.  However, you must make
--     sure that you rotate the "Wai.CSRF".'Wai.CSRF.Token' ocassionally, at
--     least each time a new user session is established, so as to avoid CSRF
--     risks.
--
-- * This 'defaultConfig' suggests you should be composing 'middleware' and
--   "Wai.CSRF".'Wai.CSRF.middleware' in this way:
--
--      @
--      "Wai.CSRF".'Wai.CSRF.middleware' /myCsrfConfig/
--         . "Wai.CryptoCookie".'middleware' /myCryptoCookieEnv/
--              :: ('Maybe' ("Wai.CSRF".'Wai.CSRF.Token', msg) -> 'Wai.Application')
--              -> 'Wai.Application'
--      @
defaultConfig
   :: (Ae.ToJSON msg, Ae.FromJSON msg)
   => Key "AEAD_AES_256_GCM_SIV"
   -- ^ Consider using 'autoKeyFileBase16' or
   -- 'readKeyFileBase16' for safely reading a 'Key' from a
   -- 'FilePath'. Alternatively, if you have the base-16 representation of the
   -- 'Key' in JSON configuration, you could use
   -- 'Data.Aeson.FromJSON'.
   -> Config Wai.CSRF.Token msg
defaultConfig :: forall msg.
(ToJSON msg, FromJSON msg) =>
Key "AEAD_AES_256_GCM_SIV" -> Config Token msg
defaultConfig Key "AEAD_AES_256_GCM_SIV"
key =
   Config
      { cookieName :: ByteString
cookieName = ByteString
"SESSION"
      , Key "AEAD_AES_256_GCM_SIV"
key :: Key "AEAD_AES_256_GCM_SIV"
key :: Key "AEAD_AES_256_GCM_SIV"
key
      , aadEncode :: Token -> ByteString
aadEncode = \(Wai.CSRF.Token SizedByteArray 32 ByteString
t) ->
         ByteString -> ByteString
BL.fromStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ SizedByteArray 32 ByteString -> ByteString
forall (n :: Nat) ba. SizedByteArray n ba -> ba
BAS.unSizedByteArray SizedByteArray 32 ByteString
t
      , msgEncode :: msg -> ByteString
msgEncode = msg -> ByteString
forall a. ToJSON a => a -> ByteString
Ae.encode
      , msgDecode :: ByteString -> Maybe msg
msgDecode = ByteString -> Maybe msg
forall a. FromJSON a => ByteString -> Maybe a
Ae.decode
      }

-- | Configuration for 'Env'.
--
-- Consider using 'defaultConfig' and updating desired fields only.
data Config (aad :: Type) (msg :: Type) = forall e.
    (Encryption e) =>
   Config
   { forall aad msg. Config aad msg -> ByteString
cookieName :: B.ByteString
   -- ^ Consider using a @\"SESSION\"@.
   , ()
key :: Key e
   -- ^ Consider using a @'Key' \"AEAD_AES_256_GCM_SIV\"@.
   , forall aad msg. Config aad msg -> aad -> ByteString
aadEncode :: aad -> BL.ByteString
   -- ^ These are the exact bytes that will be used as AEAD associated data.
   -- Consider using the raw bytes of a "Wai.CSRF".'Wai.CSRF.Token'.
   , forall aad msg. Config aad msg -> msg -> ByteString
msgEncode :: msg -> BL.ByteString
   -- ^ These are the exact bytes that will be encrypted.
   , forall aad msg. Config aad msg -> ByteString -> Maybe msg
msgDecode :: BL.ByteString -> Maybe msg
   -- ^ Undo what @msgEncode@ did, if possible.
   }

-- | Stateful encryption environment for interacting with the encrypted cookie.
--
-- It is safe to use 'Env' concurrently if necessary. Concurrency is handled
-- safely internally.
--
-- Obtain with 'newEnv'.
data Env (aad :: Type) (msg :: Type) = Env
   { forall aad msg. Env aad msg -> ByteString
cookieName :: B.ByteString
   , forall aad msg. Env aad msg -> aad -> msg -> IO ByteString
encodeEncrypt :: aad -> msg -> IO BL.ByteString
   , forall aad msg. Env aad msg -> aad -> ByteString -> Maybe msg
decryptDecode :: aad -> BL.ByteString -> Maybe msg
   }

--------------------------------------------------------------------------------

-- | Obtain a new 'Env'.
newEnv :: (MonadIO m) => Config aad msg -> m (Env aad msg)
newEnv :: forall (m :: * -> *) aad msg.
MonadIO m =>
Config aad msg -> m (Env aad msg)
newEnv c :: Config aad msg
c@Config{Key e
key :: ()
key :: Key e
key} = IO (Env aad msg) -> m (Env aad msg)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
   let dc :: Decrypt e
dc = Key e -> Decrypt e
forall k (e :: k). Encryption e => Key e -> Decrypt e
initDecrypt Key e
key
   IORef (Encrypt e)
ecRef <- Encrypt e -> IO (IORef (Encrypt e))
forall a. a -> IO (IORef a)
newIORef (Encrypt e -> IO (IORef (Encrypt e)))
-> IO (Encrypt e) -> IO (IORef (Encrypt e))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Key e -> IO (Encrypt e)
forall k (e :: k) (m :: * -> *).
(Encryption e, MonadRandom m) =>
Key e -> m (Encrypt e)
forall (m :: * -> *). MonadRandom m => Key e -> m (Encrypt e)
initEncrypt Key e
key
   Env aad msg -> IO (Env aad msg)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
      Env
         { encodeEncrypt :: aad -> msg -> IO ByteString
encodeEncrypt = \aad
aad0 msg
msg0 -> do
            let !ByteString
aad1 :: BL.ByteString = Config aad msg
c.aadEncode aad
aad0
                !ByteString
msg1 :: BL.ByteString = Config aad msg
c.msgEncode msg
msg0
            Encrypt e
ec <- IORef (Encrypt e)
-> (Encrypt e -> (Encrypt e, Encrypt e)) -> IO (Encrypt e)
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef (Encrypt e)
ecRef \Encrypt e
ec -> (Encrypt e -> Encrypt e
forall k (e :: k). Encryption e => Encrypt e -> Encrypt e
advance Encrypt e
ec, Encrypt e
ec)
            ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ Encrypt e -> ByteString -> ByteString -> ByteString
forall k (e :: k).
Encryption e =>
Encrypt e -> ByteString -> ByteString -> ByteString
encrypt Encrypt e
ec ByteString
aad1 ByteString
msg1
         , decryptDecode :: aad -> ByteString -> Maybe msg
decryptDecode = \aad
aad0 !ByteString
cry -> do
            let !aad1 :: ByteString
aad1 = Config aad msg
c.aadEncode aad
aad0
            case Decrypt e -> ByteString -> ByteString -> Either String ByteString
forall k (e :: k).
Encryption e =>
Decrypt e -> ByteString -> ByteString -> Either String ByteString
decrypt Decrypt e
dc ByteString
aad1 ByteString
cry of
               Right ByteString
msg -> Config aad msg
c.msgDecode ByteString
msg
               Either String ByteString
_ -> Maybe msg
forall a. Maybe a
Nothing
         , cookieName :: ByteString
cookieName = Config aad msg
c.cookieName
         }

-- | Transform an 'Wai.Application' so that if there is an encrypted
-- message in the incoming 'Wai.Request' cookies, it will be automatically
-- decrypted and made available to the underlying 'Wai.Application'.
--
-- The @aad@ is the AEAD associated data that came with the 'Wai.Request'.
-- Consider using 'middleware' in conjunction with
-- "Wai.CSRF".'Wai.CSRF.middleware', using "Wai.CSRF".'Wai.CSRF.Token' as
-- @aad@.
middleware
   :: Env aad msg
   -- ^ Encryption environment. Obtain with 'newEnv'.
   -> (Maybe (aad, Maybe msg) -> Wai.Application)
   -- ^ Underlying 'Wai.Application' having access to the decrypted cookie
   -- @msg@, if any.
   --
   -- Also, seeing as @msg@ being available implies @aad@ is available too, we
   -- output both values together in a manner that represents this relationship.
   -> Maybe aad
   -- ^ AEAD associated data of the incomming 'Wai.Request', if any.
   -> Wai.Application
middleware :: forall aad msg.
Env aad msg
-> (Maybe (aad, Maybe msg) -> Application)
-> Maybe aad
-> Application
middleware Env aad msg
env Maybe (aad, Maybe msg) -> Application
fapp Maybe aad
yaad Request
req Response -> IO ResponseReceived
respond = do
   let ymsg :: Maybe msg
ymsg = Env aad msg -> Request -> aad -> Maybe msg
forall aad msg. Env aad msg -> Request -> aad -> Maybe msg
msgFromRequestCookie Env aad msg
env Request
req (aad -> Maybe msg) -> Maybe aad -> Maybe msg
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Maybe aad
yaad
   Maybe (aad, Maybe msg) -> Application
fapp ((aad -> (aad, Maybe msg)) -> Maybe aad -> Maybe (aad, Maybe msg)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (,Maybe msg
ymsg) Maybe aad
yaad) Request
req Response -> IO ResponseReceived
respond

-- | Obtain the @msg@ from the 'Wai.Request' cookies.
--
-- You don't need to use this if you are using 'middleware'.
msgFromRequestCookie :: Env aad msg -> Wai.Request -> aad -> Maybe msg
msgFromRequestCookie :: forall aad msg. Env aad msg -> Request -> aad -> Maybe msg
msgFromRequestCookie Env aad msg
env Request
r aad
aad = do
   [ByteString
d64] <- [ByteString] -> Maybe [ByteString]
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ByteString] -> Maybe [ByteString])
-> [ByteString] -> Maybe [ByteString]
forall a b. (a -> b) -> a -> b
$ ByteString -> [(ByteString, ByteString)] -> [ByteString]
forall k v. Eq k => k -> [(k, v)] -> [v]
lookupMany Env aad msg
env.cookieName ([(ByteString, ByteString)] -> [ByteString])
-> [(ByteString, ByteString)] -> [ByteString]
forall a b. (a -> b) -> a -> b
$ Request -> [(ByteString, ByteString)]
requestCookies Request
r
   case Base -> ByteString -> Either String ByteString
forall bin bout.
(ByteArrayAccess bin, ByteArray bout) =>
Base -> bin -> Either String bout
BA.convertFromBase Base
BA.Base64URLUnpadded ByteString
d64 of
      Right ByteString
cry -> Env aad msg
env.decryptDecode aad
aad (ByteString -> Maybe msg) -> ByteString -> Maybe msg
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
BL.fromStrict ByteString
cry
      Left String
_ -> Maybe msg
forall a. Maybe a
Nothing

--------------------------------------------------------------------------------

-- | Construct a 'C.SetCookie' containing the encrypted @msg@.
--
-- The associated data @aad@ will not be included in this cookie, but it will
-- be taken into account for encryption and necessary for eventual decryption.
--
-- The 'C.SetCookie' has these settings, some of which could be overriden.
--
--      * Cookie name is 'Config'\'s @cookieName@.
--
--      * @HttpOnly@: Yes, and you shouldn't change this.
--
--      * @Max-Age@ and @Expires@: This cookie never expires. We recommend
--      relying on server-side expiration instead, as the lifetime of the
--      cookie could easily be extended by a legitimate but malicious client.
--      You can store a creation or expiration timestamp inside @msg@, and
--      make a decision based on that.
--
--      * @Path@: @\/@
--
--      * @SameSite@: @Lax@.
--
--      * @Secure@: Yes.
--
--      * @Domain@: Not set.
--
-- Note: If you are using "Wai.CSRF".'Wai.CSRF.Token' as @aad@, it is
-- recommended that you generate a new "Wai.CSRF".'Wai.CSRF.Token' at least
-- each time a new user session is established, but possibly more frequently,
-- and send it alongside this one (see "Wai.CSRF".'Wai.CSRF.setCookie').
setCookie :: (MonadIO m) => Env aad msg -> aad -> msg -> m WC.SetCookie
setCookie :: forall (m :: * -> *) aad msg.
MonadIO m =>
Env aad msg -> aad -> msg -> m SetCookie
setCookie Env aad msg
env aad
aad msg
msg = IO SetCookie -> m SetCookie
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO do
   ByteString
cry <- Env aad msg
env.encodeEncrypt aad
aad msg
msg
   SetCookie -> IO SetCookie
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SetCookie -> IO SetCookie) -> SetCookie -> IO SetCookie
forall a b. (a -> b) -> a -> b
$
      (Env aad msg -> SetCookie
forall aad msg. Env aad msg -> SetCookie
expireCookie Env aad msg
env)
         { WC.setCookieExpires = Nothing
         , WC.setCookieMaxAge = Nothing
         , WC.setCookieValue =
            BA.convertToBase BA.Base64URLUnpadded $ BL.toStrict cry
         }

-- | Construct a 'C.SetCookie' expiring the cookie named 'Config'\'s
-- @cookieName@.
expireCookie :: Env aad msg -> WC.SetCookie
expireCookie :: forall aad msg. Env aad msg -> SetCookie
expireCookie Env aad msg
env =
   SetCookie
WC.defaultSetCookie
      { WC.setCookieDomain = Nothing
      , WC.setCookieExpires = Just (Time.posixSecondsToUTCTime 0)
      , WC.setCookieHttpOnly = True
      , WC.setCookieMaxAge = Just (negate 1)
      , WC.setCookieName = env.cookieName
      , WC.setCookiePath = Just "/"
      , WC.setCookieSameSite = Just WC.sameSiteLax
      , WC.setCookieSecure = True
      , WC.setCookieValue = ""
      }

--------------------------------------------------------------------------------

requestCookies :: Wai.Request -> [(B.ByteString, B.ByteString)]
requestCookies :: Request -> [(ByteString, ByteString)]
requestCookies Request
r =
   ByteString -> [(ByteString, ByteString)]
WC.parseCookies (ByteString -> [(ByteString, ByteString)])
-> [ByteString] -> [(ByteString, ByteString)]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< HeaderName -> [(HeaderName, ByteString)] -> [ByteString]
forall k v. Eq k => k -> [(k, v)] -> [v]
lookupMany HeaderName
"Cookie" (Request -> [(HeaderName, ByteString)]
Wai.requestHeaders Request
r)

lookupMany :: (Eq k) => k -> [(k, v)] -> [v]
lookupMany :: forall k v. Eq k => k -> [(k, v)] -> [v]
lookupMany k
k = (k -> Bool) -> [(k, v)] -> [v]
forall k v. Eq k => (k -> Bool) -> [(k, v)] -> [v]
findMany (k -> k -> Bool
forall a. Eq a => a -> a -> Bool
== k
k)

findMany :: (Eq k) => (k -> Bool) -> [(k, v)] -> [v]
findMany :: forall k v. Eq k => (k -> Bool) -> [(k, v)] -> [v]
findMany k -> Bool
f = ((k, v) -> v) -> [(k, v)] -> [v]
forall a b. (a -> b) -> [a] -> [b]
map (k, v) -> v
forall a b. (a, b) -> b
snd ([(k, v)] -> [v]) -> ([(k, v)] -> [(k, v)]) -> [(k, v)] -> [v]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k, v) -> Bool) -> [(k, v)] -> [(k, v)]
forall a. (a -> Bool) -> [a] -> [a]
filter (\(k
a, v
_) -> k -> Bool
f k
a)