1 {-# LANGUAGE OverloadedStrings #-}
2 module FuncTorrent.Peer
8 import Prelude hiding (lookup, concat, replicate, splitAt)
11 import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton, append)
12 import Data.ByteString.Lazy (fromStrict, fromChunks)
13 import qualified Data.ByteString.Char8 as BC (replicate, pack, readInt, putStrLn)
14 import Network (connectTo, PortID(..))
15 import Data.Binary (Binary(..), decode)
16 import Data.Binary.Put (putWord32be, putWord16be, putWord8)
17 import Data.Binary.Get (getWord32be, getWord16be, getWord8, runGet)
18 import Control.Monad (replicateM, liftM, forever)
19 import Control.Applicative ((<$>), liftA3)
20 import Control.Concurrent (threadDelay)
26 data PeerState = PeerState { handle :: Handle
28 , am_interested :: Bool
29 , peer_choking :: Bool
30 , peer_interested :: Bool}
32 -- Maintain info on every piece and the current state of it.
33 -- should probably be a TVar.
34 type Pieces = [PieceData]
36 data PieceState = Pending
41 data PieceData = PieceData { index :: Int -- ^ Piece number
42 , peers :: [Peer] -- ^ list of peers who have this piece
43 , state :: PieceState } -- ^ state of the piece from download perspective.
45 -- | Peer is a PeerID, IP address, port tuple
46 data Peer = Peer ID IP Port
49 data PeerMsg = KeepAliveMsg
55 | BitFieldMsg ByteString
56 | RequestMsg Integer Integer Integer
57 | PieceMsg Integer Integer ByteString
58 | CancelMsg Integer Integer Integer
62 genHandShakeMsg :: ByteString -> String -> ByteString
63 genHandShakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID]
64 where pstrlen = singleton 19
65 pstr = BC.pack "BitTorrent protocol"
66 reserved = BC.replicate 8 '\0'
67 peerID = BC.pack peer_id
69 handShake :: Peer -> ByteString -> String -> IO Handle
70 handShake (Peer _ ip port) infoHash peerid = do
71 let hs = genHandShakeMsg infoHash peerid
72 h <- connectTo ip (PortNumber (fromIntegral port))
73 hSetBuffering h LineBuffering
76 let rlen = fromIntegral $ (unpack rlenBS) !! 0
80 instance Binary PeerMsg where
82 KeepAliveMsg -> putWord32be 0
83 ChokeMsg -> do putWord32be 1
85 UnChokeMsg -> do putWord32be 1
87 InterestedMsg -> do putWord32be 1
89 NotInterestedMsg -> do putWord32be 1
91 HaveMsg i -> do putWord32be 5
93 putWord32be (fromIntegral i)
94 BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen)
97 where bfList = unpack bf
98 bfListLen = length bfList
99 RequestMsg i o l -> do putWord32be 13
101 putWord32be (fromIntegral i)
102 putWord32be (fromIntegral o)
103 putWord32be (fromIntegral l)
104 PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
106 putWord32be (fromIntegral i)
107 putWord32be (fromIntegral o)
108 mapM_ putWord8 blockList
109 where blockList = unpack b
110 blocklen = length blockList
111 CancelMsg i o l -> do putWord32be 13
113 putWord32be (fromIntegral i)
114 putWord32be (fromIntegral o)
115 putWord32be (fromIntegral l)
116 PortMsg p -> do putWord32be 3
118 putWord16be (fromIntegral p)
124 1 -> return UnChokeMsg
125 2 -> return InterestedMsg
126 3 -> return NotInterestedMsg
127 4 -> liftM (HaveMsg . fromIntegral) getWord32be
128 5 -> liftM (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8)
129 6 -> liftA3 RequestMsg getInteger getInteger getInteger
130 where getInteger = fromIntegral <$> getWord32be
131 7 -> liftA3 PieceMsg getInteger getInteger (pack <$> replicateM (fromIntegral l - 9) getWord8)
132 where getInteger = fromIntegral <$> getWord32be
133 8 -> liftA3 CancelMsg getInteger getInteger getInteger
134 where getInteger = fromIntegral <$> getWord32be
135 9 -> liftM (PortMsg . fromIntegral) getWord16be
136 _ -> error "unknown message ID"
138 getMsg :: Handle -> IO PeerMsg
141 let lenBS = length (unpack lBS)
142 putStrLn $ "bytes read: " ++ (show lenBS)
145 then return KeepAliveMsg
147 putStrLn $ "len: " ++ (show lBS)
149 return $ decode $ fromStrict $ append lBS msg
151 bsToInt :: ByteString -> Int
152 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
154 -- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
155 -- recvMsg :: Peer -> Handle -> Msg
157 msgLoop :: Handle -> IO ()
158 msgLoop h = forever $ do
160 putStrLn $ "got a " ++ (show msg)