th7/api/websocket.go

91 lines
1.5 KiB
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
}
}
}
}