Codechange: encapsulate writing data from Packets into sockets/files/buffers to prevent packet state modifications outside of the Packet

4 files changed, 57 insertions, 7 deletions

diff --git a/src/network/core/packet.h b/src/network/core/packet.h
index 3eee0522f..b091d8a7e 100644
--- a/src/network/core/packet.h
+++ b/src/network/core/packet.h
@@ -89,6 +89,58 @@ public:
 	size_t RemainingBytesToTransfer() const;
 
 	/**
+	 * Transfer data from the packet to the given function. It starts reading at the
+	 * position the last transfer stopped.
+	 * See Packet::TransferIn for more information about transferring data to functions.
+	 * @param transfer_function The function to pass the buffer as second parameter and the
+	 *                          amount to write as third parameter. It returns the amount that
+	 *                          was written or -1 upon errors.
+	 * @param limit             The maximum amount of bytes to transfer.
+	 * @param destination       The first parameter of the transfer function.
+	 * @param args              The fourth and further parameters to the transfer function, if any.
+	 * @return The return value of the transfer_function.
+	 */
+	template <
+		typename A = size_t, ///< The type for the amount to be passed, so it can be cast to the right type.
+		typename F,          ///< The type of the function.
+		typename D,          ///< The type of the destination.
+		typename ... Args>   ///< The types of the remaining arguments to the function.
+	ssize_t TransferOutWithLimit(F transfer_function, size_t limit, D destination, Args&& ... args)
+	{
+		size_t amount = std::min(this->RemainingBytesToTransfer(), limit);
+		if (amount == 0) return 0;
+
+		assert(this->pos < this->buffer.size());
+		assert(this->pos + amount <= this->buffer.size());
+		/* Making buffer a char means casting a lot in the Recv/Send functions. */
+		const char *output_buffer = reinterpret_cast<const char*>(this->buffer + this->pos);
+		ssize_t bytes = transfer_function(destination, output_buffer, static_cast<A>(amount), std::forward<Args>(args)...);
+		if (bytes > 0) this->pos += bytes;
+		return bytes;
+	}
+
+	/**
+	 * Transfer data from the packet to the given function. It starts reading at the
+	 * position the last transfer stopped.
+	 * See Packet::TransferIn for more information about transferring data to functions.
+	 * @param transfer_function The function to pass the buffer as second parameter and the
+	 *                          amount to write as third parameter. It returns the amount that
+	 *                          was written or -1 upon errors.
+	 * @param destination       The first parameter of the transfer function.
+	 * @param args              The fourth and further parameters to the transfer function, if any.
+	 * @tparam A    The type for the amount to be passed, so it can be cast to the right type.
+	 * @tparam F    The type of the transfer_function.
+	 * @tparam D    The type of the destination.
+	 * @tparam Args The types of the remaining arguments to the function.
+	 * @return The return value of the transfer_function.
+	 */
+	template <typename A = size_t, typename F, typename D, typename ... Args>
+	ssize_t TransferOut(F transfer_function, D destination, Args&& ... args)
+	{
+		return TransferOutWithLimit<A>(transfer_function, std::numeric_limits<size_t>::max(), destination, std::forward<Args>(args)...);
+	}
+
+	/**
 	 * Transfer data from the given function into the packet. It starts writing at the
 	 * position the last transfer stopped.
 	 *
diff --git a/src/network/core/tcp.cpp b/src/network/core/tcp.cpp
index aa1e1cbed..ab18f47a8 100644
--- a/src/network/core/tcp.cpp
+++ b/src/network/core/tcp.cpp
@@ -103,7 +103,7 @@ SendPacketsState NetworkTCPSocketHandler::SendPackets(bool closing_down)
 
 	p = this->packet_queue;
 	while (p != nullptr) {
-		res = send(this->sock, (const char*)p->buffer + p->pos, p->size - p->pos, 0);
+		res = p->TransferOut<int>(send, this->sock, 0);
 		if (res == -1) {
 			int err = GET_LAST_ERROR();
 			if (err != EWOULDBLOCK) {
@@ -122,10 +122,8 @@ SendPacketsState NetworkTCPSocketHandler::SendPackets(bool closing_down)
 			return SPS_CLOSED;
 		}
 
-		p->pos += res;
-
 		/* Is this packet sent? */
-		if (p->pos == p->size) {
+		if (p->RemainingBytesToTransfer() == 0) {
 			/* Go to the next packet */
 			this->packet_queue = p->next;
 			delete p;
diff --git a/src/network/core/tcp_listen.h b/src/network/core/tcp_listen.h
index 1f073aa73..53a3d57cc 100644
--- a/src/network/core/tcp_listen.h
+++ b/src/network/core/tcp_listen.h
@@ -63,7 +63,7 @@ public:
 
 					DEBUG(net, 1, "[%s] Banned ip tried to join (%s), refused", Tsocket::GetName(), entry.c_str());
 
-					if (send(s, (const char*)p.buffer, p.size, 0) < 0) {
+					if (p.TransferOut<int>(send, s, 0) < 0) {
 						DEBUG(net, 0, "send failed with error %d", GET_LAST_ERROR());
 					}
 					closesocket(s);
@@ -80,7 +80,7 @@ public:
 				Packet p(Tfull_packet);
 				p.PrepareToSend();
 
-				if (send(s, (const char*)p.buffer, p.size, 0) < 0) {
+				if (p.TransferOut<int>(send, s, 0) < 0) {
 					DEBUG(net, 0, "send failed with error %d", GET_LAST_ERROR());
 				}
 				closesocket(s);
diff --git a/src/network/core/udp.cpp b/src/network/core/udp.cpp
index 398f53142..8e476f4e2 100644
--- a/src/network/core/udp.cpp
+++ b/src/network/core/udp.cpp
@@ -99,7 +99,7 @@ void NetworkUDPSocketHandler::SendPacket(Packet *p, NetworkAddress *recv, bool a
 		}
 
 		/* Send the buffer */
-		int res = sendto(s.second, (const char*)p->buffer, p->size, 0, (const struct sockaddr *)send.GetAddress(), send.GetAddressLength());
+		ssize_t res = p->TransferOut<int>(sendto, s.second, 0, (const struct sockaddr *)send.GetAddress(), send.GetAddressLength());
 		DEBUG(net, 7, "[udp] sendto(%s)", send.GetAddressAsString().c_str());
 
 		/* Check for any errors, but ignore it otherwise */