extend connection test to support uploading as well
This commit is contained in:
Arvid Norberg
parent
573a291495
commit
244303ea6b
|
|
@ -38,17 +38,29 @@ POSSIBILITY OF SUCH DAMAGE.
|
|||
#include "libtorrent/io.hpp"
|
||||
#include "libtorrent/torrent_info.hpp"
|
||||
#include "libtorrent/thread.hpp"
|
||||
#include "libtorrent/create_torrent.hpp"
|
||||
#include "libtorrent/hasher.hpp"
|
||||
#include <cstring>
|
||||
#include <boost/bind.hpp>
|
||||
#include <iostream>
|
||||
#include <boost/array.hpp>
|
||||
|
||||
using namespace libtorrent;
|
||||
using namespace libtorrent::detail; // for write_* and read_*
|
||||
|
||||
void generate_block(boost::uint32_t* buffer, int piece, int start, int length)
|
||||
{
|
||||
boost::uint32_t fill = (piece << 16) | (start / 0x4000);
|
||||
for (int i = 0; i < length / 4; ++i)
|
||||
{
|
||||
buffer[i] = fill;
|
||||
}
|
||||
}
|
||||
|
||||
struct peer_conn
|
||||
{
|
||||
peer_conn(io_service& ios, int num_pieces, int blocks_pp, tcp::endpoint const& ep
|
||||
, char const* ih)
|
||||
, char const* ih, bool seed_)
|
||||
: s(ios)
|
||||
, read_pos(0)
|
||||
, state(handshaking)
|
||||
|
|
@ -59,17 +71,22 @@ struct peer_conn
|
|||
, blocks_per_piece(blocks_pp)
|
||||
, info_hash(ih)
|
||||
, outstanding_requests(0)
|
||||
, seed(seed_)
|
||||
{
|
||||
// build a list of all pieces and request them all!
|
||||
for (int i = 0; i < int(pieces.size()); ++i)
|
||||
pieces[i] = i;
|
||||
std::random_shuffle(pieces.begin(), pieces.end());
|
||||
if (!seed_)
|
||||
{
|
||||
// build a list of all pieces and request them all!
|
||||
for (int i = 0; i < int(pieces.size()); ++i)
|
||||
pieces[i] = i;
|
||||
std::random_shuffle(pieces.begin(), pieces.end());
|
||||
}
|
||||
|
||||
s.async_connect(ep, boost::bind(&peer_conn::on_connect, this, _1));
|
||||
}
|
||||
|
||||
stream_socket s;
|
||||
char buffer[17*1024];
|
||||
boost::uint32_t write_buffer[17*1024/4];
|
||||
boost::uint32_t buffer[17*1024/4];
|
||||
int read_pos;
|
||||
|
||||
enum state_t
|
||||
|
|
@ -84,6 +101,8 @@ struct peer_conn
|
|||
int blocks_per_piece;
|
||||
char const* info_hash;
|
||||
int outstanding_requests;
|
||||
// if this is true, this connection is a seed
|
||||
bool seed;
|
||||
|
||||
void on_connect(error_code const& ec)
|
||||
{
|
||||
|
|
@ -101,7 +120,8 @@ struct peer_conn
|
|||
memcpy(h, handshake, sizeof(handshake));
|
||||
std::memcpy(h + 28, info_hash, 20);
|
||||
std::generate(h + 48, h + 68, &rand);
|
||||
boost::asio::async_write(s, libtorrent::asio::buffer(h, sizeof(handshake) - 1)
|
||||
// for seeds, don't send the interested message
|
||||
boost::asio::async_write(s, libtorrent::asio::buffer(h, (sizeof(handshake) - 1) - (seed ? 5 : 0))
|
||||
, boost::bind(&peer_conn::on_handshake, this, h, _1, _2));
|
||||
}
|
||||
|
||||
|
|
@ -115,7 +135,7 @@ struct peer_conn
|
|||
}
|
||||
|
||||
// read handshake
|
||||
boost::asio::async_read(s, libtorrent::asio::buffer(buffer, 68)
|
||||
boost::asio::async_read(s, libtorrent::asio::buffer((char*)buffer, 68)
|
||||
, boost::bind(&peer_conn::on_handshake2, this, _1, _2));
|
||||
}
|
||||
|
||||
|
|
@ -127,7 +147,37 @@ struct peer_conn
|
|||
return;
|
||||
}
|
||||
|
||||
work();
|
||||
if (seed)
|
||||
{
|
||||
write_have_all();
|
||||
}
|
||||
else
|
||||
{
|
||||
work_download();
|
||||
}
|
||||
}
|
||||
|
||||
void write_have_all()
|
||||
{
|
||||
// have_all and unchoke
|
||||
static char msg[] = "\0\0\0\x01\x0e\0\0\0\x01\x01";
|
||||
error_code ec;
|
||||
boost::asio::async_write(s, libtorrent::asio::buffer(msg, sizeof(msg) - 1)
|
||||
, boost::bind(&peer_conn::on_have_all_sent, this, _1, _2));
|
||||
|
||||
}
|
||||
|
||||
void on_have_all_sent(error_code const& ec, size_t bytes_transferred)
|
||||
{
|
||||
if (ec)
|
||||
{
|
||||
fprintf(stderr, "ERROR SEND HAVE ALL: %s\n", ec.message().c_str());
|
||||
return;
|
||||
}
|
||||
|
||||
// read message
|
||||
boost::asio::async_read(s, asio::buffer((char*)buffer, 4)
|
||||
, boost::bind(&peer_conn::on_msg_length, this, _1, _2));
|
||||
}
|
||||
|
||||
void write_request()
|
||||
|
|
@ -169,10 +219,10 @@ struct peer_conn
|
|||
|
||||
++outstanding_requests;
|
||||
|
||||
work();
|
||||
work_download();
|
||||
}
|
||||
|
||||
void work()
|
||||
void work_download()
|
||||
{
|
||||
if (pieces.empty() && outstanding_requests == 0)
|
||||
{
|
||||
|
|
@ -188,7 +238,7 @@ struct peer_conn
|
|||
}
|
||||
|
||||
// read message
|
||||
boost::asio::async_read(s, asio::buffer(buffer, 4)
|
||||
boost::asio::async_read(s, asio::buffer((char*)buffer, 4)
|
||||
, boost::bind(&peer_conn::on_msg_length, this, _1, _2));
|
||||
}
|
||||
|
||||
|
|
@ -199,14 +249,14 @@ struct peer_conn
|
|||
fprintf(stderr, "ERROR RECEIVE MESSAGE PREFIX: %s\n", ec.message().c_str());
|
||||
return;
|
||||
}
|
||||
char* ptr = buffer;
|
||||
char* ptr = (char*)buffer;
|
||||
unsigned int length = read_uint32(ptr);
|
||||
if (length > sizeof(buffer))
|
||||
{
|
||||
fprintf(stderr, "ERROR RECEIVE MESSAGE PREFIX: packet too big\n");
|
||||
return;
|
||||
}
|
||||
boost::asio::async_read(s, asio::buffer(buffer, length)
|
||||
boost::asio::async_read(s, asio::buffer((char*)buffer, length)
|
||||
, boost::bind(&peer_conn::on_message, this, _1, _2));
|
||||
}
|
||||
|
||||
|
|
@ -217,43 +267,154 @@ struct peer_conn
|
|||
fprintf(stderr, "ERROR RECEIVE MESSAGE: %s\n", ec.message().c_str());
|
||||
return;
|
||||
}
|
||||
char* ptr = buffer;
|
||||
char* ptr = (char*)buffer;
|
||||
int msg = read_uint8(ptr);
|
||||
if (msg == 7) --outstanding_requests;
|
||||
|
||||
work();
|
||||
if (seed)
|
||||
{
|
||||
if (msg == 6 && bytes_transferred == 13)
|
||||
{
|
||||
int piece = detail::read_int32(ptr);
|
||||
int start = detail::read_int32(ptr);
|
||||
int length = detail::read_int32(ptr);
|
||||
write_piece(piece, start, length);
|
||||
}
|
||||
else
|
||||
{
|
||||
// read another message
|
||||
boost::asio::async_read(s, asio::buffer(buffer, 4)
|
||||
, boost::bind(&peer_conn::on_msg_length, this, _1, _2));
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
if (msg == 7) --outstanding_requests;
|
||||
work_download();
|
||||
}
|
||||
}
|
||||
|
||||
void write_piece(int piece, int start, int length)
|
||||
{
|
||||
generate_block(write_buffer, piece, start, length);
|
||||
static char msg[] = " \x07"
|
||||
" " // piece
|
||||
" "; // start
|
||||
char* ptr = msg;
|
||||
write_uint32(9 + length, ptr);
|
||||
assert(length == 0x4000);
|
||||
assert(*ptr == 7);
|
||||
++ptr; // skip message id
|
||||
write_uint32(piece, ptr);
|
||||
write_uint32(start, ptr);
|
||||
boost::array<libtorrent::asio::const_buffer, 2> vec;
|
||||
vec[0] = libtorrent::asio::buffer(msg, sizeof(msg)-1);
|
||||
vec[1] = libtorrent::asio::buffer(write_buffer, length);
|
||||
boost::asio::async_write(s, vec, boost::bind(&peer_conn::on_have_all_sent, this, _1, _2));
|
||||
}
|
||||
};
|
||||
|
||||
void print_usage()
|
||||
{
|
||||
fprintf(stderr, "usage: connection_tester command ...\n\n"
|
||||
"command is one of:\n"
|
||||
" gen-torrent generate a test torrent\n"
|
||||
" upload start an uploader test\n"
|
||||
" download start a downloader test\n\n"
|
||||
"examples:\n\n"
|
||||
"connection_tester gen-torrent torrent-name\n"
|
||||
"connection_tester upload number-of-connections destination-ip destination-port torrent-file\n"
|
||||
"connection_tester download number-of-connections destination-ip destination-port torrent-file\n");
|
||||
exit(1);
|
||||
}
|
||||
|
||||
void generate_torrent(std::vector<char>& buf)
|
||||
{
|
||||
file_storage fs;
|
||||
// 1 MiB piece size
|
||||
const int piece_size = 1024 * 1024;
|
||||
// 50 GiB should be enough to not fit in physical RAM
|
||||
const int num_pieces = 1 * 1024;
|
||||
const size_type total_size = size_type(piece_size) * num_pieces;
|
||||
fs.add_file("stress_test_file", total_size);
|
||||
libtorrent::create_torrent t(fs, piece_size);
|
||||
|
||||
boost::uint32_t piece[0x4000 / 4];
|
||||
for (int i = 0; i < num_pieces; ++i)
|
||||
{
|
||||
hasher ph;
|
||||
for (int j = 0; j < piece_size; j += 0x4000)
|
||||
{
|
||||
generate_block(piece, i, j, 0x4000);
|
||||
ph.update((char*)piece, 0x4000);
|
||||
}
|
||||
t.set_hash(i, ph.final());
|
||||
}
|
||||
|
||||
std::back_insert_iterator<std::vector<char> > out(buf);
|
||||
|
||||
bencode(out, t.generate());
|
||||
}
|
||||
|
||||
int main(int argc, char* argv[])
|
||||
{
|
||||
if (argc < 5)
|
||||
if (argc <= 1) print_usage();
|
||||
|
||||
bool upload_test = false;
|
||||
bool download_test = false;
|
||||
|
||||
if (strcmp(argv[1], "gen-torrent") == 0)
|
||||
{
|
||||
fprintf(stderr, "usage: connection_tester number-of-connections destination-ip destination-port torrent-file\n");
|
||||
return 1;
|
||||
if (argc != 3) print_usage();
|
||||
|
||||
std::vector<char> tmp;
|
||||
generate_torrent(tmp);
|
||||
|
||||
FILE* output = stdout;
|
||||
if (strcmp("-", argv[2]) != 0)
|
||||
output = fopen(argv[2], "wb+");
|
||||
fwrite(&tmp[0], 1, tmp.size(), output);
|
||||
if (output != stdout)
|
||||
fclose(output);
|
||||
|
||||
return 0;
|
||||
}
|
||||
int num_connections = atoi(argv[1]);
|
||||
else if (strcmp(argv[1], "upload") == 0)
|
||||
{
|
||||
if (argc != 6) print_usage();
|
||||
upload_test = true;
|
||||
}
|
||||
else if (strcmp(argv[1], "download") == 0)
|
||||
{
|
||||
if (argc != 6) print_usage();
|
||||
download_test = true;
|
||||
}
|
||||
|
||||
if (!download_test && !upload_test) print_usage();
|
||||
|
||||
int num_connections = atoi(argv[2]);
|
||||
error_code ec;
|
||||
address_v4 addr = address_v4::from_string(argv[2], ec);
|
||||
address_v4 addr = address_v4::from_string(argv[3], ec);
|
||||
if (ec)
|
||||
{
|
||||
fprintf(stderr, "ERROR RESOLVING %s: %s\n", argv[2], ec.message().c_str());
|
||||
fprintf(stderr, "ERROR RESOLVING %s: %s\n", argv[3], ec.message().c_str());
|
||||
return 1;
|
||||
}
|
||||
int port = atoi(argv[3]);
|
||||
int port = atoi(argv[4]);
|
||||
tcp::endpoint ep(addr, port);
|
||||
torrent_info ti(argv[4], ec);
|
||||
|
||||
torrent_info ti(argv[5], ec);
|
||||
if (ec)
|
||||
{
|
||||
fprintf(stderr, "ERROR LOADING .TORRENT: %s\n", ec.message().c_str());
|
||||
return 1;
|
||||
}
|
||||
|
||||
std::list<peer_conn*> conns;
|
||||
io_service ios;
|
||||
for (int i = 0; i < num_connections; ++i)
|
||||
{
|
||||
conns.push_back(new peer_conn(ios, ti.num_pieces(), ti.piece_length() / 16 / 1024
|
||||
, ep, (char const*)&ti.info_hash()[0]));
|
||||
, ep, (char const*)&ti.info_hash()[0], upload_test));
|
||||
libtorrent::sleep(1);
|
||||
ios.poll_one(ec);
|
||||
if (ec)
|
||||
|
|
|
|||
Loading…
Reference in New Issue