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