src/DNS.hs

   1 module DNS (
   2   LookupResult,
   3   dns_tests,
   4   lookupMX',
   5   lookupNS',
   6   normalize,
   7   normalize_case,
   8   resolve_address )
   9 where
  10
  11 import Control.Applicative ((<$>))
  12 import Control.Monad (liftM)
  13 import qualified Data.ByteString.Char8 as BS (
  14   append,
  15   last,
  16   map,
  17   pack )
  18 import Data.Char (toLower)
  19 import Data.IP (IPv4)
  20 import Network.DNS (
  21   Domain,
  22   DNSFormat(..),
  23   Resolver,
  24   RDATA(..),
  25   TYPE(..),
  26   defaultResolvConf,
  27   lookupA,
  28   lookupMX,
  29   lookupNS,
  30   lookupRaw,
  31   makeResolvSeed,
  32   rdata,
  33   rrtype,
  34   withResolver )
  35 import Test.Framework (Test, testGroup)
  36 import Test.Framework.Providers.HUnit (testCase)
  37 import Test.HUnit (assertEqual)
  38 import Text.Read (readMaybe)
  39
  40 type LookupResult = (Domain, Maybe [Domain])
  41
  42 -- | Perform a query, but take the result from the authority section
  43 --   of the response rather than the answer section. Code shamelessly
  44 --   stolen from Network.DNS.lookup.
  45 lookup_authority :: Resolver -> Domain -> TYPE -> IO (Maybe [RDATA])
  46 lookup_authority rlv dom typ = (>>= toRDATA) <$> lookupRaw rlv dom typ
  47   where
  48     correct r = rrtype r == typ
  49     listToMaybe [] = Nothing
  50     listToMaybe xs = Just xs
  51     toRDATA = listToMaybe . map rdata . filter correct . authority
  52
  53 -- | Like lookupNS, except we take the result from the authority
  54 --   section of the response (as opposed to the answer section).
  55 lookupNS_authority :: Resolver -> Domain -> IO (Maybe [Domain])
  56 lookupNS_authority rlv dom = toNS <$> DNS.lookup_authority rlv dom NS
  57   where
  58     toNS = fmap (map unTag)
  59     unTag (RD_NS dm) = dm
  60     unTag _ = error "lookupNS_authority"
  61
  62 resolve_address :: String -> IO (Maybe IPv4)
  63 resolve_address s =
  64   case read_result of
  65     Just _  -> return read_result
  66     Nothing -> do
  67       default_rs <- makeResolvSeed defaultResolvConf
  68       withResolver default_rs $ \resolver -> do
  69         result <- lookupA resolver (BS.pack s)
  70         return $ case result of
  71                    Just (x:_) -> Just x
  72                    _           -> Nothing
  73   where
  74     read_result :: Maybe IPv4
  75     read_result = readMaybe s
  76
  77 lookupMX' :: Resolver -> Domain -> IO LookupResult
  78 lookupMX' resolver domain =
  79   liftM (pair_em . drop_priority) $ lookupMX resolver domain
  80   where
  81     drop_priority :: Maybe [(Domain, Int)] -> Maybe [Domain]
  82     drop_priority = fmap (map fst)
  83
  84     pair_em :: a -> (Domain, a)
  85     pair_em = (,) domain
  86
  87 -- This function keeps the domain matches with its NS records.
  88 lookupNS' :: Resolver -> Domain -> IO LookupResult
  89 lookupNS' resolver domain = do
  90   answer_result <- lookupNS resolver domain
  91   auth_result <- lookupNS_authority resolver domain
  92   liftM pair_em $ return $ combine answer_result auth_result
  93   where
  94     pair_em :: a -> (Domain, a)
  95     pair_em = (,) domain
  96
  97     combine :: (Maybe [Domain]) ->  (Maybe [Domain]) -> (Maybe [Domain])
  98     combine Nothing Nothing = Nothing
  99     combine m1 Nothing = m1
 100     combine Nothing m2 = m2
 101     combine (Just ds1) (Just ds2) = Just (ds1 ++ ds2)
 102
 103 -- | Normalize the given name by lowercasing and appending a trailing
 104 --   dot (the root) if necessary.
 105 normalize :: Domain -> Domain
 106 normalize = normalize_case . normalize_root
 107
 108
 109 normalize_root :: Domain -> Domain
 110 normalize_root d
 111   | BS.last d == '.' = d
 112   | otherwise = d `BS.append` trailing_dot
 113     where
 114       trailing_dot = BS.pack "."
 115
 116
 117 normalize_case :: Domain -> Domain
 118 normalize_case = BS.map toLower
 119
 120
 121 test_normalize_case :: Test
 122 test_normalize_case =
 123   testCase desc $
 124     assertEqual desc expected actual
 125   where
 126     desc = "normalize_case lowercases DNS names"
 127     expected = BS.pack "example.com"
 128     actual = normalize_case $ BS.pack "ExAmPlE.COM"
 129
 130 dns_tests :: Test
 131 dns_tests =
 132   testGroup "DNS Tests" [
 133     test_normalize_case ]