module Network.Wai.Session.PostgreSQL
    ( clearSession
    , dbStore
    , defaultSettings
    , fromSimpleConnection
    , purgeOldSessions
    , purger
    , ratherSecureGen
    , SimpleConnection
    , StoreSettings (..)
    , WithPostgreSQLConn (..)
    ) where

import Control.Concurrent
import Control.Concurrent.MVar
import Control.Exception.Base
import Control.Exception
import Control.Monad
import Control.Monad.IO.Class
import Data.Default
import Data.Int (Int64)
import Data.Serialize (encode, decode, Serialize)
import Data.String (fromString)
import Data.Time.Clock.POSIX (getPOSIXTime)
import Database.PostgreSQL.Simple
import Network.Wai (Request, requestHeaders)
import Network.Wai.Session
import Numeric (showHex)
import System.Entropy (getEntropy)
import Web.Cookie (parseCookies)

import qualified Data.ByteString as B
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE

-- |These settings control how the session store is behaving
data StoreSettings = StoreSettings {
    -- |The number of seconds a session is valid
    -- Seconds are counted since the session is last accessed (read or written),
    -- not since it was created.
      storeSettingsSessionTimeout :: Int64
    -- |A random session key generator. The session ID should provide
    -- sufficient entropy, and must not be predictable. It is recommended
    -- to use a cryptographically secure random number generator.
    , storeSettingsKeyGen :: IO B.ByteString
    -- |Whether to create the database table if it does not exist upon
    -- creating the session store. If set to false, the database table
    -- must exist or be created by some other means.
    , storeSettingsCreateTable :: Bool
    -- |A function that is called by to log events such as session
    -- purges or the table creation.
    , storeSettingsLog :: String -> IO ()
    -- |The number of microseconds to sleep between two runs of the
    -- old session purge worker.
    , storeSettingsPurgeInterval :: Int
    }

instance Default StoreSettings where
    def = defaultSettings

-- |By default, you pass a postgresql connection to the session store
-- when creating it. The passed connection will have to stay open
-- for the (possibly very long) existence of the session and it should
-- not be used for any other purpose during that time.
-- You can implement an instance of this class for a connection pool
-- instead, so that the session manager will not require a permanent
-- open PostgreSQL connection.
class WithPostgreSQLConn a where
    -- |Call the function (Connection -> IO b) with a valid and open
    -- PostgreSQL connection.
    withPostgreSQLConn :: a -> (Connection -> IO b) -> IO b


-- |Prepare a simple postgresql connection for use by the postgresql
-- session store. This basically wraps the connection along with a mutex
-- to ensure transactions work correctly. Connections used this way must
-- not be used anywhere else for the duration of the session store!
-- It is recommended to use a connection pool instead. To use a connection
-- pool, you simply need to implement the WithPostgreSQLConn type class.
fromSimpleConnection :: Connection -> IO SimpleConnection
fromSimpleConnection connection = do
    mvar <- newMVar ()
    return $ SimpleConnection (mvar, connection)

newtype SimpleConnection = SimpleConnection (MVar (), Connection)

instance WithPostgreSQLConn SimpleConnection where
    withPostgreSQLConn (SimpleConnection (mvar, conn)) =
        bracket (takeMVar mvar >> return conn) (\_ -> putMVar mvar ())

qryCreateTable1 :: Query
qryCreateTable1         = "CREATE TABLE wai_pg_sessions (id bigserial NOT NULL, session_key character varying NOT NULL, session_created_at bigint NOT NULL, session_last_access bigint NOT NULL, session_invalidate_key boolean NOT NULL DEFAULT false, CONSTRAINT session_pkey PRIMARY KEY (id), CONSTRAINT session_session_key_key UNIQUE (session_key)) WITH ( OIDS=FALSE );"

qryCreateTable2 :: Query
qryCreateTable2         = "CREATE TABLE wai_pg_session_data ( id bigserial NOT NULL, wai_pg_session bigint, key bytea, value bytea, CONSTRAINT wai_pg_session_data_pkey PRIMARY KEY (id), CONSTRAINT wai_pg_session_data_wai_pg_session_fkey FOREIGN KEY (wai_pg_session) REFERENCES wai_pg_sessions (id) MATCH SIMPLE ON UPDATE RESTRICT ON DELETE CASCADE, CONSTRAINT wai_pg_session_data_wai_pg_session_key_key UNIQUE (wai_pg_session, key) ) WITH (OIDS=FALSE);"

