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{"*"}

const (
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second

	// Time allowed to read the next pong message from the peer.
	pongWait = 10 * time.Second

	// Send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (pongWait * 9) / 10
)

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()
	sendPong   chan string
	send       chan []byte
	unregister chan []byte
}

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,
		sendPong:   make(chan string),
		send:       make(chan []byte),
		unregister: make(chan []byte),
	}

	conn.SetPingHandler(func(appData string) error {
		if client.OnPing != nil {
			client.OnPing()
		}
		conn.SetWriteDeadline(time.Now().Add(writeWait))
		conn.SetReadDeadline(time.Now().Add(writeWait))
		client.sendPong <- appData
		return nil
	})

	conn.SetPongHandler(func(string) error {
		conn.SetReadDeadline(time.Now().Add(pongWait))
		if client.OnPong != nil {
			client.OnPong()
		}
		return nil
	})

	// Start reading messages from client
	go client.Read()
	go client.Write()
	return client, nil
}

func (c *Client) Read() {
	if c.OnOpen != nil {
		c.OnOpen()
	}

	c.conn.SetReadDeadline(time.Now().Add(writeWait))
	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.Close(websocket.CloseNormalClosure, "Client closed")
			return
		case websocket.TextMessage:
			if isPing := c.handleJsonPing(msg); !isPing {
				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) Write() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		c.conn.Close()
	}()
	for {

		select {
		case message, ok := <-c.send:
			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if !ok {
				// The hub closed the channel.
				if err := c.conn.WriteMessage(websocket.CloseMessage, []byte{}); err != nil {
					c.handleError(err)
					return
				}
				c.handleError(fmt.Errorf("server %s closed channel", c.Id))
				return
			} else {
				if err := c.conn.WriteMessage(websocket.TextMessage, message); err != nil {
					c.handleError(err)
					return
				}
			}
		case <-ticker.C:
			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				c.handleError(err)
				return
			}

			if c.OnPing != nil {
				c.OnPing()
			}

		case message, ok := <-c.sendPong:
			if ok {
				c.conn.WriteMessage(websocket.PongMessage, []byte(message))
			}
		case message := <-c.unregister:
			c.conn.WriteMessage(websocket.CloseMessage, message)
			c.Connected = false
			close(c.sendPong)
			close(c.send)
			close(c.unregister)
			return
		}

	}
}

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(writeWait))
		// Respond with pong JSON
		select {
		case c.send <- GetPongByteSlice():
		default:
			// optional: log or handle if send buffer is full
			c.handleError(fmt.Errorf("failed to queue pong message"))
			return
		}
		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) {
	if !c.Connected {
		return
	}
	c.send <- data
}

func (c *Client) Close(code int, reason string) error {
	closeMsg := websocket.FormatCloseMessage(code, reason)
	select {
	case c.unregister <- closeMsg: // Attempt to send
	default: // If the channel is full, this runs
		return fmt.Errorf("attempt close client socket failed")
	}
	if c.OnClose != nil {
		c.OnClose(code, reason)
	}
	return nil
}

func (c *Client) handleError(err error) {
	if c.OnError != nil {
		c.OnError(err)
	}
	if err := c.Close(websocket.CloseInternalServerErr, err.Error()); err != nil {
		if c.OnError != nil {
			c.OnError(err)
		} else {
			fmt.Println("error: ", err)
		}
	}
}