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)
12 import qualified Data.ByteString.Char8 as BC (replicate, pack, length)
13 import Network (connectTo, PortID(..))
14 import Data.Binary (Binary(..), decode)
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)
22 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
27 data PeerState = PeerState { handle :: Handle
29 , am_interested :: Bool
30 , peer_choking :: Bool
31 , peer_interested :: Bool}
33 -- Maintain info on every piece and the current state of it.
34 -- should probably be a TVar.
35 type Pieces = [PieceData]
37 data PieceState = Pending
42 data PieceData = PieceData { index :: Int -- ^ Piece number
43 , peers :: [Peer] -- ^ list of peers who have this piece
44 , state :: PieceState } -- ^ state of the piece from download perspective.
46 -- | Peer is a PeerID, IP address, port tuple
47 data Peer = Peer ID IP Port
50 data PeerMsg = KeepAliveMsg
56 | BitFieldMsg ByteString
57 | RequestMsg Integer Integer Integer
58 | PieceMsg Integer Integer ByteString
59 | CancelMsg Integer Integer Integer
63 genHandShakeMsg :: ByteString -> String -> ByteString
64 genHandShakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID]
65 where pstrlen = singleton 19
66 pstr = BC.pack "BitTorrent protocol"
67 reserved = BC.replicate 8 '\0'
68 peerID = BC.pack peer_id
70 handShake :: Peer -> ByteString -> String -> IO Handle
71 handShake (Peer _ ip port) infoHash peerid = do
72 let hs = genHandShakeMsg infoHash peerid
73 h <- connectTo ip (PortNumber (fromIntegral port))
74 hSetBuffering h LineBuffering
76 rlenBS <- hGet h (length (unpack hs))
77 putStrLn $ "got handshake from peer: " ++ show rlenBS
80 instance Binary PeerMsg where
82 KeepAliveMsg -> putWord32be 0
83 ChokeMsg -> do putWord32be 1
85 UnChokeMsg -> do putWord32be 1
87 InterestedMsg -> do putWord32be 1
89 NotInterestedMsg -> do putWord32be 1
91 HaveMsg i -> do putWord32be 5
93 putWord32be (fromIntegral i)
94 BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen)
97 where bfList = unpack bf
98 bfListLen = length bfList
99 RequestMsg i o l -> do putWord32be 13
101 putWord32be (fromIntegral i)
102 putWord32be (fromIntegral o)
103 putWord32be (fromIntegral l)
104 PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
106 putWord32be (fromIntegral i)
107 putWord32be (fromIntegral o)
108 mapM_ putWord8 blockList
109 where blockList = unpack b
110 blocklen = length blockList
111 CancelMsg i o l -> do putWord32be 13
113 putWord32be (fromIntegral i)
114 putWord32be (fromIntegral o)
115 putWord32be (fromIntegral l)
116 PortMsg p -> do putWord32be 3
118 putWord16be (fromIntegral p)
124 1 -> return UnChokeMsg
125 2 -> return InterestedMsg
126 3 -> return NotInterestedMsg
127 4 -> liftM (HaveMsg . fromIntegral) getWord32be
128 5 -> liftM (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8)
129 6 -> liftA3 RequestMsg getInteger getInteger getInteger
130 where getInteger = fromIntegral <$> getWord32be
131 7 -> liftA3 PieceMsg getInteger getInteger (pack <$> replicateM (fromIntegral l - 9) getWord8)
132 where getInteger = fromIntegral <$> getWord32be
133 8 -> liftA3 CancelMsg getInteger getInteger getInteger
134 where getInteger = fromIntegral <$> getWord32be
135 9 -> liftM (PortMsg . fromIntegral) getWord16be
136 _ -> error ("unknown message ID: " ++ show msgid)
138 getMsg :: Handle -> IO PeerMsg
143 then return KeepAliveMsg
146 return $ decode $ fromStrict $ concat [lBS, msg]
149 bsToInt :: ByteString -> Int
150 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
152 bitfieldToList :: [Word8] -> [Integer]
153 bitfieldToList bs = go bs 0
156 let setBits = [pos*8 + (toInteger i) | i <- [0..8], testBit b i]
158 setBits ++ (go bs' (pos + 1))
160 -- downloadPiece :: Integer -> Handle -> IO ()
162 createDummyFile :: FilePath -> Int -> IO ()
163 createDummyFile path size = do
164 writeFile path (BC.replicate size '\0')
166 -- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
167 -- recvMsg :: Peer -> Handle -> Msg
168 msgLoop :: Handle -> ByteString -> IO ()
169 msgLoop h pieceHash =
170 let numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
174 putStrLn $ "got a " ++ show msg
176 BitFieldMsg bss -> do
177 let pieceList = bitfieldToList (unpack bss)
178 putStrLn (show pieceList)
179 -- download each of the piece in order
180 _ -> putStrLn (show msg)
182 handlePeerMsgs :: Peer -> Metainfo -> String -> (String -> IO ()) -> IO ()
183 handlePeerMsgs p m peerId logFn = do
184 h <- handShake p (infoHash m) peerId
186 msgLoop h (pieces (info m))