Implementing response cache.

Network/DNS/LookupRaw.hs

@@ -11,9 +11,14 @@ module Network.DNS.LookupRaw (
   , fromDNSFormat
   ) where
 
+import Data.ByteString.Short (toShort)
+import Data.Time (getCurrentTime, addUTCTime)
+
 import Network.DNS.IO
+import Network.DNS.Memo
 import Network.DNS.Transport
 import Network.DNS.Types
+import Network.DNS.Types.Internal
 
 import Prelude hiding (lookup)
 
@@ -32,18 +37,69 @@ lookupSection :: (DNSMessage -> [ResourceRecord])
               -> Domain
               -> TYPE
               -> IO (Either DNSError [RData])
-lookupSection section rlv dom typ = do
-    eans <- lookupRaw rlv dom typ
-    case eans of
-        Left  err -> return $ Left err
-        Right ans -> return $ fromDNSMessage ans toRData
+lookupSection section rlv dom typ = case mcacheConf of
+    Nothing -> do
+        eans <- lookupRaw rlv dom typ
+        case eans of
+          Left err  -> return $ Left err
+          Right ans -> return $ fromDNSMessage ans toRData
+    Just cacheconf -> lookupCacheSection section rlv dom typ cacheconf
   where
-    {- CNAME hack
-    dom' = if "." `isSuffixOf` dom then dom else dom ++ "."
-    correct r = rrname r == dom' && rrtype r == typ
-    -}
     correct ResourceRecord{..} = rrtype == typ
-    toRData x = map rdata . filter correct $ section x
+    toRData = map rdata . filter correct . section
+    mcacheConf = resolvCache $ resolvconf $ resolvseed rlv
+
+lookupCacheSection :: (DNSMessage -> [ResourceRecord])
+                   -> Resolver
+                   -> Domain
+                   -> TYPE
+                   -> CacheConf
+                   -> IO (Either DNSError [RData])
+lookupCacheSection section rlv dom typ cconf = do
+    mx <- lookupCacheRef (sdom,typ) cref
+    case mx of
+      Nothing -> do
+          eans <- lookupRaw rlv dom typ
+          case eans of
+            Left  err -> do
+                let v = Left err
+                insertNegative cconf cref key v
+                return v
+            Right ans -> do
+                let errs = fromDNSMessage ans toRR
+                case errs of
+                  Left  err -> do
+                      let v = Left err
+                      insertNegative cconf cref key v
+                      return v
+                  Right rss -> do
+                      let rds = map rdata rss
+                          v = Right rds
+                          ttls = map rrttl rss
+                      insertPositive cconf cref key v ttls
+                      return v
+      Just (_,x) -> return x
+  where
+    correct ResourceRecord{..} = rrtype == typ
+    toRR = filter correct . section
+    cref = cache rlv
+    sdom = toShort dom
+    key = (sdom,typ)
+
+insertPositive :: CacheConf -> CacheRef -> Key -> Entry -> [TTL] -> IO ()
+insertPositive CacheConf{..} ref k v ttls = do
+    tim <- addUTCTime life <$> getCurrentTime
+    insertCacheRef k tim v ref
+  where
+    life = fromIntegral $ case ttls of
+      []    -> minimumTTL -- fixme: what is a proper value?
+      ttl:_ -> minimumTTL `max` (maximumTTL `min` ttl)
+
+insertNegative :: CacheConf -> CacheRef -> Key -> Entry -> IO ()
+insertNegative CacheConf{..} ref k v = do
+    let life = fromIntegral negativeTTL
+    tim <- addUTCTime life <$> getCurrentTime
+    insertCacheRef k tim v ref
 
 -- | Extract necessary information from 'DNSMessage'
 fromDNSMessage :: DNSMessage -> (DNSMessage -> a) -> Either DNSError a
@@ -71,6 +127,7 @@ fromDNSFormat = fromDNSMessage
 --   from the ANSWER section of the response.
 --   See manual the manual of 'lookupRaw'
 --   to understand the concrete behavior.
+--   Cache is used if 'resolvCache' is 'Just'.
 --
 --   Example:
 --
