From 7e6e8841b0172951ada5ac7aa08dcc4bbcd58c2c Mon Sep 17 00:00:00 2001
From: muXxer <git@muxxer.de>
Date: Fri, 19 Apr 2024 10:06:26 +0200
Subject: [PATCH] Add SendRaw and BroadcastMsgRaw to websocket hub

---
 web/websockethub/client.go | 23 ++++++++++++++++++++---
 web/websockethub/hub.go    |  7 +++++++
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/web/websockethub/client.go b/web/websockethub/client.go
index 9a4f6420..65010a96 100644
--- a/web/websockethub/client.go
+++ b/web/websockethub/client.go
@@ -33,6 +33,11 @@ type WebsocketMsg struct {
 	Data []byte
 }
 
+// rawMsg is used to skip JSON marshalling while sending messages.
+type rawMsg struct {
+	Data []byte
+}
+
 // ClientID is the ID of a client.
 type ClientID uint32
 
@@ -293,6 +298,11 @@ func (c *Client) writePump() {
 		ctx, cancel := context.WithTimeout(c.ctx, writeWait)
 		defer cancel()
 
+		// check if the message is a raw message that does not need to be JSON marshaled.
+		if rawMsg, ok := msg.(*rawMsg); ok {
+			return c.conn.Write(ctx, websocket.MessageText, rawMsg.Data)
+		}
+
 		return wsjson.Write(ctx, c.conn, msg)
 	}
 
@@ -340,7 +350,8 @@ func (c *Client) writePump() {
 }
 
 // Send sends a message to the client.
-func (c *Client) Send(ctx context.Context, msg interface{}, dontDrop ...bool) error {
+// JSON marshaling is done automatically based on "json" tags.
+func (c *Client) Send(ctx context.Context, data interface{}, dontDrop ...bool) error {
 	if c.hub.Stopped() {
 		// hub was already shut down
 		return ErrWebsocketServerUnavailable
@@ -374,7 +385,7 @@ func (c *Client) Send(ctx context.Context, msg interface{}, dontDrop ...bool) er
 				return ErrClientDisconnected
 			case <-c.sendChanClosed:
 				return ErrClientDisconnected
-			case c.sendChan <- msg:
+			case c.sendChan <- data:
 				return nil
 			}
 		}
@@ -394,10 +405,16 @@ func (c *Client) Send(ctx context.Context, msg interface{}, dontDrop ...bool) er
 		return ErrClientDisconnected
 	default:
 		select {
-		case c.sendChan <- msg:
+		case c.sendChan <- data:
 			return nil
 		default:
 			return nil
 		}
 	}
 }
+
+// SendRaw sends a raw message to the client.
+// The message is not JSON marshaled.
+func (c *Client) SendRaw(ctx context.Context, data []byte, dontDrop ...bool) error {
+	return c.Send(ctx, &rawMsg{Data: data}, dontDrop...)
+}
diff --git a/web/websockethub/hub.go b/web/websockethub/hub.go
index b2c8ae72..ace7153e 100644
--- a/web/websockethub/hub.go
+++ b/web/websockethub/hub.go
@@ -105,6 +105,7 @@ func (h *Hub) Events() *Events {
 }
 
 // BroadcastMsg sends a message to all clients.
+// JSON marshaling is done automatically based on "json" tags.
 func (h *Hub) BroadcastMsg(ctx context.Context, data interface{}, dontDrop ...bool) error {
 	if h.shutdownFlag.Load() {
 		// hub was already shut down or was not started yet
@@ -157,6 +158,12 @@ func (h *Hub) BroadcastMsg(ctx context.Context, data interface{}, dontDrop ...bo
 	}
 }
 
+// BroadcastMsgRaw sends a raw message to all clients.
+// The message is not JSON marshaled.
+func (h *Hub) BroadcastMsgRaw(ctx context.Context, data []byte, dontDrop ...bool) error {
+	return h.BroadcastMsg(ctx, &rawMsg{Data: data}, dontDrop...)
+}
+
 func (h *Hub) removeClient(client *Client) {
 	delete(h.clients, client)
 	close(client.ExitSignal)