1 {-# LANGUAGE OverloadedStrings #-}
2 module FuncTorrent.Peer
8 import Prelude hiding (lookup, concat, replicate, splitAt, empty)
11 import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton)
12 import Data.ByteString.Lazy (fromStrict, fromChunks)
13 import qualified Data.ByteString.Char8 as BC (replicate, pack, length)
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)
21 import Data.Word (Word8)
27 data PeerState = PeerState { handle :: Handle
29 , am_interested :: Bool
30 , peer_choking :: Bool
31 , peer_interested :: Bool}
33 -- Maintain info on every piece and the current state of it.
34 -- should probably be a TVar.
35 type Pieces = [PieceData]
37 data PieceState = Pending
42 data PieceData = PieceData { index :: Int -- ^ Piece number
43 , peers :: [Peer] -- ^ list of peers who have this piece
44 , state :: PieceState } -- ^ state of the piece from download perspective.
46 -- | Peer is a PeerID, IP address, port tuple
47 data Peer = Peer ID IP Port
50 data PeerMsg = KeepAliveMsg
56 | BitFieldMsg ByteString
57 | RequestMsg Integer Integer Integer
58 | PieceMsg Integer Integer ByteString
59 | CancelMsg Integer Integer Integer
63 genHandShakeMsg :: ByteString -> String -> ByteString
64 genHandShakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID]
65 where pstrlen = singleton 19
66 pstr = BC.pack "BitTorrent protocol"
67 reserved = BC.replicate 8 '\0'
68 peerID = BC.pack peer_id
70 handShake :: Peer -> ByteString -> String -> IO Handle
71 handShake (Peer _ ip port) infoHash peerid = do
72 let hs = genHandShakeMsg infoHash peerid
73 h <- connectTo ip (PortNumber (fromIntegral port))
74 hSetBuffering h LineBuffering
76 rlenBS <- hGet h (length (unpack hs))
77 putStrLn $ "got handshake from peer: " ++ show rlenBS
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: " ++ show msgid)
138 getMsg :: Handle -> IO PeerMsg
143 then return KeepAliveMsg
146 return $ decode $ fromStrict $ concat [lBS, msg]
149 bsToInt :: ByteString -> Int
150 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
152 bitfieldToList :: [Word8] -> [Integer]
153 bitfieldToList bs = go bs 0
156 let setBits = [pos*8 + (toInteger i) | i <- [0..8], testBit b i]
158 setBits ++ (go bs' (pos + 1))
160 -- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
161 -- recvMsg :: Peer -> Handle -> Msg
162 msgLoop :: Handle -> ByteString -> IO ()
163 msgLoop h pieceHash =
164 let numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
168 putStrLn $ "got a " ++ show msg
170 BitFieldMsg bss -> do
171 let pieceList = bitfieldToList (unpack bss) 0
172 putStrLn (show pieceList)
173 -- doenload each of the piece
174 _ -> putStrLn (show msg)