{-# LANGUAGE DeriveDataTypeable #-} {-# LANGUAGE DeriveFoldable #-} {-# LANGUAGE DeriveFunctor #-} {-# LANGUAGE DeriveTraversable #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE LambdaCase #-} {-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE StandaloneDeriving #-} module Mainline where import Control.Applicative import Control.Arrow import Control.Concurrent.STM import Crypto.Random import Data.BEncode as BE import qualified Data.BEncode.BDict as BE ;import Data.BEncode.BDict (BKey) import Data.Bits import Data.Bits.ByteString import Data.Bool import qualified Data.ByteArray as BA ;import Data.ByteArray (ByteArrayAccess) import qualified Data.ByteString as B ;import Data.ByteString (ByteString) import Data.ByteString.Lazy (toStrict) import Data.Coerce import Data.Data import Data.Default import Data.Hashable import Data.IP import Data.List import Data.Maybe import Data.Monoid import qualified Data.Serialize as S import Data.Set (Set) import Data.Torrent import Data.Typeable import Data.Word import Network.Address (Address, fromSockAddr, setPort, sockAddrPort, toSockAddr) import Network.BitTorrent.DHT.ContactInfo as Peers import Network.BitTorrent.DHT.Token as Token import qualified Network.DHT.Routing as R ;import Network.DHT.Routing (Info, Timestamp, getTimestamp) import Network.QueryResponse import Network.Socket newtype NodeId = NodeId ByteString deriving (Eq,Ord,Show,ByteArrayAccess, BEncode, Bits) instance FiniteBits NodeId where finiteBitSize _ = 160 data NodeInfo = NodeInfo { nodeId :: NodeId , nodeIP :: IP , nodePort :: PortNumber } instance Hashable NodeInfo where hashWithSalt s ni = hashWithSalt s (nodeIP ni , nodePort ni) {-# INLINE hashWithSalt #-} {- -- | KRPC 'compact list' compatible encoding: contact information for -- nodes is encoded as a 26-byte string. Also known as "Compact node -- info" the 20-byte Node ID in network byte order has the compact -- IP-address/port info concatenated to the end. get = NodeInfo <$> (NodeId <$> S.getBytes 20 ) <*> S.get <*> S.get -} getNodeInfo4 :: S.Get NodeInfo getNodeInfo4 = NodeInfo <$> (NodeId <$> S.getBytes 20) <*> (IPv4 <$> S.get) <*> S.get putNodeInfo4 :: NodeInfo -> S.Put putNodeInfo4 (NodeInfo (NodeId nid) ip port) | IPv4 ip4 <- ip = put4 ip4 | IPv6 ip6 <- ip , Just ip4 <- un4map ip6 = put4 ip4 | otherwise = return () where put4 ip4 = S.putByteString nid >> S.put ip4 >> S.put port getNodeInfo6 :: S.Get NodeInfo getNodeInfo6 = NodeInfo <$> (NodeId <$> S.getBytes 20) <*> (IPv6 <$> S.get) <*> S.get putNodeInfo6 :: NodeInfo -> S.Put putNodeInfo6 (NodeInfo (NodeId nid) (IPv6 ip) port) = S.putByteString nid >> S.put ip >> S.put port putNodeInfo6 _ = return () nodeAddr :: NodeInfo -> SockAddr nodeAddr (NodeInfo _ ip port) = setPort port $ toSockAddr ip nodeInfo :: NodeId -> SockAddr -> Either String NodeInfo nodeInfo nid saddr | Just ip <- fromSockAddr saddr , Just port <- sockAddrPort saddr = Right $ NodeInfo nid ip port | otherwise = Left "Address family not supported." -- | Types of RPC errors. data ErrorCode -- | Some error doesn't fit in any other category. = GenericError -- | Occurs when server fail to process procedure call. | ServerError -- | Malformed packet, invalid arguments or bad token. | ProtocolError -- | Occurs when client trying to call method server don't know. | MethodUnknown deriving (Show, Read, Eq, Ord, Bounded, Typeable, Data) -- | According to the table: -- instance Enum ErrorCode where fromEnum GenericError = 201 fromEnum ServerError = 202 fromEnum ProtocolError = 203 fromEnum MethodUnknown = 204 {-# INLINE fromEnum #-} toEnum 201 = GenericError toEnum 202 = ServerError toEnum 203 = ProtocolError toEnum 204 = MethodUnknown toEnum _ = GenericError {-# INLINE toEnum #-} instance BEncode ErrorCode where toBEncode = toBEncode . fromEnum {-# INLINE toBEncode #-} fromBEncode b = toEnum <$> fromBEncode b {-# INLINE fromBEncode #-} data Error = Error { errorCode :: !ErrorCode -- ^ The type of error. , errorMessage :: !ByteString -- ^ Human-readable text message. } deriving ( Show, Eq, Ord, Typeable, Data, Read ) newtype TransactionId = TransactionId ByteString deriving (Eq, Ord, Show, BEncode) newtype Method = Method ByteString deriving (Eq, Ord, Show, BEncode) data Message a = Q { msgOrigin :: NodeId , msgID :: TransactionId , qryPayload :: a , qryMethod :: Method , qryReadOnly :: Bool } | R { msgOrigin :: NodeId , msgID :: TransactionId , rspPayload :: Either Error a , rspReflectedIP :: Maybe SockAddr } instance BE.BEncode (Message BValue) where toBEncode = encodeMessage fromBEncode = error "fromBEncode" encodeMessage (Q origin tid a meth ro) = case a of BDict args -> encodeQuery tid meth (BDict $ genericArgs origin ro `BE.union` args) _ -> encodeQuery tid meth a -- XXX: Not really a valid query. encodeMessage (R origin tid v ip) = case v of Right vals -> encodeResponse tid vals (BString . encodeAddr <$> ip) Left err -> encodeError tid err encodeAddr :: SockAddr -> ByteString encodeAddr (SockAddrInet port addr) = S.runPut (S.putWord32host addr >> S.put (fromIntegral port :: Word16)) encodeAddr (SockAddrInet6 port _ addr _) = S.runPut (S.put addr >> S.put (fromIntegral port :: Word16)) encodeAddr _ = B.empty genericArgs nodeid ro = "id" .=! nodeid .: "ro" .=? bool Nothing (Just (1 :: Int)) ro .: endDict encodeError tid (Error ecode emsg) = encodeAny tid "e" (ecode,emsg) id encodeResponse tid rvals rip = encodeAny tid "r" rvals ("ip" .=? rip .:) encodeQuery tid qmeth qargs = encodeAny tid "q" qmeth ("a" .=! qargs .:) encodeAny tid key val aux = toDict $ aux $ key .=! val .: "t" .=! tid .: "y" .=! key .: endDict parsePacket :: ByteString -> SockAddr -> Either String (Message BValue, NodeInfo) parsePacket bs addr = do pkt <- BE.decode bs ni <- nodeInfo (msgOrigin pkt) addr return (pkt, ni) encodePacket :: Message BValue -> NodeInfo -> (ByteString, SockAddr) encodePacket msg ni = ( toStrict $ BE.encode msg , nodeAddr ni ) classify :: Message BValue -> MessageClass String Method TransactionId classify (Q { msgID = tid, qryMethod = meth }) = IsQuery meth tid classify (R { msgID = tid }) = IsResponse tid encodePayload :: BEncode a => TransactionId -> NodeInfo -> NodeInfo -> a -> Message BValue encodePayload tid self dest b = R (nodeId self) tid (Right $ BE.toBEncode b) (Just $ nodeAddr dest) errorPayload :: TransactionId -> NodeInfo -> NodeInfo -> Error -> Message a errorPayload tid self dest e = R (nodeId self) tid (Left e) (Just $ nodeAddr dest) decodePayload :: BEncode a => Message BValue -> Either String a decodePayload msg = BE.fromBEncode $ qryPayload msg type Handler = MethodHandler String TransactionId NodeInfo (Message BValue) handler :: ( BEncode a , BEncode b ) => (NodeInfo -> a -> IO b) -> Maybe Handler handler f = Just $ MethodHandler decodePayload encodePayload f handlerE :: ( BEncode a , BEncode b ) => (NodeInfo -> a -> IO (Either Error b)) -> Maybe Handler handlerE f = Just $ MethodHandler decodePayload enc f where enc tid self dest (Left e) = errorPayload tid self dest e enc tid self dest (Right b) = encodePayload tid self dest b type AnnounceSet = Set (InfoHash, PortNumber) data SwarmsDatabase = SwarmsDatabase { contactInfo :: !( TVar PeerStore ) -- ^ Published by other nodes. , sessionTokens :: !( TVar SessionTokens ) -- ^ Query session IDs. , announceInfo :: !( TVar AnnounceSet ) -- ^ To publish by this node. } newSwarmsDatabase :: IO SwarmsDatabase newSwarmsDatabase = do toks <- nullSessionTokens atomically $ SwarmsDatabase <$> newTVar def <*> newTVar toks <*> newTVar def type RoutingInfo = Info NodeInfo NodeId data Routing = Routing { tentativeId :: NodeId , routing4 :: !( TVar (Maybe RoutingInfo) ) , routing6 :: !( TVar (Maybe RoutingInfo) ) } newClient :: SockAddr -> IO (Client String Method TransactionId NodeInfo (Message BValue)) newClient addr = do udp <- udpTransport addr nid <- NodeId <$> getRandomBytes 20 self <- atomically $ newTVar $ NodeInfo nid (fromMaybe (toEnum 0) $ fromSockAddr addr) (fromMaybe 0 $ sockAddrPort addr) routing <- atomically $ Routing nid <$> newTVar Nothing <*> newTVar Nothing swarms <- newSwarmsDatabase let net = onInbound (updateRouting routing) $ layerTransport parsePacket encodePacket $ udp dispatch tbl = DispatchMethods { classifyInbound = classify , lookupHandler = handlers , tableMethods = tbl } handlers :: Method -> Maybe Handler handlers ( Method "ping" ) = handler pingH handlers ( Method "find_node" ) = handler $ findNodeH routing handlers ( Method "get_peers" ) = handler $ getPeersH routing swarms handlers ( Method "announce_peer" ) = handlerE $ announceH swarms handlers ( Method meth ) = Just $ defaultHandler meth mapT = transactionMethods mapMethods gen gen :: Word16 -> (TransactionId, Word16) gen cnt = (TransactionId $ S.encode cnt, cnt+1) map_var <- atomically $ newTVar (0, mempty) return Client { clientNet = net , clientDispatcher = dispatch mapT , clientErrorReporter = ignoreErrors -- TODO , clientPending = map_var , clientAddress = atomically (readTVar self) , clientResponseId = return } defaultHandler :: ByteString -> Handler defaultHandler meth = MethodHandler decodePayload errorPayload returnError where returnError :: NodeInfo -> BValue -> IO Error returnError _ _ = return $ Error MethodUnknown ("Unknown method " <> meth) updateRouting :: Routing -> NodeInfo -> Message BValue -> IO () updateRouting routing naddr _ = do -- TODO Update kademlia table. -- TODO Update external ip address and update BEP-42 node id. return () data Ping = Ping deriving Show -- Pong is the same as Ping. type Pong = Ping pattern Pong = Ping instance BEncode Ping where toBEncode Ping = toDict endDict fromBEncode _ = pure Ping data WantIP = Want_IP4 | Want_IP6 | Want_Both deriving (Eq, Enum, Ord, Show) wantList :: WantIP -> [ByteString] wantList Want_IP4 = ["ip4"] wantList Want_IP6 = ["ip6"] wantList Want_Both = ["ip4","ip6"] instance BEncode WantIP where toBEncode w = toBEncode $ wantList w fromBEncode bval = do wants <- fromBEncode bval let _ = wants :: [ByteString] case (elem "ip4" wants, elem "ip6" wants) of (True,True) -> Right Want_Both (True,False) -> Right Want_IP4 (False,True) -> Right Want_IP6 _ -> Left "Unrecognized IP type." data FindNode = FindNode NodeId (Maybe WantIP) instance BEncode FindNode where toBEncode (FindNode nid iptyp) = toDict $ target_key .=! nid .: want_key .=? iptyp .: endDict fromBEncode = fromDict $ FindNode <$>! target_key <*>? want_key data NodeFound = NodeFound { nodes4 :: [NodeInfo] , nodes6 :: [NodeInfo] } instance BEncode NodeFound where toBEncode (NodeFound ns ns6) = toDict $ nodes_key .=? (if Prelude.null ns then Nothing else Just (S.runPut (mapM_ putNodeInfo4 ns))) .: nodes6_key .=? (if Prelude.null ns6 then Nothing else Just (S.runPut (mapM_ putNodeInfo6 ns6))) .: endDict fromBEncode bval = NodeFound <$> ns4 <*> ns6 where ns4 = fromDict (binary getNodeInfo4 nodes_key) bval ns6 = fromDict (binary getNodeInfo6 nodes6_key) bval binary :: S.Get a -> BKey -> BE.Get [a] binary get k = field (req k) >>= either (fail . format) return . S.runGet (many get) where format str = "fail to deserialize " ++ show k ++ " field: " ++ str pingH :: NodeInfo -> Ping -> IO Pong pingH _ Ping = return Pong -- | True if the argument is an IPv4-mapped address with prefix ::FFFF:0:0/96 -- as defined in RFC 4291. is4mapped :: IPv6 -> Bool is4mapped ip | [0,0,0,0,0,0xffff,_,_] <- fromIPv6 ip = True | otherwise = False un4map :: IPv6 -> Maybe IPv4 un4map ip | [0,0,0,0,0,0xffff,x,y] <- fromIPv6 ip = Just $ toIPv4 $ map (.&. 0xFF) [x `shiftR` 8, x, y `shiftR` 8, y ] | otherwise = Nothing prefer4or6 :: NodeInfo -> Maybe WantIP -> WantIP prefer4or6 addr iptyp = fromMaybe (ipFamily $ nodeIP addr) iptyp ipFamily :: IP -> WantIP ipFamily ip = case ip of IPv4 _ -> Want_IP4 IPv6 a | is4mapped a -> Want_IP4 | otherwise -> Want_IP6 findNodeH :: Routing -> NodeInfo -> FindNode -> IO NodeFound findNodeH routing addr (FindNode node iptyp) = do let preferred = prefer4or6 addr iptyp ks <- bool (return []) (go $ routing4 routing) (preferred /= Want_IP6) ks6 <- bool (return []) (go $ routing6 routing) (preferred /= Want_IP4) return $ NodeFound ks ks6 where go var = do let myid = tentativeId routing :: NodeId k = R.defaultK :: Int nobkts = R.defaultBucketCount :: Int nfo <- atomically $ readTVar var let tbl = maybe (R.nullTable myid nobkts) R.myBuckets nfo return $ R.kclosest nodeId k node tbl data GetPeers = GetPeers InfoHash (Maybe WantIP) instance BEncode GetPeers where toBEncode (GetPeers ih iptyp) = toDict $ info_hash_key .=! ih .: want_key .=? iptyp .: endDict fromBEncode = fromDict $ GetPeers <$>! info_hash_key <*>? want_key data GotPeers = GotPeers { -- | If the queried node has no peers for the infohash, returned -- the K nodes in the queried nodes routing table closest to the -- infohash supplied in the query. peers :: [PeerAddr] , nodes :: NodeFound -- | The token value is a required argument for a future -- announce_peer query. , grantedToken :: Token } -- deriving (Show, Eq, Typeable) nodeIsIPv6 :: NodeInfo -> Bool nodeIsIPv6 (NodeInfo _ (IPv6 _) _) = True nodeIsIPv6 _ = False instance BEncode GotPeers where toBEncode GotPeers { nodes = NodeFound ns4 ns6, ..} = toDict $ nodes_key .=? (if null ns4 then Nothing else Just $ S.runPut (mapM_ putNodeInfo4 ns4)) .: nodes6_key .=? (if null ns6 then Nothing else Just $ S.runPut (mapM_ putNodeInfo4 ns6)) .: token_key .=! grantedToken .: peers_key .=! map S.encode peers .: endDict fromBEncode = fromDict $ do ns4 <- fromMaybe [] <$> optional (binary getNodeInfo4 nodes_key) -- "nodes" ns6 <- fromMaybe [] <$> optional (binary getNodeInfo6 nodes6_key) -- "nodes6" tok <- field (req token_key) -- "token" ps <- fromMaybe [] <$> optional (field (req peers_key) >>= decodePeers) -- "values" pure $ GotPeers ps (NodeFound ns4 ns6) tok where decodePeers = either fail pure . mapM S.decode getPeersH :: Routing -> SwarmsDatabase -> NodeInfo -> GetPeers -> IO GotPeers getPeersH routing (SwarmsDatabase peers toks _) naddr (GetPeers ih iptyp) = do ps <- do tm <- getTimestamp atomically $ do (ps,store') <- Peers.freshPeers ih tm <$> readTVar peers writeTVar peers store' return ps -- Filter peer results to only a single address family, IPv4 or IPv6, as -- per BEP 32. let notboth = iptyp >>= \case Want_Both -> Nothing specific -> Just specific selected = prefer4or6 naddr notboth ps' = filter ( (== selected) . ipFamily . peerHost ) ps tok <- grantToken toks naddr ns <- findNodeH routing naddr (FindNode (coerce ih) iptyp) return $ GotPeers ps' ns tok -- | Announce that the peer, controlling the querying node, is -- downloading a torrent on a port. data Announce = Announce { -- | If set, the 'port' field should be ignored and the source -- port of the UDP packet should be used as the peer's port -- instead. This is useful for peers behind a NAT that may not -- know their external port, and supporting uTP, they accept -- incoming connections on the same port as the DHT port. impliedPort :: Bool -- | infohash of the torrent; , topic :: InfoHash -- | some clients announce the friendly name of the torrent here. , announcedName :: Maybe ByteString -- | the port /this/ peer is listening; , port :: PortNumber -- TODO: optional boolean "seed" key -- | received in response to a previous get_peers query. , sessionToken :: Token } deriving (Show, Eq, Typeable) peer_ip_key = "ip" peer_id_key = "peer id" peer_port_key = "port" msg_type_key = "msg_type" piece_key = "piece" total_size_key = "total_size" node_id_key :: BKey node_id_key = "id" read_only_key :: BKey read_only_key = "ro" want_key :: BKey want_key = "want" target_key :: BKey target_key = "target" nodes_key :: BKey nodes_key = "nodes" nodes6_key :: BKey nodes6_key = "nodes6" info_hash_key :: BKey info_hash_key = "info_hash" peers_key :: BKey peers_key = "values" token_key :: BKey token_key = "token" name_key :: BKey name_key = "name" port_key :: BKey port_key = "port" implied_port_key :: BKey implied_port_key = "implied_port" instance BEncode Announce where toBEncode Announce {..} = toDict $ implied_port_key .=? flagField impliedPort .: info_hash_key .=! topic .: name_key .=? announcedName .: port_key .=! port .: token_key .=! sessionToken .: endDict where flagField flag = if flag then Just (1 :: Int) else Nothing fromBEncode = fromDict $ do Announce <$> (boolField <$> optional (field (req implied_port_key))) <*>! info_hash_key <*>? name_key <*>! port_key <*>! token_key where boolField = maybe False (/= (0 :: Int)) -- | The queried node must verify that the token was previously sent -- to the same IP address as the querying node. Then the queried node -- should store the IP address of the querying node and the supplied -- port number under the infohash in its store of peer contact -- information. data Announced = Announced deriving (Show, Eq, Typeable) instance BEncode Announced where toBEncode _ = toBEncode Ping fromBEncode _ = pure Announced announceH :: SwarmsDatabase -> NodeInfo -> Announce -> IO (Either Error Announced) announceH (SwarmsDatabase peers toks _) naddr announcement = do checkToken toks naddr (sessionToken announcement) >>= bool (Left <$> return (Error ProtocolError "invalid parameter: token")) (Right <$> go) where go = atomically $ do modifyTVar' peers $ insertPeer (topic announcement) (announcedName announcement) $ PeerAddr { peerId = Nothing -- Avoid storing IPv4-mapped addresses. , peerHost = case nodeIP naddr of IPv6 ip6 | Just ip4 <- un4map ip6 -> IPv4 ip4 a -> a , peerPort = if impliedPort announcement then nodePort naddr else port announcement } return Announced