    {-# LANGUAGE RecordWildCards #-}

    module Network.DNS.Memo where

    import qualified Control.Reaper as R
    import qualified Data.ByteString as B
    import qualified Data.CaseInsensitive as CI
    import Data.Hourglass (Elapsed)
    import Data.OrdPSQ (OrdPSQ)
    import qualified Data.OrdPSQ as PSQ
    import Time.System (timeCurrent)

    import Network.DNS.Imports
    import Network.DNS.Types.Internal

    data Section = Answer | Authority deriving (Section -> Section -> Bool
    (Section -> Section -> Bool)
    -> (Section -> Section -> Bool) -> Eq Section
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: Section -> Section -> Bool
    == :: Section -> Section -> Bool
    $c/= :: Section -> Section -> Bool
    /= :: Section -> Section -> Bool
    Eq, Eq Section
    Eq Section =>
    (Section -> Section -> Ordering)
    -> (Section -> Section -> Bool)
    -> (Section -> Section -> Bool)
    -> (Section -> Section -> Bool)
    -> (Section -> Section -> Bool)
    -> (Section -> Section -> Section)
    -> (Section -> Section -> Section)
    -> Ord Section
    Section -> Section -> Bool
    Section -> Section -> Ordering
    Section -> Section -> Section
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    $ccompare :: Section -> Section -> Ordering
    compare :: Section -> Section -> Ordering
    $c< :: Section -> Section -> Bool
    < :: Section -> Section -> Bool
    $c<= :: Section -> Section -> Bool
    <= :: Section -> Section -> Bool
    $c> :: Section -> Section -> Bool
    > :: Section -> Section -> Bool
    $c>= :: Section -> Section -> Bool
    >= :: Section -> Section -> Bool
    $cmax :: Section -> Section -> Section
    max :: Section -> Section -> Section
    $cmin :: Section -> Section -> Section
    min :: Section -> Section -> Section
    Ord, Int -> Section -> ShowS
    [Section] -> ShowS
    Section -> String
    (Int -> Section -> ShowS)
    -> (Section -> String) -> ([Section] -> ShowS) -> Show Section
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> Section -> ShowS
    showsPrec :: Int -> Section -> ShowS
    $cshow :: Section -> String
    show :: Section -> String
    $cshowList :: [Section] -> ShowS
    showList :: [Section] -> ShowS
    Show)

    type Key = (ByteString
               ,TYPE)
    type Prio = Elapsed

    type Entry = Either DNSError [RData]

    type DB = OrdPSQ Key Prio Entry

    type Cache = R.Reaper DB (Key,Prio,Entry)

    newCache :: Int -> IO Cache
    newCache :: Int -> IO Cache
    newCache Int
    delay = ReaperSettings DB (Key, Prio, Entry) -> IO Cache
    forall workload item.
    ReaperSettings workload item -> IO (Reaper workload item)
    R.mkReaper ReaperSettings [Any] Any
    forall item. ReaperSettings [item] item
    R.defaultReaperSettings {
        R.reaperEmpty  = PSQ.empty
      , R.reaperCons   = \(Key
    k, Prio
    tim, Entry
    v) DB
    psq -> Key -> Prio -> Entry -> DB -> DB
    forall k p v.
    (Ord k, Ord p) =>
    k -> p -> v -> OrdPSQ k p v -> OrdPSQ k p v
    PSQ.insert Key
    k Prio
    tim Entry
    v DB
    psq
      , R.reaperAction = prune
      , R.reaperDelay  = delay * 1000000
      , R.reaperNull   = PSQ.null
      }

    lookupCache :: Key -> Cache -> IO (Maybe (Prio, Entry))
    lookupCache :: Key -> Cache -> IO (Maybe (Prio, Entry))
    lookupCache Key
    key Cache
    reaper = Key -> DB -> Maybe (Prio, Entry)
    forall k p v. Ord k => k -> OrdPSQ k p v -> Maybe (p, v)
    PSQ.lookup Key
    key (DB -> Maybe (Prio, Entry)) -> IO DB -> IO (Maybe (Prio, Entry))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Cache -> IO DB
    forall workload item. Reaper workload item -> IO workload
    R.reaperRead Cache
    reaper

    insertCache :: Key -> Prio -> Entry -> Cache -> IO ()
    insertCache :: Key -> Prio -> Entry -> Cache -> IO ()
    insertCache (ByteString
    dom,TYPE
    typ) Prio
    tim Entry
    ent0 Cache
    reaper = Cache -> (Key, Prio, Entry) -> IO ()
    forall workload item. Reaper workload item -> item -> IO ()
    R.reaperAdd Cache
    reaper (Key
    key,Prio
    tim,Entry
    ent)
      where
        key :: Key
    key = (ByteString -> ByteString
    B.copy ByteString
    dom,TYPE
    typ)
        ent :: Entry
    ent = case Entry
    ent0 of
          l :: Entry
    l@(Left DNSError
    _)  -> Entry
    l
          (Right [RData]
    rds) -> [RData] -> Entry
    forall a b. b -> Either a b
    Right ([RData] -> Entry) -> [RData] -> Entry
    forall a b. (a -> b) -> a -> b
    $ (RData -> RData) -> [RData] -> [RData]
    forall a b. (a -> b) -> [a] -> [b]
    map RData -> RData
    copy [RData]
    rds

    -- Theoretically speaking, atMostView itself is good enough for pruning.
    -- But auto-update assumes a list based db which does not provide atMost
    -- functions. So, we need to do this redundant way.
    prune :: DB -> IO (DB -> DB)
    prune :: DB -> IO (DB -> DB)
    prune DB
    oldpsq = do
        Prio
    tim <- IO Prio
    timeCurrent
        let ([(Key, Prio, Entry)]
    _, DB
    pruned) = Prio -> DB -> ([(Key, Prio, Entry)], DB)
    forall k p v.
    (Ord k, Ord p) =>
    p -> OrdPSQ k p v -> ([(k, p, v)], OrdPSQ k p v)
    PSQ.atMostView Prio
    tim DB
    oldpsq
        (DB -> DB) -> IO (DB -> DB)
    forall a. a -> IO a
    forall (m :: * -> *) a. Monad m => a -> m a
    return ((DB -> DB) -> IO (DB -> DB)) -> (DB -> DB) -> IO (DB -> DB)
    forall a b. (a -> b) -> a -> b
    $ \DB
    newpsq -> (DB -> (Key, Prio, Entry) -> DB)
    -> DB -> [(Key, Prio, Entry)] -> DB
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl' DB -> (Key, Prio, Entry) -> DB
    forall {k} {p} {v}.
    (Ord k, Ord p) =>
    OrdPSQ k p v -> (k, p, v) -> OrdPSQ k p v
    ins DB
    pruned ([(Key, Prio, Entry)] -> DB) -> [(Key, Prio, Entry)] -> DB
    forall a b. (a -> b) -> a -> b
    $ DB -> [(Key, Prio, Entry)]
    forall k p v. OrdPSQ k p v -> [(k, p, v)]
    PSQ.toList DB
    newpsq
      where
        ins :: OrdPSQ k p v -> (k, p, v) -> OrdPSQ k p v
    ins OrdPSQ k p v
    psq (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
    PSQ.insert k
    k p
    p v
    v OrdPSQ k p v
    psq

    copy :: RData -> RData
    copy :: RData -> RData
    copy r :: RData
    r@(RD_A IPv4
    _)           = RData
    r
    copy (RD_NS ByteString
    dom)          = ByteString -> RData
    RD_NS (ByteString -> RData) -> ByteString -> RData
    forall a b. (a -> b) -> a -> b
    $ ByteString -> ByteString
    B.copy ByteString
    dom
    copy (RD_CNAME ByteString
    dom)       = ByteString -> RData
    RD_CNAME (ByteString -> RData) -> ByteString -> RData
    forall a b. (a -> b) -> a -> b
    $ ByteString -> ByteString
    B.copy ByteString
    dom
    copy (RD_SOA ByteString
    mn ByteString
    mr Word32
    a Word32
    b Word32
    c Word32
    d Word32
    e) = ByteString
    -> ByteString
    -> Word32
    -> Word32
    -> Word32
    -> Word32
    -> Word32
    -> RData
    RD_SOA (ByteString -> ByteString
    B.copy ByteString
    mn) (ByteString -> ByteString
    B.copy ByteString
    mr) Word32
    a Word32
    b Word32
    c Word32
    d Word32
    e
    copy (RD_PTR ByteString
    dom)         = ByteString -> RData
    RD_PTR (ByteString -> RData) -> ByteString -> RData
    forall a b. (a -> b) -> a -> b
    $ ByteString -> ByteString
    B.copy ByteString
    dom
    copy (RD_NULL ByteString
    bytes)      = ByteString -> RData
    RD_NULL (ByteString -> RData) -> ByteString -> RData
    forall a b. (a -> b) -> a -> b
    $ ByteString -> ByteString
    B.copy ByteString
    bytes
    copy (RD_MX Word16
    prf ByteString
    dom)      = Word16 -> ByteString -> RData
    RD_MX Word16
    prf (ByteString -> RData) -> ByteString -> RData
    forall a b. (a -> b) -> a -> b
    $ ByteString -> ByteString
    B.copy ByteString
    dom
    copy (RD_TXT ByteString
    txt)         = ByteString -> RData
    RD_TXT (ByteString -> RData) -> ByteString -> RData
    forall a b. (a -> b) -> a -> b
    $ ByteString -> ByteString
    B.copy ByteString
    txt
    copy (RD_RP ByteString
    mbox ByteString
    dname)   = ByteString -> ByteString -> RData
    RD_RP (ByteString -> ByteString
    B.copy ByteString
    mbox) (ByteString -> ByteString
    B.copy ByteString
    dname)
    copy r :: RData
    r@(RD_AAAA IPv6
    _)        = RData
    r
    copy (RD_SRV Word16
    a Word16
    b Word16
    c ByteString
    dom)   = Word16 -> Word16 -> Word16 -> ByteString -> RData
    RD_SRV Word16
    a Word16
    b Word16
    c (ByteString -> RData) -> ByteString -> RData
    forall a b. (a -> b) -> a -> b
    $ ByteString -> ByteString
    B.copy ByteString
    dom
    copy (RD_DNAME ByteString
    dom)       = ByteString -> RData
    RD_DNAME (ByteString -> RData) -> ByteString -> RData
    forall a b. (a -> b) -> a -> b
    $ ByteString -> ByteString
    B.copy ByteString
    dom
    copy (RD_OPT [OData]
    od)          = [OData] -> RData
    RD_OPT ([OData] -> RData) -> [OData] -> RData
    forall a b. (a -> b) -> a -> b
    $ (OData -> OData) -> [OData] -> [OData]
    forall a b. (a -> b) -> [a] -> [b]
    map OData -> OData
    copyOData [OData]
    od
    copy (RD_DS Word16
    t Word8
    a Word8
    dt ByteString
    dv)    = Word16 -> Word8 -> Word8 -> ByteString -> RData
    RD_DS Word16
    t Word8
    a Word8
    dt (ByteString -> RData) -> ByteString -> RData
    forall a b. (a -> b) -> a -> b
    $ ByteString -> ByteString
    B.copy ByteString
    dv
    copy (RD_CDS Word16
    t Word8
    a Word8
    dt ByteString
    dv)   = Word16 -> Word8 -> Word8 -> ByteString -> RData
    RD_CDS Word16
    t Word8
    a Word8
    dt (ByteString -> RData) -> ByteString -> RData
    forall a b. (a -> b) -> a -> b
    $ ByteString -> ByteString
    B.copy ByteString
    dv
    copy (RD_NSEC ByteString
    dom [TYPE]
    ts)     = ByteString -> [TYPE] -> RData
    RD_NSEC (ByteString -> ByteString
    B.copy ByteString
    dom) [TYPE]
    ts
    copy (RD_DNSKEY Word16
    f Word8
    p Word8
    a ByteString
    k)  = Word16 -> Word8 -> Word8 -> ByteString -> RData
    RD_DNSKEY Word16
    f Word8
    p Word8
    a (ByteString -> RData) -> ByteString -> RData
    forall a b. (a -> b) -> a -> b
    $ ByteString -> ByteString
    B.copy ByteString
    k
    copy (RD_CDNSKEY Word16
    f Word8
    p Word8
    a ByteString
    k) = Word16 -> Word8 -> Word8 -> ByteString -> RData
    RD_CDNSKEY Word16
    f Word8
    p Word8
    a (ByteString -> RData) -> ByteString -> RData
    forall a b. (a -> b) -> a -> b
    $ ByteString -> ByteString
    B.copy ByteString
    k
    copy (RD_TLSA Word8
    a Word8
    b Word8
    c ByteString
    dgst) = Word8 -> Word8 -> Word8 -> ByteString -> RData
    RD_TLSA Word8
    a Word8
    b Word8
    c (ByteString -> RData) -> ByteString -> RData
    forall a b. (a -> b) -> a -> b
    $ ByteString -> ByteString
    B.copy ByteString
    dgst
    copy (RD_NSEC3 Word8
    a Word8
    b Word16
    c ByteString
    s ByteString
    h [TYPE]
    t) = Word8
    -> Word8 -> Word16 -> ByteString -> ByteString -> [TYPE] -> RData
    RD_NSEC3 Word8
    a Word8
    b Word16
    c (ByteString -> ByteString
    B.copy ByteString
    s) (ByteString -> ByteString
    B.copy ByteString
    h) [TYPE]
    t
    copy (RD_NSEC3PARAM Word8
    a Word8
    b Word16
    c ByteString
    salt) = Word8 -> Word8 -> Word16 -> ByteString -> RData
    RD_NSEC3PARAM Word8
    a Word8
    b Word16
    c (ByteString -> RData) -> ByteString -> RData
    forall a b. (a -> b) -> a -> b
    $ ByteString -> ByteString
    B.copy ByteString
    salt
    copy (RD_RRSIG RD_RRSIG
    sig)       = RD_RRSIG -> RData
    RD_RRSIG (RD_RRSIG -> RData) -> RD_RRSIG -> RData
    forall a b. (a -> b) -> a -> b
    $ RD_RRSIG -> RD_RRSIG
    copysig RD_RRSIG
    sig
      where
        copysig :: RD_RRSIG -> RD_RRSIG
    copysig s :: RD_RRSIG
    s@RDREP_RRSIG{Int64
    Word8
    Word16
    Word32
    ByteString
    TYPE
    rrsigType :: TYPE
    rrsigKeyAlg :: Word8
    rrsigNumLabels :: Word8
    rrsigTTL :: Word32
    rrsigExpiration :: Int64
    rrsigInception :: Int64
    rrsigKeyTag :: Word16
    rrsigZone :: ByteString
    rrsigValue :: ByteString
    rrsigType :: RD_RRSIG -> TYPE
    rrsigKeyAlg :: RD_RRSIG -> Word8
    rrsigNumLabels :: RD_RRSIG -> Word8
    rrsigTTL :: RD_RRSIG -> Word32
    rrsigExpiration :: RD_RRSIG -> Int64
    rrsigInception :: RD_RRSIG -> Int64
    rrsigKeyTag :: RD_RRSIG -> Word16
    rrsigZone :: RD_RRSIG -> ByteString
    rrsigValue :: RD_RRSIG -> ByteString
    ..} =
            RD_RRSIG
    s { rrsigZone = B.copy rrsigZone
              , rrsigValue = B.copy rrsigValue }
    copy (RD_CAA Word8
    f CI ByteString
    t ByteString
    v)       = Word8 -> CI ByteString -> ByteString -> RData
    RD_CAA Word8
    f (ByteString -> CI ByteString
    forall s. FoldCase s => s -> CI s
    CI.mk (ByteString -> ByteString
    B.copy (CI ByteString -> ByteString
    forall s. CI s -> s
    CI.original CI ByteString
    t))) (ByteString -> ByteString
    B.copy ByteString
    v)
    copy (UnknownRData ByteString
    is)    = ByteString -> RData
    UnknownRData (ByteString -> RData) -> ByteString -> RData
    forall a b. (a -> b) -> a -> b
    $ ByteString -> ByteString
    B.copy ByteString
    is

    copyOData :: OData -> OData
    copyOData :: OData -> OData
    copyOData (OD_ECSgeneric Word16
    family Word8
    srcBits Word8
    scpBits ByteString
    bs) =
        Word16 -> Word8 -> Word8 -> ByteString -> OData
    OD_ECSgeneric Word16
    family Word8
    srcBits Word8
    scpBits (ByteString -> OData) -> ByteString -> OData
    forall a b. (a -> b) -> a -> b
    $ ByteString -> ByteString
    B.copy ByteString
    bs
    copyOData (OD_NSID ByteString
    nsid) = ByteString -> OData
    OD_NSID (ByteString -> OData) -> ByteString -> OData
    forall a b. (a -> b) -> a -> b
    $ ByteString -> ByteString
    B.copy ByteString
    nsid
    copyOData (UnknownOData Word16
    c ByteString
    b)        = Word16 -> ByteString -> OData
    UnknownOData Word16
    c (ByteString -> OData) -> ByteString -> OData
    forall a b. (a -> b) -> a -> b
    $ ByteString -> ByteString
    B.copy ByteString
    b

    -- No copying required for the rest, but avoiding a wildcard pattern match
    -- so that if more option types are added in the future, the compiler will
    -- complain about a partial function.
    --
    copyOData o :: OData
    o@OD_ClientSubnet {} = OData
    o
    copyOData o :: OData
    o@OD_DAU {} = OData
    o
    copyOData o :: OData
    o@OD_DHU {} = OData
    o
    copyOData o :: OData
    o@OD_N3U {} = OData
    o