diff --git a/examples/connection_tester.cpp b/examples/connection_tester.cpp index 7c7ce2cfb..f03eab019 100644 --- a/examples/connection_tester.cpp +++ b/examples/connection_tester.cpp @@ -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 #include #include +#include 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 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& 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 > 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 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 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)