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

1 files changed, 49 insertions, 1 deletions

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 */