<commit_before>%%%-------------------------------------------------------------------
%%% File    : peer_communication.erl
%%% Author  : <NAME> <<EMAIL>>
%%% License : See COPYING
%%% Description : Various pieces of the peer protocol that takes a bit to
%%%   handle.
%%%
%%% Created : 26 Jan 2007 by <NAME> <<EMAIL>>
%%%-------------------------------------------------------------------
-module(peer_communication).

%% API
-export([initiate_handshake/4]).
-export([send_message/2, recv_message/1,
	 construct_bitfield/2, destruct_bitfield/2]).

-define(DEFAULT_HANDSHAKE_TIMEOUT, 120000).
-define(HANDSHAKE_SIZE, 68).
-define(PROTOCOL_STRING, "BitTorrent protocol").
-define(RESERVED_BYTES, 0:64/big).

%% Packet types
-define(CHOKE, 0:8).
-define(UNCHOKE, 1:8).
-define(INTERESTED, 2:8).
-define(NOT_INTERESTED, 3:8).
-define(HAVE, 4:8).
-define(BITFIELD, 5:8).
-define(REQUEST, 6:8).
-define(PIECE, 7:8).
-define(CANCEL, 8:8).
-define(PORT, 9:8).

%%====================================================================
%% API
%%====================================================================
%%--------------------------------------------------------------------
%% Function: recv_message(Message)
%% Description: Receive a message from a peer and decode it
%%--------------------------------------------------------------------
recv_message(Message) ->
    case Message of
	<<>> ->
	    keep_alive;
	<<?CHOKE>> ->
	    choke;
	<<?UNCHOKE>> ->
	    unchoke;
	<<?INTERESTED>> ->
	    interested;
	<<?NOT_INTERESTED>> ->
	    not_interested;
	<<?HAVE, PieceNum:32/big>> ->
	    {have, PieceNum};
	<<?BITFIELD, BitField/binary>> ->
	    {bitfield, BitField};
	<<?REQUEST, Index:32/big, Begin:32/big, Len:32/big>> ->
	    {request, Index, Begin, Len};
	<<?PIECE, Index:32/big, Begin:32/big, Data/binary>> ->
	    {piece, Index, Begin, Data};
	<<?CANCEL, Index:32/big, Begin:32/big, Len:32/big>> ->
	    {cancel, Index, Begin, Len};
	<<?PORT, Port:16/big>> ->
	    {port, Port}
    end.

%%--------------------------------------------------------------------
%% Function: send_message(Socket, Message)
%% Description: Send a message on a socket
%%--------------------------------------------------------------------
send_message(Socket, Message) ->
    Datagram =
	case Message of
	    keep_alive ->
		<<>>;
	    choke ->
		<<?CHOKE>>;
	    unchoke ->
		<<?UNCHOKE>>;
	    interested ->
		<<?INTERESTED>>;
	    not_interested ->
		<<?NOT_INTERESTED>>;
	    {have, PieceNum} ->
		<<?HAVE, PieceNum:32/big>>;
	    {bitfield, BitField} ->
		<<?BITFIELD, BitField/binary>>;
	    {request, Index, Begin, Len} ->
		<<?REQUEST, Index:32/big, Begin:32/big, Len:32/big>>;
	    {piece, Index, Begin, Data} ->
		<<?PIECE,
		 Index:32/big, Begin:32/big, Data/binary>>;
	    {cancel, Index, Begin, Len} ->
		<<?CANCEL, Index:32/big, Begin:32/big, Len:32/big>>;
	    {port, PortNum} ->
		<<?PORT, PortNum:16/big>>
        end,
    gen_tcp:send(Socket, Datagram).

%%--------------------------------------------------------------------
%% Function: initiate_handshake(socket(), peer_id(), peer_id(),
%%                              info_hash()) ->
%%                                         {ok, protocol_version()} |
%%                                              {error, Reason}
%% Description: Handshake with a peer where we have initiated with him.
%%  This call is used if we are the initiator of a torrent handshake as
%%  we then know the peer_id completely.
%%--------------------------------------------------------------------
initiate_handshake(Socket, PeerId, MyPeerId, InfoHash) ->
    PSSize = length(?PROTOCOL_STRING),
    BinPeerId = list_to_binary(PeerId),
    % Since we are the initiator, send out this handshake
    Header = <<PSSize:8, ?PROTOCOL_STRING, ?RESERVED_BYTES>>,
    ok = gen_tcp:send(Socket, Header),
    ok = gen_tcp:send(Socket, InfoHash),
    ok = gen_tcp:send(Socket, MyPeerId),
    % Now, we wait for his handshake to arrive on the socket
    % Since we are the initiator, he is requested to fire off everything
    % to us.
    case gen_tcp:recv(Socket, ?HANDSHAKE_SIZE, ?DEFAULT_HANDSHAKE_TIMEOUT) of
	{ok, Packet} ->
	    <<PSL:8/integer, ?PROTOCOL_STRING, ReservedBytes:8/binary,
	      IH:20/binary, PI:20/binary>> = Packet,
	    if
		PSL /= PSSize ->
		    {error, packet_size_mismatch};
		IH /= InfoHash ->
		    {error, infohash_mismatch};
		PI /= BinPeerId ->
		    {error, peer_id_mismatch};
		true ->
		    {ok, ReservedBytes}
	    end;
	{error, X} ->
	    {error, X}
    end.

