import Prelude hiding (lookup, concat, replicate, splitAt, take, drop)
+import Control.Concurrent.MVar (MVar, readMVar, putMVar, takeMVar)
import Control.Monad.State
-import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, drop, empty)
+import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, drop, empty, singleton)
import Data.Bits
import Data.Word (Word8)
-import Data.Map ((!), adjust)
+import Data.Map (Map, (!), adjust, fromList, insert)
import Network (connectTo, PortID(..))
import System.IO (Handle, BufferMode(..), hSetBuffering, hClose)
+import FuncTorrent.Bencode(BVal(..), encode, decode, decodeWithLeftOvers)
import FuncTorrent.Metainfo (Metainfo(..))
import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
import FuncTorrent.Utils (splitNum, verifyHash)
, heChoking :: Bool
, heInterested :: Bool}
+data InfoPieceMap = InfoPieceMap { infoLength :: Integer
+ , infoMap :: Map Integer (Maybe ByteString) }
+
+newtype InfoState = InfoState (MVar InfoPieceMap)
+
havePiece :: PieceMap -> Integer -> Bool
havePiece pm index =
dlstate (pm ! index) == Have
hSetBuffering h LineBuffering
return h
+
doHandshake :: Bool -> Handle -> Peer -> ByteString -> String -> IO ()
doHandshake True h p infohash peerid = do
let hs = genHandshakeMsg infohash peerid
hPut h hs
putStrLn $ "--> handhake to peer: " ++ show p
- _ <- hGet h (length (unpack hs))
+ hsMsg <- hGet h (length (unpack hs))
putStrLn $ "<-- handshake from peer: " ++ show p
return ()
+ -- if doesPeerSupportExtendedMsg hsMsg
+ -- then
+ -- return doExtendedHandshake h
+ -- else
+ -- return Nothing
doHandshake False h p infohash peerid = do
let hs = genHandshakeMsg infohash peerid
putStrLn "waiting for a handshake"
- hsMsg <- hGet h (length (unpack hs))
+ -- read 28 bytes. '19' ++ 'BitTorrent Protocol' ++ 8 reserved bytes
+ hsMsg <- hGet h 28
putStrLn $ "<-- handshake from peer: " ++ show p
let rxInfoHash = take 20 $ drop 28 hsMsg
if rxInfoHash /= infohash
else do
_ <- hPut h hs
putStrLn $ "--> handhake to peer: " ++ show p
- return ()
+ -- if doesPeerSupportExtendedMsg hsMsg
+ -- then do
+ -- doExtendedHandshake h
+ -- else
+ -- return Nothing
+
bitfieldToList :: [Word8] -> [Integer]
bitfieldToList bs = go bs 0
At this point, we have the infodict.
--)
+-}
+{-
+data InfoPieceMap = { infoLength :: Integer
+ , infoMap :: Map Integer (Maybe ByteString)
+ }
+
+newtype InfoState = InfoState (MVar InfoPieceMap)
+
+-}
+
+
+metadataMsgLoop :: Handle -> InfoState -> IO ()
+metadataMsgLoop h (InfoState st) = do
+ infoState <- readMVar st
+ let metadataLen = infoLength infoState
+ -- send the handshake msg
+ metadata = encode (metadataMsg metadataLen)
+ sendMsg h (ExtendedMsg 0 metadata)
+ -- recv return msg from the peer. Will have 'metadata_size'
+ msg <- getMsg h
+ case msg of
+ ExtendedMsg 0 rBs -> do
+ -- decode rBs
+ let (Right (Bdict msgMap)) = decode rBs
+ (Bdict mVal) = msgMap ! "m" -- which is another dict
+ (Bint metadata_msgID) = mVal ! "ut_metadata"
+ (Bint metadata_size) = msgMap ! "metadata_size"
+ -- divide metadata_size into 16384 sized pieces, find number of pieces
+ (q, r) = metadata_size `divMod` 16384
+ -- pNumLengthPairs = zip [0..q-1] (take q (repeat 16384)) ++ (q, r)
+ -- TODO: corner case where infodict size is a multiple of 16384
+ -- and start sending request msg for each.
+ if metadataLen == 0
+ then -- We don't have any piece. Send request msg for all pieces.
+ mapM_ (\n -> do
+ sendMsg h (ExtendedMsg metadata_msgID (encode (requestMsg n)))
+ dataOrRejectMsg <- getMsg h
+ case dataOrRejectMsg of
+ ExtendedMsg 3 payload -> do
+ -- bencoded dict followed by XXXXXX
+ infoState <- takeMVar st
+ let (Right (Bdict bval, pieceData)) = decodeWithLeftOvers payload
+ (Bint pieceIndex) = bval ! "piece"
+ payloadLen = length (unpack pieceData)
+ infoMapVal = infoMap infoState
+ putMVar st infoState {
+ infoMap = insert pieceIndex (Just payload) infoMapVal }
+ )
+ [0..q]
+ else
+ return () -- TODO: reject for now
+ where
+ metadataMsg 0 = Bdict (fromList [("m", Bdict (fromList [("ut_metadata", (Bint 3))]))])
+ metadataMsg l = Bdict (fromList [("m", Bdict (fromList [("ut_metadata", (Bint 3))])),
+ ("metadata_size", (Bint l))])
+ requestMsg i = Bdict (fromList [("msg_type", (Bint 0)), ("piece", (Bint i))])
+ rejectmsg i = Bdict (fromList [("msg_type", (Bint 2)), ("piece", (Bint i))])
+
+doesPeerSupportExtendedMsg :: ByteString -> Bool
+doesPeerSupportExtendedMsg bs = take 1 (drop 5 bs) == singleton 0x10
| CancelMsg Integer Integer Integer
| PortMsg Port
| ExtendedMsg Integer ByteString
- deriving (Show)
+ deriving (Show)
+
+data ExtMetadataMsg = Request Integer
+ | Data Integer Integer
+ | Reject Integer
+ deriving (Eq, Show)
instance Binary PeerMsg where
put msg = case msg of
PortMsg p -> do putWord32be 3
putWord8 9
putWord16be (fromIntegral p)
- ExtendedHandshakeMsg t b-> do putWord32be msgLen
- putWord8 20
- putWord8 t -- 0 => handshake msg
- -- actual extension msg follows
- mapM_ putWord8 blockList
- where blockList = unpack b
- blockLen = length blockList
-
+ ExtendedMsg t b-> do putWord32be (fromIntegral blockLen)
+ putWord8 20
+ putWord8 (fromIntegral t) -- 0 => handshake msg
+ -- actual extension msg follows
+ mapM_ putWord8 blockList
+ where blockList = unpack b
+ blockLen = length blockList
where putIndexOffsetLength i o l = do
putWord32be (fromIntegral i)
genHandshakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved1, reserved2, reserved3, infoHash, peerID]
where pstrlen = singleton 19
pstr = BC.pack "BitTorrent protocol"
- reserved1 = BC.replicate 4 '\0'
+ reserved1 = BC.replicate 5 '\0'
reserved2 = singleton 0x10 -- support extension protocol
- reserved3 = BC.replicate 3 '\0'
+ reserved3 = BC.replicate 2 '\0'
peerID = BC.pack peer_id
bsToInt :: ByteString -> Int