@@ -122,6 +179,8 @@ lookupAuth = lookupSection authority
 --    The first received response is accepted even if
 --    it is an error.
 --
+--  Cache is not used even if 'resolvCache' is 'Just'.
+--
 --   The example code:
 --
 --   @

Network/DNS/Memo.hs

@@ -0,0 +1,33 @@
+module Network.DNS.Memo where
+
+import Control.Applicative ((<$>))
+import Data.ByteString.Short (ShortByteString)
+import Data.IORef (IORef, newIORef, readIORef, atomicModifyIORef', IORef)
+import Data.OrdPSQ (OrdPSQ)
+import qualified Data.OrdPSQ as PSQ
+import Data.Time (UTCTime)
+
+import Network.DNS.Types
+
+type Key = (ShortByteString -- avoiding memory fragmentation
+           ,TYPE)
+type Prio = UTCTime
+
+type Entry = Either DNSError [RData]
+
+type PSQ = OrdPSQ
+newtype CacheRef = CacheRef (IORef (PSQ Key Prio Entry))
+
+newCacheRef :: IO CacheRef
+newCacheRef = CacheRef <$> newIORef PSQ.empty
+
+lookupCacheRef :: Key -> CacheRef -> IO (Maybe (Prio, Entry))
+lookupCacheRef key (CacheRef ref) = PSQ.lookup key <$> readIORef ref
+
+insertCacheRef :: Key -> Prio -> Entry -> CacheRef -> IO ()
+insertCacheRef key tim ent (CacheRef ref) =
+    atomicModifyIORef' ref $ \q -> (PSQ.insert key tim ent q, ())
+
+pruneCacheRef :: Prio -> CacheRef -> IO ()
+pruneCacheRef tim (CacheRef ref) =
+    atomicModifyIORef' ref $ \p -> (snd (PSQ.atMostView tim p), ())

Network/DNS/Resolver.hs

@@ -14,8 +14,15 @@ module Network.DNS.Resolver (
   , resolvRetry
   , resolvEDNS
   , resolvConcurrent
-  -- ** Related types
+  , resolvCache
+  -- ** Specifying DNS servers
   , FileOrNumericHost(..)
+  -- ** Configuring cache
+  , CacheConf
+  , defaultCacheConf
+  , minimumTTL
+  , maximumTTL
+  , negativeTTL
   -- * Intermediate data type for resolver
   , ResolvSeed
   , makeResolvSeed
@@ -49,6 +56,7 @@ import Network.BSD (getProtocolNumber)
 import Network.DNS.Transport
 import Network.DNS.Types
 import Network.DNS.Types.Internal
+import Network.DNS.Memo
 import Network.Socket (AddrInfoFlag(..), AddrInfo(..), PortNumber(..), HostName, SocketType(Datagram), getAddrInfo, defaultHints)
 import Prelude hiding (lookup)
 
@@ -138,7 +146,8 @@ makeResolver seed = do
   let n = NE.length $ nameservers seed
   refs <- replicateM n (C.drgNew >>= I.newIORef)
   let gens = NE.fromList $ map getRandom refs
-  return $ Resolver seed gens
+  cacheref <- newCacheRef
+  return $ Resolver seed gens cacheref
 
 getRandom :: IORef C.ChaChaDRG -> IO Word16
 getRandom ref = I.atomicModifyIORef' ref $ \gen ->

Network/DNS/Types.hs

@@ -127,7 +127,7 @@ type Mailbox = ByteString
 newtype TYPE = TYPE {
     -- | From type to number.
     fromTYPE :: Word16
-  } deriving Eq
+  } deriving (Eq, Ord)
 
 -- https://www.iana.org/assignments/dns-parameters/dns-parameters.xhtml#dns-parameters-4
 
@@ -257,7 +257,7 @@ data TYPE = A          -- ^ IPv4 address
           | ANY        -- ^ A request for all records the server/cache
                        --   has available
           | UnknownTYPE Word16  -- ^ Unknown type
-          deriving (Eq, Show, Read)
+          deriving (Eq, Ord, Show, Read)
 
 -- | From type to number.
 fromTYPE :: TYPE -> Word16

Network/DNS/Types/Internal.hs

@@ -4,6 +4,7 @@ import Data.List.NonEmpty (NonEmpty(..))
 import Network.Socket (AddrInfo(..), PortNumber(..), HostName)
 import Data.Word (Word16)
 
+import Network.DNS.Memo
 import Network.DNS.Types
 
 ----------------------------------------------------------------
@@ -22,6 +23,27 @@ data FileOrNumericHost = RCFilePath FilePath -- ^ A path for \"resolv.conf\"
                        | RCHostPort HostName PortNumber -- ^ A numeric IP address and port number. /Warning/: host names are invalid.
                        deriving Show
 
+----------------------------------------------------------------
+
+-- | Cache configuration for responses.
+data CacheConf = CacheConf {
+    -- | If RR's TTL is lower than this value, this value is used instead.
+    minimumTTL  :: TTL
+    -- | If RR's TTL is higher than this value, this value is used instead.
+  , maximumTTL  :: TTL
+    -- | TTL in the case of 'DNSError'.
+  , negativeTTL :: TTL
+  } deriving Show
+
+-- | Default cache configuration.
+--
+-- >>> defaultCacheConf
+-- CacheConf {minimumTTL = 60, maximumTTL = 300, negativeTTL = 300}
+defaultCacheConf :: CacheConf
+defaultCacheConf = CacheConf 60 300 300
+
+----------------------------------------------------------------
+
 -- | Type for resolver configuration.
 --  Use 'defaultResolvConf' to create a new value.
 --
@@ -40,6 +62,10 @@ data FileOrNumericHost = RCFilePath FilePath -- ^ A path for \"resolv.conf\"
 --  An example to disable EDNS0 with a 1,280-bytes buffer:
 --
 --  >>> let conf = defaultResolvConf { resolvEDNS = [fromEDNS0 defaultEDNS0 { udpSize = 1280 }] }
+--
+--  An example to enable cache:
+--
+--  >>> let conf = defaultResolvConf { resolvCache = Just defaultCacheConf }
 data ResolvConf = ResolvConf {
    -- | Server information.
     resolvInfo       :: FileOrNumericHost
@@ -51,6 +77,8 @@ data ResolvConf = ResolvConf {
   , resolvEDNS       :: [ResourceRecord]
    -- | Concurrent queries if multiple DNS servers are specified.
   , resolvConcurrent :: Bool
+   -- | Cache configuration
+  , resolvCache      :: Maybe CacheConf
 } deriving Show
 
 -- | Return a default 'ResolvConf':
@@ -67,6 +95,7 @@ defaultResolvConf = ResolvConf {
   , resolvRetry      = 3
   , resolvEDNS       = [fromEDNS0 defaultEDNS0]
   , resolvConcurrent = False
+  , resolvCache      = Nothing
 }
 
 ----------------------------------------------------------------
@@ -86,8 +115,9 @@ data ResolvSeed = ResolvSeed {
 
 -- | Abstract data type of DNS Resolver.
 --   This includes newly seeded identifier generators for all
---   specified DNS servers.
+--   specified DNS servers and a cache database.
 data Resolver = Resolver {
     resolvseed :: ResolvSeed
   , genIds     :: NonEmpty (IO Word16)
+  , cache      :: CacheRef
 }

dns.cabal

@@ -27,14 +27,15 @@ Library
                         Network.DNS.Decode
                         Network.DNS.IO
   Other-Modules:        Network.DNS.Decode.Internal
+                        Network.DNS.Memo
                         Network.DNS.StateBinary
                         Network.DNS.Transport
                         Network.DNS.Types.Internal
   if impl(ghc < 8)
     Build-Depends:      semigroups
   Build-Depends:        base >= 4 && < 5
-                      , attoparsec
                       , async
+                      , attoparsec
                       , base64-bytestring
                       , binary
                       , bytestring
@@ -45,7 +46,9 @@ Library
                       , iproute >= 1.3.2
                       , mtl
                       , network >= 2.3
+                      , psqueues
                       , safe == 0.3.*
+                      , time
   if os(windows)
     Build-Depends:    split
     C-Sources:        cbits/dns.c