module Network.KRPC.Protocol
(
KError(..)
, ErrorCode
, errorCode
, mkKError
, KQuery(queryMethod, queryArgs)
, MethodName
, ParamName
, kquery
, KResponse(respVals)
, ValName
, kresponse
, sendMessage
, recvResponse
, KRemote
, KRemoteAddr
, 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
data KError
= GenericError { errorMessage :: !ByteString }
| ServerError { errorMessage :: !ByteString }
| ProtocolError { errorMessage :: !ByteString }
| MethodUnknown { errorMessage :: !ByteString }
deriving (Show, Read, Eq, Ord, Typeable)
instance BEncode KError where
toBEncode e = toDict $
"e" .=! (errorCode e, errorMessage e)
.: "y" .=! ("e" :: ByteString)
.: endDict
fromBEncode be @ (BDict d)
| BE.lookup "y" d == Just (BString "e")
= (`fromDict` be) $ do
uncurry mkKError <$>! "e"
fromBEncode _ = decodingError "KError"
type ErrorCode = Int
errorCode :: KError -> ErrorCode
errorCode (GenericError _) = 201
errorCode (ServerError _) = 202
errorCode (ProtocolError _) = 203
errorCode (MethodUnknown _) = 204
mkKError :: ErrorCode -> ByteString -> KError
mkKError 201 = GenericError
mkKError 202 = ServerError
mkKError 203 = ProtocolError
mkKError 204 = MethodUnknown
mkKError _ = GenericError
serverError :: SomeException -> KError
serverError = ServerError . BC.pack . show
type MethodName = ByteString
type ParamName = ByteString
data KQuery = KQuery {
queryMethod :: !MethodName
, queryArgs :: BDict
} deriving (Show, Read, Eq, Ord, Typeable)
instance BEncode KQuery where
toBEncode (KQuery m args) = toDict $
"a" .=! BDict args
.: "q" .=! m
.: "y" .=! ("q" :: ByteString)
.: endDict
fromBEncode bv @ (BDict d)
| BE.lookup "y" d == Just (BString "q") = (`fromDict` bv) $ do
a <- field (req "a")
q <- field (req "q")
return $! KQuery q a
fromBEncode _ = decodingError "KQuery"
kquery :: MethodName -> BDict -> KQuery
kquery = KQuery
type ValName = ByteString
newtype KResponse = KResponse { respVals :: BDict }
deriving (Show, Read, Eq, Ord, Typeable)
instance BEncode KResponse where
toBEncode (KResponse vals) = toDict $
"r" .=! vals
.: "y" .=! ("r" :: ByteString)
.: endDict
fromBEncode bv @ (BDict d)
| BE.lookup "y" d == Just (BString "r") = (`fromDict` bv) $ do
KResponse <$>! "r"
fromBEncode _ = decodingError "KDict"
kresponse :: BDict -> KResponse
kresponse = KResponse
type KRemoteAddr = SockAddr
type KRemote = Socket
sockAddrFamily :: SockAddr -> Family
sockAddrFamily (SockAddrInet _ _ ) = AF_INET
sockAddrFamily (SockAddrInet6 _ _ _ _) = AF_INET6
sockAddrFamily (SockAddrUnix _ ) = AF_UNIX
withRemote :: (MonadBaseControl IO m, MonadIO m) => (KRemote -> m a) -> m a
withRemote = bracket (liftIO (socket AF_INET6 Datagram defaultProtocol))
(liftIO . sClose)
maxMsgSize :: Int
maxMsgSize = 64 * 1024
sendMessage :: BEncode msg => msg -> KRemoteAddr -> KRemote -> IO ()
sendMessage msg addr sock = sendManyTo sock (LB.toChunks (encode msg)) addr
recvResponse :: KRemote -> 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)
remoteServer :: (MonadBaseControl IO remote, MonadIO remote)
=> KRemoteAddr
-> (KRemoteAddr -> 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))