%%====================================================================
%% Internal functions
%%====================================================================

%%--------------------------------------------------------------------
%% Function: construct_bitfield
%% Description: Construct a BitField for sending to the peer
%%--------------------------------------------------------------------
construct_bitfield(Size, PieceSet) ->
    PadBits = 8 - (Size rem 8),
    Bits = lists:append(
	     [utils:list_tabulate(
		Size,
		fun(N) ->
			case sets:is_element(N, PieceSet) of
			    true -> 1;
			    false -> 0
			end
		end),
	      utils:list_tabulate(
	       PadBits,
	       fun(_N) -> 0 end)]),
    0 = length(Bits) rem 8,
    list_to_binary(build_bytes(Bits)).
<commit_msg>Begin generalizing handshake code.<commit_after>%%%-------------------------------------------------------------------
%%% File    : peer_communication.erl
%%% Author  : <NAME> <<EMAIL>>
%%% License : See COPYING
%%% Description : Various pieces of the peer protocol that takes a bit to
%%%   handle.
%%%
%%% Created : 26 Jan 2007 by <NAME> <<EMAIL>>
%%%-------------------------------------------------------------------
-module(peer_communication).

%% API
-export([initiate_handshake/4, recieve_handshake/1,
	 complete_handshake_header/3]).
-export([send_message/2, recv_message/1,
	 construct_bitfield/2, destruct_bitfield/2]).

-define(DEFAULT_HANDSHAKE_TIMEOUT, 120000).
-define(HANDSHAKE_SIZE, 68).
-define(PROTOCOL_STRING, "BitTorrent protocol").
-define(RESERVED_BYTES, 0:64/big).

%% Packet types
-define(CHOKE, 0:8).
-define(UNCHOKE, 1:8).
-define(INTERESTED, 2:8).
-define(NOT_INTERESTED, 3:8).
-define(HAVE, 4:8).
-define(BITFIELD, 5:8).
-define(REQUEST, 6:8).
-define(PIECE, 7:8).
-define(CANCEL, 8:8).
-define(PORT, 9:8).

%%====================================================================
%% API
%%====================================================================
%%--------------------------------------------------------------------
%% Function: recv_message(Message) -> keep_alive | choke | unchoke |
%%   interested | not_interested | {have, integer()} | ...
%% Description: Receive a message from a peer and decode it
%%--------------------------------------------------------------------
recv_message(Message) ->
    case Message of
	<<>> ->
	    keep_alive;
	<<?CHOKE>> ->
	    choke;
	<<?UNCHOKE>> ->
	    unchoke;
	<<?INTERESTED>> ->
	    interested;
	<<?NOT_INTERESTED>> ->
	    not_interested;
	<<?HAVE, PieceNum:32/big>> ->
	    {have, PieceNum};
	<<?BITFIELD, BitField/binary>> ->
	    {bitfield, BitField};
	<<?REQUEST, Index:32/big, Begin:32/big, Len:32/big>> ->
	    {request, Index, Begin, Len};
	<<?PIECE, Index:32/big, Begin:32/big, Data/binary>> ->
	    {piece, Index, Begin, Data};
	<<?CANCEL, Index:32/big, Begin:32/big, Len:32/big>> ->
	    {cancel, Index, Begin, Len};
	<<?PORT, Port:16/big>> ->
	    {port, Port}
    end.

%%--------------------------------------------------------------------
%% Function: send_message(Socket, Message)
%% Description: Send a message on a socket
%%--------------------------------------------------------------------
send_message(Socket, Message) ->
    Datagram =
	case Message of
	    keep_alive ->
		<<>>;
	    choke ->
		<<?CHOKE>>;
	    unchoke ->
		<<?UNCHOKE>>;
	    interested ->
		<<?INTERESTED>>;
	    not_interested ->
		<<?NOT_INTERESTED>>;
	    {have, PieceNum} ->
		<<?HAVE, PieceNum:32/big>>;
	    {bitfield, BitField} ->
		<<?BITFIELD, BitField/binary>>;
	    {request, Index, Begin, Len} ->
		<<?REQUEST, Index:32/big, Begin:32/big, Len:32/big>>;
	    {piece, Index, Begin, Data} ->
		<<?PIECE,
		 Index:32/big, Begin:32/big, Data/binary>>;
	    {cancel, Index, Begin, Len} ->
		<<?CANCEL, Index:32/big, Begin:32/big, Len:32/big>>;
	    {port, PortNum} ->
		<<?PORT, PortNum:16/big>>
        end,
    gen_tcp:send(Socket, Datagram).

