ChronoTomato/internal/websocket/connect.go

package websocket

import (
	"time"

	"github.com/gorilla/websocket"

	ChronoTomato "git.smsvc.net/pomodoro/ChronoTomato/pkg/models"
)

type Client ChronoTomato.GoTomatoClient // New websocket client

// Connects to websocket
func Connect(url string) Client {
	conn, _, err := websocket.DefaultDialer.Dial(url, nil)

	return Client{Conn: conn, Server: url, LastErr: err}
}

func (c Client) Connected() bool {
	if c.Conn == nil {
		return false
	}

	// Set a pong handler to signal when a pong is received
	c.Conn.SetPongHandler(func(s string) error {
		c.Conn.SetReadDeadline(time.Now().Add(10 * time.Second))
		return nil
	})

	// Send a ping
	err := c.Conn.WriteMessage(websocket.PingMessage, nil)
	if err != nil {
		return false
	}

	return true
}

// Disconnects from websocket
func (c Client) Disconnect() {
	// Cleanly close the connection by sending a close message and then
	// waiting (with timeout) for the server to close the connection.
	if c.Connected() {
		c.Conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
		select {
		case <-Done:
		case <-time.After(time.Second):
		}
	}
	return
}