lib/etorrent-1.0/src/etorrent_peer_communication.erl

   1 %%%-------------------------------------------------------------------
   2 %%% File    : peer_communication.erl
   3 %%% Author  : Jesper Louis Andersen <jlouis@succubus>
   4 %%% License : See COPYING
   5 %%% Description : Various pieces of the peer protocol that takes a bit to
   6 %%%   handle.
   7 %%%
   8 %%% Created : 26 Jan 2007 by Jesper Louis Andersen <jlouis@succubus>
   9 %%%-------------------------------------------------------------------
  10 -module(etorrent_peer_communication).
  11
  12 %% API
  13 -export([initiate_handshake/3, receive_handshake/1,
  14          complete_handshake/3]).
  15 -export([send_message/3, recv_message/2,
  16          construct_bitfield/2, destruct_bitfield/2]).
  17
  18 -define(DEFAULT_HANDSHAKE_TIMEOUT, 120000).
  19 -define(HANDSHAKE_SIZE, 68).
  20 -define(PROTOCOL_STRING, "BitTorrent protocol").
  21
  22 %% Extensions
  23 -define(EXT_BASIS, 0). % The protocol basis
  24 -define(EXT_FAST,  4). % The Fast Extension
  25
  26 %% Packet types
  27 -define(CHOKE, 0:8).
  28 -define(UNCHOKE, 1:8).
  29 -define(INTERESTED, 2:8).
  30 -define(NOT_INTERESTED, 3:8).
  31 -define(HAVE, 4:8).
  32 -define(BITFIELD, 5:8).
  33 -define(REQUEST, 6:8).
  34 -define(PIECE, 7:8).
  35 -define(CANCEL, 8:8).
  36 -define(PORT, 9:8).
  37
  38 %% FAST EXTENSION Packet types
  39 -define(SUGGEST, 13:8).
  40 -define(HAVE_ALL, 14:8).
  41 -define(HAVE_NONE, 15:8).
  42 -define(REJECT_REQUEST, 16:8).
  43 -define(ALLOWED_FAST, 17:8).
  44
  45 %%====================================================================
  46 %% API
  47 %%====================================================================
  48 %%--------------------------------------------------------------------
  49 %% Function: recv_message(Message) -> keep_alive | choke | unchoke |
  50 %%   interested | not_interested | {have, integer()} | ...
  51 %% Description: Receive a message from a peer and decode it
  52 %%--------------------------------------------------------------------
  53 recv_message(Rate, Message) ->
  54     MSize = size(Message),
  55     Decoded =
  56         case Message of
  57             <<>> ->
  58                 keep_alive;
  59             <<?CHOKE>> ->
  60                 choke;
  61             <<?UNCHOKE>> ->
  62                 unchoke;
  63             <<?INTERESTED>> ->
  64                 interested;
  65             <<?NOT_INTERESTED>> ->
  66                 not_interested;
  67             <<?HAVE, PieceNum:32/big>> ->
  68                 {have, PieceNum};
  69             <<?BITFIELD, BitField/binary>> ->
  70                 {bitfield, BitField};
  71             <<?REQUEST, Index:32/big, Begin:32/big, Len:32/big>> ->
  72                 {request, Index, Begin, Len};
  73             <<?PIECE, Index:32/big, Begin:32/big, Data/binary>> ->
  74                 {piece, Index, Begin, Data};
  75             <<?CANCEL, Index:32/big, Begin:32/big, Len:32/big>> ->
  76                 {cancel, Index, Begin, Len};
  77             <<?PORT, Port:16/big>> ->
  78                 {port, Port};
  79             %% FAST EXTENSION MESSAGES
  80             <<?SUGGEST, Index:32/big>> ->
  81                 {suggest, Index};
  82             <<?HAVE_ALL>> ->
  83                 have_all;
  84             <<?HAVE_NONE>> ->
  85                 have_none;
  86             <<?REJECT_REQUEST, Index:32, Offset:32, Len:32>> ->
  87                 {reject_request, Index, Offset, Len};
  88             <<?ALLOWED_FAST, FastSet/binary>> ->
  89                 {allowed_fast, decode_allowed_fast(FastSet)}
  90         end,
  91     {Decoded, etorrent_rate:update(Rate, MSize), MSize}.
  92
  93 %%--------------------------------------------------------------------
  94 %% Function: send_message(Socket, Message)
  95 %% Description: Send a message on a socket
  96 %%--------------------------------------------------------------------
  97 send_message(Rate, Socket, Message) ->
  98     Datagram =
  99         case Message of
 100             keep_alive ->
 101                 <<>>;
 102             choke ->
 103                 <<?CHOKE>>;
 104             unchoke ->
 105                 <<?UNCHOKE>>;
 106             interested ->
 107                 <<?INTERESTED>>;
 108             not_interested ->
 109                 <<?NOT_INTERESTED>>;
 110             {have, PieceNum} ->
 111                 <<?HAVE, PieceNum:32/big>>;
 112             {bitfield, BitField} ->
 113                 <<?BITFIELD, BitField/binary>>;
 114             {request, Index, Begin, Len} ->
 115                 <<?REQUEST, Index:32/big, Begin:32/big, Len:32/big>>;
 116             {piece, Index, Begin, Data} ->
 117                 <<?PIECE,
 118                  Index:32/big, Begin:32/big, Data/binary>>;
 119             {cancel, Index, Begin, Len} ->
 120                 <<?CANCEL, Index:32/big, Begin:32/big, Len:32/big>>;
 121             {port, PortNum} ->
 122                 <<?PORT, PortNum:16/big>>;
 123             %% FAST EXTENSION
 124             {suggest, Index} ->
 125                 <<?SUGGEST, Index:32>>;
 126             have_all ->
 127                 <<?HAVE_ALL>>;
 128             have_none ->
 129                 <<?HAVE_NONE>>;
 130             {reject_request, Index, Offset, Len} ->
 131                 <<?REJECT_REQUEST, Index, Offset, Len>>;
 132             {allowed_fast, FastSet} ->
 133                 BinFastSet = encode_fastset(FastSet),
 134                 <<?ALLOWED_FAST, BinFastSet>>
 135         end,
 136     Sz = size(Datagram),
 137     Res = gen_tcp:send(Socket, <<Sz:32/big, Datagram/binary>>),
 138     {Res, etorrent_rate:update(Rate, Sz), Sz}.
 139
 140 %%--------------------------------------------------------------------
 141 %% Function: receive_handshake(Socket) -> {ok, protocol_version,
 142 %%                                             remote_peer_id()} |
 143 %%                                       {ok, proto_version(),
 144 %%                                            info_hash(),
 145 %%                                            remote_peer_id()} |
 146 %%                                        {error, Reason}
 147 %% Description: Receive a handshake from another peer. In the receive,
 148 %%  we don't send the info_hash, but expect the initiator to send what
 149 %%  he thinks is the correct hash. For the return value, see the
 150 %%  function receive_header()
 151 %%--------------------------------------------------------------------
 152 receive_handshake(Socket) ->
 153     Header = build_peer_protocol_header(),
 154     case gen_tcp:send(Socket, Header) of
 155         ok ->
 156             receive_header(Socket, await);
 157         {error, X}  ->
 158             {error, X}
 159     end.
 160
 161 %%--------------------------------------------------------------------
 162 %% Function: initiate_handshake(socket(), peer_id(), info_hash()) ->
 163 %%                                         {ok, protocol_version()} |
 164 %%                                              {error, Reason}
 165 %% Description: Handshake with a peer where we have initiated with him.
 166 %%  This call is used if we are the initiator of a torrent handshake as
 167 %%  we then know the peer_id completely.
 168 %%--------------------------------------------------------------------
 169 initiate_handshake(Socket, LocalPeerId, InfoHash) ->
 170     % Since we are the initiator, send out this handshake
 171     Header = build_peer_protocol_header(),
 172     try
 173         ok = gen_tcp:send(Socket, Header),
 174         ok = gen_tcp:send(Socket, InfoHash),
 175         ok = gen_tcp:send(Socket, LocalPeerId),
 176         receive_header(Socket, InfoHash)
 177     catch
 178         error:_ -> {error, stop}
 179     end.
 180
 181 %%--------------------------------------------------------------------
 182 %% Function: complete_handshake/3
 183 %% Args: Socket ::= socket()
 184 %%       InfoHash ::= binary()
 185 %%       LocalPeerId ::= binary()
 186 %% Description: Complete a handshake.
 187 %%--------------------------------------------------------------------
 188 complete_handshake(Socket, InfoHash, LocalPeerId) ->
 189     Header = build_peer_protocol_header(),
 190     try
 191         ok = gen_tcp:send(Socket, Header),
 192         ok = gen_tcp:send(Socket, InfoHash),
 193         ok = gen_tcp:send(Socket, LocalPeerId),
 194         ok
 195     catch
 196         error:_ -> {error, stop}
 197     end.
 198
 199
 200 %%====================================================================
 201 %% Internal functions
 202 %%====================================================================
 203
 204 %%--------------------------------------------------------------------
 205 %% Function: build_peer_protocol_header() -> binary()
 206 %% Description: Returns the Peer Protocol header.
 207 %%--------------------------------------------------------------------
 208 build_peer_protocol_header() ->
 209     PSSize = length(?PROTOCOL_STRING),
 210     ReservedBytes = protocol_capabilities(),
 211     <<PSSize:8, ?PROTOCOL_STRING, ReservedBytes/binary>>.
 212
 213 protocol_capabilities() ->
 214     ProtoSpec = lists:sum([%?EXT_FAST,
 215                            ?EXT_BASIS]),
 216     <<ProtoSpec:64/big>>.
 217
 218 %%--------------------------------------------------------------------
 219 %% Function: receive_header(socket()) -> {ok, proto_version(),
 220 %%                                            remote_peer_id()} |
 221 %%                                       {ok, proto_version(),
 222 %%                                            info_hash(),
 223 %%                                            remote_peer_id()} |
 224 %%                                       {error, Reason}
 225 %% Description: Receive the full header from a peer. The function
 226 %% returns either with an error or successfully with a
 227 %% protocol_version string, the infohash the remote sent us and his
 228 %% peer_id.
 229 %% --------------------------------------------------------------------
 230 receive_header(Socket, InfoHash) ->
 231     %% Last thing we do on the socket, catch an error here.
 232     case gen_tcp:recv(Socket, ?HANDSHAKE_SIZE, ?DEFAULT_HANDSHAKE_TIMEOUT) of
 233         %% Fail if the header length is wrong
 234         {ok, <<PSL:8/integer, ?PROTOCOL_STRING, _:8/binary,
 235                _IH:20/binary, _PI:20/binary>>}
 236           when PSL /= length(?PROTOCOL_STRING) ->
 237             {error, packet_size_mismatch};
 238         %% If the infohash is await, return the infohash along.
 239         {ok, <<_PSL:8/integer, ?PROTOCOL_STRING, ReservedBytes:8/binary,
 240                IH:20/binary, PI:20/binary>>}
 241           when InfoHash =:= await ->
 242             {ok, ReservedBytes, IH, PI};
 243         %% Infohash mismatches. Error it.
 244         {ok, <<_PSL:8/integer, ?PROTOCOL_STRING, _ReservedBytes:8/binary,
 245                IH:20/binary, _PI:20/binary>>}
 246           when IH /= InfoHash ->
 247             {error, infohash_mismatch};
 248         %% Everything ok
 249         {ok, <<_PSL:8/integer, ?PROTOCOL_STRING, ReservedBytes:8/binary,
 250                _IH:20/binary, PI:20/binary>>} ->
 251             {ok, ReservedBytes, PI};
 252         %% This is not even a header!
 253         {ok, X} when is_binary(X) ->
 254             {error, {bad_header, X}};
 255         %% Propagate errors upwards, most importantly, {error, closed}
 256         {error, Reason} ->
 257             {error, Reason}
 258     end.
 259
 260 %%--------------------------------------------------------------------
 261 %% Function: construct_bitfield
 262 %% Description: Construct a BitField for sending to the peer
 263 %%--------------------------------------------------------------------
 264 construct_bitfield(Size, PieceSet) ->
 265     PadBits = 8 - (Size rem 8),
 266     F = fun(N) ->
 267                 case gb_sets:is_element(N, PieceSet) of
 268                     true -> 1;
 269                     false -> 0
 270                 end
 271         end,
 272     Bits = lists:append([F(N) || N <- lists:seq(0, Size-1)],
 273                         [0 || _N <- lists:seq(1,PadBits)]),
 274     0 = length(Bits) rem 8,
 275     list_to_binary(build_bytes(Bits)).
 276
 277 build_bytes(BitField) ->
 278     build_bytes(BitField, []).
 279
 280 build_bytes([], Acc) ->
 281     lists:reverse(Acc);
 282 build_bytes(L, Acc) ->
 283     {Byte, Rest} = lists:split(8, L),
 284     build_bytes(Rest, [bytify(Byte) | Acc]).
 285
 286 bytify([B1, B2, B3, B4, B5, B6, B7, B8]) ->
 287     <<B1:1/integer, B2:1/integer, B3:1/integer, B4:1/integer,
 288       B5:1/integer, B6:1/integer, B7:1/integer, B8:1/integer>>.
 289
 290 destruct_bitfield(Size, BinaryLump) ->
 291     ByteList = binary_to_list(BinaryLump),
 292     Numbers = decode_bytes(0, ByteList),
 293     PieceSet = gb_sets:from_list(lists:flatten(Numbers)),
 294     case max_element(PieceSet) < Size of
 295         true ->
 296             {ok, PieceSet};
 297         false ->
 298             {error, bitfield_had_wrong_padding}
 299     end.
 300
 301 max_element(Set) ->
 302     gb_sets:fold(fun(E, Max) ->
 303                          case E > Max of
 304                              true ->
 305                                  E;
 306                              false ->
 307                                  Max
 308                          end
 309                  end, 0, Set).
 310
 311 decode_byte(B, Add) ->
 312     <<B1:1/integer, B2:1/integer, B3:1/integer, B4:1/integer,
 313       B5:1/integer, B6:1/integer, B7:1/integer, B8:1/integer>> = <<B>>,
 314     Bytes = [{B1, 0}, {B2, 1}, {B3, 2}, {B4, 3},
 315              {B5, 4}, {B6, 5}, {B7, 6}, {B8, 7}],
 316     [N+Add || {K, N} <- Bytes, K =:= 1].
 317
 318 decode_bytes(_SoFar, []) -> [];
 319 decode_bytes(SoFar, [B | Rest]) ->
 320     [decode_byte(B, SoFar) | decode_bytes(SoFar + 8, Rest)].
 321
 322
 323 decode_allowed_fast(<<>>) -> [];
 324 decode_allowed_fast(<<Index:32, Rest/binary>>) ->
 325     [Index | decode_allowed_fast(Rest)].
 326
 327 encode_fastset([]) -> <<>>;
 328 encode_fastset([Idx | Rest]) ->
 329     R = encode_fastset(Rest),
 330     <<R/binary, Idx:32>>.