-
Notifications
You must be signed in to change notification settings - Fork 35
/
gateway.go
376 lines (317 loc) · 10.3 KB
/
gateway.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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
package voice
import (
"bytes"
"context"
"errors"
"fmt"
"io"
"log/slog"
"sync"
"syscall"
"time"
"github.com/disgoorg/json"
"github.com/disgoorg/snowflake/v2"
"github.com/gorilla/websocket"
"github.com/disgoorg/disgo/discord"
)
var (
// ErrGatewayNotConnected is returned when the gateway is not connected and a message is attempted to be sent.
ErrGatewayNotConnected = fmt.Errorf("voice gateway not connected")
// ErrGatewayAlreadyConnected is returned when the gateway is already connected and a connection is attempted to be opened.
ErrGatewayAlreadyConnected = fmt.Errorf("voice gateway already connected")
)
// GatewayVersion is the version of the voice gateway we are using.
const GatewayVersion = 4
// Status returns the current status of the gateway.
type Status int
const (
StatusUnconnected Status = iota
StatusConnecting
StatusWaitingForHello
StatusIdentifying
StatusResuming
StatusWaitingForReady
StatusReady
StatusDisconnected
)
type (
// EventHandlerFunc is a function that handles a voice gateway event.
EventHandlerFunc func(opCode Opcode, data GatewayMessageData)
// CloseHandlerFunc is a function that handles a voice gateway close.
CloseHandlerFunc func(gateway Gateway, err error)
// GatewayCreateFunc is a function that creates a new voice gateway.
GatewayCreateFunc func(eventHandlerFunc EventHandlerFunc, closeHandlerFunc CloseHandlerFunc, opts ...GatewayConfigOpt) Gateway
// StateProviderFunc is a function that provides the current conn state of the voice gateway.
StateProviderFunc func() State
)
// State is the current state of the voice conn.
type State struct {
GuildID snowflake.ID
UserID snowflake.ID
ChannelID *snowflake.ID
SessionID string
Token string
Endpoint string
}
// Gateway is a websocket connection to the Discord voice gateway.
type Gateway interface {
// SSRC returns the SSRC of the current voice connection.
SSRC() uint32
// Latency returns the current latency of the voice gateway connection.
Latency() time.Duration
// Open opens a new websocket connection to the voice gateway.
Open(ctx context.Context, state State) error
// Close closes the websocket connection to the voice gateway.
Close()
// CloseWithCode closes the websocket connection to the voice gateway with a specific close code.
CloseWithCode(code int, message string)
// Send sends a message to the voice gateway.
Send(ctx context.Context, opCode Opcode, data GatewayMessageData) error
}
// NewGateway creates a new voice Gateway.
func NewGateway(eventHandlerFunc EventHandlerFunc, closeHandlerFunc CloseHandlerFunc, opts ...GatewayConfigOpt) Gateway {
config := DefaultGatewayConfig()
config.Apply(opts)
config.Logger = config.Logger.With(slog.String("name", "voice_conn_gateway"))
return &gatewayImpl{
config: *config,
eventHandlerFunc: eventHandlerFunc,
closeHandlerFunc: closeHandlerFunc,
}
}
type gatewayImpl struct {
config GatewayConfig
eventHandlerFunc EventHandlerFunc
closeHandlerFunc CloseHandlerFunc
ssrc uint32
state State
conn *websocket.Conn
connMu sync.Mutex
status Status
heartbeatTicker *time.Ticker
heartbeatInterval time.Duration
lastHeartbeatSent time.Time
lastHeartbeatReceived time.Time
lastNonce int64
}
func (g *gatewayImpl) SSRC() uint32 {
return g.ssrc
}
func (g *gatewayImpl) Open(ctx context.Context, state State) error {
g.config.Logger.Debug("opening voice gateway connection")
g.state = state
g.connMu.Lock()
defer g.connMu.Unlock()
if g.conn != nil {
return ErrGatewayAlreadyConnected
}
g.status = StatusConnecting
gatewayURL := fmt.Sprintf("wss://%s?v=%d", state.Endpoint, GatewayVersion)
g.config.Logger.Debug("connecting to voice gateway at", slog.String("url", gatewayURL))
g.lastHeartbeatSent = time.Now().UTC()
conn, rs, err := g.config.Dialer.DialContext(ctx, gatewayURL, nil)
if err != nil {
g.Close()
defer rs.Body.Close()
return fmt.Errorf("error connecting to voice gateway: %w", err)
}
conn.SetCloseHandler(func(code int, text string) error {
return nil
})
g.conn = conn
g.status = StatusWaitingForHello
go g.listen(g.conn)
return nil
}
func (g *gatewayImpl) Close() {
g.CloseWithCode(websocket.CloseNormalClosure, "Shutting down")
}
func (g *gatewayImpl) CloseWithCode(code int, message string) {
if g.heartbeatTicker != nil {
g.config.Logger.Debug("closing heartbeat goroutines")
g.heartbeatTicker.Stop()
g.heartbeatTicker = nil
}
g.connMu.Lock()
defer g.connMu.Unlock()
if g.conn != nil {
g.config.Logger.Debug("closing voice gateway connection", slog.Int("code", code), slog.String("message", message))
if err := g.conn.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(code, message)); err != nil && !errors.Is(err, websocket.ErrCloseSent) {
g.config.Logger.Debug("error writing close code", slog.Any("err", err))
}
_ = g.conn.Close()
g.conn = nil
// clear resume data as we closed gracefully
if code == websocket.CloseNormalClosure || code == websocket.CloseGoingAway {
g.ssrc = 0
}
}
}
func (g *gatewayImpl) heartbeat() {
g.heartbeatTicker = time.NewTicker(g.heartbeatInterval)
defer g.heartbeatTicker.Stop()
defer g.config.Logger.Debug("exiting voice heartbeat goroutine")
for range g.heartbeatTicker.C {
g.sendHeartbeat()
}
}
func (g *gatewayImpl) sendHeartbeat() {
g.lastNonce = time.Now().UnixMilli()
ctx, cancel := context.WithTimeout(context.Background(), g.heartbeatInterval)
defer cancel()
if err := g.Send(ctx, OpcodeHeartbeat, GatewayMessageDataHeartbeat(g.lastNonce)); err != nil {
if !errors.Is(err, ErrGatewayNotConnected) || errors.Is(err, syscall.EPIPE) {
return
}
g.config.Logger.Error("failed to send heartbeat", slog.Any("err", err))
g.CloseWithCode(websocket.CloseServiceRestart, "heartbeat timeout")
go g.reconnect()
return
}
g.lastHeartbeatSent = time.Now().UTC()
}
func (g *gatewayImpl) listen(conn *websocket.Conn) {
defer g.config.Logger.Debug("exiting listen goroutine")
loop:
for {
_, reader, err := conn.NextReader()
if err != nil {
g.connMu.Lock()
sameConn := g.conn == conn
g.connMu.Unlock()
// if sameConn is false, it means the connection has been closed by the user, and we can just exit
if !sameConn {
return
}
reconnect := true
var closeError *websocket.CloseError
if errors.As(err, &closeError) {
closeCode := GatewayCloseEventCodeByCode(closeError.Code)
reconnect = closeCode.Reconnect
}
g.CloseWithCode(websocket.CloseServiceRestart, "listen error")
if g.config.AutoReconnect && reconnect {
go g.reconnect()
} else if g.closeHandlerFunc != nil {
go g.closeHandlerFunc(g, err)
}
break loop
}
message, err := g.parseMessage(reader)
if err != nil {
g.config.Logger.Error("error while parsing voice gateway event", slog.Any("err", err))
continue
}
switch d := message.D.(type) {
case GatewayMessageDataHello:
g.status = StatusWaitingForReady
g.lastHeartbeatReceived = time.Now().UTC()
g.heartbeatInterval = time.Duration(d.HeartbeatInterval) * time.Millisecond
go g.heartbeat()
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
if g.ssrc == 0 {
g.status = StatusIdentifying
err = g.Send(ctx, OpcodeIdentify, GatewayMessageDataIdentify{
GuildID: g.state.GuildID,
UserID: g.state.UserID,
SessionID: g.state.SessionID,
Token: g.state.Token,
})
} else {
g.status = StatusResuming
err = g.Send(ctx, OpcodeResume, GatewayMessageDataResume{
GuildID: g.state.GuildID,
SessionID: g.state.SessionID,
Token: g.state.Token,
})
}
cancel()
if err != nil {
g.CloseWithCode(websocket.CloseServiceRestart, "failed to send identify or resume")
go g.reconnect()
return
}
case GatewayMessageDataReady:
g.status = StatusReady
g.ssrc = d.SSRC
case GatewayMessageDataHeartbeatACK:
if int64(d) != g.lastNonce {
g.config.Logger.Error("received heartbeat ack with nonce", slog.Int64("nonce", int64(d)), slog.Int64("last_nonce", g.lastNonce))
go g.reconnect()
break loop
}
g.lastHeartbeatReceived = time.Now().UTC()
}
g.eventHandlerFunc(message.Op, message.D)
}
}
func (g *gatewayImpl) Latency() time.Duration {
return g.lastHeartbeatReceived.Sub(g.lastHeartbeatSent)
}
func (g *gatewayImpl) Send(ctx context.Context, op Opcode, d GatewayMessageData) error {
data, err := json.Marshal(GatewayMessage{
Op: op,
D: d,
})
if err != nil {
return fmt.Errorf("failed to marshal voice gateway message: %w", err)
}
return g.send(ctx, websocket.TextMessage, data)
}
func (g *gatewayImpl) send(ctx context.Context, messageType int, data []byte) error {
g.connMu.Lock()
defer g.connMu.Unlock()
if g.conn == nil {
return ErrGatewayNotConnected
}
g.config.Logger.Debug("sending message to voice gateway", slog.String("data", string(data)))
deadline, ok := ctx.Deadline()
if ok {
if err := g.conn.SetWriteDeadline(deadline); err != nil {
return err
}
}
if err := g.conn.WriteMessage(messageType, data); err != nil {
return fmt.Errorf("failed to send message to voice gateway: %w", err)
}
return nil
}
func (g *gatewayImpl) reconnectTry(ctx context.Context, try int) error {
delay := time.Duration(try) * 2 * time.Second
if delay > 30*time.Second {
delay = 30 * time.Second
}
timer := time.NewTimer(time.Duration(try) * delay)
defer timer.Stop()
select {
case <-ctx.Done():
timer.Stop()
return ctx.Err()
case <-timer.C:
}
g.config.Logger.Debug("reconnecting voice gateway")
if err := g.Open(ctx, g.state); err != nil {
if errors.Is(err, discord.ErrGatewayAlreadyConnected) {
return err
}
g.config.Logger.Error("failed to reconnect voice gateway", slog.Any("err", err))
g.status = StatusDisconnected
return g.reconnectTry(ctx, try+1)
}
return nil
}
func (g *gatewayImpl) reconnect() {
if err := g.reconnectTry(context.Background(), 0); err != nil {
g.config.Logger.Error("failed to reopen voice gateway", slog.Any("err", err))
}
}
func (g *gatewayImpl) parseMessage(r io.Reader) (GatewayMessage, error) {
buff := &bytes.Buffer{}
data, _ := io.ReadAll(io.TeeReader(r, buff))
g.config.Logger.Debug("received message from voice gateway", slog.String("data", string(data)))
var message GatewayMessage
if err := json.NewDecoder(buff).Decode(&message); err != nil {
return GatewayMessage{}, err
}
return message, nil
}