2 files changed, 66 insertions, 81 deletions

diff --git a/src/Network/KRPC.hs b/src/Network/KRPC.hs
index d295a965..802cddc5 100644
--- a/src/Network/KRPC.hs
+++ b/src/Network/KRPC.hs
@@ -114,17 +114,21 @@ module Network.KRPC
        , server
        ) where
 
-import Control.Exception
+import Control.Applicative
+import Control.Exception.Lifted as Lifted
+import Control.Monad
 import Control.Monad.Trans.Control
 import Control.Monad.IO.Class
 import Data.BEncode as BE
 import Data.ByteString.Char8 as BC
+import Data.ByteString.Lazy as BL
 import Data.List as L
 import Data.Monoid
 import Data.String
 import Data.Typeable
 import Network
 import Network.Socket
+import Network.Socket.ByteString as BS
 
 import Network.KRPC.Protocol
 
@@ -166,6 +170,32 @@ showsMethod (Method name) =
     paramsTy = typeOf (impossible :: a)
     valuesTy = typeOf (impossible :: b)
 
+{-----------------------------------------------------------------------
+--  Client
+-----------------------------------------------------------------------}
+
+sendMessage :: BEncode msg => msg -> SockAddr -> Socket -> IO ()
+sendMessage msg addr sock = sendManyTo sock (BL.toChunks (encode msg)) addr
+{-# INLINE sendMessage #-}
+
+maxMsgSize :: Int
+--maxMsgSize = 512 -- release: size of payload of one udp packet
+maxMsgSize = 64 * 1024 -- bench: max UDP MTU
+{-# INLINE maxMsgSize #-}
+
+recvResponse :: Socket -> IO (Either KError KResponse)
+recvResponse sock = do
+  (raw, _) <- BS.recvFrom sock maxMsgSize
+  return $ case decode raw of
+    Right resp -> Right resp
+    Left decE -> Left $ case decode raw of
+      Right kerror -> kerror
+      _ -> ProtocolError (BC.pack decE)
+
+withRemote :: (MonadBaseControl IO m, MonadIO m) => (Socket -> m a) -> m a
+withRemote = bracket (liftIO (socket AF_INET6 Datagram defaultProtocol))
+                     (liftIO .  sClose)
+{-# SPECIALIZE withRemote :: (Socket -> IO a) -> IO a #-}
 
 getResult :: BEncode result => Socket -> IO result
 getResult sock = do
@@ -183,6 +213,10 @@ call addr arg = liftIO $ withRemote $ \sock -> do
   where
     Method name = method :: Method req resp
 
+{-----------------------------------------------------------------------
+--  Server
+-----------------------------------------------------------------------}
+
 type HandlerBody remote = SockAddr -> KQuery -> remote (Either KError KResponse)
 
 -- | Procedure signature and implementation binded up.
@@ -204,6 +238,36 @@ handler body = (name, newbody)
           r <- body addr a
           return (Right (KResponse (toBEncode r)))
 
+sockAddrFamily :: SockAddr -> Family
+sockAddrFamily (SockAddrInet  _ _    ) = AF_INET
+sockAddrFamily (SockAddrInet6 _ _ _ _) = AF_INET6
+sockAddrFamily (SockAddrUnix  _      ) = AF_UNIX
+
+-- | Run server using a given port. Method invocation should be done manually.
+remoteServer :: (MonadBaseControl IO remote, MonadIO remote)
+             => SockAddr -- ^ Port number to listen.
+             -> (SockAddr -> KQuery -> remote (Either KError KResponse))
+             -> remote ()
+remoteServer servAddr action = bracket (liftIO bindServ) (liftIO . sClose) loop
+  where
+    bindServ = do
+        let family = sockAddrFamily servAddr
+        sock <- socket family Datagram defaultProtocol
+        when (family == AF_INET6) $ do
+          setSocketOption sock IPv6Only 0
+        bindSocket sock servAddr
+        return sock
+
+    loop sock = forever $ do
+        (bs, addr) <- liftIO $ BS.recvFrom sock maxMsgSize
+        reply <- handleMsg bs addr
+        liftIO $ sendMessage reply addr sock
+      where
+        handleMsg bs addr = case decode bs of
+          Right query -> (either toBEncode toBEncode <$> action addr query)
+                        `Lifted.catch` (return . toBEncode . serverError)
+          Left decodeE   -> return $ toBEncode (ProtocolError (BC.pack decodeE))
+
 -- | Run RPC server on specified port by using list of handlers.
 --   Server will dispatch procedure specified by callee, but note that
 --   it will not create new thread for each connection.
diff --git a/src/Network/KRPC/Protocol.hs b/src/Network/KRPC/Protocol.hs
index 5b072d79..55bbdf4e 100644
--- a/src/Network/KRPC/Protocol.hs
+++ b/src/Network/KRPC/Protocol.hs
@@ -21,45 +21,23 @@
 module Network.KRPC.Protocol
        ( -- * Error
          KError(..)
-       , ErrorCode
-       , errorCode
-       , mkKError
+       , serverError
 
          -- * Query
        , KQuery(..)
        , MethodName
-       , ParamName
 
          -- * Response
        , KResponse(..)
-       , ValName
-
-       , sendMessage
-       , recvResponse
-
-         -- * Remote
-       , withRemote
-       , remoteServer
        ) where
 
-import Control.Applicative
 import Control.Exception.Lifted as Lifted
-import Control.Monad
-import Control.Monad.IO.Class
-import Control.Monad.Trans.Control
-
 import Data.BEncode as BE
 import Data.BEncode.BDict as BE
-import Data.BEncode.Types as BE
 import Data.ByteString as B
 import Data.ByteString.Char8 as BC
-import qualified Data.ByteString.Lazy as LB
 import Data.Typeable
 
-import Network.Socket hiding (recvFrom)
-import Network.Socket.ByteString
-
-
 -- | Errors used to signal that some error occurred while processing a
 -- procedure call. Error may be send only from server to client but
 -- not in the opposite direction.
@@ -120,9 +98,7 @@ mkKError _   = GenericError
 serverError :: SomeException -> KError
 serverError = ServerError . BC.pack . show
 
-
 type MethodName = ByteString
-type ParamName  = ByteString
 
 -- | Query used to signal that caller want to make procedure call to
 -- callee and pass arguments in. Therefore query may be only sent from
@@ -155,8 +131,6 @@ instance BEncode KQuery where
 
   fromBEncode _ = decodingError "KQuery"
 
-type ValName = ByteString
-
 -- | KResponse used to signal that callee successufully process a
 -- procedure call and to return values from procedure. KResponse should
 -- not be sent if error occurred during RPC. Thus KResponse may be only
@@ -183,56 +157,3 @@ instance BEncode KResponse where
       KResponse <$>! "r"
 
   fromBEncode _ = decodingError "KDict"
-
-sockAddrFamily :: SockAddr -> Family
-sockAddrFamily (SockAddrInet  _ _    ) = AF_INET
-sockAddrFamily (SockAddrInet6 _ _ _ _) = AF_INET6
-sockAddrFamily (SockAddrUnix  _      ) = AF_UNIX
-
-withRemote :: (MonadBaseControl IO m, MonadIO m) => (Socket -> m a) -> m a
-withRemote = bracket (liftIO (socket AF_INET6 Datagram defaultProtocol))
-                     (liftIO .  sClose)
-{-# SPECIALIZE withRemote :: (Socket -> IO a) -> IO a #-}
-
-maxMsgSize :: Int
---maxMsgSize = 512 -- release: size of payload of one udp packet
-maxMsgSize = 64 * 1024 -- bench: max UDP MTU
-{-# INLINE maxMsgSize #-}
-
-sendMessage :: BEncode msg => msg -> SockAddr -> Socket -> IO ()
-sendMessage msg addr sock = sendManyTo sock (LB.toChunks (encode msg)) addr
-{-# INLINE sendMessage #-}
-
-recvResponse :: Socket -> IO (Either KError KResponse)
-recvResponse sock = do
-  (raw, _) <- recvFrom sock maxMsgSize
-  return $ case decode raw of
-    Right resp -> Right resp
-    Left decE -> Left $ case decode raw of
-      Right kerror -> kerror
-      _ -> ProtocolError (BC.pack decE)
-
--- | Run server using a given port. Method invocation should be done manually.
-remoteServer :: (MonadBaseControl IO remote, MonadIO remote)
-             => SockAddr -- ^ Port number to listen.
-             -> (SockAddr -> KQuery -> remote (Either KError KResponse))
-             -> remote ()
-remoteServer servAddr action = bracket (liftIO bindServ) (liftIO . sClose) loop
-  where
-    bindServ = do
-        let family = sockAddrFamily servAddr
-        sock <- socket family Datagram defaultProtocol
-        when (family == AF_INET6) $ do
-          setSocketOption sock IPv6Only 0
-        bindSocket sock servAddr
-        return sock
-
-    loop sock = forever $ do
-        (bs, addr) <- liftIO $ recvFrom sock maxMsgSize
-        reply <- handleMsg bs addr
-        liftIO $ sendMessage reply addr sock
-      where
-        handleMsg bs addr = case decode bs of
-          Right query -> (either toBEncode toBEncode <$> action addr query)
-                        `Lifted.catch` (return . toBEncode . serverError)
-          Left decodeE   -> return $ toBEncode (ProtocolError (BC.pack decodeE))