Codechange: encapsulate reading data from sockets into Packets to prevent packet state modifications outside of the Packet

4 files changed, 75 insertions, 15 deletions

diff --git a/src/network/core/packet.cpp b/src/network/core/packet.cpp
index 54f5a79e1..5cb6716bc 100644
--- a/src/network/core/packet.cpp
+++ b/src/network/core/packet.cpp
@@ -17,17 +17,24 @@
 #include "../../safeguards.h"
 
 /**
- * Create a packet that is used to read from a network socket
- * @param cs the socket handler associated with the socket we are reading from
+ * Create a packet that is used to read from a network socket.
+ * @param cs                The socket handler associated with the socket we are reading from.
+ * @param initial_read_size The initial amount of data to transfer from the socket into the
+ *                          packet. This defaults to just the required bytes to determine the
+ *                          packet's size. That default is the wanted for streams such as TCP
+ *                          as you do not want to read data of the next packet yet. For UDP
+ *                          you need to read the whole packet at once otherwise you might
+ *                          loose some the data of the packet, so there you pass the maximum
+ *                          size for the packet you expect from the network.
  */
-Packet::Packet(NetworkSocketHandler *cs)
+Packet::Packet(NetworkSocketHandler *cs, size_t initial_read_size)
 {
 	assert(cs != nullptr);
 
 	this->cs     = cs;
 	this->next   = nullptr;
 	this->pos    = 0; // We start reading from here
-	this->size   = 0;
+	this->size   = static_cast<int>(initial_read_size);
 	this->buffer = MallocT<byte>(SEND_MTU);
 }
 
@@ -336,3 +343,12 @@ void Packet::Recv_string(char *buffer, size_t size, StringValidationSettings set
 
 	str_validate(bufp, last, settings);
 }
+
+/**
+ * Get the amount of bytes that are still available for the Transfer functions.
+ * @return The number of bytes that still have to be transfered.
+ */
+size_t Packet::RemainingBytesToTransfer() const
+{
+	return this->size - this->pos;
+}
diff --git a/src/network/core/packet.h b/src/network/core/packet.h
index 901d3f593..3eee0522f 100644
--- a/src/network/core/packet.h
+++ b/src/network/core/packet.h
@@ -12,9 +12,11 @@
 #ifndef NETWORK_CORE_PACKET_H
 #define NETWORK_CORE_PACKET_H
 
+#include "os_abstraction.h"
 #include "config.h"
 #include "core.h"
 #include "../../string_type.h"
+#include <functional>
 
 typedef uint16 PacketSize; ///< Size of the whole packet.
 typedef uint8  PacketType; ///< Identifier for the packet
@@ -56,7 +58,7 @@ private:
 	NetworkSocketHandler *cs;
 
 public:
-	Packet(NetworkSocketHandler *cs);
+	Packet(NetworkSocketHandler *cs, size_t initial_read_size = sizeof(PacketSize));
 	Packet(PacketType type);
 	~Packet();
 
@@ -83,6 +85,52 @@ public:
 	uint32 Recv_uint32();
 	uint64 Recv_uint64();
 	void   Recv_string(char *buffer, size_t size, StringValidationSettings settings = SVS_REPLACE_WITH_QUESTION_MARK);
