82 lines
1.6 KiB
Go
82 lines
1.6 KiB
Go
package dart
|
|
|
|
import (
|
|
"fmt"
|
|
"net/http"
|
|
"sync"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
type WebsocketManager struct {
|
|
conns map[string]map[*websocket.Conn]bool
|
|
mutex sync.RWMutex
|
|
}
|
|
|
|
func (wm *WebsocketManager) createHandler() Handler {
|
|
return func(ctx *Context) error {
|
|
upgrader := websocket.Upgrader{
|
|
CheckOrigin: func(r *http.Request) bool { return true },
|
|
}
|
|
conn, err := upgrader.Upgrade(ctx.response, ctx.request, nil)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
wm.register(ctx.request.URL.Path, conn)
|
|
|
|
for {
|
|
if _, _, err := conn.NextReader(); err != nil {
|
|
break
|
|
}
|
|
}
|
|
|
|
wm.unregister(ctx.request.URL.Path, conn)
|
|
conn.Close()
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func (wm *WebsocketManager) register(path string, conn *websocket.Conn) {
|
|
wm.mutex.Lock()
|
|
defer wm.mutex.Unlock()
|
|
|
|
if _, exists := wm.conns[path]; !exists {
|
|
wm.conns[path] = make(map[*websocket.Conn]bool)
|
|
}
|
|
wm.conns[path][conn] = true
|
|
}
|
|
|
|
func (wm *WebsocketManager) unregister(path string, conn *websocket.Conn) {
|
|
wm.mutex.Lock()
|
|
defer wm.mutex.Unlock()
|
|
|
|
if conns, exists := wm.conns[path]; exists {
|
|
delete(conns, conn)
|
|
if len(conns) == 0 {
|
|
delete(wm.conns, path)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (wm *WebsocketManager) Broadcast(path string, msg []byte) error {
|
|
wm.mutex.RLock()
|
|
conns, exists := wm.conns[path]
|
|
wm.mutex.RUnlock()
|
|
|
|
if !exists {
|
|
return fmt.Errorf("there is no websocket on path '%s'", path)
|
|
}
|
|
|
|
for conn := range conns {
|
|
go func(c *websocket.Conn) {
|
|
if err := c.WriteMessage(websocket.TextMessage, msg); err != nil {
|
|
wm.unregister(path, conn)
|
|
conn.Close()
|
|
}
|
|
}(conn)
|
|
}
|
|
|
|
return nil
|
|
}
|