Supporting server side and brushing up.

4 files changed, 62 insertions, 38 deletions

diff --git a/Network/DNS.hs b/Network/DNS.hs
index acf56c6..73b44db 100644
--- a/Network/DNS.hs
+++ b/Network/DNS.hs
@@ -1,10 +1,9 @@
 {-|
   Thread-safe DNS library written in Haskell.
 
-  Currently, only resolver side is supported. This code is written in
-  Haskell, not using FFI.
+  This code is written in Haskell, not using FFI.
 
-  Sample code:
+  Sample code for DNS lookup:
 
 @
     import qualified Network.DNS as DNS (lookup)
@@ -18,11 +17,19 @@
 -}
 
 module Network.DNS (
+  -- * High level
     module Network.DNS.Lookup
   , module Network.DNS.Resolver
   , module Network.DNS.Types
+  -- * Low level
+  , module Network.DNS.Decode
+  , module Network.DNS.Encode
   ) where
 
 import Network.DNS.Lookup
 import Network.DNS.Resolver
 import Network.DNS.Types
+import Network.DNS.Decode
+import Network.DNS.Encode
+
+
diff --git a/Network/DNS/Response.hs b/Network/DNS/Decode.hs
index 523335b..c84a06d 100644
--- a/Network/DNS/Response.hs
+++ b/Network/DNS/Decode.hs
@@ -1,33 +1,46 @@
 {-# LANGUAGE OverloadedStrings #-}
 
-module Network.DNS.Response (responseIter, parseResponse, runDNSFormat, runDNSFormat_) where
+module Network.DNS.Decode (
+    receive
+  , decode
+  ) where
 
 import Control.Applicative
 import Control.Monad
 import Data.Bits
+import Data.ByteString (ByteString)
 import qualified Data.ByteString.Char8 as BS
+import qualified Data.ByteString.Lazy as BL
+import Data.Enumerator (Enumerator, run_, ($$))
 import Data.IP
 import Data.Maybe
+import Network
 import Network.DNS.Internal
 import Network.DNS.StateBinary
-import Data.Enumerator (Enumerator, Iteratee, run_, ($$))
-import Data.ByteString (ByteString)
-import qualified Data.ByteString.Lazy as BL
+import Network.Socket.Enumerator
 
-runDNSFormat :: BL.ByteString -> Either String (DNSFormat, PState)
-runDNSFormat = runSGet decodeResponse
+----------------------------------------------------------------
 
-runDNSFormat_ :: BL.ByteString -> Either String DNSFormat
-runDNSFormat_ bs = fst <$> runDNSFormat bs
+{-| Receiving DNS data from 'Socket' and parse it.
+    The second argument is a buffer size for the socket.
+-}
+receive :: Socket -> Integer -> IO DNSFormat
+receive sock bufsize = receiveDNSFormat responseEnum
+  where
+    responseEnum = enumSocket bufsize sock
+
+{-| Parsing DNS data.
+-}
+decode :: BL.ByteString -> Either String DNSFormat
+decode bs = fst <$> runSGet decodeResponse bs
 
-responseIter :: Monad m => Iteratee ByteString m (DNSFormat, PState)
-responseIter = iterSGet decodeResponse
+----------------------------------------------------------------
 
-parseResponse :: (Functor m, Monad m)
-              => Enumerator ByteString m (a,b)
-              -> Iteratee ByteString m (a,b)
-              -> m a
-parseResponse enum iter = fst <$> run_ (enum $$ iter)
+receiveDNSFormat :: Enumerator ByteString IO (DNSFormat, PState)
+                 -> IO DNSFormat
+receiveDNSFormat enum = fst <$> run_ (enum $$ iter)
+  where
+    iter = iterSGet decodeResponse
 
 ----------------------------------------------------------------
 
diff --git a/Network/DNS/Query.hs b/Network/DNS/Encode.hs
index bdf86bc..6116a04 100644
--- a/Network/DNS/Query.hs
+++ b/Network/DNS/Encode.hs
@@ -1,5 +1,5 @@
 {-# LANGUAGE RecordWildCards #-}
-module Network.DNS.Query (composeQuery, composeDNSFormat) where
+module Network.DNS.Encode (encode) where
 
 import qualified Data.ByteString.Lazy.Char8 as BL (ByteString)
 import qualified Data.ByteString.Char8 as BS (length, null, break, drop)
@@ -16,20 +16,10 @@ import Data.IP
 
 ----------------------------------------------------------------
 
-composeDNSFormat :: DNSFormat -> BL.ByteString
-composeDNSFormat fmt = runSPut (encodeDNSFormat fmt)
-
-composeQuery :: Int -> [Question] -> BL.ByteString
-composeQuery idt qs = composeDNSFormat qry
-  where
-    hdr = header defaultQuery
-    qry = defaultQuery {
-        header = hdr {
-           identifier = idt
-         , qdCount = length qs
-         }
-      , question = qs
-      }
+{-| Composing DNS data.
+-}
+encode :: DNSFormat -> BL.ByteString
+encode fmt = runSPut (encodeDNSFormat fmt)
 
 ----------------------------------------------------------------
 
diff --git a/Network/DNS/Resolver.hs b/Network/DNS/Resolver.hs
index b3182f6..44aee7d 100644
--- a/Network/DNS/Resolver.hs
+++ b/Network/DNS/Resolver.hs
@@ -28,19 +28,20 @@ module Network.DNS.Resolver (
 
 import Control.Applicative
 import Control.Exception
+import qualified Data.ByteString.Lazy as BL
 import Data.Char
 import Data.Int
 import Data.List hiding (find, lookup)
 import Network.BSD
-import Network.DNS.Query
-import Network.DNS.Response
+import Network.DNS.Decode
+import Network.DNS.Encode
+import Network.DNS.Internal
 import Network.DNS.Types
 import Network.Socket hiding (send, sendTo, recv, recvFrom)
 import Network.Socket.ByteString.Lazy
 import Prelude hiding (lookup)
 import System.Random
 import System.Timeout
-import Network.Socket.Enumerator
 
 ----------------------------------------------------------------
 
@@ -170,8 +171,7 @@ lookupRaw :: Resolver -> Domain -> TYPE -> IO (Maybe DNSFormat)
 lookupRaw rlv dom typ = do
     seqno <- genId rlv
     sendAll sock (composeQuery seqno [q])
-    let responseEnum = enumSocket bufsize sock
-    (>>= check seqno) <$> timeout tm (parseResponse responseEnum responseIter)
+    (>>= check seqno) <$> timeout tm (receive sock bufsize)
   where
     sock = dnsSock rlv
     bufsize = dnsBufsize rlv
@@ -182,3 +182,17 @@ lookupRaw rlv dom typ = do
         if identifier hdr == seqno && anCount hdr /= 0
             then Just res
             else Nothing
+
+----------------------------------------------------------------
+
+composeQuery :: Int -> [Question] -> BL.ByteString
+composeQuery idt qs = encode qry
+  where
+    hdr = header defaultQuery
+    qry = defaultQuery {
+        header = hdr {
+           identifier = idt
+         , qdCount = length qs
+         }
+      , question = qs
+      }