diff --git a/src/net/packetsocket.cpp b/src/net/packetsocket.cpp index 7095b28f2907f647e85ec211e562ace4b763313c..80c7a67b57e9c3de0c1c31f6b3a312407ce5f90f 100644 --- a/src/net/packetsocket.cpp +++ b/src/net/packetsocket.cpp @@ -15,6 +15,7 @@ namespace net PacketSocket::PacketSocket(SocketDevice *sock) : TrafficShapedSocket(sock) , ctrl_packets_sent(0) + , pending_upload_data_bytes(0) , uploaded_data_bytes(0) { } @@ -22,6 +23,7 @@ PacketSocket::PacketSocket(SocketDevice *sock) PacketSocket::PacketSocket(int fd, int ip_version) : TrafficShapedSocket(fd, ip_version) , ctrl_packets_sent(0) + , pending_upload_data_bytes(0) , uploaded_data_bytes(0) { } @@ -29,6 +31,7 @@ PacketSocket::PacketSocket(int fd, int ip_version) PacketSocket::PacketSocket(bool tcp, int ip_version) : TrafficShapedSocket(tcp, ip_version) , ctrl_packets_sent(0) + , pending_upload_data_bytes(0) , uploaded_data_bytes(0) { } @@ -82,6 +85,7 @@ Uint32 PacketSocket::write(Uint32 max, bt::TimeStamp now) QMutexLocker locker(&mutex); if (curr_packet->getType() == PIECE) { up_speed->onData(ret, now); + pending_upload_data_bytes -= ret; uploaded_data_bytes += ret; } } else @@ -113,10 +117,12 @@ Uint32 PacketSocket::write(Uint32 max, bt::TimeStamp now) void PacketSocket::addPacket(Packet::Ptr packet) { + Q_ASSERT(!packet->sending()); QMutexLocker locker(&mutex); - if (packet->getType() == PIECE) + if (packet->getType() == PIECE) { data_packets.push_back(packet); - else + pending_upload_data_bytes += packet->getDataLength(); + } else control_packets.push_back(packet); // tell upload thread we have data ready should it be sleeping net::SocketMonitor::instance().signalPacketReady(); @@ -151,7 +157,7 @@ void PacketSocket::clearPieces(bool reject) if (p->getType() == bt::PIECE && !p->sending() && curr_packet != p) { if (reject) addPacket(Packet::Ptr(p->makeRejectOfPiece())); - + pending_upload_data_bytes -= p->getDataLength(); i = data_packets.erase(i); } else { i++; @@ -166,6 +172,7 @@ void PacketSocket::doNotSendPiece(const bt::Request &req, bool reject) while (i != data_packets.end()) { Packet::Ptr p = *i; if (p->isPiece(req) && !p->sending() && p != curr_packet) { + pending_upload_data_bytes -= p->getDataLength(); i = data_packets.erase(i); if (reject) { // queue a reject packet @@ -183,4 +190,10 @@ Uint32 PacketSocket::numPendingPieceUploads() const return data_packets.size(); } +Uint32 PacketSocket::numPendingPieceUploadBytes() const +{ + QMutexLocker locker(&mutex); + return pending_upload_data_bytes; +} + } diff --git a/src/net/packetsocket.h b/src/net/packetsocket.h index bbcb9e5e476d69b3b0bf341dc830663b9e1c9e45..8d06bbc5f4cf202fb8a09c7867f38f75ed5f4768 100644 --- a/src/net/packetsocket.h +++ b/src/net/packetsocket.h @@ -61,6 +61,9 @@ public: /// Get the number of pending piece uploads Uint32 numPendingPieceUploads() const; + /// Get the number of pending piece upload bytes (including message headers) + Uint32 numPendingPieceUploadBytes() const; + protected: /** * Preprocess the packet data, before it is sent. Default implementation does nothing. @@ -77,6 +80,7 @@ protected: bt::Packet::Ptr curr_packet; Uint32 ctrl_packets_sent; + Uint32 pending_upload_data_bytes; Uint32 uploaded_data_bytes; }; diff --git a/src/peer/peer.cpp b/src/peer/peer.cpp index 0f0fffda3ef755b1a654f6823cc18c0e6eff1ff0..09a7b9d923f2c7125b332318242397d85d636244 100644 --- a/src/peer/peer.cpp +++ b/src/peer/peer.cpp @@ -32,6 +32,8 @@ using namespace net; namespace bt { static const int MAX_METADATA_SIZE = 100 * 1024 * 1024; // maximum allowed metadata_size (up to 100 MiB) +static const unsigned MAX_PENDING_UPLOAD_BLOCKS = 512; // allow up to 8 MiB of 16KiB blocks +static const Uint32 MAX_PENDING_UPLOAD_BYTES = MAX_PENDING_UPLOAD_BLOCKS * (13 + 16384); static Uint32 peer_id_counter = 1; bool Peer::resolve_hostname = true; @@ -818,6 +820,12 @@ bool Peer::sendChunk(Uint32 index, Uint32 begin, Uint32 len, Chunk *ch) Out(SYS_CON | LOG_NOTICE) << "\tPiece : begin = " << begin << " len = " << len << endl; return false; } + if (sock->numPendingPieceUploads() >= MAX_PENDING_UPLOAD_BLOCKS || + sock->numPendingPieceUploadBytes() + 13 + len > MAX_PENDING_UPLOAD_BYTES) + { + Out(SYS_CON | LOG_NOTICE) << "Warning : rejecting piece request due to limit on pending uploads" << endl; + return false; + } /* Out(SYS_CON|LOG_DEBUG) << QString("Uploading %1 %2 %3 %4 %5") * .arg(index).arg(begin).arg(len).arg((quint64)ch,0,16).arg((quint64)ch->getData(),0,16) * << endl;;