8 files changed, 29 insertions, 25 deletions

diff --git a/README.md b/README.md
index ccbd6789..189bda04 100644
--- a/README.md
+++ b/README.md
@@ -13,6 +13,10 @@ language, thus it's hard to shoot yourself in the foot accidently.
 See bittorrent DHT [specification][spec] for detailed protocol
 description.
 
+### Example
+
+TODO
+
 #### Modules
 
 * Remote.KRPC — simple interface which reduce all RPC related stuff to
diff --git a/bench/Main.hs b/bench/Main.hs
index fdf76cc2..024d4d93 100644
--- a/bench/Main.hs
+++ b/bench/Main.hs
@@ -6,10 +6,11 @@ import Data.ByteString (ByteString)
 import qualified Data.ByteString as B
 import Criterion.Main
 import Network.KRPC
+import Network.Socket
 
 
 addr :: RemoteAddr
-addr = (0, 6000)
+addr = SockAddrInet 6000 0
 
 echo :: Method ByteString ByteString
 echo = method "echo" ["x"] ["x"]
diff --git a/bench/Server.hs b/bench/Server.hs
index 444362c1..ef20c08a 100644
--- a/bench/Server.hs
+++ b/bench/Server.hs
@@ -3,10 +3,11 @@ module Main (main) where
 
 import Data.ByteString (ByteString)
 import Network.KRPC
+import Network.Socket
 
 
 echo :: Method ByteString ByteString
 echo = method "echo" ["x"] ["x"]
 
 main :: IO ()
-main = server 6000 [ echo ==> return ]
+main = server (SockAddrInet 6000 0) [ echo ==> return ]
diff --git a/krpc.cabal b/krpc.cabal
index d6995410..9929cea7 100644
--- a/krpc.cabal
+++ b/krpc.cabal
@@ -69,6 +69,7 @@ test-suite test-client
 
                      , bencoding
                      , krpc
+                     , network
 
                      , HUnit
                      , test-framework
@@ -84,6 +85,7 @@ executable test-server
                      , bytestring
                      , bencoding
                      , krpc
+                     , network
 
 executable bench-server
   default-language:    Haskell2010
@@ -92,6 +94,7 @@ executable bench-server
   build-depends:       base        == 4.*
                      , bytestring
                      , krpc
+                     , network
   ghc-options:         -fforce-recomp
 
 benchmark bench-client
@@ -103,4 +106,5 @@ benchmark bench-client
                      , bytestring
                      , criterion
                      , krpc
+                     , network
   ghc-options:         -O2 -fforce-recomp
\ No newline at end of file
diff --git a/src/Network/KRPC.hs b/src/Network/KRPC.hs
index e667853a..3c9f9bee 100644
--- a/src/Network/KRPC.hs
+++ b/src/Network/KRPC.hs
@@ -97,7 +97,8 @@
 module Network.KRPC
        ( -- * Method
          Method(..)
-       , method, idM
+       , method
+       , idM
 
          -- * Client
        , RemoteAddr
@@ -349,11 +350,11 @@ infix 1 ==>@
 --   it will not create new thread for each connection.
 --
 server :: (MonadBaseControl IO remote, MonadIO remote)
-       => PortNumber              -- ^ Port used to accept incoming connections.
+       => KRemoteAddr             -- ^ Port used to accept incoming connections.
        -> [MethodHandler remote]  -- ^ Method table.
        -> remote ()
-server servport handlers = do
-    remoteServer servport $ \addr q -> do
+server servAddr handlers = do
+    remoteServer servAddr $ \addr q -> do
       case dispatch (queryMethod q) of
         Nothing -> return $ Left $ MethodUnknown (queryMethod q)
         Just  m -> m addr q
diff --git a/src/Network/KRPC/Protocol.hs b/src/Network/KRPC/Protocol.hs
index ad1dabca..2d905f06 100644
--- a/src/Network/KRPC/Protocol.hs
+++ b/src/Network/KRPC/Protocol.hs
@@ -202,10 +202,7 @@ kresponse :: [(ValName, BValue)] -> KResponse
 kresponse = KResponse . M.fromList
 {-# INLINE kresponse #-}
 
-
-
-type KRemoteAddr = (HostAddress, PortNumber)
-
+type KRemoteAddr = SockAddr
 type KRemote = Socket
 
 withRemote :: (MonadBaseControl IO m, MonadIO m) => (KRemote -> m a) -> m a
@@ -224,8 +221,7 @@ maxMsgSize = 64 * 1024 -- max udp size
 
 -- TODO eliminate toStrict
 sendMessage :: BEncode msg => msg -> KRemoteAddr -> KRemote -> IO ()
-sendMessage msg (host, port) sock =
-  sendAllTo sock (LB.toStrict (encoded msg)) (SockAddrInet port host)
+sendMessage msg addr sock = sendAllTo sock (LB.toStrict (encoded msg)) addr
 {-# INLINE sendMessage #-}
 
 recvResponse :: KRemote -> IO (Either KError KResponse)
@@ -239,26 +235,21 @@ recvResponse sock = do
 
 -- | Run server using a given port. Method invocation should be done manually.
 remoteServer :: (MonadBaseControl IO remote, MonadIO remote)
-             => PortNumber -- ^ Port number to listen.
+             => KRemoteAddr -- ^ Port number to listen.
              -> (KRemoteAddr -> KQuery -> remote (Either KError KResponse))
              -- ^ Handler.
              -> remote ()
-remoteServer servport action = bracket (liftIO bindServ) (liftIO . sClose) loop
+remoteServer servAddr action = bracket (liftIO bindServ) (liftIO . sClose) loop
   where
     bindServ = do
      sock <- socket AF_INET Datagram defaultProtocol
-     bindSocket sock (SockAddrInet servport iNADDR_ANY)
+     bindSocket sock servAddr
      return sock
 
     loop sock = forever $ do
-      (bs, addr) <- liftIO $ recvFrom sock maxMsgSize
-      case addr of
-        SockAddrInet port host -> do
-           let kaddr = (host, port)
-           reply <- handleMsg bs kaddr
-           liftIO $ sendMessage reply kaddr sock
-        _ -> return ()
-
+        (bs, addr) <- liftIO $ recvFrom sock maxMsgSize
+        reply <- handleMsg bs addr
+        liftIO $ sendMessage reply addr sock
       where
         handleMsg bs addr = case decoded bs of
           Right query -> (either toBEncode toBEncode <$> action addr query)
diff --git a/tests/Client.hs b/tests/Client.hs
index db7a3219..cda01631 100644
--- a/tests/Client.hs
+++ b/tests/Client.hs
@@ -15,11 +15,12 @@ import Test.Framework
 import Test.Framework.Providers.HUnit
 
 import Network.KRPC
+import Network.Socket
 import Shared
 
 
 addr :: RemoteAddr
-addr = (0, 6000)
+addr = SockAddrInet 6000 0
 
 withServ :: FilePath -> IO () -> IO ()
 withServ serv_path = bracket up terminateProcess . const
diff --git a/tests/Server.hs b/tests/Server.hs
index 9e70b70b..b4b34891 100644
--- a/tests/Server.hs
+++ b/tests/Server.hs
@@ -3,11 +3,12 @@ module Main (main) where
 
 import Data.BEncode
 import Network.KRPC
+import Network.Socket
 import Shared
 
 
 main :: IO ()
-main = server 6000
+main = server (SockAddrInet 6000 0)
   [ unitM ==> return
   , echoM ==> return
   , echoBytes ==> return