/
tracker.go
95 lines (81 loc) · 2 KB
/
tracker.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
package pufferpanel
import (
"encoding/json"
"github.com/gorilla/websocket"
"github.com/pufferpanel/pufferpanel/v3/logging"
"github.com/pufferpanel/pufferpanel/v3/messages"
"io"
"sync"
)
type Tracker struct {
sockets []*Socket
locker sync.Mutex
}
func CreateTracker() *Tracker {
return &Tracker{sockets: make([]*Socket, 0)}
}
func (ws *Tracker) Register(conn *Socket) {
ws.locker.Lock()
defer ws.locker.Unlock()
ws.sockets = append(ws.sockets, conn)
}
func (ws *Tracker) WriteMessage(msg messages.Message) error {
d, err := json.Marshal(&messages.Transmission{Message: msg, Type: msg.Key()})
if err != nil {
return err
}
ws.locker.Lock()
defer ws.locker.Unlock()
for i := 0; i < len(ws.sockets); i++ {
go func(conn *Socket, data []byte) {
_, err := conn.Write(data)
if err != nil {
logging.Debug.Printf("websocket encountered error, dropping (%s)", err.Error())
ws.locker.Lock()
defer ws.locker.Unlock()
for i, k := range ws.sockets {
if k == conn {
ws.sockets[i] = ws.sockets[len(ws.sockets)-1]
ws.sockets[len(ws.sockets)-1] = nil
ws.sockets = ws.sockets[:len(ws.sockets)-1]
break
}
}
}
}(ws.sockets[i], d)
}
return nil
}
func (ws *Tracker) Write(source []byte) (n int, e error) {
packet := messages.Console{Logs: source}
e = ws.WriteMessage(packet)
n = len(source)
return
}
func Create(ws *websocket.Conn) *Socket {
return &Socket{conn: ws}
}
type Socket struct {
conn *websocket.Conn
locker sync.Mutex
io.WriteCloser
}
func (s *Socket) WriteMessage(msg messages.Message) error {
return s.WriteJSON(messages.Transmission{Type: msg.Key(), Message: msg})
}
func (s *Socket) Write(data []byte) (int, error) {
s.locker.Lock()
defer s.locker.Unlock()
return len(data), s.conn.WriteMessage(websocket.TextMessage, data)
}
func (s *Socket) WriteJSON(data interface{}) error {
d, err := json.Marshal(data)
if err != nil {
return err
}
_, err = s.Write(d)
return err
}
func (s *Socket) Close() error {
return s.conn.Close()
}