From f9610a6d357f2ed79f2f18baa477244970f8b40f Mon Sep 17 00:00:00 2001
From: Michael Orlitzky <michael@orlitzky.com>
Date: Tue, 3 Sep 2013 11:19:15 -0400
Subject: [PATCH] Add a "timeout" command-line argument.

---
 doc/man1/haeredes.1 |  2 ++
 haeredes.cabal      |  3 ++-
 src/CommandLine.hs  | 21 +++++++++++++++++++++
 src/Main.hs         | 30 +++++++++++++++++-------------
 src/Timeout.hs      | 18 ++++++++++++++++++
 5 files changed, 60 insertions(+), 14 deletions(-)
 create mode 100644 src/Timeout.hs

diff --git a/doc/man1/haeredes.1 b/doc/man1/haeredes.1
index 569897e..900b4c6 100644
--- a/doc/man1/haeredes.1
+++ b/doc/man1/haeredes.1
@@ -100,6 +100,8 @@ probably just lead to false positives.
 .IP \fB\-\-server\fR,\ \fB-s\fR
 Use the given DNS server rather than the resolvers listed in
 /etc/resolv.conf. Either an IP address or a hostname will work.
+.IP \fB\-\-timeout\fR,\ \fB-t\fR
+The number of seconds to wait for an answer from DNS (default: 15).
 .SH EXAMPLES
 
 .IP \[bu] 2
diff --git a/haeredes.cabal b/haeredes.cabal
index 631f5eb..5b7969a 100644
--- a/haeredes.cabal
+++ b/haeredes.cabal
@@ -1,5 +1,5 @@
 name:           haeredes
-version:        0.1
+version:        0.2
 cabal-version:  >= 1.8
 author:         Michael Orlitzky
 maintainer:	Michael Orlitzky <michael@orlitzky.com>
@@ -104,6 +104,7 @@ executable haeredes
     CommandLine
     DNS
     ExitCodes
+    Timeout
 
   ghc-options:
     -Wall
diff --git a/src/CommandLine.hs b/src/CommandLine.hs
index 2118f59..b94e378 100644
--- a/src/CommandLine.hs
+++ b/src/CommandLine.hs
@@ -30,6 +30,8 @@ import System.Console.CmdArgs (
 import Paths_haeredes (version)
 import Data.Version (showVersion)
 
+import Timeout (Timeout(..))
+
 -- | Description of the 'NS' mode.
 ns_description :: String
 ns_description =
@@ -56,15 +58,24 @@ server_help =
   "IP address or hostname of server to query " ++
   "(will use resolv.conf if not specified)"
 
+-- | Help string for the --timeout flag.
+timeout_help :: String
+timeout_help =
+  "Query timeout, in seconds (default: " ++ defstr ++ ")"
+  where
+    defstr = show $ seconds (def :: Timeout)
+
 -- | The Args type represents the possible command-line options. The
 --   duplication here seems necessary; CmdArgs' magic requires us to
 --   define some things explicitly.
 data Args =
   NS { no_append_root :: Bool,
        server :: Maybe String,
+       timeout :: Timeout,
        delegates :: [String] } |
   MX { no_append_root :: Bool,
        server :: Maybe String,
+       timeout :: Timeout,
        delegates :: [String] }
   deriving (Data, Show, Typeable)
 
@@ -94,6 +105,11 @@ arg_spec =
                  &= typ "HOST"
                  &= help server_help,
 
+           timeout = def
+                   &= groupname "Common flags"
+                   &= typ "SECONDS"
+                   &= help timeout_help,
+
            delegates = def
                     &= args
                     &= typ "DELEGATES" }
@@ -111,6 +127,11 @@ arg_spec =
                  &= typ "IP"
                  &= help server_help,
 
+           timeout = def
+                   &= groupname "Common flags"
+                   &= typ "SECONDS"
+                   &= help timeout_help,
+
            delegates = def
                     &= args
                     &= typ "DELEGATES" }
diff --git a/src/Main.hs b/src/Main.hs
index 6b276b7..4e0da6b 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -14,7 +14,7 @@ import Data.String.Utils (join)
 import Network.DNS (
   Domain,
   FileOrNumericHost(RCHostName),
-  ResolvConf(resolvInfo),
+  ResolvConf(resolvInfo, resolvTimeout),
   defaultResolvConf,
   makeResolvSeed,
   withResolver )
@@ -31,6 +31,7 @@ import DNS (
   normalize_case,
   resolve_address )
 import ExitCodes (exit_bad_server)
+import Timeout (Timeout(..))
 
 
 -- | Report results for this LookupResult. If there's a Nothing in the
@@ -83,18 +84,21 @@ main = do
   let nrml_domains   = map normalize_function raw_domains
   let nrml_delegates = map normalize_function raw_delegates
 
-  rc <- case (server cfg) of
-          Nothing -> return defaultResolvConf
-          Just s -> do
-            s' <- resolve_address s
-            case s' of
-              Nothing -> do
-                hPutStrLn stderr ("Bad DNS server or lookup error: " ++ s)
-                exitWith (ExitFailure exit_bad_server)
-              Just s'' ->
-                return $ defaultResolvConf { resolvInfo =
-                                               RCHostName (show s'') }
-
+  rc' <- case (server cfg) of
+           Nothing -> return defaultResolvConf
+           Just s -> do
+             s' <- resolve_address s
+             case s' of
+               Nothing -> do
+                 hPutStrLn stderr ("Bad DNS server or lookup error: " ++ s)
+                 exitWith (ExitFailure exit_bad_server)
+               Just s'' ->
+                 return $ defaultResolvConf { resolvInfo =
+                                                RCHostName (show s'') }
+
+  -- Set the timeout from the command line. The resolvTimeout field is
+  -- in microseconds, so we multiply by one million.
+  let rc = rc' { resolvTimeout = 1000 * 1000 * (seconds $ timeout cfg) }
   rs <- makeResolvSeed rc
 
   let lookup_function = case cfg of
diff --git a/src/Timeout.hs b/src/Timeout.hs
new file mode 100644
index 0000000..6f0f3d4
--- /dev/null
+++ b/src/Timeout.hs
@@ -0,0 +1,18 @@
+{-# LANGUAGE DeriveDataTypeable #-}
+
+-- | A newtype around the DNS query timeout, in seconds. Needed for
+--   the 'Default' instance.
+--
+module Timeout
+where
+
+import Data.Data (Data)
+import System.Console.CmdArgs.Default (Default(..))
+import Data.Typeable (Typeable)
+
+newtype Timeout =
+  Timeout { seconds :: Int }
+  deriving (Data, Show, Typeable)
+
+instance Default Timeout where
+  def = Timeout 15
-- 
2.43.2