introducing ResolvSeed.

4 files changed, 57 insertions, 39 deletions

diff --git a/Network/DNS.hs b/Network/DNS.hs
index 5d70f2e..db03c2d 100644
--- a/Network/DNS.hs
+++ b/Network/DNS.hs
@@ -11,14 +11,19 @@
     import qualified Network.DNS as DNS (lookup)
     import Network.DNS hiding (lookup)
     main :: IO ()
-    main = makeDefaultResolver >>= DNS.lookup \"www.example.com\" A >>= print
+    main = do
+        rs <- makeDefaultResolvSeed
+        withResolver rs $ \\resolver -> do
+            DNS.lookup resolver \"www.example.com\" A >>= print
 @
 -}
 
 module Network.DNS (
-    module Network.DNS.Types
+    module Network.DNS.Lookup
   , module Network.DNS.Resolver
+  , module Network.DNS.Types
   ) where
 
-import Network.DNS.Types
+import Network.DNS.Lookup
 import Network.DNS.Resolver
+import Network.DNS.Types
diff --git a/Network/DNS/Query.hs b/Network/DNS/Query.hs
index db30bd8..66bfb31 100644
--- a/Network/DNS/Query.hs
+++ b/Network/DNS/Query.hs
@@ -49,7 +49,7 @@ encodeFlags _ = put16 0x0100 -- xxx
 
 encodeQuestion :: [Question] -> SPut
 encodeQuestion qs = do
-  let q = head qs 
+  let q = head qs
       dom = qname q
       typ = qtype q
   encodeDomain dom
diff --git a/Network/DNS/Resolver.hs b/Network/DNS/Resolver.hs
index 6ae7d50..e959bcb 100644
--- a/Network/DNS/Resolver.hs
+++ b/Network/DNS/Resolver.hs
@@ -1,13 +1,15 @@
 {-|
-  APIs of DNS Resolver.
+  DNS Resolver and lookup functions.
 -}
 
 module Network.DNS.Resolver (
-    Resolver, makeResolver, makeDefaultResolver
+    ResolvSeed, makeResolvSeed, makeDefaultResolvSeed
+  , Resolver, withResolver
   , lookup, lookupRaw
   ) where
 
 import Control.Applicative
+import Control.Exception
 import Data.List hiding (find, lookup)
 import Data.Int
 import Network.DNS.Types
@@ -22,11 +24,15 @@ import Prelude hiding (lookup)
 ----------------------------------------------------------------
 
 {-|
-  Abstract data type of DNS Resolver
+  Abstract data type of DNS Resolver seed
 -}
+data ResolvSeed = ResolvSeed {
+    addrInfo :: AddrInfo
+}
+
 data Resolver = Resolver {
-    genId    :: IO Int
-  , addrInfo :: AddrInfo
+    genId   :: IO Int
+  , dnsSock :: Socket
 }
 
 ----------------------------------------------------------------
@@ -40,21 +46,19 @@ dnsBufferSize = 512
 ----------------------------------------------------------------
 
 {-|
-  Making Resolver from an IP address of a DNS cache server.
+  Making 'ResolvSeed' from an IP address of a DNS cache server.
 -}
-makeResolver :: HostName -> IO Resolver
-makeResolver addr = do
-    ai <- makeAddrInfo addr
-    return $ Resolver { genId = getRandom, addrInfo = ai }
+makeResolvSeed :: HostName -> IO ResolvSeed
+makeResolvSeed addr = ResolvSeed <$> makeAddrInfo addr
 
 {-|
-  Making Resolver from \"/etc/resolv.conf\".
+  Making 'ResolvSeed' from \"/etc/resolv.conf\".
 -}
-makeDefaultResolver :: IO Resolver
-makeDefaultResolver = do
-  cs <- readFile resolvConf
-  let l:_ = filter ("nameserver" `isPrefixOf`) $ lines cs
-  makeResolver $ drop 11 l
+makeDefaultResolvSeed :: IO ResolvSeed
+makeDefaultResolvSeed = toAddr <$> readFile resolvConf >>= makeResolvSeed
+  where
+    toAddr cs = let l:_ = filter ("nameserver" `isPrefixOf`) $ lines cs
+                in drop 11 l
 
 ----------------------------------------------------------------
 
@@ -74,15 +78,26 @@ makeAddrInfo addr = do
 
 ----------------------------------------------------------------
 
+withResolver :: ResolvSeed -> (Resolver -> IO ()) -> IO ()
+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
+  func resolv `finally` sClose sock
+
+----------------------------------------------------------------
+
 {-|
   Looking up resource records of a domain.
 -}
-lookup :: Domain -> TYPE -> Resolver -> IO (Maybe [RDATA])
-lookup dom typ rlv = do
-    idnt <- genId rlv
-    res <- lookupRaw' dom typ rlv idnt
+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 == idnt && anCount hdr /= 0
+    if identifier hdr == seqno && anCount hdr /= 0
        then return . listToMaybe . map rdata . filter correct $ answer res
        else return Nothing
   where
@@ -99,16 +114,14 @@ lookup dom typ rlv = do
 {-|
   Looking up a domain and returning an entire DNS Response.
 -}
-lookupRaw :: Domain -> TYPE -> Resolver -> IO DNSFormat
-lookupRaw dom typ rlv = genId rlv >>= lookupRaw' dom typ rlv
-
-lookupRaw' :: Domain -> TYPE -> Resolver -> Int -> IO DNSFormat
-lookupRaw' dom typ rlv idnt = do
-  let ai = addrInfo rlv
-      q = makeQuestion dom typ
-  sock <- socket (addrFamily ai) (addrSocketType ai) (addrProtocol ai)
-  connect sock (addrAddress ai)
-  sendAll sock (composeQuery idnt [q])
-  fmt <- parseResponse <$> recv sock dnsBufferSize
-  sClose sock
-  return fmt
+lookupRaw :: Resolver -> Domain -> TYPE -> IO 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
diff --git a/Network/DNS/StateBinary.hs b/Network/DNS/StateBinary.hs
index 3b89665..4241364 100644
--- a/Network/DNS/StateBinary.hs
+++ b/Network/DNS/StateBinary.hs
@@ -12,7 +12,7 @@ import Network.DNS.Types
 import Prelude hiding (lookup)
 
 ----------------------------------------------------------------
-  
+
 type SGet = StateT PState Get
 
 type PState = IntMap Domain