cda8b96274414d6bd90b6dd71bf8d5f1f08e2309
[haeredes.git] / 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 ]