-- | The implementation of "Network.TLS.SessionManager", exposed for
-- testing.  This module may change without notice.
module Network.TLS.SessionManager.Internal where

import Control.Exception (assert)
import Control.Reaper
import Data.ByteArray (convert)
import Data.ByteString (ByteString)
import Data.ByteString.Short (ShortByteString)
import qualified Data.ByteString.Short as SBS
import Data.IORef
import Data.OrdPSQ (OrdPSQ)
import qualified Data.OrdPSQ as Q
import Network.TLS
import qualified Network.TLS.Internal as I
import qualified System.Clock as C

import Network.TLS.Imports

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

-- | Configuration for session managers.
data Config = Config
    { Config -> Int
ticketLifetime :: Int
    -- ^ Ticket lifetime in seconds.  It should match 'serverTicketLifetime'
    -- on a server.  A TLS 1.3 session is kept no longer than the lifetime of
    -- its ticket.
    , Config -> Int
pruningDelay :: Int
    -- ^ Pruning delay in seconds. This is set to 'reaperDelay'.
    , Config -> Int
dbMaxSize :: Int
    -- ^ The limit size of session data entries.
    }

-- | ticketLifetime: 2 hours (7200 seconds), pruningDelay: 10 minutes (600 seconds), dbMaxSize: 1000 entries.
defaultConfig :: Config
defaultConfig :: Config
defaultConfig =
    Config
        { ticketLifetime :: Int
ticketLifetime = Int
7200
        , pruningDelay :: Int
pruningDelay = Int
600
        , dbMaxSize :: Int
dbMaxSize = Int
1000
        }

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

toKey :: ByteString -> SessionIDCopy
toKey :: ByteString -> SessionIDCopy
toKey = ByteString -> SessionIDCopy
SBS.toShort

toValue :: SessionData -> SessionDataCopy
toValue :: SessionData -> SessionDataCopy
toValue SessionData
sd =
    SessionData -> SessionDataCopy
SessionDataCopy (SessionData -> SessionDataCopy) -> SessionData -> SessionDataCopy
forall a b. (a -> b) -> a -> b
$
        SessionData
sd
            { sessionSecret = convert $ sessionSecret sd
            , sessionALPN = convert <$> sessionALPN sd
            }

fromValue :: SessionDataCopy -> SessionData
fromValue :: SessionDataCopy -> SessionData
fromValue (SessionDataCopy SessionData
sd) =
    SessionData
sd
        { sessionSecret = convert $ sessionSecret sd
        , sessionALPN = convert <$> sessionALPN sd
        }

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

type SessionIDCopy = ShortByteString
newtype SessionDataCopy = SessionDataCopy SessionData
    deriving (Int -> SessionDataCopy -> ShowS
[SessionDataCopy] -> ShowS
SessionDataCopy -> String
(Int -> SessionDataCopy -> ShowS)
-> (SessionDataCopy -> String)
-> ([SessionDataCopy] -> ShowS)
-> Show SessionDataCopy
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SessionDataCopy -> ShowS
showsPrec :: Int -> SessionDataCopy -> ShowS
$cshow :: SessionDataCopy -> String
show :: SessionDataCopy -> String
$cshowList :: [SessionDataCopy] -> ShowS
showList :: [SessionDataCopy] -> ShowS
Show, SessionDataCopy -> SessionDataCopy -> Bool
(SessionDataCopy -> SessionDataCopy -> Bool)
-> (SessionDataCopy -> SessionDataCopy -> Bool)
-> Eq SessionDataCopy
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SessionDataCopy -> SessionDataCopy -> Bool
== :: SessionDataCopy -> SessionDataCopy -> Bool
$c/= :: SessionDataCopy -> SessionDataCopy -> Bool
/= :: SessionDataCopy -> SessionDataCopy -> Bool
Eq)

type Sec = Int64
type Value = (SessionDataCopy, IORef Availability)
type Sessions = OrdPSQ SessionIDCopy Sec Value

-- | The sessions, and the keys deleted since the reaper last took them away
-- to prune.  While it prunes, the workload left in their place is empty, so
-- a deletion there is kept to be applied to the pruned sessions too.
data DB = DB Sessions [SessionIDCopy]

data Item = Add SessionIDCopy Sec Value | Del SessionIDCopy

data Use = SingleUse | MultipleUse
data Availability = Fresh | Used

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

-- | Creating an in-memory session manager.
newSessionManager :: Config -> IO SessionManager
newSessionManager :: Config -> IO SessionManager
newSessionManager Config
conf = do
    let lifetime :: Sec