%%--------------------------------------------------------------------
%% Function: recieve_handshake(Socket) -> {ok, protocol_version,
%%                                             info_hash(),
%%                                             remote_peer_id()} |
%%                                        {error, Reason}
%% Description: Recieve a handshake from another peer. In the recieve,
%%  we don't send the info_hash, but expect the initiator to send what
%%  he thinks is the correct hash. For the return value, see the
%%  function recieve_header()
%%--------------------------------------------------------------------
recieve_handshake(Socket) ->
    Header = build_peer_protocol_header(),
    ok = gen_tcp:send(Socket, Header),
    recieve_header(Socket).

%%--------------------------------------------------------------------
%% Function: complete_handshake_header(socket(),
%%   info_hash(), peer_id()) -> ok
%% Description: Complete the handshake with the peer in the other end.
%%--------------------------------------------------------------------
complete_handshake_header(Socket, InfoHash, LocalPeerId) ->
    ok = gen_tcp:send(Socket, InfoHash),
    ok = gen_tcp:send(Socket, LocalPeerId),
    ok.

%%--------------------------------------------------------------------
%% Function: initiate_handshake(socket(), peer_id(), peer_id(),
%%                              info_hash()) ->
%%                                         {ok, protocol_version()} |
%%                                              {error, Reason}
%% Description: Handshake with a peer where we have initiated with him.
%%  This call is used if we are the initiator of a torrent handshake as
%%  we then know the peer_id completely.
%%--------------------------------------------------------------------
initiate_handshake(Socket, RemotePeerId, LocalPeerId, InfoHash) ->
    BinPeerId = list_to_binary(RemotePeerId),
    % Since we are the initiator, send out this handshake
    Header = build_peer_protocol_header(),
    ok = gen_tcp:send(Socket, Header),
    complete_handshake_header(Socket, InfoHash, LocalPeerId),
    % Now, we wait for his handshake to arrive on the socket
    % Since we are the initiator, he is requested to fire off everything
    % to us.
    case recieve_header(Socket) of
	{ok, ReservedBytes, IH, PI} ->
	    if
		IH /= InfoHash ->
		    {error, infohash_mismatch};
		PI /= BinPeerId ->
		    {error, peer_id_mismatch};
		true ->
		    {ok, ReservedBytes}
	    end;
	{error, X} ->
	    {error, X}
    end.

%%====================================================================
%% Internal functions
%%====================================================================

%%--------------------------------------------------------------------
%% Function: build_peer_protocol_header() -> binary()
%% Description: Returns the Peer Protocol header.
%%--------------------------------------------------------------------
build_peer_protocol_header() ->
    PSSize = length(?PROTOCOL_STRING),
    <<PSSize:8, ?PROTOCOL_STRING, ?RESERVED_BYTES>>.

%%--------------------------------------------------------------------
%% Function: recieve_header(socket()) -> {ok, proto_version(),
%%                                            info_hash(),
%%                                            remote_peer_id()} |
%%                                       {error, Reason}
%% Description: Recieve the full header from a peer. The function
%% returns either with an error or successfully with a
%% protocol_version string, the infohash the remote sent us and his
%% peer_id.
%% --------------------------------------------------------------------
recieve_header(Socket) ->
    case gen_tcp:recv(Socket, ?HANDSHAKE_SIZE, ?DEFAULT_HANDSHAKE_TIMEOUT) of
	{ok, Packet} ->
	    <<PSL:8/integer, ?PROTOCOL_STRING, ReservedBytes:8/binary,
	     IH:20/binary, PI:20/binary>> = Packet,
	    if
		PSL /= length(?PROTOCOL_STRING) ->
		    {error, packet_size_mismatch};
		true ->
		    {ok, ReservedBytes, IH, PI}
	    end;
	{error, X} ->
	    {error, X}
    end.

%%--------------------------------------------------------------------
%% Function: construct_bitfield
%% Description: Construct a BitField for sending to the peer
%%--------------------------------------------------------------------
construct_bitfield(Size, PieceSet) ->
    PadBits = 8 - (Size rem 8),
    Bits = lists:append(
	     [utils:list_tabulate(
		Size,
		fun(N) ->
			case sets:is_element(N, PieceSet) of
			    true -> 1;
			    false -> 0
			end
		end),
	      utils:list_tabulate(
	       PadBits,
	       fun(_N) -> 0 end)]),
    0 = length(Bits) rem 8,
    list_to_binary(build_bytes(Bits)).

{'fraction_non_alphanumeric': 0.31585775335775335, 'fraction_numerical': 0.014804639804639804, 'mean_word_length': 1.8489130434782608, 'pattern_counts': {'":': 0, '<': 111, '<?xml version=': 0, '>': 197, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6976296', 'n_tokens_mistral': 4028, 'n_tokens_neox': 3700, 'n_words': 893}