src/Main.hs

   1 module Main (main)
   2 where
   3
   4 import Control.Concurrent.ParallelIO.Global (
   5   parallelInterleaved,
   6   stopGlobalPool )
   7 import Control.Monad (unless)
   8 import qualified Data.ByteString.Char8 as BS (
   9   getContents,
  10   pack,
  11   words )
  12 import Data.List ( (\\), intercalate, sort )
  13 import Network.DNS (
  14   DNSError(NameError),
  15   Domain,
  16   FileOrNumericHost(RCHostName),
  17   ResolvConf(resolvInfo, resolvTimeout),
  18   defaultResolvConf,
  19   makeResolvSeed,
  20   normalize,
  21   normalizeCase,
  22   withResolver )
  23 import System.Exit (ExitCode(ExitFailure), exitWith)
  24 import System.IO (hPutStrLn, stderr)
  25
  26 import CommandLine (
  27    Args(NS,MX,delegates,no_append_root,server,timeout),
  28    get_args)
  29 import DNS (
  30   LookupResult,
  31   lookupMX',
  32   lookupNS',
  33   resolve_address )
  34 import ExitCodes (exit_bad_server)
  35 import Timeout (Timeout(seconds))
  36
  37
  38 -- | Given a list of delegates, report results for this
  39 --   'LookupResult'.
  40 --
  41 --   If there's an empty list in the second component, there were no
  42 --   query results, so we report that the domain was not delegated. If
  43 --   there were some results and there are leftovers (after removing
  44 --   the delegates), we report those as well.
  45 --
  46 --   Before processing, all names are normalized using the supplied
  47 --   function @normalize_function@. Ideally this should be the same
  48 --   function applied to the user-input names.
  49 --
  50 report :: [Domain] -- ^ The list of @delgts@
  51        -> (Domain -> Domain) -- ^ Domain name normalization function,
  52                              --   @normalize_function@.
  53        -> LookupResult
  54        -> IO ()
  55
  56 -- If the lookup succeeded but there were no results, report that the
  57 -- domain is not delegated. Note that the behavior of the DNS library
  58 -- changed with regard to this at some point: we used to get back
  59 -- a "success," but with an empty list of results. Now a NameError
  60 -- (which is not actually an error!) is returned.
  61 report _ _ (d, Left NameError) =
  62   putStrLn $ "Domain " ++ (show d) ++ " not delegated."
  63
  64 -- If the lookup resulted in some other DNS error, we just ignore the
  65 -- whole thing.
  66 report _ _ (_, Left _) = return ()
  67
  68 -- Otherwise, subtract our delegates from the list of results and
  69 -- report the leftovers.
  70 report delgts normalize_function (d, Right raw_hosts) = do
  71   let nrml_hosts = map normalize_function raw_hosts
  72   -- Sort the leftovers so that we can test the expected output.
  73   let leftovers = sort (nrml_hosts \\ delgts)
  74   unless (null leftovers) $
  75     putStrLn $ "Domain " ++
  76                (show d) ++
  77                " delegates somewhere else: " ++
  78                (join " " (map show leftovers))
  79   where
  80     -- Create one big string by joining together a list of smaller
  81     -- strings and placing a delimiter between them.
  82     join :: String -> [String] -> String
  83     join delimiter strings = intercalate delimiter strings
  84
  85
  86 main :: IO ()
  87 main = do
  88   cfg <- get_args
  89
  90   -- This reads stdin.
  91   input <- BS.getContents
  92
  93   -- Split the input on any whitespace characters.
  94   let raw_domains = BS.words input
  95
  96   -- Convert these to ByteStrings.
  97   let raw_delegates = map BS.pack (delegates cfg)
  98
  99   let normalize_function =
 100         if (no_append_root cfg)
 101         then normalizeCase
 102         else normalize
 103
 104   -- Normalize the given names and delegates
 105   let nrml_domains   = map normalize_function raw_domains
 106   let nrml_delegates = map normalize_function raw_delegates
 107
 108   rc' <- case (server cfg) of
 109            Nothing -> return defaultResolvConf
 110            Just s -> do
 111              s' <- resolve_address s
 112              case s' of
 113                Left err -> do
 114                  let errmsg = show err
 115                  hPutStrLn stderr ("Bad DNS server or lookup error: " ++ errmsg)
 116                  exitWith (ExitFailure exit_bad_server)
 117                Right []  -> do
 118                  hPutStrLn stderr ("Hostname " ++ s ++ " has no 'A' records.")
 119                  exitWith (ExitFailure exit_bad_server)
 120                Right (srv:_) ->
 121                  return $ defaultResolvConf { resolvInfo =
 122                                                 RCHostName (show srv) }
 123
 124   -- Set the timeout from the command line. The resolvTimeout field is
 125   -- in microseconds, so we multiply by one million.
 126   let a_million = 1000 * 1000 :: Int
 127   let rc = rc' { resolvTimeout = a_million * seconds (timeout cfg) }
 128   rs <- makeResolvSeed rc
 129
 130   let lookup_function = case cfg of
 131                           NS{} -> lookupNS'
 132                           MX{} -> lookupMX'
 133
 134   let lookup' d = withResolver rs $ \resolver ->
 135                     lookup_function resolver d
 136
 137   -- Construct a list of [IO whatever]. The withResolver calls
 138   -- are the ones that should be run in parallel.
 139   let actions = map lookup' nrml_domains
 140
 141   -- Run the lookup actions in parallel.
 142   results <- parallelInterleaved actions
 143
 144   -- Output the results.
 145   mapM_ (report nrml_delegates normalize_function) results
 146
 147   stopGlobalPool