dnscache 0.1.2 → 1.0.1
raw patch · 3 files changed
+217/−103 lines, 3 filesdep ~contstuff
Dependency ranges changed: contstuff
Files
- MassDns.hs +6/−3
- Network/DnsCache.hs +209/−98
- dnscache.cabal +2/−2
MassDns.hs view
@@ -7,6 +7,7 @@ module Main where +import Control.ContStuff import Network.DnsCache import System.Environment import System.IO@@ -20,9 +21,11 @@ printDns :: Int -> MassType -> [String] -> IO () printDns numThreads mtype domains =- seq mtype . withDnsCache numThreads 300 $ \dns ->- massLookupReport dns mtype domains $ \res -> do- let dom = massDomain res+ seq mtype .+ withDnsStateT numThreads 300 .+ massLookupReport mtype domains $ \res -> do+ let dom = massDomain res+ liftIO $ do printEntry dom "A" $ map show (massA res) printEntry dom "AAAA" $ map show (massAAAA res) printEntry dom "MX" $ massMX res
Network/DnsCache.hs view
@@ -7,13 +7,19 @@ -- -- This module implements an asynchronous, caching DNS resolver. -{-# LANGUAGE GADTs, TypeFamilies #-}+{-# LANGUAGE FlexibleInstances, GADTs, RankNTypes, TypeFamilies #-} module Network.DnsCache ( -- * DNS cache DnsCache,+ createDnsCache,+ freeDnsCache, withDnsCache, + -- * Monadic interface+ DnsMonad(..),+ withDnsStateT,+ -- * DNS lookup resolveA, resolveAAAA,@@ -30,17 +36,18 @@ ) where -import qualified Control.Exception as E import qualified Data.Map as M import Control.Concurrent import Control.ContStuff import Data.Char+import Data.IORef import Data.IP import Data.List import Data.Map (Map) import Data.Ord import Data.Time.Clock import Network.DNS hiding (lookup)+import System.IO.Unsafe -- | Command to a DNS cache.@@ -57,10 +64,11 @@ data DnsCache = DnsCache {- dnsAVar :: MVar (CacheCmd IPv4),- dnsAaaaVar :: MVar (CacheCmd IPv6),- dnsMxVar :: MVar (CacheCmd Domain),- dnsResVar :: MVar ResolverCmd+ dnsAVar :: MVar (CacheCmd IPv4),+ dnsAaaaVar :: MVar (CacheCmd IPv6),+ dnsMxVar :: MVar (CacheCmd Domain),+ dnsResVar :: MVar ResolverCmd,+ dnsNumThreads :: Int } @@ -106,6 +114,24 @@ type Waiter res = Maybe [res] -> IO () +-- | Monads, which contain a DNS cache.++class MonadIO m => DnsMonad m where+ -- | Get the current DNS cache.+ getDnsCache :: m DnsCache++instance MonadIO m => DnsMonad (StateT r DnsCache m) where+ getDnsCache = get++instance DnsMonad IO where+ getDnsCache = readIORef globalCache++{-# NOINLINE globalCache #-}+globalCache :: IORef DnsCache+globalCache =+ unsafePerformIO (createDnsCache 1 10 >>= newIORef)++ -- | DNS cache thread. dnsCache :: NominalDiffTime -> MVar (CacheCmd res) -> IO ()@@ -151,106 +177,64 @@ modify (M.insert dom (Resolved curTime res)) --- | Perform a mass lookup.+-- | Resolve an A record. -massLookup :: DnsCache -> MassType -> [Domain] -> IO (Map Domain MassResult)-massLookup dns mtype domains = do- let numDomains = length domains- resultVar <- newEmptyMVar- massResolver dns mtype domains resultVar- execStateT M.empty . replicateM_ numDomains $ do- m@(MassResult dom _ _ _) <- liftIO $ takeMVar resultVar- modify (M.insert dom m)+resolveA :: DnsMonad m => Domain -> m (Maybe [IPv4])+resolveA dom = getDnsCache >>= resolveA_ dom --- | Perform a mass lookup with report function.+-- | Resolve an A record. -massLookupReport ::- MonadIO m =>- DnsCache -> MassType -> [Domain] -> (MassResult -> m ()) -> m ()-massLookupReport dns mtype domains rep = do- let numDomains = length domains- resultVar <- liftIO newEmptyMVar- liftIO $ massResolver dns mtype domains resultVar- replicateM_ numDomains $ liftIO (takeMVar resultVar) >>= rep+resolveA_ :: MonadIO m => Domain -> DnsCache -> m (Maybe [IPv4])+resolveA_ dom cfg =+ liftIO $ do+ let cacheVar = dnsAVar cfg+ let resVar = dnsResVar cfg + answerVar <- newEmptyMVar+ putMVar resVar (Resolve dom cacheVar lookupA (putMVar answerVar))+ takeMVar answerVar --- | Mass resolver. -massResolver :: DnsCache -> MassType -> [Domain] -> MVar MassResult -> IO ()-massResolver dns mtype domains resultVar = do- forM_ domains $ \domain ->- forkIO $ do- case mtype of- MassA -> do- resVar <- newEmptyMVar- forkIO $ resolveA dns domain >>= putMVar resVar . maybe [] id- res <- takeMVar resVar- putMVar resultVar (MassResult domain res [] [])-- MassAAAA -> do- resVar <- newEmptyMVar- forkIO $ resolveAAAA dns domain >>= putMVar resVar . maybe [] id- res <- takeMVar resVar- putMVar resultVar (MassResult domain [] res [])-- MassMX -> do- resVar <- newEmptyMVar- forkIO $ resolveMX dns domain >>= putMVar resVar . maybe [] id- res <- takeMVar resVar- putMVar resultVar (MassResult domain [] [] res)-- MassAll -> do- aVar <- newEmptyMVar- aaaaVar <- newEmptyMVar- mxVar <- newEmptyMVar-- forkIO $ resolveA dns domain >>= putMVar aVar . maybe [] id- forkIO $ resolveAAAA dns domain >>= putMVar aaaaVar . maybe [] id- forkIO $ resolveMX dns domain >>= putMVar mxVar . maybe [] id+-- | Resolve an AAAA record. - a <- takeMVar aVar- aaaa <- takeMVar aaaaVar- mx <- takeMVar mxVar- putMVar resultVar (MassResult domain a aaaa mx)+resolveAAAA :: DnsMonad m => Domain -> m (Maybe [IPv6])+resolveAAAA dom = getDnsCache >>= resolveAAAA_ dom --- | Resolve an A record.+-- | Resolve an AAAA record. -resolveA :: DnsCache -> Domain -> IO (Maybe [IPv4])-resolveA cfg dom = do- let cacheVar = dnsAVar cfg- let resVar = dnsResVar cfg+resolveAAAA_ :: MonadIO m => Domain -> DnsCache -> m (Maybe [IPv6])+resolveAAAA_ dom cfg =+ liftIO $ do+ let cacheVar = dnsAaaaVar cfg+ let resVar = dnsResVar cfg - answerVar <- newEmptyMVar- putMVar resVar (Resolve dom cacheVar lookupA (putMVar answerVar))- takeMVar answerVar+ answerVar <- newEmptyMVar+ putMVar resVar (Resolve dom cacheVar lookupAAAA (putMVar answerVar))+ takeMVar answerVar --- | Resolve an AAAA record.--resolveAAAA :: DnsCache -> Domain -> IO (Maybe [IPv6])-resolveAAAA cfg dom = do- let cacheVar = dnsAaaaVar cfg- let resVar = dnsResVar cfg+-- | Resolve an MX record. - answerVar <- newEmptyMVar- putMVar resVar (Resolve dom cacheVar lookupAAAA (putMVar answerVar))- takeMVar answerVar+resolveMX :: DnsMonad m => Domain -> m (Maybe [Domain])+resolveMX dom = getDnsCache >>= resolveMX_ dom -- | Resolve an MX record. -resolveMX :: DnsCache -> Domain -> IO (Maybe [Domain])-resolveMX cfg dom = do- let cacheVar = dnsMxVar cfg- let resVar = dnsResVar cfg+resolveMX_ :: MonadIO m => Domain -> DnsCache -> m (Maybe [Domain])+resolveMX_ dom cfg =+ liftIO $ do+ let cacheVar = dnsMxVar cfg+ let resVar = dnsResVar cfg - let mxLookup r d = fmap (map fst . sortBy (comparing snd)) <$> lookupMX r d+ let mxLookup r d = fmap (map fst . sortBy (comparing snd))+ <$> lookupMX r d - answerVar <- newEmptyMVar- putMVar resVar (Resolve dom cacheVar mxLookup (putMVar answerVar))- takeMVar answerVar+ answerVar <- newEmptyMVar+ putMVar resVar (Resolve dom cacheVar mxLookup (putMVar answerVar))+ takeMVar answerVar -- | DNS resolver thread.@@ -287,29 +271,48 @@ ResolverQuit c -> liftIO c >> abort () +-- ================== --+-- DNS cache creation --+-- ================== --++ -- | Start a DNS cache with the given number of resolver threads and the -- given cache timeout. -withDnsCache :: Int -> NominalDiffTime -> (DnsCache -> IO a) -> IO a-withDnsCache numThreads timeout comp = do- resolverVar <- newEmptyMVar- aVar <- newEmptyMVar- aaaaVar <- newEmptyMVar- mxVar <- newEmptyMVar+createDnsCache :: MonadIO m => Int -> NominalDiffTime -> m DnsCache+createDnsCache numThreads timeout = do+ resolverVar <- liftIO newEmptyMVar+ aVar <- liftIO newEmptyMVar+ aaaaVar <- liftIO newEmptyMVar+ mxVar <- liftIO newEmptyMVar - replicateM_ numThreads . forkIO $ resolver resolverVar- forkIO $ dnsCache timeout aVar- forkIO $ dnsCache timeout aaaaVar- forkIO $ dnsCache timeout mxVar+ liftIO $ do+ replicateM_ numThreads . forkIO $ resolver resolverVar+ forkIO $ dnsCache timeout aVar+ forkIO $ dnsCache timeout aaaaVar+ forkIO $ dnsCache timeout mxVar - let cfg = DnsCache {+ let dns = DnsCache { dnsAVar = aVar, dnsAaaaVar = aaaaVar, dnsMxVar = mxVar,- dnsResVar = resolverVar+ dnsResVar = resolverVar,+ dnsNumThreads = numThreads }+ return dns - E.finally (comp cfg) $ do++-- | Free existing DNS cache.++freeDnsCache :: MonadIO m => DnsCache -> m ()+freeDnsCache dns =+ liftIO $ do+ let aVar = dnsAVar dns+ let aaaaVar = dnsAaaaVar dns+ let mxVar = dnsMxVar dns+ let resolverVar = dnsResVar dns+ let numThreads = dnsNumThreads dns+ quitVar <- newEmptyMVar replicateM_ numThreads $ do putMVar resolverVar (ResolverQuit $ putMVar quitVar ())@@ -319,3 +322,111 @@ putMVar aaaaVar (CacheQuit (putMVar quitVar ())) putMVar mxVar (CacheQuit (putMVar quitVar ())) replicateM_ 3 $ takeMVar quitVar+++-- | Convenient wrapper around 'createDnsCache' and 'freeDnsCache'.++withDnsCache :: (HasExceptions m, MonadIO m) => Int -> NominalDiffTime -> (DnsCache -> m a) -> m a+withDnsCache numThreads timeout comp =+ bracket (createDnsCache numThreads timeout)+ freeDnsCache+ comp+++-- ==================== --+-- Running computations --+-- ==================== --++-- | Run a 'StateT' computation with a DNS cache.++withDnsStateT ::+ (Applicative m, HasExceptions m, MonadIO m) =>+ Int -> NominalDiffTime -> StateT a DnsCache m a -> m a+withDnsStateT numThreads timeout comp =+ withDnsCache numThreads timeout (\dns -> evalStateT dns comp)+++-- =================== --+-- DNS mass resolution --+-- =================== --+++-- | Perform a mass lookup.++massLookup :: DnsMonad m => MassType -> [Domain] -> m (Map Domain MassResult)+massLookup mtype domains = do+ getDnsCache >>= massLookup_ mtype domains+++-- | Perform a mass lookup.++massLookup_ :: MonadIO m => MassType -> [Domain] -> DnsCache -> m (Map Domain MassResult)+massLookup_ mtype domains dns =+ liftIO $ do+ let numDomains = length domains+ resultVar <- newEmptyMVar+ massResolver dns mtype domains resultVar+ execStateT M.empty . replicateM_ numDomains $ do+ m@(MassResult dom _ _ _) <- liftIO $ takeMVar resultVar+ modify (M.insert dom m)+++-- | Perform a mass lookup with report function.++massLookupReport ::+ DnsMonad m =>+ MassType -> [Domain] -> (MassResult -> m ()) -> m ()+massLookupReport mtype domains rep =+ getDnsCache >>= massLookupReport_ mtype domains rep+++-- | Perform a mass lookup with report function.++massLookupReport_ ::+ MonadIO m =>+ MassType -> [Domain] -> (MassResult -> m ()) -> DnsCache -> m ()+massLookupReport_ mtype domains rep dns = do+ let numDomains = length domains+ resultVar <- liftIO newEmptyMVar+ liftIO $ massResolver dns mtype domains resultVar+ replicateM_ numDomains $ liftIO (takeMVar resultVar) >>= rep+++-- | Mass resolver.++massResolver :: DnsCache -> MassType -> [Domain] -> MVar MassResult -> IO ()+massResolver dns mtype domains resultVar = do+ forM_ domains $ \domain ->+ forkIO $ do+ case mtype of+ MassA -> do+ resVar <- newEmptyMVar+ forkIO $ resolveA_ domain dns >>= putMVar resVar . maybe [] id+ res <- takeMVar resVar+ putMVar resultVar (MassResult domain res [] [])++ MassAAAA -> do+ resVar <- newEmptyMVar+ forkIO $ resolveAAAA_ domain dns >>= putMVar resVar . maybe [] id+ res <- takeMVar resVar+ putMVar resultVar (MassResult domain [] res [])++ MassMX -> do+ resVar <- newEmptyMVar+ forkIO $ resolveMX_ domain dns >>= putMVar resVar . maybe [] id+ res <- takeMVar resVar+ putMVar resultVar (MassResult domain [] [] res)++ MassAll -> do+ aVar <- newEmptyMVar+ aaaaVar <- newEmptyMVar+ mxVar <- newEmptyMVar++ forkIO $ resolveA_ domain dns >>= putMVar aVar . maybe [] id+ forkIO $ resolveAAAA_ domain dns >>= putMVar aaaaVar . maybe [] id+ forkIO $ resolveMX_ domain dns >>= putMVar mxVar . maybe [] id++ a <- takeMVar aVar+ aaaa <- takeMVar aaaaVar+ mx <- takeMVar mxVar+ putMVar resultVar (MassResult domain a aaaa mx)
dnscache.cabal view
@@ -1,5 +1,5 @@ Name: dnscache-Version: 0.1.2+Version: 1.0.1 Category: Network Synopsis: Caching DNS resolver library and mass DNS resolver utility Maintainer: Ertugrul Söylemez <es@ertes.de>@@ -19,7 +19,7 @@ Build-depends: base >= 4 && <= 5, containers >= 0.3.0.0,- contstuff >= 1.0.0,+ contstuff >= 1.0.1, dns >= 0.1.0, iproute >= 1.2.0, time >= 1.1.4