1 {-# LANGUAGE OverloadedStrings #-}
2 module FuncTorrent.Peer
7 import Prelude hiding (lookup, concat, replicate, splitAt, take)
9 import System.IO (Handle, BufferMode(..), hSetBuffering)
10 import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, empty)
11 import qualified Data.ByteString.Char8 as BC (length)
12 import Network (connectTo, PortID(..))
13 import Control.Monad (liftM, forM)
15 import Data.Word (Word8)
16 import Data.Map (Map, fromList, toList, (!), mapWithKey, adjust)
17 import qualified Crypto.Hash.SHA1 as SHA1 (hash)
19 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
20 import FuncTorrent.Utils (splitN, splitNum)
21 import FuncTorrent.Fileops (createDummyFile, writeFileAtOffset)
22 import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
24 -- PeerState is a misnomer
25 data PeerState = PeerState { handle :: Handle
28 , meInterested :: Bool
30 , heInterested :: Bool}
32 data PieceDlState = Pending
37 -- todo - map with index to a new data structure (peers who have that piece amd state)
38 data PieceData = PieceData { peers :: [Peer] -- ^ list of peers who have this piece
39 , state :: PieceDlState -- ^ state of the piece from download perspective.
40 , hash :: ByteString -- ^ piece hash
41 , len :: Integer } -- ^ piece length
43 -- which piece is with which peers
44 type PieceMap = Map Integer PieceData
47 -- Make the initial Piece map, with the assumption that no peer has the
48 -- piece and that every piece is pending download.
49 mkPieceMap :: Integer -> ByteString -> [Integer] -> PieceMap
50 mkPieceMap numPieces pieceHash pLengths = fromList kvs
51 where kvs = [(i, PieceData { peers = []
55 | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
56 hashes = splitN 20 pieceHash
58 havePiece :: PieceMap -> Integer -> Bool
60 state (pm ! index) == Have
62 connectToPeer :: Peer -> IO Handle
63 connectToPeer (Peer _ ip port) = do
64 h <- connectTo ip (PortNumber (fromIntegral port))
65 hSetBuffering h LineBuffering
68 doHandshake :: Handle -> Peer -> ByteString -> String -> IO ()
69 doHandshake h peer infoHash peerid = do
70 let hs = genHandshakeMsg infoHash peerid
72 putStrLn $ "--> handhake to peer: " ++ show peer
73 _ <- hGet h (length (unpack hs))
74 putStrLn $ "<-- handshake from peer: " ++ show peer
77 bitfieldToList :: [Word8] -> [Integer]
78 bitfieldToList bs = go bs 0
81 let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
83 setBits ++ go bs' (pos + 1)
85 -- recvMsg :: Peer -> Handle -> Msg
86 msgLoop :: PeerState -> PieceMap -> IO ()
87 msgLoop pState@(PeerState { meInterested = False, heChoking = True }) pieceStatus =
89 -- if me NOT Interested and she is Choking, tell her that
92 sendMsg h InterestedMsg
93 putStrLn $ "--> InterestedMsg to peer: " ++ show (peer pState)
94 msgLoop (pState { meInterested = True }) pieceStatus
95 msgLoop pState@(PeerState { meInterested = True, heChoking = False }) pieceStatus =
96 -- if me Interested and she not Choking, send her a request
98 case pickPiece pieceStatus of
99 Nothing -> putStrLn "Nothing to download"
101 let pLen = len (pieceStatus ! workPiece)
102 putStrLn $ "piece length = " ++ show pLen
103 pBS <- downloadPiece (handle pState) workPiece pLen
104 if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
106 putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS))
108 let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
109 putStrLn $ "Write into file at offset: " ++ show fileOffset
110 writeFileAtOffset "/tmp/download.file" fileOffset pBS
111 msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus)
112 msgLoop pState pieceStatus = do
113 msg <- getMsg (handle pState)
114 putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
117 sendMsg (handle pState) KeepAliveMsg
118 putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState)
119 msgLoop pState pieceStatus
120 BitFieldMsg bss -> do
121 let pieceList = bitfieldToList (unpack bss)
122 pieceStatus' = updatePieceAvailability pieceStatus (peer pState) pieceList
123 putStrLn $ show (length pieceList) ++ " Pieces"
124 -- for each pieceIndex in pieceList, make an entry in the pieceStatus
125 -- map with pieceIndex as the key and modify the value to add the peer.
126 -- download each of the piece in order
127 msgLoop pState pieceStatus'
129 msgLoop (pState { heChoking = False }) pieceStatus
131 msgLoop pState pieceStatus
133 -- simple algorithm to pick piece.
134 -- pick the first piece from 0 that is not downloaded yet.
135 pickPiece :: PieceMap -> Maybe Integer
137 let pieceList = toList m
138 allPending = filter (\(_, v) -> state v == Pending) pieceList
144 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
145 updatePieceAvailability pieceStatus p pieceList =
146 mapWithKey (\k pd -> if k `elem` pieceList
147 then (pd { peers = p : peers pd })
150 handlePeerMsgs :: Peer -> Metainfo -> String -> IO ()
151 handlePeerMsgs p m peerId = do
153 doHandshake h p (infoHash m) peerId
154 let state = PeerState { handle = h
156 , heInterested = False
158 , meInterested = False
160 pieceHash = pieces (info m)
161 numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
162 pLen = pieceLength (info m)
163 fileLen = lengthInBytes (info m)
164 pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
165 createDummyFile "/tmp/download.file" (fromIntegral fileLen)
166 msgLoop state pieceStatus
168 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
169 downloadPiece h index pieceLength = do
170 let chunks = splitNum pieceLength 16384
171 liftM concat $ forM (zip [0..] chunks) (\(i, pLen) -> do
172 sendMsg h (RequestMsg index (i*pLen) pLen)
173 putStrLn $ "--> " ++ "RequestMsg for Piece "
174 ++ show index ++ ", part: " ++ show i ++ " of length: "
178 PieceMsg index begin block -> do
179 putStrLn $ " <-- PieceMsg for Piece: "
185 putStrLn "ignoring irrelevant msg"
188 verifyHash :: ByteString -> ByteString -> Bool
189 verifyHash bs pieceHash =
190 take 20 (SHA1.hash bs) == pieceHash