SSX3LobbyServer/lib/aries/MessageIo.hpp

#include <aries/Message.hpp>
#include <aries/Tags.hpp>
#include <base/types.hpp>
#include <boost/asio/read.hpp>
#include <boost/asio/write.hpp>
#include <exception>
#include <impl/asio_config.hpp>

#include "boost/asio/buffer.hpp"

namespace ls::aries {

	// Mostly to be conservative. I don't think the game will really care but bleh
	constexpr static auto MAX_TAGFIELD_SIZE_IN_MB = 1;
	constexpr static auto MAX_TAGFIELD_SIZE_IN_BYTES = MAX_TAGFIELD_SIZE_IN_MB * (1024 * 1024);


	namespace errors {

		struct TagPayloadTooLarge : std::exception {
			TagPayloadTooLarge(u32 size)
				: payloadSize(size) {
				whatStr = std::format("Tag field data size is over {} MB (Max is {}MB).", (static_cast<u32>(payloadSize) / 1024 / 1024), MAX_TAGFIELD_SIZE_IN_MB);
			}

			const char* what() const noexcept override {
				return whatStr.c_str();
			}

		   private:
			u32 payloadSize;
			std::string whatStr;
		};

	} // namespace errors

	/// Reads an Aries message from an Boost.Asio async read stream.
	/// Returns the raw Aries message header and the tag field buffer.
	template <class AsyncReadStream>
	base::Awaitable<RawAriesMessage> AsyncReadAriesMessage(AsyncReadStream& stream) {
		RawAriesMessage res;

		// Read the header first
		co_await asio::async_read(stream, asio::buffer(&res.header, sizeof(res.header)), asio::deferred);

		auto realPayloadSize = res.header.messageSize - sizeof(res.header);

		// Read tag payload (if there is one)
		if(res.header.messageSize != sizeof(res.header)) {
			// Sanity check. I don't expect game payloads to ever reach this large, but who knows.
			if(realPayloadSize > MAX_TAGFIELD_SIZE_IN_BYTES)
				throw errors::TagPayloadTooLarge(realPayloadSize);

			res.tagFields.resize(realPayloadSize);

			co_await asio::async_read(stream, asio::buffer(res.tagFields), asio::deferred);
		}

		co_return res;
	}

	template <class AsyncReadStream>
	base::Awaitable<void> AsyncWriteAriesMessage(AsyncReadStream& stream, const RawAriesMessage& message) {
		auto realTagFieldSize = message.header.messageSize - sizeof(message.header);

		// Make sure *we* won't write a message the official Aries protocol
		// won't like (even though it'd probably just crash, it's nice for us to do this.)
		if(message.header.messageSize != sizeof(message.header)) {
			// Sanity check. I don't expect game payloads to ever reach this large, but who knows.
			if(realTagFieldSize > MAX_TAGFIELD_SIZE_IN_BYTES)
				throw errors::TagPayloadTooLarge(realTagFieldSize);
		}

		// Our buffer list. We pass this to asio::async_write so we only actually have to perform
		// one (scatter-gather) I/O operation in this function
		std::array<asio::const_buffer, 2> buffers = {
			asio::buffer(&message.header, sizeof(message.header)),
			asio::buffer(message.tagFields, realTagFieldSize)
		};

		co_await asio::async_write(stream, buffers, asio::deferred);
		co_return;
	}

} // namespace ls::aries