diff --git a/MassDns.hs b/MassDns.hs --- a/MassDns.hs +++ b/MassDns.hs @@ -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 diff --git a/Network/DnsCache.hs b/Network/DnsCache.hs --- a/Network/DnsCache.hs +++ b/Network/DnsCache.hs @@ -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) diff --git a/dnscache.cabal b/dnscache.cabal --- a/dnscache.cabal +++ b/dnscache.cabal @@ -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 @@ -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