Merge pull request #186 from disgoorg/patch/resume-url

disgoorg · Aug 29, 2022 · 6ea9ab4 · 6ea9ab4
2 parents cba399e + bcb3760
commit 6ea9ab4
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 7 deletions.
diff --git a/gateway/gateway_config.go b/gateway/gateway_config.go
@@ -18,6 +18,7 @@ func DefaultConfig() *Config {
 		ShardCount:        1,
 		AutoReconnect:     true,
 		MaxReconnectTries: 10,
+		EnableResumeURL:   true,
 	}
 }
 
@@ -32,10 +33,12 @@ type Config struct {
 	ShardID                   int
 	ShardCount                int
 	SessionID                 *string
+	ResumeGatewayURL          *string
 	LastSequenceReceived      *int
 	AutoReconnect             bool
 	MaxReconnectTries         int
 	EnableRawEvents           bool
+	EnableResumeURL           bool
 	RateLimiter               RateLimiter
 	RateRateLimiterConfigOpts []RateLimiterConfigOpt
 	Presence                  *MessageDataPresenceUpdate
@@ -155,6 +158,13 @@ func WithEnableRawEvents(enableRawEventEvents bool) ConfigOpt {
 	}
 }
 
+// WithEnableResumeURL enables/disables usage of resume URLs sent by Discord.
+func WithEnableResumeURL(enableResumeURL bool) ConfigOpt {
+	return func(config *Config) {
+		config.EnableResumeURL = enableResumeURL
+	}
+}
+
 // WithRateLimiter sets the grate.RateLimiter for the Gateway.
 func WithRateLimiter(rateLimiter RateLimiter) ConfigOpt {
 	return func(config *Config) {

diff --git a/gateway/gateway_events.go b/gateway/gateway_events.go
@@ -16,12 +16,13 @@ type EventData interface {
 
 // EventReady is the event sent by discord when you successfully Identify
 type EventReady struct {
-	Version     int                        `json:"v"`
-	User        discord.OAuth2User         `json:"user"`
-	Guilds      []discord.UnavailableGuild `json:"guilds"`
-	SessionID   string                     `json:"session_id"`
-	Shard       []int                      `json:"shard,omitempty"`
-	Application discord.PartialApplication `json:"application"`
+	Version          int                        `json:"v"`
+	User             discord.OAuth2User         `json:"user"`
+	Guilds           []discord.UnavailableGuild `json:"guilds"`
+	SessionID        string                     `json:"session_id"`
+	ResumeGatewayURL string                     `json:"resume_gateway_url"`
+	Shard            []int                      `json:"shard,omitempty"`
+	Application      discord.PartialApplication `json:"application"`
 }
 
 func (EventReady) messageData() {}

diff --git a/gateway/gateway_impl.go b/gateway/gateway_impl.go
@@ -94,7 +94,11 @@ func (g *gatewayImpl) Open(ctx context.Context) error {
 	}
 	g.status = StatusConnecting
 
-	gatewayURL := fmt.Sprintf("%s?v=%d&encoding=json", g.config.URL, Version)
+	wsURL := g.config.URL
+	if g.config.ResumeGatewayURL != nil && g.config.EnableResumeURL {
+		wsURL = *g.config.ResumeGatewayURL
+	}
+	gatewayURL := fmt.Sprintf("%s?v=%d&encoding=json", wsURL, Version)
 	g.lastHeartbeatSent = time.Now().UTC()
 	conn, rs, err := g.config.Dialer.DialContext(ctx, gatewayURL, nil)
 	if err != nil {
@@ -156,6 +160,7 @@ func (g *gatewayImpl) CloseWithCode(ctx context.Context, code int, message strin
 		// clear resume data as we closed gracefully
 		if code == websocket.CloseNormalClosure || code == websocket.CloseGoingAway {
 			g.config.SessionID = nil
+			g.config.ResumeGatewayURL = nil
 			g.config.LastSequenceReceived = nil
 		}
 	}
@@ -330,6 +335,7 @@ loop:
 					g.config.Logger.Error(g.formatLogs("invalid sequence provided. reconnecting..."))
 					g.config.LastSequenceReceived = nil
 					g.config.SessionID = nil
+					g.config.ResumeGatewayURL = nil
 				} else {
 					message := g.formatLogsf("gateway close received, reconnect: %t, code: %d, error: %s", g.config.AutoReconnect && reconnect, closeError.Code, closeError.Text)
 					if reconnect {
@@ -390,6 +396,7 @@ loop:
 			// get session id here
 			if readyEvent, ok := data.(EventReady); ok {
 				g.config.SessionID = &readyEvent.SessionID
+				g.config.ResumeGatewayURL = &readyEvent.ResumeGatewayURL
 				g.status = StatusReady
 				g.config.Logger.Debug(g.formatLogs("ready event received"))
 			}
@@ -424,6 +431,7 @@ loop:
 				// clear resume info
 				g.config.SessionID = nil
 				g.config.LastSequenceReceived = nil
+				g.config.ResumeGatewayURL = nil
 			}
 
 			g.CloseWithCode(context.TODO(), code, "invalid session")