src/DNS.hs

   1 module DNS (
   2   LookupResult,
   3   dns_properties,
   4   dns_tests,
   5   lookupMX',
   6   lookupNS',
   7   normalize,
   8   normalize_case,
   9   resolve_address )
  10 where
  11
  12 import Control.Applicative ((<$>))
  13 import Control.Monad (liftM)
  14 import qualified Data.ByteString.Char8 as BS (
  15   append,
  16   last,
  17   map,
  18   null,
  19   pack )
  20 import Data.Char (toLower)
  21 import Data.IP (IPv4)
  22 import Network.DNS (
  23   Domain,
  24   DNSFormat(..),
  25   Resolver,
  26   RDATA(..),
  27   TYPE(..),
  28   defaultResolvConf,
  29   lookupA,
  30   lookupMX,
  31   lookupNS,
  32   lookupRaw,
  33   makeResolvSeed,
  34   rdata,
  35   rrtype,
  36   withResolver )
  37 import Test.Framework (Test, testGroup)
  38 import Test.Framework.Providers.HUnit (testCase)
  39 import Test.Framework.Providers.QuickCheck2 (testProperty)
  40 import Test.HUnit (assertEqual)
  41 import Text.Read (readMaybe)
  42
  43 type LookupResult = (Domain, Maybe [Domain])
  44
  45 -- | Perform a query, but take the result from the authority section
  46 --   of the response rather than the answer section. Code shamelessly
  47 --   stolen from Network.DNS.lookup.
  48 lookup_authority :: Resolver -> Domain -> TYPE -> IO (Maybe [RDATA])
  49 lookup_authority rlv dom typ = (>>= toRDATA) <$> lookupRaw rlv dom typ
  50   where
  51     correct r = rrtype r == typ
  52     listToMaybe [] = Nothing
  53     listToMaybe xs = Just xs
  54     toRDATA = listToMaybe . map rdata . filter correct . authority
  55
  56 -- | Like lookupNS, except we take the result from the authority
  57 --   section of the response (as opposed to the answer section).
  58 lookupNS_authority :: Resolver -> Domain -> IO (Maybe [Domain])
  59 lookupNS_authority rlv dom = toNS <$> DNS.lookup_authority rlv dom NS
  60   where
  61     toNS = fmap (map unTag)
  62     unTag (RD_NS dm) = dm
  63     unTag _ = error "lookupNS_authority"
  64
  65
  66 -- | Takes a String representing either a hostname or an IP
  67 --   address. If a hostname was supplied, it is resolved to either an
  68 --   IPv4 or Nothing. If an IP address is supplied, it is returned as an
  69 --   IPv4.
  70 --
  71 --   Examples:
  72 --
  73 --   >>> resolve_address "example.com"
  74 --   Just 93.184.216.119
  75 --   >>> resolve_address "93.184.216.119"
  76 --   Just 93.184.216.119
  77 --
  78 resolve_address :: String -> IO (Maybe IPv4)
  79 resolve_address s =
  80   case read_result of
  81     Just _  -> return read_result
  82     Nothing -> do
  83       default_rs <- makeResolvSeed defaultResolvConf
  84       withResolver default_rs $ \resolver -> do
  85         result <- lookupA resolver (BS.pack s)
  86         return $ case result of
  87                    Just (x:_) -> Just x
  88                    _           -> Nothing
  89   where
  90     read_result :: Maybe IPv4
  91     read_result = readMaybe s
  92
  93
  94 -- | This calls lookupMX, and returns the result as the second
  95 --   component of a tuple whose first component is the domain name
  96 --   itself.
  97 --
  98 --   Examples:
  99 --
 100 --   The example domain, example.com, has no MX record.
 101 --
 102 --   >>> rs <- makeResolvSeed defaultResolvConf
 103 --   >>> let domain = BS.pack "example.com."
 104 --   >>> withResolver rs $ \resolver -> lookupMX' resolver domain
 105 --   ("example.com.",Nothing)
 106 --
 107 lookupMX' :: Resolver -> Domain -> IO LookupResult
 108 lookupMX' resolver domain =
 109   liftM (pair_em . drop_priority) $ lookupMX resolver domain
 110   where
 111     drop_priority :: Maybe [(Domain, Int)] -> Maybe [Domain]
 112     drop_priority = fmap (map fst)
 113
 114     pair_em :: a -> (Domain, a)
 115     pair_em = (,) domain
 116
 117
 118 -- | This calls lookupNS, and returns the result as the second
 119 --   component of a tuple whose first component is the domain name
 120 --   itself.
 121 --
 122 --   Examples:
 123 --
 124 --   The example domain, example.com, does have NS records, but the
 125 --   order in which they are returned is variable, so we have to sort
 126 --   them to get a reliable result.
 127 --
 128 --   >>> import Data.List (sort)
 129 --   >>> let sort_snd (x,y) = (x, sort <$> y)
 130 --   >>> rs <- makeResolvSeed defaultResolvConf
 131 --   >>> let domain = BS.pack "example.com."
 132 --   >>> withResolver rs $ \resolver -> sort_snd <$> lookupNS' resolver domain
 133 --   ("example.com.",Just ["a.iana-servers.net.","b.iana-servers.net."])
 134 --
 135 lookupNS' :: Resolver -> Domain -> IO LookupResult
 136 lookupNS' resolver domain = do
 137   answer_result <- lookupNS resolver domain
 138   auth_result <- lookupNS_authority resolver domain
 139   liftM pair_em $ return $ combine answer_result auth_result
 140   where
 141     pair_em :: a -> (Domain, a)
 142     pair_em = (,) domain
 143
 144     combine :: (Maybe [Domain]) ->  (Maybe [Domain]) -> (Maybe [Domain])
 145     combine Nothing Nothing = Nothing
 146     combine m1 Nothing = m1
 147     combine Nothing m2 = m2
 148     combine (Just ds1) (Just ds2) = Just (ds1 ++ ds2)
 149
 150 -- | Perform both normalize_case and normalize_root.
 151 normalize :: Domain -> Domain
 152 normalize = normalize_case . normalize_root
 153
 154 -- | Normalize the given name by appending a trailing dot (the DNS
 155 --   root) if necessary.
 156 normalize_root :: Domain -> Domain
 157 normalize_root d
 158   | BS.null d = BS.pack "."
 159   | BS.last d == '.' = d
 160   | otherwise = d `BS.append` trailing_dot
 161     where
 162       trailing_dot = BS.pack "."
 163
 164
 165 -- | Normalize the given name by lowercasing it.
 166 normalize_case :: Domain -> Domain
 167 normalize_case = BS.map toLower
 168
 169
 170 test_normalize_case :: Test
 171 test_normalize_case =
 172   testCase desc $
 173     assertEqual desc expected actual
 174   where
 175     desc = "normalize_case lowercases DNS names"
 176     expected = BS.pack "example.com"
 177     actual = normalize_case $ BS.pack "ExAmPlE.COM"
 178
 179 prop_normalize_case_idempotent :: String -> Bool
 180 prop_normalize_case_idempotent s =
 181   (normalize_case . normalize_case) bs  == normalize_case bs
 182   where
 183     bs = BS.pack s
 184
 185 test_normalize_root_adds_dot :: Test
 186 test_normalize_root_adds_dot =
 187   testCase desc $
 188     assertEqual desc expected actual
 189   where
 190     desc = "normalize_root adds a trailing dot"
 191     expected = BS.pack "example.com."
 192     actual = normalize_root $ BS.pack "example.com"
 193
 194 prop_normalize_root_idempotent :: String -> Bool
 195 prop_normalize_root_idempotent s =
 196   (normalize_root . normalize_root) bs  == normalize_root bs
 197   where
 198     bs = BS.pack s
 199
 200 dns_tests :: Test
 201 dns_tests =
 202   testGroup "DNS Tests" [
 203     test_normalize_case,
 204     test_normalize_root_adds_dot ]
 205
 206 dns_properties :: Test
 207 dns_properties =
 208   testGroup "DNS Properties" [
 209     testProperty
 210       "normalize_case is idempotent"
 211       prop_normalize_case_idempotent,
 212     testProperty
 213       "normalize_root is idempotent"
 214       prop_normalize_root_idempotent ]