refactor(network): clean up packet bundling

src/network/connection.cpp

@@ -67,7 +67,9 @@ namespace Network {
 
 bool Connection::beginPacket() {
 	if (m_IsBundle) {
-		m_BundlePacketPosition = 0;
+		m_BundledPacketsPosition = 0;
+		m_BundlePartFinished = false;
+
 		return true;
 	}
 
@@ -84,24 +86,23 @@ bool Connection::beginPacket() {
 
 bool Connection::endPacket() {
 	if (m_IsBundle) {
-		uint32_t innerPacketSize = m_BundlePacketPosition;
-
-		MUST_TRANSFER_BOOL((innerPacketSize > 0));
+		m_BundlePartFinished = true;
 
-		m_IsBundle = false;
+		auto innerPacketSize = m_BundledPacketsPosition;
 
-		if (m_BundlePacketInnerCount == 0) {
-			sendPacketType(PACKET_BUNDLE);
-			sendPacketNumber();
-		}
-		sendShort(innerPacketSize);
-		sendBytes(m_Packet, innerPacketSize);
+		MUST_TRANSFER_BOOL(sendU16(innerPacketSize));
+		MUST_TRANSFER_BOOL(sendBytes(m_BundledPacket, innerPacketSize));
 
 		m_BundlePacketInnerCount++;
-		m_IsBundle = true;
+		m_BundledPacketsPosition = 0;
+
 		return true;
 	}
 
+	if (m_BundleFinished) {
+		m_BundleFinished = false;
+	}
+
 	int r = m_UDP.endPacket();
 	if (r == 0) {
 		// This is usually just `ERR_ABRT` but the UDP client doesn't expose
@@ -114,40 +115,49 @@ bool Connection::endPacket() {
 }
 
 bool Connection::beginBundle() {
+	memset(m_BundledPacket, 0, sizeof(m_BundledPacket));
+	m_BundlePacketInnerCount = 0;
+
 	MUST_TRANSFER_BOOL(m_ServerFeatures.has(ServerFeatures::PROTOCOL_BUNDLE_SUPPORT));
 	MUST_TRANSFER_BOOL(m_Connected);
 	MUST_TRANSFER_BOOL(!m_IsBundle);
 	MUST_TRANSFER_BOOL(beginPacket());
 
+	MUST_TRANSFER_BOOL(sendPacketType(PACKET_BUNDLE));
+	MUST_TRANSFER_BOOL(sendPacketNumber());
+
 	m_IsBundle = true;
-	m_BundlePacketInnerCount = 0;
+	m_BundleFinished = false;
+	m_BundlePartFinished = false;
 	return true;
 }
 
 bool Connection::endBundle() {
 	MUST_TRANSFER_BOOL(m_IsBundle);
 
 	m_IsBundle = false;
-
-	MUST_TRANSFER_BOOL((m_BundlePacketInnerCount > 0));
+	m_BundleFinished = true;
+	m_BundlePartFinished = false;
 
 	return endPacket();
 }
 
 size_t Connection::write(const uint8_t* buffer, size_t size) {
-	if (m_IsBundle) {
-		if (m_BundlePacketPosition + size > sizeof(m_Packet)) {
+	if (m_IsBundle && !m_BundlePartFinished) {
+		if (m_BundledPacketsPosition + size > sizeof(m_BundledPacket)) {
+			m_Logger.error("Bundled packet too large");
 			return 0;
 		}
-		memcpy(m_Packet + m_BundlePacketPosition, buffer, size);
-		m_BundlePacketPosition += size;
+
+		memcpy(m_BundledPacket + m_BundledPacketsPosition, buffer, size);
+		m_BundledPacketsPosition += size;
+
 		return size;
 	}
+
 	return m_UDP.write(buffer, size);
 }
 
-size_t Connection::write(uint8_t byte) { return write(&byte, 1); }
-
 bool Connection::sendFloat(float f) {
 	convert_to_chars(f, m_Buf);
 
@@ -156,7 +166,7 @@ bool Connection::sendFloat(float f) {
 
 bool Connection::sendByte(uint8_t c) { return write(&c, 1) != 0; }
 
-bool Connection::sendShort(int16_t i) {
+bool Connection::sendU16(uint16_t i) {
 	convert_to_chars(i, m_Buf);
 
 	return write(m_Buf, sizeof(i)) != 0;
@@ -197,13 +207,7 @@ bool Connection::sendShortString(const char* str) {
 	return true;
 }
 
-bool Connection::sendPacketType(uint8_t type) {
-	MUST_TRANSFER_BOOL(sendByte(0));
-	MUST_TRANSFER_BOOL(sendByte(0));
-	MUST_TRANSFER_BOOL(sendByte(0));
-
-	return sendByte(type);
-}
+bool Connection::sendPacketType(uint8_t type) { return sendInt(type); }
 
 bool Connection::sendLongString(const char* str) {
 	int size = strlen(str);
@@ -618,13 +622,14 @@ void Connection::update() {
 	std::vector<Sensor*>& sensors = sensorManager.getSensors();
 
 	updateSensorState(sensors);
-	maybeRequestFeatureFlags();
 
 	if (!m_Connected) {
 		searchForServer();
 		return;
 	}
 
+	maybeRequestFeatureFlags();
+
 	if (m_LastPacketTimestamp + TIMEOUT < millis()) {
 		statusManager.setStatus(SlimeVR::Status::SERVER_CONNECTING, true);
 

src/network/connection.h

@@ -26,11 +26,11 @@
 #include <Arduino.h>
 #include <WiFiUdp.h>
 
+#include "featureflags.h"
 #include "globals.h"
 #include "quat.h"
 #include "sensors/sensor.h"
 #include "wifihandler.h"
-#include "featureflags.h"
 
 namespace SlimeVR {
 namespace Network {
@@ -107,28 +107,25 @@ class Connection {
 	);
 #endif
 
-	const ServerFeatures& getServerFeatureFlags() {
-		return m_ServerFeatures;
-	}
+	const ServerFeatures& getServerFeatureFlags() { return m_ServerFeatures; }
 
 	bool beginBundle();
 	bool endBundle();
 
 private:
-	void updateSensorState(std::vector<Sensor *> & sensors);
+	void updateSensorState(std::vector<Sensor*>& sensors);
 	void maybeRequestFeatureFlags();
 
 	bool beginPacket();
 	bool endPacket();
 
-	size_t write(const uint8_t *buffer, size_t size);
-	size_t write(uint8_t byte);
+	size_t write(const uint8_t* buffer, size_t size);
 
 	bool sendPacketType(uint8_t type);
 	bool sendPacketNumber();
 	bool sendFloat(float f);
 	bool sendByte(uint8_t c);
-	bool sendShort(int16_t i);
+	bool sendU16(uint16_t c);
 	bool sendInt(int32_t i);
 	bool sendU64(uint64_t l);
 	bool sendBytes(const uint8_t* c, size_t length);
@@ -152,7 +149,11 @@ class Connection {
 	SlimeVR::Logging::Logger m_Logger = SlimeVR::Logging::Logger("UDPConnection");
 
 	WiFiUDP m_UDP;
-	unsigned char m_Packet[128];  // buffer for incoming packets
+	/*
+	  The current incoming packet that is being handled
+	  TODO: remove this from the class and make it a local variable
+	*/
+	uint8_t m_Packet[128];
 	uint64_t m_PacketNumber = 0;
 
 	int m_ServerPort = 6969;
@@ -168,7 +169,13 @@ class Connection {
 	ServerFeatures m_ServerFeatures{};
 
 	bool m_IsBundle = false;
-	uint16_t m_BundlePacketPosition = 0;
+	bool m_BundleFinished = false;
+	bool m_BundlePartFinished = false;
+	/* The packet that is currently being bundled */
+	uint8_t m_BundledPacket[128];
+	/* Position in `m_BundledPacket` */
+	uint16_t m_BundledPacketsPosition = 0;
+	/* Count of packets that are in the currently forming bundle */
 	uint16_t m_BundlePacketInnerCount = 0;
 
 	unsigned char m_Buf[8];