th7/api/websocket.go

package api

import (
	"fmt"
	"log"
	"net/http"
	"strings"

	"github.com/gorilla/websocket"
)

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

type MessageType = int

const (
	ALL MessageType = iota
	CHANNELS
	CONFIG
)

type Message struct {
	messageType MessageType
	data        []byte
}

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

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

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

func (man *Manager) BroadcastAll(msg []byte) {
	for client := range man.clients {
		select {
		case client.send <- msg:
		default:
			man.unregister <- client
		}
	}
}

func (man *Manager) BroadcastChannels(msg []byte) {
	for client := range man.clients {
		if !client.all && !client.channels {
			continue
		}

		select {
		case client.send <- msg:
		default:
			man.unregister <- client
		}
	}
}

func (man *Manager) BroadcastConfig(msg []byte) {
	for client := range man.clients {
		if !client.all && !client.config {
			continue
		}

		select {
		case client.send <- msg:
		default:
			man.unregister <- 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 {
				fmt.Println("Unregistering:", conn)
				close(conn.send)
				delete(man.clients, conn)
			}

		// broadcast a message to all connected clients
		case message := <-man.broadcast:
			switch message.messageType {
			case ALL:
				man.BroadcastAll(message.data)
			case CHANNELS:
				man.BroadcastChannels(message.data)
			case CONFIG:
				man.BroadcastConfig(message.data)
			}

		}

	}
}

func (c *Client) read() {

	defer func() {
		c.manager.unregister <- c
		_ = c.socket.Close()
	}()

	for {
		_, message, err := c.socket.ReadMessage()
		if err != nil {
			c.manager.unregister <- c
			_ = c.socket.Close()
			break
		}

		switch strings.ToLower(string(message)) {
		case "all":
			c.all = true
		case "channels":
			c.channels = true
		case "config":
			c.config = true
		default:
			log.Println("received: ", message)
		}
	}
}

func (c *Client) write() {

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

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

			err := c.socket.WriteMessage(websocket.TextMessage, message)
			if err != nil {
				fmt.Println(err)
				return
			}
		}
	}

}