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