1 {-# LANGUAGE OverloadedStrings #-}
2 module FuncTorrent.Peer
7 import Prelude hiding (lookup, concat, replicate, splitAt, empty, writeFile)
9 import System.IO (Handle, BufferMode(..), hSetBuffering)
10 import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton, writeFile)
11 import Data.ByteString.Lazy (fromStrict, fromChunks, toStrict)
12 import qualified Data.ByteString.Char8 as BC (replicate, pack, length)
13 import Network (connectTo, PortID(..))
14 import Data.Binary (Binary(..), decode, encode)
15 import Data.Binary.Put (putWord32be, putWord16be, putWord8)
16 import Data.Binary.Get (getWord32be, getWord16be, getWord8, runGet)
17 import Control.Monad (replicateM, liftM, forever)
18 import Control.Applicative ((<$>), liftA3)
20 import Data.Word (Word8)
21 import Data.Map (Map(..), fromList, (!))
23 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
24 import FuncTorrent.Utils (splitN)
30 -- PeerState is a misnomer
31 data PeerState = PeerState { handle :: Handle
34 , meInterested :: Bool
36 , heInterested :: Bool}
38 -- Maintain info on every piece and the current state of it.
39 -- should probably be a TVar.
40 type Pieces = [PieceData]
42 data PieceDlState = Pending
47 -- todo - map with index to a new data structure (peers who have that piece amd state)
48 data PieceData = PieceData { peers :: [Peer] -- ^ list of peers who have this piece
49 , state :: PieceDlState -- ^ state of the piece from download perspective.
50 , hash :: ByteString } -- ^ piece hash
52 -- which piece is with which peers
53 type PieceMap = Map Integer PieceData
55 -- | Peer is a PeerID, IP address, port tuple
56 data Peer = Peer ID IP Port
59 data PeerMsg = KeepAliveMsg
65 | BitFieldMsg ByteString
66 | RequestMsg Integer Integer Integer
67 | PieceMsg Integer Integer ByteString
68 | CancelMsg Integer Integer Integer
72 -- Make the initial Piece map, with the assumption that no peer has the
73 -- piece and that every piece is pending download.
74 mkPieceMap :: Integer -> ByteString -> PieceMap
75 mkPieceMap numPieces pieceHash = fromList kvs
76 where kvs = [(i, PieceData { peers = []
78 , hash = h }) | (i, h) <- zip [0..numPieces] hashes]
79 hashes = splitN (fromIntegral numPieces) pieceHash
81 havePiece :: PieceMap -> Integer -> Bool
83 state (pm ! index) == Have
85 genHandShakeMsg :: ByteString -> String -> ByteString
86 genHandShakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID]
87 where pstrlen = singleton 19
88 pstr = BC.pack "BitTorrent protocol"
89 reserved = BC.replicate 8 '\0'
90 peerID = BC.pack peer_id
92 handShake :: Peer -> ByteString -> String -> IO Handle
93 handShake (Peer _ ip port) infoHash peerid = do
94 let hs = genHandShakeMsg infoHash peerid
95 h <- connectTo ip (PortNumber (fromIntegral port))
96 hSetBuffering h LineBuffering
98 rlenBS <- hGet h (length (unpack hs))
99 putStrLn $ "got handshake from peer: " ++ show rlenBS
102 instance Binary PeerMsg where
103 put msg = case msg of
104 KeepAliveMsg -> putWord32be 0
105 ChokeMsg -> do putWord32be 1
107 UnChokeMsg -> do putWord32be 1
109 InterestedMsg -> do putWord32be 1
111 NotInterestedMsg -> do putWord32be 1
113 HaveMsg i -> do putWord32be 5
115 putWord32be (fromIntegral i)
116 BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen)
118 mapM_ putWord8 bfList
119 where bfList = unpack bf
120 bfListLen = length bfList
121 RequestMsg i o l -> do putWord32be 13
123 putWord32be (fromIntegral i)
124 putWord32be (fromIntegral o)
125 putWord32be (fromIntegral l)
126 PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
128 putWord32be (fromIntegral i)
129 putWord32be (fromIntegral o)
130 mapM_ putWord8 blockList
131 where blockList = unpack b
132 blocklen = length blockList
133 CancelMsg i o l -> do putWord32be 13
135 putWord32be (fromIntegral i)
136 putWord32be (fromIntegral o)
137 putWord32be (fromIntegral l)
138 PortMsg p -> do putWord32be 3
140 putWord16be (fromIntegral p)
146 1 -> return UnChokeMsg
147 2 -> return InterestedMsg
148 3 -> return NotInterestedMsg
149 4 -> liftM (HaveMsg . fromIntegral) getWord32be
150 5 -> liftM (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8)
151 6 -> liftA3 RequestMsg getInteger getInteger getInteger
152 where getInteger = fromIntegral <$> getWord32be
153 7 -> liftA3 PieceMsg getInteger getInteger (pack <$> replicateM (fromIntegral l - 9) getWord8)
154 where getInteger = fromIntegral <$> getWord32be
155 8 -> liftA3 CancelMsg getInteger getInteger getInteger
156 where getInteger = fromIntegral <$> getWord32be
157 9 -> liftM (PortMsg . fromIntegral) getWord16be
158 _ -> error ("unknown message ID: " ++ show msgid)
160 getMsg :: Handle -> IO PeerMsg
165 then return KeepAliveMsg
168 return $ decode $ fromStrict $ concat [lBS, msg]
170 sendMsg :: Handle -> PeerMsg -> IO ()
172 let bsMsg = toStrict $ encode msg
176 bsToInt :: ByteString -> Int
177 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
179 bitfieldToList :: [Word8] -> [Integer]
180 bitfieldToList bs = go bs 0
183 let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
185 setBits ++ go bs' (pos + 1)
187 -- downloadPiece :: Integer -> Handle -> IO ()
189 createDummyFile :: FilePath -> Int -> IO ()
190 createDummyFile path size =
191 writeFile path (BC.replicate size '\0')
193 -- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
194 -- recvMsg :: Peer -> Handle -> Msg
195 msgLoop :: PeerState -> ByteString -> IO ()
196 msgLoop state pieceHash =
197 let numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
198 pieceStatus = mkPieceMap numPieces pieceHash
201 -- if meInterested and he NOT Choking, pick a piece to download
202 -- and send a requestmsg.
203 msg <- getMsg (handle state)
204 putStrLn $ "got a " ++ show msg
206 BitFieldMsg bss -> do
207 let pieceList = bitfieldToList (unpack bss)
209 -- for each pieceIndex in pieceList, make an entry in the pieceStatus
210 -- map with pieceIndex as the key and modify the value to add the peer.
212 -- download each of the piece in order
215 msgLoop (state {heChoking = False}) pieceHash
218 handlePeerMsgs :: Peer -> Metainfo -> String -> (String -> IO ()) -> IO ()
219 handlePeerMsgs p m peerId logFn = do
220 h <- handShake p (infoHash m) peerId
222 let state = PeerState { handle = h
224 , heInterested = False
226 , meInterested = True
227 , meChoking = False }
228 msgLoop state (pieces (info m))