lifetime = Int -> Sec
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Sec) -> Int -> Sec
forall a b. (a -> b) -> a -> b
$ Config -> Int
ticketLifetime Config
conf
        maxsiz :: Int
maxsiz = Config -> Int
dbMaxSize Config
conf
    reaper <-
        ReaperSettings DB Item -> IO (Reaper DB Item)
forall workload item.
ReaperSettings workload item -> IO (Reaper workload item)
mkReaper
            ReaperSettings [ZonkAny 0] (ZonkAny 0)
forall item. ReaperSettings [item] item
defaultReaperSettings
                { reaperEmpty = DB Q.empty []
                , reaperCons = cons maxsiz
                , reaperAction = clean maxsiz
                , reaperNull = \(DB Sessions
q [SessionIDCopy]
_) -> Sessions -> Bool
forall k p v. OrdPSQ k p v -> Bool
Q.null Sessions
q
                , reaperDelay = pruningDelay conf * 1000000
                , reaperThreadName = "TLS session manager"
                }
    return $
        noSessionManager
            { sessionResume = resume reaper MultipleUse
            , sessionResumeOnlyOnce = resume reaper SingleUse
            , sessionEstablish = \ByteString
x SessionData
y -> Reaper DB Item -> Sec -> ByteString -> SessionData -> IO ()
establish Reaper DB Item
reaper Sec
lifetime ByteString
x SessionData
y IO () -> IO (Maybe ByteString) -> IO (Maybe ByteString)
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Maybe ByteString -> IO (Maybe ByteString)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ByteString
forall a. Maybe a
Nothing
            , sessionInvalidate = invalidate reaper
            , sessionUseTicket = False
            }

cons :: Int -> Item -> DB -> DB
cons :: Int -> Item -> DB -> DB
cons Int
lim (Add SessionIDCopy
k Sec
t Value
v) (DB Sessions
q [SessionIDCopy]
ds)
    | Int
lim Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = Sessions -> [SessionIDCopy] -> DB
DB Sessions
forall k p v. OrdPSQ k p v
Q.empty [SessionIDCopy]
ds
    -- Replacing a session makes no room for it.
    | Sessions -> Int
forall k p v. OrdPSQ k p v -> Int
Q.size Sessions
q Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
lim Bool -> Bool -> Bool
&& Bool -> Bool
not (SessionIDCopy -> Sessions -> Bool
forall k p v. Ord k => k -> OrdPSQ k p v -> Bool
Q.member SessionIDCopy
k Sessions
q) = case Sessions -> Maybe (SessionIDCopy, Sec, Value, Sessions)
forall k p v.
(Ord k, Ord p) =>
OrdPSQ k p v -> Maybe (k, p, v, OrdPSQ k p v)
Q.minView Sessions
q of
        Maybe (SessionIDCopy, Sec, Value, Sessions)
