th7/api/websocket.go

package api

import (
	"net/http"

	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

type Client struct {
	socket  *websocket.Conn
	send    chan []byte
	manager *Manager
}

type Manager struct {
	clients    map[*Client]bool
	broadcast  chan []byte
	register   chan *Client
	unregister chan *Client
}

func NewManager() *Manager {
	return &Manager{
		clients:    make(map[*Client]bool),
		broadcast:  make(chan []byte),
		register:   make(chan *Client),
		unregister: make(chan *Client),
	}
}

func (man *Manager) start() {
	for {
		select {
		// new connection
		case conn := <-man.register:
			man.clients[conn] = true

		// connection closed
		case conn := <-man.unregister:
			if _, ok := man.clients[conn]; ok {
				close(conn.send)
				delete(man.clients, conn)
			}

		// broadcast a message to all connected clients
		case message := <-man.broadcast:
			for client := range man.clients {
				select {
				case client.send <- message:
				default:
					man.unregister <- client
				}
			}
		}

	}
}

func (c *Client) writer() {

	defer func() {
		_ = c.socket.Close()
	}()

	for {
		select {
		case message, ok := <-c.send:
			if !ok {
				_ = c.socket.WriteMessage(websocket.CloseMessage, []byte{})
				c.manager.unregister <- c
				return
			}

			err := c.socket.WriteMessage(websocket.TextMessage, message)
			if err != nil {
				c.manager.unregister <- c
				return
			}
		}
	}

}