From d16870464640fcd94d318a919dbf10eec0f96ee3 Mon Sep 17 00:00:00 2001
From: Kazu Yamamoto <kazu@iij.ad.jp>
Date: Mon, 24 Oct 2011 16:02:33 +0900
Subject: pull request from yihuang. Applying by hand.

---
 SimpleServer.hs | 40 +++++++++++++++++-----------------------
 1 file changed, 17 insertions(+), 23 deletions(-)

diff --git a/SimpleServer.hs b/SimpleServer.hs
index 8e12a29..6cd1c11 100644
--- a/SimpleServer.hs
+++ b/SimpleServer.hs
@@ -6,9 +6,9 @@ import Control.Monad
 import qualified Data.ByteString as S
 import Data.ByteString.Lazy hiding (putStrLn, filter, length)
 import Data.Default
+import Data.IP
 import Data.Maybe
 import Data.Monoid
-import Debug.Trace
 import Network.BSD
 import Network.DNS hiding (lookup)
 import Network.Socket hiding (recvFrom)
@@ -20,6 +20,7 @@ data Conf = Conf {
     bufSize :: Int
   , timeOut :: Int
   , realDNS :: HostName
+  , hosts   :: [(Domain, IPv4)]
 }
 
 instance Default Conf where
@@ -27,6 +28,7 @@ instance Default Conf where
         bufSize = 512
       , timeOut = 3 * 1000 * 1000
       , realDNS = "192.168.1.1"
+      , hosts   = [("localhost", "127.0.0.1")]
     }
 
 timeout' :: String -> Int -> IO a -> IO (Maybe a)
@@ -37,11 +39,10 @@ timeout' msg tm io = do
 
 proxyRequest :: Conf -> ResolvConf -> DNSFormat -> IO (Maybe DNSFormat)
 proxyRequest Conf{..} rc req = do
-    let
-      worker Resolver{..} = do
-        let packet = mconcat . toChunks $ encode req
-        sendAll dnsSock packet
-        receive dnsSock dnsBufsize
+    let worker Resolver{..} = do
+            let packet = mconcat . toChunks $ encode req
+            sendAll dnsSock packet
+            receive dnsSock dnsBufsize
     rs <- makeResolvSeed rc
     withResolver rs $ \r ->
         (>>= check) <$> timeout' "proxy timeout" timeOut (worker r)
@@ -51,38 +52,31 @@ proxyRequest Conf{..} rc req = do
     check rsp = let hdr = header rsp
                 in  if identifier hdr == ident
                         then Just rsp
-                        else trace "identifier not match" Nothing
+                        else Nothing
 
-{--
- - TBD
- --}
 handleRequest :: Conf -> ResolvConf -> DNSFormat -> IO (Maybe DNSFormat)
-handleRequest conf rc req = maybe (proxyRequest conf rc req) (trace "return A record" $ return . Just) mresponse
+handleRequest conf@Conf{hosts=hosts} rc req =
+    maybe
+      (proxyRequest conf rc req)
+      (return . Just)
+      lookupHosts
   where
     filterA = filter ((==A) . qtype)
-    mresponse = do
-        let ident = identifier . header $ req
+    ident = identifier . header $ req
+    lookupHosts = do
         q <- listToMaybe . filterA . question $ req
-        let dom = qname q
-        ip <- lookup dom hosts
+        ip <- lookup (qname q) hosts
         return $ responseA ident q ip
-    hosts = [ ("proxy.com.", "127.0.0.1")
-            --, ("*.proxy.com", "127.0.0.1")
-            ]
 
 handlePacket :: Conf -> Socket -> SockAddr -> S.ByteString -> IO ()
 handlePacket conf@Conf{..} sock addr bs = case decode (fromChunks [bs]) of
     Right req -> do
-        print req
         let rc = defaultResolvConf { resolvInfo = RCHostName realDNS }
         mrsp <- handleRequest conf rc req
-        print mrsp
         case mrsp of
             Just rsp ->
                 let packet = mconcat . toChunks $ encode rsp
-                in  timeout' "send timeout" timeOut (sendAllTo sock packet addr) >>
-                    print (S.length packet) >>
-                    return ()
+                in void $ timeout' "send timeout" timeOut (sendAllTo sock packet addr)
             Nothing -> return ()
     Left msg -> putStrLn msg
 
-- 
cgit v1.2.3