module ADNS.Cache
(
DnsCache
, withDnsCache
, stopDnsCache
, resolveA
, showHostAddress
) where
import ADNS hiding (HostName, HostAddress)
import ADNS.Base
import ADNS.Endian
import Data.List
import Data.Word
import qualified Data.HashMap.Strict as HM
import qualified Data.Text as T
import Control.Monad
import Data.Bits
import Control.Concurrent
import qualified Control.Concurrent.MSem as MSem
import Data.Time.Clock.POSIX
import Network.Socket (HostName, HostAddress, inet_ntoa)
data DnsCache
= DnsCache
{ resolver :: Resolver
, sem :: MSem.MSem Int
, cache :: MVar (HM.HashMap T.Text (POSIXTime,
Either T.Text (Queue HostAddress)))
}
data Queue a = Queue ![a] ![a]
deriving Show
listQ :: [a] -> Queue a
listQ l = Queue l []
rotQ :: Queue a -> Maybe (a, Queue a)
rotQ (Queue [] []) = Nothing
rotQ (Queue (x:xs) ys) = Just (x, Queue xs (x:ys))
rotQ (Queue [] ys) = rotQ (Queue (reverse ys) [])
maxQueries :: Int
maxQueries = 30
withDnsCache :: (DnsCache -> IO a) -> IO a
withDnsCache act =
initResolver [NoErrPrint, NoServerWarn, NoSigPipe] $ \ r -> do
c <- newMVar HM.empty
s <- MSem.new maxQueries
act (DnsCache r s c)
stopDnsCache :: DnsCache -> IO ()
stopDnsCache d =
replicateM_ maxQueries $ MSem.wait (sem d)
resolveA :: DnsCache -> HostName -> IO (Either String HostAddress)
resolveA d@(DnsCache {..}) domain
| isIPAddr domain = return $ Right $ ipToWord32 domain
| otherwise = do
t <- getPOSIXTime
mbr <- modifyMVar cache $ \ c ->
case HM.lookup key c of
Just (expT, a)
| t < expT -> fmap (\(m,r) -> (m, Just r)) $
insRot False c expT a
_ -> return (c, Nothing)
case mbr of
Just r -> return r
_ -> do
ra <- MSem.with sem $ resolveA' d [] domain
modifyMVar cache $ \ m ->
case ra of
Left e -> insRot True m (t+60) $ Left $ T.pack e
Right (max (t+60) -> et, rras) ->
insRot True m et $ Right $ listQ [a | RRAddr a <- rras]
where key = T.pack domain
err False m _ e = return (m, Left $ T.unpack e)
err True m t e = return (HM.insert key (t, Left e) m, Left $ T.unpack e)
insRot f m t (Left e) = err f m t e
insRot f m t (Right q)
| Just (a, q') <- rotQ q =
return (HM.insert key (t, Right q') m,
Right a)
| otherwise =
err f m t "No RRAddr???"
resolveA' :: DnsCache -> [HostName] -> HostName
-> IO (Either String (POSIXTime, [RRAddr]))
resolveA' d@(DnsCache {..}) parents x
| length parents > 20 =
return $ Left $ "Too many CNAMEs " ++ show (x : parents)
| x `elem` parents =
return $ Left $ "CNAME loop " ++ show x ++ " already in "
++ show parents
| otherwise = do
Answer {..} <- resolver x A
[QuoteOk_Query, QuoteOk_AnsHost] >>= takeMVar
case () of
_ | status == sOK -> do
return $ Right (realToFrac expires, [ a | RRA a <- rrs ])
| status == sPROHIBITEDCNAME
, Just cn <- cname -> resolveA' d (x:parents) cn
_ -> do
e <- adnsStrerror status
s <- adnsErrAbbrev status
return $ Left $ e ++ " (" ++ s ++ ")"
isIPAddr :: HostName -> Bool
isIPAddr hn = length groups == 4 && all ip groups
where groups = splitPoints hn
ip x = length x <= 3 && all (\ e -> e >= '0' && e <= '9') x &&
read x <= (255 :: Int)
splitPoints :: HostName -> [String]
splitPoints = filter (/= ".") . groupBy (\ a b -> a /= '.' && b /= '.')
ipToWord32 :: String -> Word32
ipToWord32 ip =
let [b4,b3,b2,b1] = (map read $ splitPoints ip) :: [Word32]
mk x1 x2 x3 x4 =
(x1 `shiftL` 24) + (x2 `shiftL` 16) + (x3 `shiftL` 8) + x4
in
case endian of
BigEndian -> mk b4 b3 b2 b1
LittleEndian -> mk b1 b2 b3 b4
PDPEndian -> mk b2 b1 b4 b3
showHostAddress :: HostAddress -> String
showHostAddress = show . RRAddr
_test :: IO ()
_test = withDnsCache $ \ c -> do
let r hn = do
h <- resolveA c hn
putStrLn $ hn ++ ": " ++ either id showHostAddress h
rn n = replicateM_ n . r
r "127.0.0.1"
r "qwerqwer"
r "qwer.google.com"
rn 10 "www.huffingtonpost.com"
rn 10 "twitter.com"
rn 10 "wordpress.com"
rn 10 "feedburner.com"
rn 10 "feeds.feedburner.com"