we can specify timeout.

2 files changed, 79 insertions, 45 deletions

diff --git a/Network/DNS.hs b/Network/DNS.hs
index 735db1f..c8a4c83 100644
--- a/Network/DNS.hs
+++ b/Network/DNS.hs
@@ -12,7 +12,7 @@
     import Network.DNS hiding (lookup)
     main :: IO ()
     main = do
-        rs <- makeDefaultResolvSeed
+        rs <- makeResolvSeed defaultResolvConf
         withResolver rs $ \\resolver -> do
             DNS.lookup resolver \"www.example.com\" A >>= print
 @
diff --git a/Network/DNS/Resolver.hs b/Network/DNS/Resolver.hs
index 54a6411..41cabc7 100644
--- a/Network/DNS/Resolver.hs
+++ b/Network/DNS/Resolver.hs
@@ -3,23 +3,55 @@
 -}
 
 module Network.DNS.Resolver (
-    ResolvSeed, makeResolvSeed, makeDefaultResolvSeed
+    FileOrNumericHost(..), ResolvConf(..), defaultResolvConf
+  , ResolvSeed, makeResolvSeed
   , Resolver, withResolver
   , lookup, lookupRaw
   ) where
 
 import Control.Applicative
 import Control.Exception
-import Data.List hiding (find, lookup)
 import Data.Int
-import Network.DNS.Types
+import Data.List hiding (find, lookup)
+import Network.BSD
 import Network.DNS.Query
 import Network.DNS.Response
-import Random
-import Network.BSD
+import Network.DNS.Types
 import Network.Socket hiding (send, sendTo, recv, recvFrom)
 import Network.Socket.ByteString.Lazy
 import Prelude hiding (lookup)
+import Random
+import System.Timeout
+
+----------------------------------------------------------------
+
+{-|
+  Union type for 'FilePath' and 'HostName'. Specify 'FilePath' to
+  \"resolv.conf\" or numeric IP address in 'String' form.
+-}
+data FileOrNumericHost = RCFilePath FilePath | RCHostName HostName
+
+{-|
+  Type for resolver configuration
+-}
+data ResolvConf = ResolvConf {
+    resolvInfo :: FileOrNumericHost
+  , resolvTimeout :: Int
+  , resolvBufsize :: Int64
+}
+
+{-|
+  Default 'ResolvConf'.
+  'resolvInfo' is 'RCFilePath' \"\/etc\/resolv.conf\".
+  'resolvTimeout' is 3,000,000 micro seconds.
+  'resolvBufsize' is 512.
+-}
+defaultResolvConf :: ResolvConf
+defaultResolvConf = ResolvConf {
+    resolvInfo = RCFilePath "/etc/resolv.conf"
+  , resolvTimeout = 3 * 1000 * 1000
+  , resolvBufsize = 512
+}
 
 ----------------------------------------------------------------
 
@@ -28,43 +60,36 @@ import Prelude hiding (lookup)
 -}
 data ResolvSeed = ResolvSeed {
     addrInfo :: AddrInfo
+  , rsTimeout :: Int
+  , rsBufsize :: Int64
 }
 
+{-|
+  Abstract data type of DNS Resolver
+-}
 data Resolver = Resolver {
     genId   :: IO Int
   , dnsSock :: Socket
+  , dnsTimeout :: Int
+  , dnsBufsize :: Int64
 }
 
 ----------------------------------------------------------------
 
-resolvConf :: String
-resolvConf = "/etc/resolv.conf"
-
-dnsBufferSize :: Int64
-dnsBufferSize = 512
-
-----------------------------------------------------------------
-
 {-|
   Making 'ResolvSeed' from an IP address of a DNS cache server.
 -}
