/
websocket.go
120 lines (108 loc) · 3.16 KB
/
websocket.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
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
package common
import (
"github.com/gorilla/websocket"
"github.com/nccgroup/tracy/api/types"
"github.com/nccgroup/tracy/log"
)
var (
subscribers map[int]*subscriber
updateChan chan interface{}
addSubChan chan *subscriber
changeSubChan chan []int
removeSubChan chan int
)
// subscriber is a struct used to keep track of client connections to the router.
// A subscriber has a key index, the tracer it is currently listening to and the
// websocket connection that manages it.
type subscriber struct {
KeyChan chan int
Sock *websocket.Conn
Tracer uint
}
// AddSubscriber takes a websocket connection and adds it to the list of subscribers.
// New events that happen in package common get pushed these events.
func AddSubscriber(conn *websocket.Conn) int {
c := make(chan int, 1)
addSubChan <- &subscriber{c, conn, 0}
return <-c
}
// RemoveSubscriber removes the websocket from the list of subscribers.
func RemoveSubscriber(key int) {
removeSubChan <- key
}
// ChangeTracer changes the tracer to send event updates for.
func ChangeTracer(key, tracer int) {
changeSubChan <- []int{key, tracer}
}
// UpdateSubscribers sends an update to all the subscribers.
func UpdateSubscribers(update interface{}) {
updateChan <- update
}
// Router is a single goroutine used to serialize notifications to each of the
// connected websocket clients. It handles adding and removing a subscriber, changing
// what tracer the subscriber is listening to and and updates to that tracer.
func router() {
id := 0
for {
select {
case change := <-changeSubChan:
subscribers[change[0]].Tracer = uint(change[1])
case add := <-addSubChan:
subscribers[id] = add
add.KeyChan <- id
id++
case remove := <-removeSubChan:
delete(subscribers, remove)
case update := <-updateChan:
updateRouter(update)
}
}
}
func updateRouter(update interface{}) {
//TODO: It'd be good to keep the types of subscribers in case we
// need to send specific messages to the extension or a particular
// type of extension
for _, sub := range subscribers {
switch u := update.(type) {
case types.Tracer:
if err := sub.Sock.WriteJSON(types.TracerWebSocket{u}); err != nil {
log.Error.Print(err)
continue
}
case types.Request:
if err := sub.Sock.WriteJSON(types.RequestWebSocket{u}); err != nil {
log.Error.Print(err)
continue
}
case types.TracerEvent:
// Only send event updates for the subscribed tracer.
if u.TracerID == sub.Tracer {
if err := sub.Sock.WriteJSON(types.TracerEventsWebSocket{u}); err != nil {
log.Error.Print(err)
continue
}
}
case types.Notification:
if err := sub.Sock.WriteJSON(types.NotificationWebSocket{u}); err != nil {
log.Error.Print(err)
continue
}
case types.Reproduction:
if err := sub.Sock.WriteJSON(types.ReproductionWebSocket{u}); err != nil {
log.Error.Print(err)
continue
}
default:
log.Error.Printf("not sure what it was: %T", u)
continue
}
}
}
func init() {
updateChan = make(chan interface{}, 10)
addSubChan = make(chan *subscriber, 10)
removeSubChan = make(chan int, 10)
changeSubChan = make(chan []int, 10)
subscribers = make(map[int]*subscriber, 25)
go router()
}