support upload-only extension message

include/libtorrent/bt_peer_connection.hpp

@@ -103,6 +103,8 @@ namespace libtorrent
 
 		void start();
 
+		enum { upload_only_msg = 2 };
+
 		~bt_peer_connection();
 		
 #ifndef TORRENT_DISABLE_ENCRYPTION
@@ -203,8 +205,8 @@ namespace libtorrent
 		void write_handshake();
 #ifndef TORRENT_DISABLE_EXTENSIONS
 		void write_extensions();
+		void write_upload_only();
 #endif
-		void write_chat_message(const std::string& msg);
 		void write_metadata(std::pair<int, int> req);
 		void write_metadata_request(std::pair<int, int> req);
 		void write_keepalive();
@@ -354,14 +356,18 @@ private:
 		std::vector<range> m_payloads;
 
 #ifndef TORRENT_DISABLE_EXTENSIONS
+		// the message ID for upload only message
+		// 0 if not supported
+		int m_upload_only_id;
+
+		char m_reserved_bits[8];
 		// this is set to true if the handshake from
 		// the peer indicated that it supports the
 		// extension protocol
-		bool m_supports_extensions;
-		char m_reserved_bits[8];
+		bool m_supports_extensions:1;
 #endif
-		bool m_supports_dht_port;
-		bool m_supports_fast;
+		bool m_supports_dht_port:1;
+		bool m_supports_fast:1;
 
 #ifndef TORRENT_DISABLE_ENCRYPTION
 		// this is set to true after the encryption method has been

include/libtorrent/peer_connection.hpp

@@ -262,12 +262,7 @@ namespace libtorrent
 
 		bool is_seed() const;
 
-		void set_upload_only(bool u)
-		{
-			m_upload_only = u;
-			disconnect_if_redundant();
-		}
-
+		void set_upload_only(bool u);
 		bool upload_only() const { return m_upload_only; }
 
 		// will send a keep-alive message to the peer

include/libtorrent/torrent.hpp

@@ -163,8 +163,15 @@ namespace libtorrent
 		void start_announcing();
 		void stop_announcing();
 
+		void send_upload_only();
+
 		void set_upload_mode(bool b);
 		bool upload_mode() const { return m_upload_mode; }
+		bool is_upload_only() const
+		{
+			return (((is_finished() && !super_seeding()) || upload_mode())
+				&& !m_ses.settings().lazy_bitfields);
+		}
 
 		int seed_rank(session_settings const& s) const;
 

src/bt_peer_connection.cpp

@@ -97,6 +97,7 @@ namespace libtorrent
 			, peerinfo)
 		, m_state(read_protocol_identifier)
 #ifndef TORRENT_DISABLE_EXTENSIONS
+		, m_upload_only_id(0)
 		, m_supports_extensions(false)
 #endif
 		, m_supports_dht_port(false)
@@ -1433,6 +1434,13 @@ namespace libtorrent
 			return;
 		}
 
+		if (extended_id == upload_only_msg)
+		{
+			if (!packet_finished()) return;
+			set_upload_only(detail::read_uint8(recv_buffer.begin));
+			return;
+		}
+
 #ifndef TORRENT_DISABLE_EXTENSIONS
 		for (extension_list_t::iterator i = m_extensions.begin()
 			, end(m_extensions.end()); i != end; ++i)
@@ -1484,6 +1492,10 @@ namespace libtorrent
 		if (is_disconnecting()) return;
 #endif
 
+		// upload_only
+		if (lazy_entry const* m = root.dict_find_dict("m"))
+			m_upload_only_id = m->dict_find_int_value("upload_only", 0);
+
 		// there is supposed to be a remote listen port
 		int listen_port = root.dict_find_int_value("p");
 		if (listen_port > 0 && peer_info_struct() != 0)
@@ -1501,7 +1513,7 @@ namespace libtorrent
 		int reqq = root.dict_find_int_value("reqq");
 		if (reqq > 0) m_max_out_request_queue = reqq;
 
-		if (root.dict_find_int_value("upload_only"))
+		if (root.dict_find_int_value("upload_only", 0))
 			set_upload_only(true);
 
 		std::string myip = root.dict_find_string_value("yourip");
