/Node.hs

  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)