]> git.rkrishnan.org Git - functorrent.git/blob - src/FuncTorrent/PeerMsgs.hs
quickcheck: size limit randomly generated Benc structures
[functorrent.git] / src / FuncTorrent / PeerMsgs.hs
1 {-
2  - Copyright (C) 2015-2016 Ramakrishnan Muthukrishnan <ram@rkrishnan.org>
3  -
4  - This file is part of FuncTorrent.
5  -
6  - FuncTorrent is free software; you can redistribute it and/or modify
7  - it under the terms of the GNU General Public License as published by
8  - the Free Software Foundation; either version 3 of the License, or
9  - (at your option) any later version.
10  -
11  - FuncTorrent is distributed in the hope that it will be useful,
12  - but WITHOUT ANY WARRANTY; without even the implied warranty of
13  - MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
14  - GNU General Public License for more details.
15  -
16  - You should have received a copy of the GNU General Public License
17  - along with FuncTorrent; if not,  see <http://www.gnu.org/licenses/>
18  -}
19
20 {-# LANGUAGE OverloadedStrings #-}
21
22 module FuncTorrent.PeerMsgs
23        (genHandshakeMsg,
24         sendMsg,
25         getMsg,
26         Peer(..),
27         makePeer,
28         PeerMsg(..)
29        ) where
30
31 import Prelude hiding (lookup, concat, replicate, splitAt, take)
32
33 import System.IO (Handle)
34 import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton)
35 import Data.ByteString.Lazy (fromStrict, fromChunks, toStrict)
36 import Data.ByteString.Char8 as BC (splitAt)
37 import qualified Data.ByteString.Char8 as BC (replicate, pack)
38 import Control.Monad (replicateM)
39 import Control.Applicative (liftA3)
40
41 import Data.Binary (Binary(..), decode, encode)
42 import Data.Binary.Put (putWord32be, putWord16be, putWord8)
43 import Data.Binary.Get (getWord32be, getWord16be, getWord8, runGet)
44
45 import FuncTorrent.Utils (toIP, toPort)
46
47 -- | Peer is a IP address, port tuple
48 data Peer = Peer IP Port
49           deriving (Show, Eq)
50
51 type ID = String
52 type IP = String
53 type Port = Integer
54
55 data PeerMsg = KeepAliveMsg
56              | ChokeMsg
57              | UnChokeMsg
58              | InterestedMsg
59              | NotInterestedMsg
60              | HaveMsg Integer
61              | BitFieldMsg ByteString
62              | RequestMsg Integer Integer Integer
63              | PieceMsg Integer Integer ByteString
64              | CancelMsg Integer Integer Integer
65              | PortMsg Port
66              | ExtendedMsg Integer ByteString
67   deriving (Show)
68
69 data ExtMetadataMsg = Request Integer
70                     | Data Integer Integer
71                     | Reject Integer
72   deriving (Eq, Show)
73
74 instance Binary PeerMsg where
75   put msg = case msg of
76              KeepAliveMsg -> putWord32be 0
77              ChokeMsg -> do putWord32be 1
78                             putWord8 0
79              UnChokeMsg -> do putWord32be 1
80                               putWord8 1
81              InterestedMsg -> do putWord32be 1
82                                  putWord8 2
83              NotInterestedMsg -> do putWord32be 1
84                                     putWord8 3
85              HaveMsg i -> do putWord32be 5
86                              putWord8 4
87                              putWord32be (fromIntegral i)
88              BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen)
89                                   putWord8 5
90                                   mapM_ putWord8 bfList
91                                     where bfList = unpack bf
92                                           bfListLen = length bfList
93              RequestMsg i o l -> do putWord32be 13
94                                     putWord8 6
95                                     putIndexOffsetLength i o l
96              PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
97                                   putWord8 7
98                                   putWord32be (fromIntegral i)
99                                   putWord32be (fromIntegral o)
100                                   mapM_ putWord8 blockList
101                                     where blockList = unpack b
102                                           blocklen = length blockList
103              CancelMsg i o l -> do putWord32be 13
104                                    putWord8 8
105                                    putIndexOffsetLength i o l
106              PortMsg p -> do putWord32be 3
107                              putWord8 9
108                              putWord16be (fromIntegral p)
109              ExtendedMsg t b-> do putWord32be (fromIntegral blockLen)
110                                   putWord8 20
111                                   putWord8 (fromIntegral t) -- 0 => handshake msg
112                                   -- actual extension msg follows
113                                   mapM_ putWord8 blockList
114                                     where blockList = unpack b
115                                           blockLen  = length blockList
116
117     where putIndexOffsetLength i o l = do
118             putWord32be (fromIntegral i)
119             putWord32be (fromIntegral o)
120             putWord32be (fromIntegral l)
121             
122   get = do
123     l <- getWord32be
124     msgid <- getWord8
125     case msgid of
126      0 -> return ChokeMsg
127      1 -> return UnChokeMsg
128      2 -> return InterestedMsg
129      3 -> return NotInterestedMsg
130      4 -> fmap (HaveMsg . fromIntegral) getWord32be
131      5 -> fmap (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8)
132      6 -> liftA3 RequestMsg getInteger getInteger getInteger
133        where getInteger = fromIntegral <$> getWord32be
134      7 -> liftA3 PieceMsg getInteger getInteger (pack  <$> replicateM (fromIntegral l - 9) getWord8)
135        where getInteger = fromIntegral <$> getWord32be
136      8 -> liftA3 CancelMsg getInteger getInteger getInteger
137        where getInteger = fromIntegral <$> getWord32be
138      9 -> fmap (PortMsg . fromIntegral) getWord16be
139      _ -> error ("unknown message ID: " ++ show msgid)
140
141 getMsg :: Handle -> IO PeerMsg
142 getMsg h = do
143   lBS <- hGet h 4
144   let l = bsToInt lBS
145   if l == 0
146     then return KeepAliveMsg
147     else do
148     msg <- hGet h l
149     return $ decode $ fromStrict $ concat [lBS, msg]
150
151 sendMsg :: Handle -> PeerMsg -> IO ()
152 sendMsg h msg = hPut h bsMsg
153   where bsMsg = toStrict $ encode msg
154
155 genHandshakeMsg :: ByteString -> String -> ByteString
156 genHandshakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved1, reserved2, reserved3, infoHash, peerID]
157   where pstrlen = singleton 19
158         pstr = BC.pack "BitTorrent protocol"
159         reserved1 = BC.replicate 5 '\0'
160         reserved2 = singleton 0x10 -- support extension protocol
161         reserved3 = BC.replicate 2 '\0'
162         peerID = BC.pack peer_id
163
164 bsToInt :: ByteString -> Int
165 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
166
167 makePeer :: ByteString -> Peer
168 makePeer peer = Peer (toIP ip') (toPort port')
169   where (ip', port') = splitAt 4 peer