Add a doctest suite.
[haeredes.git] / 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 ]