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)
|
||||
, 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)
|
||||
|
|
Loading…
Reference in New Issue