tecamino-dbm/websocket/models/client.go

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())
}