Nothing -> Bool -> DB -> DB
forall a. (?callStack::CallStack) => Bool -> a -> a
assert Bool
False (DB -> DB) -> DB -> DB
forall a b. (a -> b) -> a -> b
$ Sessions -> [SessionIDCopy] -> DB
DB (SessionIDCopy -> Sec -> Value -> Sessions -> Sessions
forall k p v.
(Ord k, Ord p) =>
k -> p -> v -> OrdPSQ k p v -> OrdPSQ k p v
Q.insert SessionIDCopy
k Sec
t Value
v Sessions
forall k p v. OrdPSQ k p v
Q.empty) [SessionIDCopy]
ds
        Just (SessionIDCopy
_, Sec
_, Value
_, Sessions
q') -> Sessions -> [SessionIDCopy] -> DB
DB (SessionIDCopy -> Sec -> Value -> Sessions -> Sessions
forall k p v.
(Ord k, Ord p) =>
k -> p -> v -> OrdPSQ k p v -> OrdPSQ k p v
Q.insert SessionIDCopy
k Sec
t Value
v Sessions
q') [SessionIDCopy]
ds
    | Bool
otherwise = Sessions -> [SessionIDCopy] -> DB
DB (SessionIDCopy -> Sec -> Value -> Sessions -> Sessions
forall k p v.
(Ord k, Ord p) =>
k -> p -> v -> OrdPSQ k p v -> OrdPSQ k p v
Q.insert SessionIDCopy
k Sec
t Value
v Sessions
q) [SessionIDCopy]
ds
-- At most as many keys as there can be sessions are kept.
cons Int
lim (Del SessionIDCopy
k) (DB Sessions
q [SessionIDCopy]
ds) = Sessions -> [SessionIDCopy] -> DB
DB (SessionIDCopy -> Sessions -> Sessions
forall k p v. (Ord k, Ord p) => k -> OrdPSQ k p v -> OrdPSQ k p v
Q.delete SessionIDCopy
k Sessions
q) (Int -> [SessionIDCopy] -> [SessionIDCopy]
forall a. Int -> [a] -> [a]
take Int
lim (SessionIDCopy
k SessionIDCopy -> [SessionIDCopy] -> [SessionIDCopy]
forall a. a -> [a] -> [a]
: [SessionIDCopy]
ds))

clean :: Int -> DB -> IO (DB -> DB)
clean :: Int -> DB -> IO (DB -> DB)
clean Int
lim (DB Sessions
oldq [SessionIDCopy]
_) = do
    currentTime <- TimeSpec -> Sec
C.sec (TimeSpec -> Sec) -> IO TimeSpec -> IO Sec
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Clock -> IO TimeSpec
C.getTime Clock
C.Monotonic
    let pruned = ([(SessionIDCopy, Sec, Value)], Sessions) -> Sessions
forall a b. (a, b) -> b
snd (([(SessionIDCopy, Sec, Value)], Sessions) -> Sessions)
-> ([(SessionIDCopy, Sec, Value)], Sessions) -> Sessions
forall a b. (a -> b) -> a -> b
$ Sec -> Sessions -> ([(SessionIDCopy, Sec, Value)], Sessions)
forall k p v.
(Ord k, Ord p) =>
p -> OrdPSQ k p v -> ([(k, p, v)], OrdPSQ k p v)
Q.atMostView Sec
currentTime Sessions
oldq
    return $ merge pruned
  where
    ins :: OrdPSQ k p v -> (k, p, v) -> OrdPSQ k p v
ins OrdPSQ k p v
q (k
k, p
p, v
v) = k -> p -> v -> OrdPSQ k p v -> OrdPSQ k p v
forall k p v.
(Ord k, Ord p) =>
k -> p -> v -> OrdPSQ k p v -> OrdPSQ k p v
Q.insert k
k p
p v
v OrdPSQ k p v
q
    -- There is not 'merge' API.
    -- We hope that the new sessions are fewer than the pruned ones.
    merge :: Sessions -> DB -> DB
merge Sessions
pruned (DB Sessions
newq [SessionIDCopy]
ds) = Sessions -> [SessionIDCopy] -> DB
DB (Sessions -> Sessions
forall {a} {b} {c}. (Ord a, Ord b) => OrdPSQ a b c -> OrdPSQ a b c
limit (Sessions -> Sessions) -> Sessions -> Sessions
forall a b. (a -> b) -> a -> b
$ (Sessions -> (SessionIDCopy, Sec, Value) -> Sessions)
-> Sessions -> [(SessionIDCopy, Sec, Value)] -> Sessions
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Sessions -> (SessionIDCopy, Sec, Value) -> Sessions
forall {k} {p} {v}.
(Ord k, Ord p) =>
OrdPSQ k p v -> (k, p, v) -> OrdPSQ k p v
ins ((SessionIDCopy -> Sessions -> Sessions)
-> Sessions -> [SessionIDCopy] -> Sessions
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr SessionIDCopy -> Sessions -> Sessions
forall k p v. (Ord k, Ord p) => k -> OrdPSQ k p v -> OrdPSQ k p v
Q.delete Sessions
pruned [SessionIDCopy]
ds) [(SessionIDCopy, Sec, Value)]
entries) []
      where
        entries :: [(SessionIDCopy, Sec, Value)]
entries = Sessions -> [(SessionIDCopy, Sec, Value)]
forall k p v. OrdPSQ k p v -> [(k, p, v)]
Q.toList Sessions
newq
    -- The size was checked against the new sessions alone.
    limit :: OrdPSQ a b c -> OrdPSQ a b c
limit OrdPSQ a b c
q
        | Int
lim Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = OrdPSQ a b c
forall k p v. OrdPSQ k p v
Q.empty
        | OrdPSQ a b c -> Int
forall k p v. OrdPSQ k p v -> Int
Q.size OrdPSQ a b c
q Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
lim = OrdPSQ a b c
-> ((a, b, c, OrdPSQ a b c) -> OrdPSQ a b c)
-> Maybe (a, b, c, OrdPSQ a b c)
-> OrdPSQ a b c
forall b a. b -> (a -> b) -> Maybe a -> b
maybe OrdPSQ a b c
q (\(a
_, b
_, c
_, OrdPSQ a b c
q') -> OrdPSQ a b c -> OrdPSQ a b c
limit OrdPSQ a b c
q') (Maybe (a, b, c, OrdPSQ a b c) -> OrdPSQ a b c)
-> Maybe (a, b, c, OrdPSQ a b c) -> OrdPSQ a b c
forall a b. (a -> b) -> a -> b
$ OrdPSQ a b c -> Maybe (a, b, c, OrdPSQ a b c)
forall k p v.
(Ord k, Ord p) =>
OrdPSQ k p v -> Maybe (k, p, v, OrdPSQ k p v)
Q.minView OrdPSQ a b c
q
        | Bool
otherwise = OrdPSQ a b c
q

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

establish
    :: Reaper DB Item
    -> Sec
    -> SessionID
    -> SessionData
    -> IO ()
establish :: Reaper DB Item -> Sec -> ByteString -> SessionData -> IO ()
establish Reaper DB Item
reaper Sec
lifetime ByteString
k SessionData
sd = do
    ref <- Availability -> IO (IORef Availability)
forall a. a -> IO (IORef a)
newIORef Availability
Fresh
    p <- (+ lifetime') . C.sec <$> C.getTime C.Monotonic
    let v = (SessionDataCopy
sd', IORef Availability
ref)
    reaperAdd reaper (Add k' p v)
  where
    k' :: SessionIDCopy
k' = ByteString -> SessionIDCopy
toKey ByteString
k
    sd' :: SessionDataCopy
sd' = SessionData -> SessionDataCopy
toValue SessionData
sd
    -- A TLS 1.3 ticket is not kept past the lifetime the server gave it.
    lifetime' :: Sec
lifetime' = case SessionData -> Maybe TLS13TicketInfo
sessionTicketInfo SessionData
sd of
        Just TLS13TicketInfo
tinfo -> Sec -> Sec -> Sec
forall a. Ord a => a -> a -> a
min Sec
lifetime (Second -> Sec
forall a b. (Integral a, Num b) => a -> b
fromIntegral (TLS13TicketInfo -> Second
I.lifetime TLS13TicketInfo
tinfo))
        Maybe TLS13TicketInfo
Nothing -> Sec
lifetime

resume
    :: Reaper DB Item
    -> Use
    -> SessionID
    -> IO (Maybe SessionData)
resume :: Reaper DB Item -> Use -> ByteString -> IO (Maybe SessionData)
resume Reaper DB Item
reaper Use
use ByteString
k = do
    DB q _ <- Reaper DB Item -> IO DB
forall workload item. Reaper workload item -> IO workload
reaperRead Reaper DB Item
reaper
    now <- C.sec <$> C.getTime C.Monotonic
    case Q.lookup k' q of
        -- An expired session stays until the next pruning; it is not used.
        Just (Sec
p, (SessionDataCopy
sd, IORef Availability
ref))
            | Sec
p Sec -> Sec -> Bool
forall a. Ord a => a -> a -> Bool
> Sec
now -> case Use
use of
                Use
SingleUse -> do
                    available <- IORef Availability
-> (Availability -> (Availability, Bool)) -> IO Bool
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Availability
ref Availability -> (Availability, Bool)
check
                    reaperAdd reaper (Del k')
                    return $ if available then Just (fromValue sd) else Nothing
                Use
MultipleUse -> Maybe SessionData -> IO (Maybe SessionData)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe SessionData -> IO (Maybe SessionData))
-> Maybe SessionData -> IO (Maybe SessionData)
forall a b. (a -> b) -> a -> b
$ SessionData -> Maybe SessionData
forall a. a -> Maybe a
Just (SessionDataCopy -> SessionData
fromValue SessionDataCopy
sd)
        Maybe (Sec, Value)
_ -> Maybe SessionData -> IO (Maybe SessionData)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe SessionData
forall a. Maybe a
Nothing
  where
    check :: Availability -> (Availability, Bool)
check Availability
Fresh = (Availability
Used, Bool
True)
    check Availability
Used = (Availability
Used, Bool
False)
    k' :: SessionIDCopy
k' = ByteString -> SessionIDCopy
toKey ByteString
k

invalidate
    :: Reaper DB Item
    -> SessionID
    -> IO ()
-- Not looked up first: while the reaper prunes, the session is not seen.
invalidate :: Reaper DB Item -> ByteString -> IO ()
invalidate Reaper DB Item
reaper ByteString
k = Reaper DB Item -> Item -> IO ()
forall workload item. Reaper workload item -> item -> IO ()
reaperAdd Reaper DB Item
reaper (SessionIDCopy -> Item
Del (ByteString -> SessionIDCopy
toKey ByteString
k))