+
+	size_t RemainingBytesToTransfer() const;
+
+	/**
+	 * Transfer data from the given function into the packet. It starts writing at the
+	 * position the last transfer stopped.
+	 *
+	 * Examples of functions that can be used to transfer data into a packet are TCP's
+	 * recv and UDP's recvfrom functions. They will directly write their data into the
+	 * packet without an intermediate buffer.
+	 * Examples of functions that can be used to transfer data from a packet are TCP's
+	 * send and UDP's sendto functions. They will directly read the data from the packet's
+	 * buffer without an intermediate buffer.
+	 * These are functions are special in a sense as even though the packet can send or
+	 * receive an amount of data, those functions can say they only processed a smaller
+	 * amount, so special handling is required to keep the position pointers correct.
+	 * Most of these transfer functions are in the form function(source, buffer, amount, ...),
+	 * so the template of this function will assume that as the base parameter order.
+	 *
+	 * This will attempt to write all the remaining bytes into the packet. It updates the
+	 * position based on how many bytes were actually written by the called transfer_function.
+	 * @param transfer_function The function to pass the buffer as second parameter and the
+	 *                          amount to read as third parameter. It returns the amount that
+	 *                          was read or -1 upon errors.
+	 * @param source            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 S    The type of the source.
+	 * @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 S, typename ... Args>
+	ssize_t TransferIn(F transfer_function, S source, Args&& ... args)
+	{
+		size_t amount = this->RemainingBytesToTransfer();
+		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. */
+		char *input_buffer = reinterpret_cast<char*>(this->buffer + this->pos);
+		ssize_t bytes = transfer_function(source, input_buffer, static_cast<A>(amount), std::forward<Args>(args)...);
+		if (bytes > 0) this->pos += bytes;
+		return bytes;
+	}
 };
 
 #endif /* NETWORK_CORE_PACKET_H */
diff --git a/src/network/core/tcp.cpp b/src/network/core/tcp.cpp
index 1461a9298..aa1e1cbed 100644
--- a/src/network/core/tcp.cpp
+++ b/src/network/core/tcp.cpp
@@ -156,9 +156,8 @@ Packet *NetworkTCPSocketHandler::ReceivePacket()
 
 	/* Read packet size */
 	if (!p->HasPacketSizeData()) {
-		while (!p->HasPacketSizeData()) {
-		/* Read the size of the packet */
-			res = recv(this->sock, (char*)p->buffer + p->pos, sizeof(PacketSize) - p->pos, 0);
+		while (p->RemainingBytesToTransfer() != 0) {
+			res = p->TransferIn<int>(recv, this->sock, 0);
 			if (res == -1) {
 				int err = GET_LAST_ERROR();
 				if (err != EWOULDBLOCK) {
@@ -175,7 +174,6 @@ Packet *NetworkTCPSocketHandler::ReceivePacket()
 				this->CloseConnection();
 				return nullptr;
 			}
-			p->pos += res;
 		}
 
 		/* Parse the size in the received packet and if not valid, close the connection. */
@@ -186,8 +184,8 @@ Packet *NetworkTCPSocketHandler::ReceivePacket()
 	}
 
 	/* Read rest of packet */
-	while (p->pos < p->size) {
-		res = recv(this->sock, (char*)p->buffer + p->pos, p->size - p->pos, 0);
+	while (p->RemainingBytesToTransfer() != 0) {
+		res = p->TransferIn<int>(recv, this->sock, 0);
 		if (res == -1) {
 			int err = GET_LAST_ERROR();
 			if (err != EWOULDBLOCK) {
@@ -204,8 +202,6 @@ Packet *NetworkTCPSocketHandler::ReceivePacket()
 			this->CloseConnection();
 			return nullptr;
 		}
-
-		p->pos += res;
 	}
 
 	/* Prepare for receiving a new packet */
diff --git a/src/network/core/udp.cpp b/src/network/core/udp.cpp
index 614ae8c3b..398f53142 100644
--- a/src/network/core/udp.cpp
+++ b/src/network/core/udp.cpp
@@ -119,12 +119,12 @@ void NetworkUDPSocketHandler::ReceivePackets()
 			struct sockaddr_storage client_addr;
 			memset(&client_addr, 0, sizeof(client_addr));
 
-			Packet p(this);
+			Packet p(this, SEND_MTU);
 			socklen_t client_len = sizeof(client_addr);
 
 			/* Try to receive anything */
 			SetNonBlocking(s.second); // Some OSes seem to lose the non-blocking status of the socket
-			int nbytes = recvfrom(s.second, (char*)p.buffer, SEND_MTU, 0, (struct sockaddr *)&client_addr, &client_len);
+			ssize_t nbytes = p.TransferIn<int>(recvfrom, s.second, 0, (struct sockaddr *)&client_addr, &client_len);
 
 			/* Did we get the bytes for the base header of the packet? */
 			if (nbytes <= 0) break;    // No data, i.e. no packet