-
Notifications
You must be signed in to change notification settings - Fork 35
/
websocket.go
120 lines (102 loc) · 2.28 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 gold
import (
"io"
"log"
"net/http"
"strings"
"sync"
"golang.org/x/net/websocket"
)
// type wsConn struct {
// subbed bool
// uuid string
// }
var (
websocketSubs = map[string]map[*websocket.Conn]string{}
websocketSubsL = new(sync.RWMutex)
)
func onDeleteURI(uri string) {
websocketPublish(uri)
}
func onUpdateURI(uri string) {
websocketPublish(uri)
}
// Handles each websocket connection
func websocketHandler(ws *websocket.Conn) {
// @@TODO switch to server logging
// log.Println("opened via:", ws.RemoteAddr())
uris := map[string]bool{}
message := ""
for {
err := websocket.Message.Receive(ws, &message)
if err == io.EOF {
break
}
if err != nil {
log.Println(err)
break
}
argv := strings.Split(message, " ")
if len(argv) < 2 {
argv = append(argv, "")
}
cmd, uri := argv[0], argv[1]
switch cmd {
case "ping":
websocket.Message.Send(ws, "pong")
case "sub":
uris[uri] = true
websocketSubsL.Lock()
if _, ex := websocketSubs[uri]; !ex {
websocketSubs[uri] = map[*websocket.Conn]string{}
}
websocketSubs[uri][ws] = NewUUID()
websocketSubsL.Unlock()
websocket.Message.Send(ws, "ack "+websocketSubs[uri][ws])
default:
log.Println("invalid message:", message)
}
}
websocketSubsL.Lock()
for k := range uris {
delete(websocketSubs[k], ws)
}
websocketSubsL.Unlock()
// @@TODO switch to server logging
// log.Println("closed via:", ws.RemoteAddr())
}
func websocketPublish(uri string, uuid ...string) {
websocketSubsL.RLock()
subs := websocketSubs[uri]
websocketSubsL.RUnlock()
for k := range subs {
uuidMatch := true
// log.Println(uuid)
if len(uuid) > 0 && subs[k] != uuid[0] {
uuidMatch = false
}
if uuidMatch {
err := websocket.Message.Send(k, "pub "+uri)
if err != nil {
log.Println(err)
}
}
}
}
// Converts an HTTP request to a websocket server
func websocketServe(w http.ResponseWriter, req *http.Request) {
websocket.Handler(websocketHandler).ServeHTTP(w, req)
}
// Checks whether an HTTP request looks like websocket
func websocketUpgrade(r *http.Request) bool {
if r == nil {
return false
}
if strings.ToLower(r.Header.Get("Connection")) != "upgrade" {
return false
}
if strings.ToLower(r.Header.Get("Upgrade")) != "websocket" {
return false
}
return true
}