module Network.KRPC.Protocol
(
KError(..)
, ErrorCode
, errorCode
, mkKError
, KQuery(queryMethod, queryArgs)
, MethodName
, ParamName
, kquery
, KResponse(respVals)
, ValName
, kresponse
, sendMessage
, recvResponse
, KRemote
, KRemoteAddr
, withRemote
, remoteServer
, encode
, encoded
, decode
, decoded
, toBEncode
, fromBEncode
) 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
import Data.ByteString as B
import Data.ByteString.Char8 as BC
import qualified Data.ByteString.Lazy as LB
import Data.Map as M
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)
instance BEncode KError where
toBEncode e = fromAscAssocs
[ "e" --> (errorCode e, errorMessage e)
, "y" --> ("e" :: ByteString)
]
fromBEncode (BDict d)
| M.lookup "y" d == Just (BString "e")
= uncurry mkKError <$> d >-- "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 :: Map ParamName BValue
} deriving (Show, Read, Eq, Ord)
instance BEncode KQuery where
toBEncode (KQuery m args) = fromAscAssocs
[ "a" --> BDict args
, "q" --> m
, "y" --> ("q" :: ByteString)
]
fromBEncode (BDict d)
| M.lookup "y" d == Just (BString "q") =
KQuery <$> d >-- "q"
<*> d >-- "a"
fromBEncode _ = decodingError "KQuery"
kquery :: MethodName -> [(ParamName, BValue)] -> KQuery
kquery name args = KQuery name (M.fromList args)
type ValName = ByteString
newtype KResponse = KResponse { respVals :: BDict }
deriving (Show, Read, Eq, Ord)
instance BEncode KResponse where
toBEncode (KResponse vals) = fromAscAssocs
[ "r" --> vals
, "y" --> ("r" :: ByteString)
]
fromBEncode (BDict d)
| M.lookup "y" d == Just (BString "r") =
KResponse <$> d >-- "r"
fromBEncode _ = decodingError "KDict"
kresponse :: [(ValName, BValue)] -> KResponse
kresponse = KResponse . M.fromList
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 (encoded msg)) addr
recvResponse :: KRemote -> IO (Either KError KResponse)
recvResponse sock = do
(raw, _) <- recvFrom sock maxMsgSize
return $ case decoded raw of
Right resp -> Right resp
Left decE -> Left $ case decoded 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 decoded bs of
Right query -> (either toBEncode toBEncode <$> action addr query)
`Lifted.catch` (return . toBEncode . serverError)
Left decodeE -> return $ toBEncode (ProtocolError (BC.pack decodeE))