qryCreateSession :: Query
qryCreateSession        = "INSERT INTO wai_pg_sessions (session_key, session_created_at, session_last_access) VALUES (?,?,?) RETURNING id"

qryCreateSessionEntry :: Query
qryCreateSessionEntry   = "INSERT INTO wai_pg_session_data (wai_pg_session,key,value) VALUES (?,?,?)"

qryUpdateSession :: Query
qryUpdateSession        = "UPDATE wai_pg_sessions SET session_last_access=? WHERE id=?"

qryUpdateSessionEntry :: Query
qryUpdateSessionEntry   = "UPDATE wai_pg_session_data SET value=? WHERE wai_pg_session=? AND key=?"

qryLookupSession :: Query
qryLookupSession        = "SELECT id FROM wai_pg_sessions WHERE session_key=? AND session_last_access>=?"

qryLookupSession' :: Query
qryLookupSession'       = "UPDATE wai_pg_sessions SET session_last_access=? WHERE id=?"

qryLookupSession'' :: Query
qryLookupSession''      = "SELECT value FROM wai_pg_session_data WHERE wai_pg_session=? AND key=?"

qryLookupSession''' :: Query
qryLookupSession'''     = "SELECT id FROM wai_pg_session_data WHERE wai_pg_session=? AND key=?"

qryPurgeOldSessions :: Query
qryPurgeOldSessions     = "DELETE FROM wai_pg_sessions WHERE session_last_access<?"

qryCheckNewKey :: Query
qryCheckNewKey          = "SELECT session_invalidate_key FROM wai_pg_sessions WHERE session_key=?"

qryInvalidateSess1 :: Query
qryInvalidateSess1      = "UPDATE wai_pg_sessions SET session_invalidate_key=TRUE WHERE session_key=?"

qryInvalidateSess2 :: Query
qryInvalidateSess2      = "DELETE FROM wai_pg_session_data WHERE wai_pg_session=(SELECT id FROM wai_pg_sessions WHERE session_key=?)"

qryUpdateKey :: Query
qryUpdateKey            = "UPDATE wai_pg_sessions SET session_key=?,session_invalidate_key=FALSE WHERE session_key=?"

-- |Create a new postgresql backed wai session store.
dbStore :: (WithPostgreSQLConn a, Serialize k, Eq k, Serialize v, MonadIO m) => a -> StoreSettings -> IO (SessionStore m k v)
dbStore pool stos = do
    when (storeSettingsCreateTable stos) $
        withPostgreSQLConn pool $ \ conn ->
            unerror $ do
                void $ execute_ conn qryCreateTable1
                void $ execute_ conn qryCreateTable2
                storeSettingsLog stos "Created tables."
    return $ dbStore' pool stos

-- |Delete expired sessions from the database.
purgeOldSessions :: WithPostgreSQLConn a => a -> StoreSettings -> IO Int64
purgeOldSessions pool stos = do
    curtime <- round <$> liftIO getPOSIXTime
    count <- withPostgreSQLConn pool $ \ conn ->
        execute conn qryPurgeOldSessions (Only (curtime - storeSettingsSessionTimeout stos))
    storeSettingsLog stos $ "Purged " ++ show count ++ " session(s)."
    return count

-- |Run a thread using forkIO that runs periodically to
-- purge old sessions.
purger :: WithPostgreSQLConn a => a -> StoreSettings -> IO ThreadId
purger pool stos = forkIO . forever . unerror $ do
    purgeOldSessions pool stos
    threadDelay $ storeSettingsPurgeInterval stos

-- |Create default settings using a session timeout of
-- one hour, a cryptographically secure session id generator
-- using 24 bytes of entropy and putStrLn to log events
-- to stdout.
defaultSettings :: StoreSettings
defaultSettings = StoreSettings
    { storeSettingsSessionTimeout=3600
    , storeSettingsKeyGen=ratherSecureGen 24
    , storeSettingsCreateTable=True
    , storeSettingsLog=putStrLn
    , storeSettingsPurgeInterval=600000000
    }

-- |Generate a session ID with n bytes of entropy
ratherSecureGen :: Int -> IO B.ByteString
ratherSecureGen n = TE.encodeUtf8 . prettyPrint <$> getEntropy n