@@ -1589,6 +1601,22 @@ namespace libtorrent
 		return packet_finished();
 	}
 
+#ifndef TORRENT_DISABLE_EXTENSIONS
+	void bt_peer_connection::write_upload_only()
+	{
+		INVARIANT_CHECK;
+		
+		boost::shared_ptr<torrent> t = associated_torrent().lock();
+		if (m_upload_only_id == 0) return;
+
+		char msg[7] = {0, 0, 0, 3, msg_extended};
+		char* ptr = msg + 5;
+		detail::write_uint8(m_upload_only_id, ptr);
+		detail::write_uint8(t->is_upload_only(), ptr);
+		send_buffer(msg, sizeof(msg));
+	}
+#endif
+
 	void bt_peer_connection::write_keepalive()
 	{
 		INVARIANT_CHECK;
@@ -1793,10 +1821,8 @@ namespace libtorrent
 		TORRENT_ASSERT(m_supports_extensions);
 		TORRENT_ASSERT(m_sent_handshake);
 
-		entry handshake(entry::dictionary_t);
-		entry extension_list(entry::dictionary_t);
-
-		handshake["m"] = extension_list;
+		entry handshake;
+		entry::dictionary_type& m = handshake["m"].dict();
 
 		// only send the port in case we bade the connection
 		// on incoming connections the other end already knows
@@ -1811,9 +1837,11 @@ namespace libtorrent
 		boost::shared_ptr<torrent> t = associated_torrent().lock();
 		TORRENT_ASSERT(t);
 
+		m["upload_only"] = upload_only_msg;
+
 		// if we're using lazy bitfields or if we're super seeding, don't say
 		// we're upload only, since it might make peers disconnect
-		if (t->is_finished() && !t->super_seeding() && !m_ses.settings().lazy_bitfields)
+		if (t->is_upload_only())
 			handshake["upload_only"] = 1;
 
 		tcp::endpoint ep = m_ses.get_ipv6_interface();

src/peer_connection.cpp

@@ -4967,5 +4967,14 @@ namespace libtorrent
 		boost::shared_ptr<torrent> t = m_torrent.lock();
 		return m_num_pieces == (int)m_have_piece.size() && m_num_pieces > 0 && t && t->valid_metadata();
 	}
+
+	void peer_connection::set_upload_only(bool u)
+	{
+		m_upload_only = u;
+		boost::shared_ptr<torrent> t = associated_torrent().lock();
+		t->get_policy().set_seed(m_peer_info, u);
+		disconnect_if_redundant();
+	}
+
 }
 

src/torrent.cpp

@@ -423,12 +423,27 @@ namespace libtorrent
 		}
 	}
 
+	void torrent::send_upload_only()
+	{
+#ifndef TORRENT_DISABLE_EXTENSIONS
+		for (std::set<peer_connection*>::iterator i = m_connections.begin()
+			, end(m_connections.end()); i != end; ++i)
+		{
+			bt_peer_connection* p = dynamic_cast<bt_peer_connection*>(*i);
+			if (p == 0) continue;
+			p->write_upload_only();
+		}
+#endif
+	}
+
 	void torrent::set_upload_mode(bool b)
 	{
 		if (b == m_upload_mode) return;
 
 		m_upload_mode = b;
 
+		send_upload_only();
+
 		if (m_upload_mode)
 		{
 			// clear request queues of all peers
@@ -4253,6 +4268,8 @@ namespace libtorrent
 		// to make sure we're cleared the piece picker
 		if (is_seed()) completed();
 
+		send_upload_only();
+
 		// disconnect all seeds
 		// TODO: should disconnect all peers that have the pieces we have
 		// not just seeds
@@ -4294,6 +4311,8 @@ namespace libtorrent
 		set_state(torrent_status::downloading);
 		set_queue_position((std::numeric_limits<int>::max)());
 		m_policy.recalculate_connect_candidates();
+
+		send_upload_only();
 	}
 
 	// called when torrent is complete (all pieces downloaded)