tecamino-dbm/websocket/models/client.go

package models

import (
	"fmt"
	"log"
	"net/http"
	"slices"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

var Origins []string = []string{"*"}

var Broadcast Clients = make(Clients)

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

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		if len(Origins) == 0 {
			return false
		} else if Origins[0] == "*" {
			return true
		}
		return slices.Contains(Origins, r.Header.Get("Origin"))
	},
	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,
		send:       make(chan []byte, 512),
		unregister: make(chan []byte, 256),
	}

	Broadcast[client.Id] = client

	conn.SetPingHandler(func(appData string) error {
		if client.OnPing != nil {
			client.OnPing()
		}
		conn.SetWriteDeadline(time.Now().Add(writeWait))
		conn.SetReadDeadline(time.Now().Add(writeWait))
		if err := client.conn.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(pongWait)); err != nil {
			client.OnError(err)
		}
		return nil
	})

	conn.SetPongHandler(func(appData 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()
	go client.PingInterval(pingPeriod)
	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 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() {
	defer 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("ping")); 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 message := <-c.unregister:
			c.conn.WriteMessage(websocket.CloseMessage, message)
			c.Connected = false
			close(c.send)
			delete(Broadcast, c.Id)
			close(c.unregister)
			return
		}

	}
}

func (c *Client) PingInterval(interval time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for range ticker.C {
		if c.OnPing != nil {
			c.OnPing()
		}

		if err := c.conn.WriteControl(websocket.PingMessage, []byte("ping"), time.Now().Add(pongWait)); err != nil {
			c.OnError(err)
			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)
		}
	}
}