-makeResolvSeed :: HostName -> IO ResolvSeed
-makeResolvSeed addr = ResolvSeed <$> makeAddrInfo addr
-
-{-|
-  Making 'ResolvSeed' from \"/etc/resolv.conf\".
--}
-makeDefaultResolvSeed :: IO ResolvSeed
-makeDefaultResolvSeed = toAddr <$> readFile resolvConf >>= makeResolvSeed
+makeResolvSeed :: ResolvConf -> IO ResolvSeed
+makeResolvSeed conf = ResolvSeed <$> addr
+                                 <*> pure (resolvTimeout conf)
+                                 <*> pure (resolvBufsize conf)
   where
+    addr = case resolvInfo conf of
+        RCHostName numhost -> makeAddrInfo numhost
+        RCFilePath file -> toAddr <$> readFile file >>= makeAddrInfo
     toAddr cs = let l:_ = filter ("nameserver" `isPrefixOf`) $ lines cs
                 in drop 11 l
 
-----------------------------------------------------------------
-
-getRandom :: IO Int
-getRandom = getStdRandom (randomR (0,65535))
-
 makeAddrInfo :: HostName -> IO AddrInfo
 makeAddrInfo addr = do
     proto <- getProtocolNumber "udp"
@@ -88,9 +113,17 @@ withResolver seed func = do
   let ai = addrInfo seed
   sock <- socket (addrFamily ai) (addrSocketType ai) (addrProtocol ai)
   connect sock (addrAddress ai)
-  let resolv = Resolver getRandom sock
+  let resolv = Resolver {
+          genId = getRandom
+        , dnsSock = sock
+        , dnsTimeout = rsTimeout seed
+        , dnsBufsize = rsBufsize seed
+        }
   func resolv `finally` sClose sock
 
+getRandom :: IO Int
+getRandom = getStdRandom (randomR (0,65535))
+
 ----------------------------------------------------------------
 
 {-|
@@ -98,13 +131,8 @@ withResolver seed func = do
 -}
 lookup :: Resolver -> Domain -> TYPE -> IO (Maybe [RDATA])
 lookup rlv dom typ = do
-    let sock = dnsSock rlv
-    seqno <- genId rlv
-    res <- lookupRaw' sock seqno dom typ
-    let hdr = header res
-    if identifier hdr == seqno && anCount hdr /= 0
-       then return . listToMaybe . map rdata . filter correct $ answer res
-       else return Nothing
+    mres <- lookupRaw rlv dom typ
+    return (mres >>= toRDATA)
   where
     {- CNAME hack
     dom' = if "." `isSuffixOf` dom
@@ -115,18 +143,24 @@ lookup rlv dom typ = do
     correct r = rrtype r == typ
     listToMaybe [] = Nothing
     listToMaybe xs = Just xs
+    toRDATA = listToMaybe . map rdata . filter correct . answer
 
 {-|
   Looking up a domain and returning an entire DNS Response.
 -}
-lookupRaw :: Resolver -> Domain -> TYPE -> IO DNSFormat
+lookupRaw :: Resolver -> Domain -> TYPE -> IO (Maybe DNSFormat)
 lookupRaw rlv dom typ = do
-   let sock = dnsSock rlv
-   seqno <- genId rlv
-   lookupRaw' sock seqno dom typ
-
-lookupRaw' :: Socket -> Int -> Domain -> TYPE -> IO DNSFormat
-lookupRaw' sock seqno dom typ = do
-  let q = makeQuestion dom typ
-  sendAll sock (composeQuery seqno [q])
-  parseResponse <$> recv sock dnsBufferSize
+    seqno <- genId rlv
+    sendAll sock (composeQuery seqno [q])
+    mres <- timeout tm (parseResponse <$> recv sock bufsize)
+    return (mres >>= check seqno)
+  where
+    sock = dnsSock rlv
+    bufsize = dnsBufsize rlv
+    tm = dnsTimeout rlv
+    q = makeQuestion dom typ
+    check seqno res = do
+        let hdr = header res
+        if identifier hdr == seqno && anCount hdr /= 0
+            then Just res
+            else Nothing