prettyPrint :: B.ByteString -> T.Text
prettyPrint = T.pack . concatMap (`showHex` "") . B.unpack

dbStore' :: (WithPostgreSQLConn a, Serialize k, Eq k, Serialize v, MonadIO m) => a -> StoreSettings -> SessionStore m k v
dbStore' pool stos Nothing = do
    newKey <- storeSettingsKeyGen stos
    curtime <- round <$> liftIO getPOSIXTime
    sessionPgId <- withPostgreSQLConn pool $ \ conn -> do
        [Only res] <- query conn qryCreateSession (newKey, curtime :: Int64, curtime) :: IO [Only Int64]
        return (res :: Int64)
    backend pool stos newKey sessionPgId
dbStore' pool stos (Just key) = do
    curtime <- round <$> liftIO getPOSIXTime
    res <- withPostgreSQLConn pool $ \ conn ->
        query conn qryLookupSession (key, curtime - storeSettingsSessionTimeout stos) :: IO [Only Int64]
    case res of
        [Only sessionPgId]  -> backend pool stos key sessionPgId
        _                   -> dbStore' pool stos Nothing

-- |This function can be called to invalidate a session and enforce creating
-- a new one with a new session ID. It should be called *before* any calls
-- to sessionStore are made. It needs to be passed a request and the cookie
-- name explicitly due to the limited nature of the Network.Wai.Session
-- interface.
-- Sessions should be cleared when a login is performed, to prevent certain
-- kinds of session hijacking attacks.
clearSession :: (WithPostgreSQLConn a) => a -> B.ByteString -> Request -> IO ()
clearSession pool cookieName req = do
    let map         = [] :: [(k, v)]
        map'        = "" -- encode map
        cookies     = parseCookies <$> lookup (fromString "Cookie") (requestHeaders req)
        Just key    = lookup cookieName =<< cookies
    withPostgreSQLConn pool $ \ conn ->
        withTransaction conn $ do
            void $ execute conn qryInvalidateSess1 (Only key)
            void $ execute conn qryInvalidateSess2 (Only key)

backend :: (WithPostgreSQLConn a, Serialize k, Eq k, Serialize v, MonadIO m) => a -> StoreSettings -> B.ByteString -> Int64 -> IO (Session m k v, IO B.ByteString)
backend pool stos key sessionPgId =
    return ( (
        reader pool key sessionPgId
      , writer pool key sessionPgId ), withPostgreSQLConn pool $ \conn -> do
        [Only shouldNewKey] <- query conn qryCheckNewKey (Only key)
        if shouldNewKey then do
            newKey' <- storeSettingsKeyGen stos
            execute conn qryUpdateKey (newKey', key)
            return newKey'
        else
            return key
      )


reader :: (WithPostgreSQLConn a, Serialize k, Eq k, Serialize v, MonadIO m) => a -> B.ByteString -> Int64 -> k -> m (Maybe v)
reader pool key sessionPgId k = do
    curtime <- round <$> liftIO getPOSIXTime
    res <- liftIO $ withPostgreSQLConn pool $ \conn -> do
        void $ execute conn qryLookupSession' (curtime :: Int64, sessionPgId)
        query conn qryLookupSession'' (sessionPgId, Binary $ encode k)
    case res of
        [Only value]    -> case decode (fromBinary value) of
            Right value'    -> return $ Just value'
            Left error      -> return Nothing
        []              -> return Nothing

writer :: (WithPostgreSQLConn a, Serialize k, Eq k, Serialize v, MonadIO m) => a -> B.ByteString -> Int64 -> k -> v -> m ()
writer pool key sessionPgId k v = do
    curtime <- round <$> liftIO getPOSIXTime
    let k' = Binary $ encode k
        v' = Binary $ encode v
    liftIO $ withPostgreSQLConn pool $ \conn ->
        withTransaction conn $ do
            res <- query conn qryLookupSession''' (sessionPgId, k') :: IO [Only Int64]
            case res of
                [Only id]   -> void $ execute conn qryUpdateSessionEntry (v', sessionPgId, k')
                _           -> void $ execute conn qryCreateSessionEntry (sessionPgId, k', v')
            void $ execute conn qryUpdateSession (curtime :: Int64, sessionPgId)

ignoreSqlError :: SqlError -> IO ()
ignoreSqlError _ = pure ()

unerror :: IO a -> IO ()
unerror action = void action `catch` ignoreSqlError