handlePeerMsgs
) where
-import Prelude hiding (lookup, concat, replicate, splitAt, empty, writeFile)
+import Prelude hiding (lookup, concat, replicate, splitAt, take)
import System.IO (Handle, BufferMode(..), hSetBuffering)
-import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton, writeFile)
+import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton, take, empty)
import Data.ByteString.Lazy (fromStrict, fromChunks, toStrict)
import qualified Data.ByteString.Char8 as BC (replicate, pack, length)
import Network (connectTo, PortID(..))
import Data.Binary (Binary(..), decode, encode)
import Data.Binary.Put (putWord32be, putWord16be, putWord8)
import Data.Binary.Get (getWord32be, getWord16be, getWord8, runGet)
-import Control.Monad (replicateM, liftM, forever)
+import Control.Monad (replicateM, liftM, forM)
import Control.Applicative ((<$>), liftA3)
import Data.Bits
import Data.Word (Word8)
-import Data.Map (Map(..), fromList, (!))
+import Data.Map (Map, fromList, toList, (!), mapWithKey, adjust)
+import qualified Crypto.Hash.SHA1 as SHA1 (hash)
import FuncTorrent.Metainfo (Info(..), Metainfo(..))
-import FuncTorrent.Utils (splitN)
+import FuncTorrent.Utils (splitN, splitNum)
+import FuncTorrent.Fileops (createDummyFile, writeFileAtOffset)
type ID = String
type IP = String
, heChoking :: Bool
, heInterested :: Bool}
--- Maintain info on every piece and the current state of it.
--- should probably be a TVar.
-type Pieces = [PieceData]
-
data PieceDlState = Pending
| InProgress
| Have
-- todo - map with index to a new data structure (peers who have that piece amd state)
data PieceData = PieceData { peers :: [Peer] -- ^ list of peers who have this piece
, state :: PieceDlState -- ^ state of the piece from download perspective.
- , hash :: ByteString } -- ^ piece hash
+ , hash :: ByteString -- ^ piece hash
+ , len :: Integer } -- ^ piece length
-- which piece is with which peers
type PieceMap = Map Integer PieceData
-- Make the initial Piece map, with the assumption that no peer has the
-- piece and that every piece is pending download.
-mkPieceMap :: Integer -> ByteString -> PieceMap
-mkPieceMap numPieces pieceHash = fromList kvs
+mkPieceMap :: Integer -> ByteString -> [Integer] -> PieceMap
+mkPieceMap numPieces pieceHash pLengths = fromList kvs
where kvs = [(i, PieceData { peers = []
, state = Pending
- , hash = h }) | (i, h) <- zip [0..numPieces] hashes]
- hashes = splitN (fromIntegral numPieces) pieceHash
+ , hash = h
+ , len = pLen })
+ | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
+ hashes = splitN 20 pieceHash
havePiece :: PieceMap -> Integer -> Bool
havePiece pm index =
reserved = BC.replicate 8 '\0'
peerID = BC.pack peer_id
-handShake :: Peer -> ByteString -> String -> IO Handle
-handShake (Peer _ ip port) infoHash peerid = do
- let hs = genHandShakeMsg infoHash peerid
+connectToPeer :: Peer -> IO Handle
+connectToPeer peer@(Peer _ ip port) = do
h <- connectTo ip (PortNumber (fromIntegral port))
hSetBuffering h LineBuffering
- hPut h hs
- rlenBS <- hGet h (length (unpack hs))
- putStrLn $ "got handshake from peer: " ++ show rlenBS
return h
+doHandShake :: Handle -> Peer -> ByteString -> String -> IO ()
+doHandShake h peer infoHash peerid = do
+ let hs = genHandShakeMsg infoHash peerid
+ hPut h hs
+ putStrLn $ "--> handhake to peer: " ++ show peer
+ _ <- hGet h (length (unpack hs))
+ putStrLn $ "<-- handshake from peer: " ++ show peer
+ return ()
+
instance Binary PeerMsg where
put msg = case msg of
KeepAliveMsg -> putWord32be 0
return $ decode $ fromStrict $ concat [lBS, msg]
sendMsg :: Handle -> PeerMsg -> IO ()
-sendMsg h msg =
- let bsMsg = toStrict $ encode msg
- in
- hPut h bsMsg
+sendMsg h msg = hPut h bsMsg
+ where bsMsg = toStrict $ encode msg
bsToInt :: ByteString -> Int
bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
in
setBits ++ go bs' (pos + 1)
--- downloadPiece :: Integer -> Handle -> IO ()
-
-createDummyFile :: FilePath -> Int -> IO ()
-createDummyFile path size =
- writeFile path (BC.replicate size '\0')
-
--- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
-- recvMsg :: Peer -> Handle -> Msg
-msgLoop :: PeerState -> ByteString -> IO ()
-msgLoop state pieceHash =
- let numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
- pieceStatus = mkPieceMap numPieces pieceHash
+msgLoop :: PeerState -> PieceMap -> IO ()
+msgLoop pState pieceStatus | not (meInterested pState) && heChoking pState = do
+ -- if me NOT Interested and she is Choking, tell her that
+ -- I am interested.
+ let h = handle pState
+ sendMsg h InterestedMsg
+ putStrLn $ "--> InterestedMsg to peer: " ++ show (peer pState)
+ msgLoop (pState { meInterested = True }) pieceStatus
+ | meInterested pState && not (heChoking pState) =
+ -- if me Interested and she not Choking, send her a request
+ -- for a piece.
+ case pickPiece pieceStatus of
+ Nothing -> putStrLn "Nothing to download"
+ Just workPiece -> do
+ let pLen = len (pieceStatus ! workPiece)
+ putStrLn $ "piece length = " ++ show pLen
+ pBS <- downloadPiece (handle pState) workPiece pLen
+ if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
+ then
+ putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS))
+ else do
+ let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
+ putStrLn $ "Write into file at offset: " ++ show fileOffset
+ writeFileAtOffset "/tmp/download.file" fileOffset pBS
+ msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus)
+ | otherwise = do
+ msg <- getMsg (handle pState)
+ putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
+ case msg of
+ KeepAliveMsg -> do
+ sendMsg (handle pState) KeepAliveMsg
+ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState)
+ msgLoop pState pieceStatus
+ BitFieldMsg bss -> do
+ let pieceList = bitfieldToList (unpack bss)
+ pieceStatus' = updatePieceAvailability pieceStatus (peer pState) pieceList
+ putStrLn $ show (length pieceList) ++ " Pieces"
+ -- for each pieceIndex in pieceList, make an entry in the pieceStatus
+ -- map with pieceIndex as the key and modify the value to add the peer.
+ -- download each of the piece in order
+ msgLoop pState pieceStatus'
+ UnChokeMsg ->
+ msgLoop (pState { heChoking = False }) pieceStatus
+ _ ->
+ msgLoop pState pieceStatus
+
+-- simple algorithm to pick piece.
+-- pick the first piece from 0 that is not downloaded yet.
+pickPiece :: PieceMap -> Maybe Integer
+pickPiece m =
+ let pieceList = toList m
+ allPending = filter (\(_, v) -> state v == Pending) pieceList
in
- forever $ do
- -- if meInterested and he NOT Choking, pick a piece to download
- -- and send a requestmsg.
- msg <- getMsg (handle state)
- putStrLn $ "got a " ++ show msg
- case msg of
- BitFieldMsg bss -> do
- let pieceList = bitfieldToList (unpack bss)
- print pieceList
- -- for each pieceIndex in pieceList, make an entry in the pieceStatus
- -- map with pieceIndex as the key and modify the value to add the peer.
-
- -- download each of the piece in order
- UnChokeMsg -> do
- print msg
- msgLoop (state {heChoking = False}) pieceHash
- _ -> print msg
-
-handlePeerMsgs :: Peer -> Metainfo -> String -> (String -> IO ()) -> IO ()
-handlePeerMsgs p m peerId logFn = do
- h <- handShake p (infoHash m) peerId
- logFn "handShake"
+ case allPending of
+ [] -> Nothing
+ ((i, _):_) -> Just i
+
+updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
+updatePieceAvailability pieceStatus p pieceList =
+ mapWithKey (\k pd -> if k `elem` pieceList
+ then (pd { peers = p : peers pd })
+ else pd) pieceStatus
+
+handlePeerMsgs :: Peer -> Metainfo -> String -> IO ()
+handlePeerMsgs p m peerId = do
+ h <- connectToPeer p
+ doHandShake h p (infoHash m) peerId
let state = PeerState { handle = h
, peer = p
, heInterested = False
, heChoking = True
- , meInterested = True
- , meChoking = False }
- msgLoop state (pieces (info m))
+ , meInterested = False
+ , meChoking = True }
+ pieceHash = pieces (info m)
+ numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
+ pLen = pieceLength (info m)
+ fileLen = lengthInBytes (info m)
+ pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
+ createDummyFile "/tmp/download.file" (fromIntegral fileLen)
+ msgLoop state pieceStatus
+downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
+downloadPiece h index pieceLength = do
+ let chunks = splitNum pieceLength 16384
+ liftM concat $ forM (zip [0..] chunks) (\(i, pLen) -> do
+ sendMsg h (RequestMsg index (i*pLen) pLen)
+ putStrLn $ "--> " ++ "RequestMsg for Piece "
+ ++ show index ++ ", part: " ++ show i ++ " of length: "
+ ++ show pLen
+ msg <- getMsg h
+ case msg of
+ PieceMsg index begin block -> do
+ putStrLn $ " <-- PieceMsg for Piece: "
+ ++ show index
+ ++ ", offset: "
+ ++ show begin
+ return block
+ _ -> do
+ putStrLn "ignoring irrelevant msg"
+ return empty)
+
+verifyHash :: ByteString -> ByteString -> Bool
+verifyHash bs pieceHash =
+ take 20 (SHA1.hash bs) == pieceHash
projects / functorrent.git / blobdiff