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 { peer :: Peer
33 , meInterested :: Bool
35 , heInterested :: Bool}
37 -- Maintain info on every piece and the current state of it.
38 -- should probably be a TVar.
39 type Pieces = [PieceData]
41 data PieceDlState = Pending
46 -- todo - map with index to a new data structure (peers who have that piece amd state)
47 data PieceData = PieceData { peers :: [Peer] -- ^ list of peers who have this piece
48 , state :: PieceDlState -- ^ state of the piece from download perspective.
49 , hash :: ByteString } -- ^ piece hash
51 -- which piece is with which peers
52 type PieceMap = Map Integer PieceData
54 -- | Peer is a PeerID, IP address, port tuple
55 data Peer = Peer ID IP Port
58 data PeerMsg = KeepAliveMsg
64 | BitFieldMsg ByteString
65 | RequestMsg Integer Integer Integer
66 | PieceMsg Integer Integer ByteString
67 | CancelMsg Integer Integer Integer
71 -- Make the initial Piece map, with the assumption that no peer has the
72 -- piece and that every piece is pending download.
73 mkPieceMap :: Integer -> ByteString -> PieceMap
74 mkPieceMap numPieces pieceHash = fromList kvs
75 where kvs = [(i, PieceData { peers = []
77 , hash = h }) | (i, h) <- zip [0..numPieces] hashes]
78 hashes = splitN (fromIntegral numPieces) pieceHash
80 havePiece :: PieceMap -> Integer -> Bool
82 state (pm ! index) == Have
84 genHandShakeMsg :: ByteString -> String -> ByteString
85 genHandShakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID]
86 where pstrlen = singleton 19
87 pstr = BC.pack "BitTorrent protocol"
88 reserved = BC.replicate 8 '\0'
89 peerID = BC.pack peer_id
91 handShake :: Peer -> ByteString -> String -> IO Handle
92 handShake (Peer _ ip port) infoHash peerid = do
93 let hs = genHandShakeMsg infoHash peerid
94 h <- connectTo ip (PortNumber (fromIntegral port))
95 hSetBuffering h LineBuffering
97 rlenBS <- hGet h (length (unpack hs))
98 putStrLn $ "got handshake from peer: " ++ show rlenBS
101 instance Binary PeerMsg where
102 put msg = case msg of
103 KeepAliveMsg -> putWord32be 0
104 ChokeMsg -> do putWord32be 1
106 UnChokeMsg -> do putWord32be 1
108 InterestedMsg -> do putWord32be 1
110 NotInterestedMsg -> do putWord32be 1
112 HaveMsg i -> do putWord32be 5
114 putWord32be (fromIntegral i)
115 BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen)
117 mapM_ putWord8 bfList
118 where bfList = unpack bf
119 bfListLen = length bfList
120 RequestMsg i o l -> do putWord32be 13
122 putWord32be (fromIntegral i)
123 putWord32be (fromIntegral o)
124 putWord32be (fromIntegral l)
125 PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
127 putWord32be (fromIntegral i)
128 putWord32be (fromIntegral o)
129 mapM_ putWord8 blockList
130 where blockList = unpack b
131 blocklen = length blockList
132 CancelMsg i o l -> do putWord32be 13
134 putWord32be (fromIntegral i)
135 putWord32be (fromIntegral o)
136 putWord32be (fromIntegral l)
137 PortMsg p -> do putWord32be 3
139 putWord16be (fromIntegral p)
145 1 -> return UnChokeMsg
146 2 -> return InterestedMsg
147 3 -> return NotInterestedMsg
148 4 -> liftM (HaveMsg . fromIntegral) getWord32be
149 5 -> liftM (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8)
150 6 -> liftA3 RequestMsg getInteger getInteger getInteger
151 where getInteger = fromIntegral <$> getWord32be
152 7 -> liftA3 PieceMsg getInteger getInteger (pack <$> replicateM (fromIntegral l - 9) getWord8)
153 where getInteger = fromIntegral <$> getWord32be
154 8 -> liftA3 CancelMsg getInteger getInteger getInteger
155 where getInteger = fromIntegral <$> getWord32be
156 9 -> liftM (PortMsg . fromIntegral) getWord16be
157 _ -> error ("unknown message ID: " ++ show msgid)
159 getMsg :: Handle -> IO PeerMsg
164 then return KeepAliveMsg
167 return $ decode $ fromStrict $ concat [lBS, msg]
169 sendMsg :: Handle -> PeerMsg -> IO ()
171 let bsMsg = toStrict $ encode msg
175 bsToInt :: ByteString -> Int
176 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
178 bitfieldToList :: [Word8] -> [Integer]
179 bitfieldToList bs = go bs 0
182 let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
184 setBits ++ go bs' (pos + 1)
186 -- downloadPiece :: Integer -> Handle -> IO ()
188 createDummyFile :: FilePath -> Int -> IO ()
189 createDummyFile path size =
190 writeFile path (BC.replicate size '\0')
192 -- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
193 -- recvMsg :: Peer -> Handle -> Msg
194 msgLoop :: Handle -> ByteString -> PeerState -> IO ()
195 msgLoop h pieceHash state =
196 let numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
197 pieceStatus = mkPieceMap numPieces pieceHash
201 putStrLn $ "got a " ++ show msg
203 BitFieldMsg bss -> do
204 let pieceList = bitfieldToList (unpack bss)
206 -- for each pieceIndex in pieceList, make an entry in the pieceStatus
207 -- map with pieceIndex as the key and modify the value to add the peer.
209 -- download each of the piece in order
212 msgLoop h pieceHash (state {heChoking = False})
215 handlePeerMsgs :: Peer -> Metainfo -> String -> (String -> IO ()) -> IO ()
216 handlePeerMsgs p m peerId logFn = do
217 h <- handShake p (infoHash m) peerId
219 let state = PeerState { peer = p
220 , heInterested = False
222 , meInterested = True
223 , meChoking = False }
224 msgLoop h (pieces (info m)) state