forked from premiere/premiere-libtorrent
improve connection tester to support upload and download mode
This commit is contained in:
parent
244303ea6b
commit
7bd0573927
|
@ -64,23 +64,15 @@ struct peer_conn
|
||||||
: s(ios)
|
: s(ios)
|
||||||
, read_pos(0)
|
, read_pos(0)
|
||||||
, state(handshaking)
|
, state(handshaking)
|
||||||
// don't request anything from the last piece
|
|
||||||
// to keep things simple
|
|
||||||
, pieces(num_pieces - 1)
|
|
||||||
, block(0)
|
, block(0)
|
||||||
, blocks_per_piece(blocks_pp)
|
, blocks_per_piece(blocks_pp)
|
||||||
, info_hash(ih)
|
, info_hash(ih)
|
||||||
, outstanding_requests(0)
|
, outstanding_requests(0)
|
||||||
, seed(seed_)
|
, seed(seed_)
|
||||||
|
, blocks_received(0)
|
||||||
|
, num_pieces(num_pieces)
|
||||||
{
|
{
|
||||||
if (!seed_)
|
pieces.reserve(num_pieces);
|
||||||
{
|
|
||||||
// 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));
|
s.async_connect(ep, boost::bind(&peer_conn::on_connect, this, _1));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -103,6 +95,8 @@ struct peer_conn
|
||||||
int outstanding_requests;
|
int outstanding_requests;
|
||||||
// if this is true, this connection is a seed
|
// if this is true, this connection is a seed
|
||||||
bool seed;
|
bool seed;
|
||||||
|
int blocks_received;
|
||||||
|
int num_pieces;
|
||||||
|
|
||||||
void on_connect(error_code const& ec)
|
void on_connect(error_code const& ec)
|
||||||
{
|
{
|
||||||
|
@ -224,7 +218,9 @@ struct peer_conn
|
||||||
|
|
||||||
void work_download()
|
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");
|
fprintf(stderr, "COMPLETED DOWNLOAD\n");
|
||||||
return;
|
return;
|
||||||
|
@ -270,6 +266,8 @@ struct peer_conn
|
||||||
char* ptr = (char*)buffer;
|
char* ptr = (char*)buffer;
|
||||||
int msg = read_uint8(ptr);
|
int msg = read_uint8(ptr);
|
||||||
|
|
||||||
|
//printf("msg: %d len: %d\n", msg, int(bytes_transferred));
|
||||||
|
|
||||||
if (seed)
|
if (seed)
|
||||||
{
|
{
|
||||||
if (msg == 6 && bytes_transferred == 13)
|
if (msg == 6 && bytes_transferred == 13)
|
||||||
|
@ -288,7 +286,29 @@ struct peer_conn
|
||||||
}
|
}
|
||||||
else
|
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();
|
work_download();
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
@ -318,12 +338,21 @@ void print_usage()
|
||||||
fprintf(stderr, "usage: connection_tester command ...\n\n"
|
fprintf(stderr, "usage: connection_tester command ...\n\n"
|
||||||
"command is one of:\n"
|
"command is one of:\n"
|
||||||
" gen-torrent generate a test torrent\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"
|
" 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"
|
"examples:\n\n"
|
||||||
"connection_tester gen-torrent torrent-name\n"
|
"connection_tester gen-torrent test.torrent\n"
|
||||||
"connection_tester upload number-of-connections destination-ip destination-port torrent-file\n"
|
"connection_tester upload 200 127.0.0.1 6881 test.torrent\n"
|
||||||
"connection_tester download number-of-connections destination-ip destination-port torrent-file\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);
|
exit(1);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -359,8 +388,7 @@ int main(int argc, char* argv[])
|
||||||
{
|
{
|
||||||
if (argc <= 1) print_usage();
|
if (argc <= 1) print_usage();
|
||||||
|
|
||||||
bool upload_test = false;
|
enum { none, upload_test, download_test, dual_test } test_mode = none;
|
||||||
bool download_test = false;
|
|
||||||
|
|
||||||
if (strcmp(argv[1], "gen-torrent") == 0)
|
if (strcmp(argv[1], "gen-torrent") == 0)
|
||||||
{
|
{
|
||||||
|
@ -381,12 +409,17 @@ int main(int argc, char* argv[])
|
||||||
else if (strcmp(argv[1], "upload") == 0)
|
else if (strcmp(argv[1], "upload") == 0)
|
||||||
{
|
{
|
||||||
if (argc != 6) print_usage();
|
if (argc != 6) print_usage();
|
||||||
upload_test = true;
|
test_mode = upload_test;
|
||||||
}
|
}
|
||||||
else if (strcmp(argv[1], "download") == 0)
|
else if (strcmp(argv[1], "download") == 0)
|
||||||
{
|
{
|
||||||
if (argc != 6) print_usage();
|
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();
|
if (!download_test && !upload_test) print_usage();
|
||||||
|
@ -413,8 +446,11 @@ int main(int argc, char* argv[])
|
||||||
io_service ios;
|
io_service ios;
|
||||||
for (int i = 0; i < num_connections; ++i)
|
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
|
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);
|
libtorrent::sleep(1);
|
||||||
ios.poll_one(ec);
|
ios.poll_one(ec);
|
||||||
if (ec)
|
if (ec)
|
||||||
|
|
Loading…
Reference in New Issue