Introduce a NormalDomain newtype to ensure comparisons are made safely.

[list-remote-forwards.git] / src / DNS.hs

diff --git a/src/DNS.hs b/src/DNS.hs
index c319a9fbb387f66ae7a047dee33d119c6cbc5428..556c68675eecb4e4a6293117143e5358255f28bc 100644 (file)
--- a/src/DNS.hs
+++ b/src/DNS.hs
@@ -1,14 +1,66 @@
-module DNS ( lookup_mxs )
+module DNS (
+  MxSetMap,
+  NormalDomain,
+  mx_set_map,
+  normalize_string )
 where
 
+import qualified Data.ByteString.Char8 as BS ( pack, unpack )
+import Data.List ( nub )
+import Data.Map ( Map )
+import qualified Data.Map as Map ( fromList )
+import Data.Set ( Set )
+import qualified Data.Set as Set ( fromList )
 import Network.DNS (
   Domain,
   defaultResolvConf,
   lookupMX,
   makeResolvSeed,
+  normalize,
   withResolver )
 
--- Slow since we create the resolver every time.
+-- | A type-safe wrapper around a domain name (represented as a
+--   string) that ensures we've created it by calling
+--   'normalize_string'. This prevents us from making
+--   comparisons on un-normalized 'Domain's or 'String's.
+--
+newtype NormalDomain =
+  NormalDomain String
+  deriving ( Eq, Ord, Show )
+
+
+-- | A set of mail exchanger names, represented as 'String's. The use
+--   of 'NormalDomain' prevents us from constructing a set of names
+--   that aren't normalized first.
+--
+type MxSet = Set NormalDomain
+
+
+-- | A map from domain names (represented as 'String's) to sets of
+--   mail exchanger names (also represented as 'String's). The use of
+--   'NormalDomain' in the key prevents us from using keys that aren't
+--   normalized; this is important because we'll be using them for
+--   lookups and want e.g. \"foo.com\" and \"FOO.com\" to look up the
+--   same MX records.
+--
+type MxSetMap = Map NormalDomain MxSet
+
+
+-- | Normalize a domain name string by converting to a 'Domain',
+--   calling 'normalize', and then converting back.
+--
+--   ==== __Examples__
+--
+--   >>> normalize_string "ExAMplE.com"
+--   NormalDomain "example.com."
+--
+normalize_string :: String -> NormalDomain
+normalize_string = NormalDomain . BS.unpack . normalize . BS.pack
+
+
+-- | Retrieve all MX records for the given domain. This is somewhat
+--   inefficient, since we create the resolver every time.
+--
 lookup_mxs :: Domain -> IO [Domain]
 lookup_mxs domain = do
   default_rs <- makeResolvSeed defaultResolvConf
@@ -17,3 +69,51 @@ lookup_mxs domain = do
     return $ case mxs of
                Left  _     -> []
                Right pairs -> map fst pairs
+
+
+-- | Takes a list of domain names represented as 'String's and
+--   constructs a map from domain names to sets of mail exchangers
+--   (for those domain names) also represented as 'String's.
+--
+--   During construction, we have to switch to the DNS internal
+--   representation of a 'Domain' which uses ByteStrings, but before
+--   we return the map to the client, we want everything to be in
+--   terms of standard 'String's for comparison purposes.
+--
+--   The list of domains is normalized and de-duped before lookups are
+--   performed to avoid doing lookups twice for identical domains.
+--
+mx_set_map :: [String] -> IO MxSetMap
+mx_set_map domains = do
+  -- Construct a list of pairs.
+  pairs <- mapM make_pair unique_domains
+
+  -- And make a map from the pairs.
+  return $ Map.fromList pairs
+
+  where
+    -- Convert, normalize, and de-dupe the @domains@.
+    unique_domains :: [Domain]
+    unique_domains = nub $ map (normalize . BS.pack) domains
+
+    -- | Convert a string domain name into a pair containing the
+    --   domain name in the first component and a set of its mail
+    --   exchangers (as strings) in the second component.
+    --
+    make_pair :: Domain -> IO (NormalDomain, Set NormalDomain)
+    make_pair domain = do
+      -- Lookup the @domain@'s MX records.
+      mx_list <- lookup_mxs domain
+
+      -- Now convert the MX records *back* to strings, and then to
+      -- NormalDomains
+      let normal_mx_list = map (normalize_string . BS.unpack) mx_list
+
+      -- Convert the list into a set...
+      let normal_mx_set = Set.fromList normal_mx_list
+
+      -- The lookup key.
+      let normal_domain = normalize_string $ BS.unpack domain
+
+      -- Finally, construct the pair and return it.
+      return (normal_domain, normal_mx_set)