improve connection tester to support upload and download mode

This commit is contained in:
Arvid Norberg 2011-04-22 07:38:46 +00:00
parent 244303ea6b
commit 7bd0573927
1 changed files with 58 additions and 22 deletions

View File

@ -64,23 +64,15 @@ struct peer_conn
: s(ios)
, read_pos(0)
, state(handshaking)
// don't request anything from the last piece
// to keep things simple
, pieces(num_pieces - 1)
, block(0)
, blocks_per_piece(blocks_pp)
, info_hash(ih)
, outstanding_requests(0)
, seed(seed_)
, blocks_received(0)
, num_pieces(num_pieces)
{
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());
}
pieces.reserve(num_pieces);
s.async_connect(ep, boost::bind(&peer_conn::on_connect, this, _1));
}
@ -103,6 +95,8 @@ struct peer_conn
int outstanding_requests;
// if this is true, this connection is a seed
bool seed;
int blocks_received;
int num_pieces;
void on_connect(error_code const& ec)
{
@ -224,7 +218,9 @@ struct peer_conn
void work_download()
{
if (pieces.empty() && outstanding_requests == 0)
if (pieces.empty()
&& outstanding_requests == 0
&& blocks_received >= num_pieces * blocks_per_piece)
{
fprintf(stderr, "COMPLETED DOWNLOAD\n");
return;
@ -270,6 +266,8 @@ struct peer_conn
char* ptr = (char*)buffer;
int msg = read_uint8(ptr);
//printf("msg: %d len: %d\n", msg, int(bytes_transferred));
if (seed)
{
if (msg == 6 && bytes_transferred == 13)
@ -288,7 +286,29 @@ struct peer_conn
}
else
{
if (msg == 7) --outstanding_requests;
if (msg == 0xe) // have_all
{
// build a list of all pieces and request them all!
pieces.resize(num_pieces);
for (int i = 0; i < int(pieces.size()); ++i)
pieces[i] = i;
std::random_shuffle(pieces.begin(), pieces.end());
}
else if (msg == 4) // have
{
int piece = detail::read_int32(ptr);
if (pieces.empty()) pieces.push_back(piece);
else pieces.insert(pieces.begin() + (rand() % pieces.size()), piece);
}
else if (msg == 5)
{
// todo: support bitfield
}
else if (msg == 7)
{
++blocks_received;
--outstanding_requests;
}
work_download();
}
}
@ -318,12 +338,21 @@ void print_usage()
fprintf(stderr, "usage: connection_tester command ...\n\n"
"command is one of:\n"
" gen-torrent generate a test torrent\n"
" this command takes one extra argument, specifying the file to save\n"
" the .torrent file to\n\n"
" upload start an uploader test\n"
" download start a downloader test\n\n"
" download start a downloader test\n"
" dual start a download and upload test\n"
" these commands set takes 4 additional arguments\n"
" 1. num-connections - the number of connections to make to the target\n"
" 2. destination-IP - the IP address of the target\n"
" 3. destination-port - the port the target listens on\n"
" 4. torrent-file - the torrent file previously generated by gen-torrent\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");
"connection_tester gen-torrent test.torrent\n"
"connection_tester upload 200 127.0.0.1 6881 test.torrent\n"
"connection_tester download 200 127.0.0.1 6881 test.torrent\n"
"connection_tester dual 200 127.0.0.1 6881 test.torrent\n");
exit(1);
}
@ -359,8 +388,7 @@ int main(int argc, char* argv[])
{
if (argc <= 1) print_usage();
bool upload_test = false;
bool download_test = false;
enum { none, upload_test, download_test, dual_test } test_mode = none;
if (strcmp(argv[1], "gen-torrent") == 0)
{
@ -381,12 +409,17 @@ int main(int argc, char* argv[])
else if (strcmp(argv[1], "upload") == 0)
{
if (argc != 6) print_usage();
upload_test = true;
test_mode = upload_test;
}
else if (strcmp(argv[1], "download") == 0)
{
if (argc != 6) print_usage();
download_test = true;
test_mode = download_test;
}
else if (strcmp(argv[1], "dual") == 0)
{
if (argc != 6) print_usage();
test_mode = dual_test;
}
if (!download_test && !upload_test) print_usage();
@ -413,8 +446,11 @@ int main(int argc, char* argv[])
io_service ios;
for (int i = 0; i < num_connections; ++i)
{
bool seed = false;
if (test_mode == upload_test) seed = true;
else if (test_mode == dual_test) seed = (i & 1);
conns.push_back(new peer_conn(ios, ti.num_pieces(), ti.piece_length() / 16 / 1024
, ep, (char const*)&ti.info_hash()[0], upload_test));
, ep, (char const*)&ti.info_hash()[0], seed));
libtorrent::sleep(1);
ios.poll_one(ec);
if (ec)