/Node.hs
Haskell | 122 lines | 105 code | 15 blank | 2 comment | 0 complexity | 11c4a8a45cb495947ec01ba12aba977d MD5 | raw file
1{-# LANGUAGE FlexibleInstances #-} 2module Node where 3 4import Control.Concurrent 5import qualified Data.ByteString.Lazy.Char8 as B8 6import qualified Data.ByteString.Char8 as SB8 7import Network.Socket hiding (send, sendTo, recv, recvFrom) 8import Network.Socket.ByteString 9import Control.Concurrent.Chan 10import Control.Monad 11import Data.IORef 12import qualified Network.Libev as Ev 13import Data.Bits ((.&.)) 14 15import InState 16import KRPC 17import BEncoding (BValue) 18 19 20data NodeState = State { stSock :: Socket, 21 stQueryHandler :: QueryHandler, 22 stReplyHandler :: ReplyHandler, 23 stLastT :: T 24 } 25type Node = IORef NodeState 26type QueryHandler = SockAddr -> BValue -> Query -> IO (Either Error Reply) 27type ReplyHandler = SockAddr -> BValue -> Reply -> IO () 28 29 30new :: Ev.EvLoopPtr -> Int -> IO Node 31new evLoop port 32 = do sock <- socket AF_INET Datagram defaultProtocol 33 bindSocket sock (SockAddrInet (fromIntegral port) 0) 34 node <- newIORef State { stSock = sock, 35 stQueryHandler = nullQueryHandler, 36 stReplyHandler = nullReplyHandler, 37 stLastT = T B8.empty 38 } 39 let callback evLoop evIo evType 40 = when (evType .&. Ev.ev_read /= 0) $ 41 runOnce node 42 fd = fromIntegral $ fdSocket sock 43 44 evIoCb <- Ev.mkIoCallback callback 45 evIo <- Ev.mkEvIo 46 Ev.evIoInit evIo evIoCb fd Ev.ev_read 47 Ev.evIoStart evLoop evIo 48 return node 49 50nullQueryHandler _ _ _ 51 = return $ Left $ Error 201 $ B8.pack "Not implemented" 52nullReplyHandler _ _ _ 53 = return () 54 55setQueryHandler :: QueryHandler -> Node -> IO () 56setQueryHandler cb = inState $ \st -> st { stQueryHandler = cb } 57setReplyHandler :: ReplyHandler -> Node -> IO () 58setReplyHandler cb = inState $ \st -> st { stReplyHandler = cb } 59 60run :: Node -> IO () 61run node = do runOnce node 62 run node 63 64runOnce :: Node -> IO () 65runOnce node = do sock <- stSock `liftM` readIORef node 66 (buf, addr) <- recvFrom sock 1024 67 let handle st = catch (handlePacket st buf addr) $ \e -> 68 do putStrLn $ "Error handling packet: " ++ show e 69 return st 70 inState handle node 71 72 73handlePacket :: NodeState -> SB8.ByteString -> SockAddr -> IO NodeState 74handlePacket st buf addr 75 = do let errorOrPkt = decodePacket buf 76 case errorOrPkt of 77 Right (bval, pkt) -> 78 do --putStrLn $ "Received from " ++ show addr ++ ": " ++ show pkt 79 let isQuery = case pkt of 80 (QPacket _ _) -> True 81 _ -> False 82 case isQuery of 83 False -> 84 case pkt of 85 RPacket _ reply -> 86 do catch (stReplyHandler st addr bval reply) print 87 return st 88 _ -> return st 89 True -> 90 do let QPacket t qry = pkt 91 qRes <- stQueryHandler st addr bval qry 92 let pkt = case qRes of 93 Left e -> EPacket t e 94 Right r -> RPacket t r 95 buf = SB8.concat $ B8.toChunks $ encodePacket pkt 96 --putStrLn $ "Replying " ++ show pkt ++ " to " ++ show addr 97 catch (sendTo (stSock st) buf addr >> 98 return () 99 ) $ \_ -> return () 100 return st 101 Left e -> do putStrLn e 102 return st 103 104sendQueryNoWait :: SockAddr -> Query -> Node -> IO () 105sendQueryNoWait addr qry 106 = inState $ \st -> 107 do let t = tSucc $ stLastT st 108 pkt = QPacket t qry 109 buf = SB8.concat $ B8.toChunks $ encodePacket pkt 110 putStrLn $ "Sending to " ++ show addr ++ ": " ++ show pkt 111 catch (sendTo (stSock st) buf addr >> 112 return () 113 ) $ \a -> return () 114 return st { stLastT = t } 115 116getAddrs :: String -> String -> IO [SockAddr] 117getAddrs name serv 118 = do putStrLn $ "Resolving " ++ show (name, serv) 119 liftM (map addrAddress . 120 filter ((== AF_INET) . addrFamily)) $ 121 getAddrInfo (Just defaultHints { addrFamily = AF_INET }) 122 (Just name) (Just serv)