modify websocketserver to broker with callback functions

websocket/models/client.go

@@ -0,0 +1,164 @@
+package models
+
+import (
+	"encoding/json"
+	"fmt"
+	"log"
+	"net/http"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
+)
+
+var Origins []string = []string{"*"}
+
+type Client struct {
+	Id        string
+	Connected bool            `json:"connected"`
+	Conn      *websocket.Conn `json:"-"`
+	OnOpen    func()
+	OnMessage func(data []byte)
+	OnClose   func(code int, reason string)
+	OnError   func(err error)
+	OnPing    func()
+	OnPong    func()
+	timeout   time.Duration
+}
+
+var upgrader = websocket.Upgrader{
+	CheckOrigin: func(r *http.Request) bool {
+		if len(Origins) == 0 {
+			return false
+		}
+		if Origins[0] == "*" {
+			return true
+		}
+		origin := r.Header.Get("Origin")
+		for _, o := range Origins {
+			if o == origin {
+				return true
+			}
+		}
+		return false
+	},
+	EnableCompression: false,
+}
+
+func ConnectNewClient(id string, c *gin.Context) (*Client, error) {
+	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
+	if err != nil {
+		return nil, fmt.Errorf("websocket upgrade error: %w", err)
+	}
+
+	client := &Client{
+		Id:        id,
+		Connected: true,
+		Conn:      conn,
+		timeout:   5,
+	}
+
+	conn.SetPingHandler(func(appData string) error {
+		if client.OnPing != nil {
+			client.OnPing()
+		}
+		conn.SetWriteDeadline(time.Now().Add(client.timeout))
+		conn.SetReadDeadline(time.Now().Add(client.timeout))
+		return conn.WriteMessage(websocket.PongMessage, []byte(appData))
+	})
+
+	conn.SetPongHandler(func(appData string) error {
+		if client.OnPong != nil {
+			client.OnPong()
+		}
+		conn.SetReadDeadline(time.Now().Add(client.timeout))
+		return nil
+	})
+
+	// Start reading messages from client
+	go client.Listen(7)
+	return client, nil
+}
+
+func (c *Client) Listen(timeout uint) {
+	if timeout > 0 {
+		c.timeout = time.Duration(timeout) * time.Second
+	}
+
+	if c.OnOpen != nil {
+		c.OnOpen()
+	}
+
+	c.Conn.SetReadDeadline(time.Now().Add(c.timeout))
+	for c.Connected {
+		msgType, msg, err := c.Conn.ReadMessage()
+		if err != nil {
+			c.handleError(fmt.Errorf("read error (id:%s): %w", c.Id, err))
+			return
+		}
+		switch msgType {
+		case websocket.CloseMessage:
+			c.handleClose(1000, "Client closed")
+			return
+		case websocket.TextMessage:
+			if isPing := c.handleJsonPing(msg); isPing {
+				continue
+			}
+			if c.OnMessage != nil {
+				c.OnMessage(msg)
+			} else {
+				log.Printf("Received message but no handler set (id:%s): %s", c.Id, string(msg))
+			}
+		default:
+			log.Printf("Unhandled message type %d (id:%s)", msgType, c.Id)
+		}
+	}
+}
+
+func (c *Client) handleJsonPing(msg []byte) (isPing bool) {
+	var wsMsg WSMessage
+	err := json.Unmarshal(msg, &wsMsg)
+	if err == nil && wsMsg.IsPing() {
+		c.Conn.SetReadDeadline(time.Now().Add(c.timeout))
+		// Respond with pong JSON
+		c.Conn.SetWriteDeadline(time.Now().Add(c.timeout))
+		err = c.Conn.WriteMessage(websocket.TextMessage, GetPongByteSlice())
+		if err != nil {
+			c.handleError(fmt.Errorf("write pong error: %w", err))
+			return
+		}
+		if c.OnPing != nil {
+			c.OnPing()
+		}
+		isPing = true
+	}
+	return
+}
+
+func (c *Client) SendResponse(data []byte, timeout uint) error {
+	c.Conn.SetWriteDeadline(time.Now().Add(time.Duration(timeout) * time.Second))
+	return c.Conn.WriteMessage(websocket.TextMessage, data)
+}
+
+func (c *Client) Close(code int, reason string) {
+	c.handleClose(code, reason)
+}
+
+func (c *Client) handleClose(code int, text string) {
+	if !c.Connected {
+		return
+	}
+	c.Connected = false
+	if c.OnClose != nil {
+		c.OnClose(code, text)
+	}
+	c.Conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(code, text))
+	c.Conn.Close()
+}
+
+func (c *Client) handleError(err error) {
+	if c.OnError != nil {
+		c.OnError(err)
+	}
+	c.Close(websocket.CloseInternalServerErr, err.Error())
+}