Refactored websocket message.

1 files changed, 120 insertions, 36 deletions

diff --git a/zenhttp/include/zenhttp/websocket.h b/zenhttp/include/zenhttp/websocket.h
index 1ab9b4804..336d98b42 100644
--- a/zenhttp/include/zenhttp/websocket.h
+++ b/zenhttp/include/zenhttp/websocket.h
@@ -1,10 +1,13 @@
 // Copyright Epic Games, Inc. All Rights Reserved.
 
+#include <zencore/compactbinarypackage.h>
 #include <zencore/memory.h>
 
 #include <compare>
 #include <functional>
+#include <future>
 #include <memory>
+#include <optional>
 
 #pragma once
 
@@ -14,10 +17,11 @@ class io_context;
 
 namespace zen {
 
-class CbPackage;
-class CbObject;
 class BinaryWriter;
 
+/**
+ * A unique socket ID.
+ */
 class WebSocketId
 {
 	static std::atomic_uint32_t NextId;
@@ -37,45 +41,132 @@ private:
 	uint32_t m_Value{};
 };
 
+/**
+ * Type of web socket message.
+ */
+enum class WebSocketMessageType : uint8_t
+{
+	kInvalid,
+	kNotification,
+	kRequest,
+	kResponse
+};
+
+/**
+ * Web socket message.
+ */
+class WebSocketMessage
+{
+	struct Header
+	{
+		static constexpr uint32_t HeaderMagic = 0x7a776d68;	 // zwmh
+
+		uint64_t			 MessageSize{};
+		uint32_t			 Magic{HeaderMagic};
+		uint32_t			 CorrelationId{};
+		uint32_t			 Crc32{};
+		WebSocketMessageType MessageType{};
+		uint8_t				 Reserved[3] = {0};
+
+		bool IsValid() const;
+	};
+
+	static_assert(sizeof Header == 24);
+
+	static std::atomic_uint32_t NextCorrelationId;
+
+public:
+	static constexpr size_t HeaderSize = sizeof(Header);
+
+	WebSocketMessage() = default;
+
+	WebSocketId			 SocketId() const { return m_SocketId; }
+	void				 SetSocketId(WebSocketId Id) { m_SocketId = Id; }
+	void				 SetMessageType(WebSocketMessageType MessageType);
+	WebSocketMessageType MessageType() const { return m_Header.MessageType; }
+	uint64_t			 MessageSize() const { return m_Header.MessageSize; }
+	void				 SetCorrelationId(uint32_t Id) { m_Header.CorrelationId = Id; }
+	uint32_t			 CorrelationId() const { return m_Header.CorrelationId; }
+
+	const CbPackage& Body() const { return m_Body.value(); }
+	void			 SetBody(CbPackage&& Body);
+	void			 SetBody(CbObject&& Body);
+	bool			 HasBody() const { return m_Body.has_value(); }
+
+	void Save(BinaryWriter& Writer);
+	bool TryLoadHeader(MemoryView Memory);
+
+	bool IsValid() const { return m_Header.MessageType != WebSocketMessageType::kInvalid; }
+
+private:
+	Header					 m_Header{};
+	WebSocketId				 m_SocketId{};
+	std::optional<CbPackage> m_Body;
+};
+
 class WebSocketServer;
 
+/**
+ * Base class for handling web socket requests and notifications from connected client(s).
+ */
 class WebSocketService
 {
 public:
 	virtual ~WebSocketService() = default;
 
-	virtual bool HandleMessage(WebSocketId Id, const CbPackage& Msg) = 0;
-	void		 Configure(WebSocketServer& Server);
+	void Configure(WebSocketServer& Server);
+
+	virtual bool HandleRequest(const WebSocketMessage&) { ZEN_ASSERT(false); }
+	virtual void HandleNotification(const WebSocketMessage&) { ZEN_ASSERT(false); }
 
 protected:
 	WebSocketService() = default;
 
-	void PublishMessage(WebSocketId Id, CbPackage&& Msg);
-	void PublishMessage(WebSocketId Id, CbObject&& Msg);
+	virtual void RegisterHandlers(WebSocketServer& Server) = 0;
+
+	WebSocketServer& SocketServer()
+	{
+		ZEN_ASSERT(m_SocketServer);
+		return *m_SocketServer;
+	}
 
 private:
-	WebSocketServer* m_Server = nullptr;
+	WebSocketServer* m_SocketServer{};
 };
 
+/**
+ * Server options.
+ */
 struct WebSocketServerOptions
 {
 	uint16_t Port		 = 8848;
 	uint32_t ThreadCount = 1;
 };
 
+/**
+ * The web socket server manages client connections and routing of requests and notifications.
+ */
 class WebSocketServer
 {
 public:
 	virtual ~WebSocketServer() = default;
 
-	virtual void RegisterService(WebSocketService& Service)		 = 0;
-	virtual bool Run(const WebSocketServerOptions& Options)		 = 0;
-	virtual void Shutdown()										 = 0;
-	virtual void PublishMessage(WebSocketId Id, CbPackage&& Msg) = 0;
+	virtual bool Run(const WebSocketServerOptions& Options) = 0;
+	virtual void Shutdown()									= 0;
+
+	virtual void RegisterService(WebSocketService& Service)									  = 0;
+	virtual void RegisterNotificationHandler(std::string_view Key, WebSocketService& Service) = 0;
+	virtual void RegisterRequestHandler(std::string_view Key, WebSocketService& Service)	  = 0;
+
+	virtual void SendNotification(WebSocketMessage&& Notification) = 0;
+	virtual void SendResponse(WebSocketMessage&& Response)		   = 0;
 
 	static std::unique_ptr<WebSocketServer> Create();
 };
 
+/**
+ * The state of the web socket.
+ */
 enum class WebSocketState : uint32_t
 {
 	kNone,
@@ -85,6 +176,9 @@ enum class WebSocketState : uint32_t
 	kError
 };
 
+/**
+ * Type of web socket client event.
+ */
 enum class WebSocketEvent : uint32_t
 {
 	kConnected,
@@ -92,6 +186,9 @@ enum class WebSocketEvent : uint32_t
 	kError
 };
 
+/**
+ * Web socket client connection info.
+ */
 struct WebSocketConnectInfo
 {
 	std::string				 Host;
@@ -101,40 +198,27 @@ struct WebSocketConnectInfo
 	uint16_t				 Version{13};
 };
 
+/**
+ * A connection to a web socket server for sending requests and listening for notifications.
+ */
 class WebSocketClient
 {
 public:
-	using EventCallback	  = std::function<void()>;
-	using MessageCallback = std::function<void(CbPackage&)>;
+	using EventCallback		   = std::function<void()>;
+	using NotificationCallback = std::function<void(WebSocketMessage&&)>;
 
 	virtual ~WebSocketClient() = default;
 
-	virtual bool		   Connect(const WebSocketConnectInfo& Info) = 0;
-	virtual void		   Disconnect()								 = 0;
-	virtual bool		   IsConnected() const						 = 0;
-	virtual WebSocketState State() const							 = 0;
-	virtual void		   SendMsg(CbObject&& Msg)					 = 0;
-	virtual void		   SendMsg(CbPackage&& Msg)					 = 0;
+	virtual std::future<bool> Connect(const WebSocketConnectInfo& Info) = 0;
+	virtual void			  Disconnect()								= 0;
+	virtual bool			  IsConnected() const						= 0;
+	virtual WebSocketState	  State() const								= 0;
 
-	virtual void On(WebSocketEvent Evt, EventCallback&& Cb) = 0;
-	virtual void OnMessage(MessageCallback&& Cb)			= 0;
+	virtual std::future<WebSocketMessage> SendRequest(WebSocketMessage&& Request)		  = 0;
+	virtual void						  OnNotification(NotificationCallback&& Cb)		  = 0;
+	virtual void						  OnEvent(WebSocketEvent Evt, EventCallback&& Cb) = 0;
 
 	static std::shared_ptr<WebSocketClient> Create(asio::io_context& IoCtx);
 };
 
-struct WebSocketMessageHeader
-{
-	static constexpr uint32_t ExpectedMagic = 0x7a776d68;  // zwmh
-
-	uint32_t Magic{};
-	uint64_t ContentLength{};
-	uint32_t Crc32{};
-
-	bool IsValid() const;
-
-	static bool Read(MemoryView Memory, WebSocketMessageHeader& OutHeader);
-
-	static void Write(BinaryWriter& Writer, const CbPackage& Msg);
-};
-
 }  // namespace zen