From 607250ae55e71a85a808a78f3ca0a0451006b4ca Mon Sep 17 00:00:00 2001
From: Maximilian von Lindern <48887425+mavolin@users.noreply.github.com>
Date: Fri, 30 Oct 2020 21:41:04 +0100
Subject: [PATCH] Gateway: Added reconnect timeout; fixed UA (#154)

* Gateway: use gateway version 8

* API: remove old v0.0.1 version tag

* Discord: fix typos

* Gateway: add timeout

* Gateway: revert to returning errors on ReconnectCtx
---
 api/api.go         |  2 +-
 discord/channel.go |  2 +-
 gateway/gateway.go | 43 +++++++++++++++++++++++++++++++++----------
 3 files changed, 35 insertions(+), 12 deletions(-)

diff --git a/api/api.go b/api/api.go
index 83ec57e..4903636 100644
--- a/api/api.go
+++ b/api/api.go
@@ -21,7 +21,7 @@ var (
 	EndpointGatewayBot = EndpointGateway + "/bot"
 )
 
-var UserAgent = "DiscordBot (https://github.com/diamondburned/arikawa/v2, v0.0.1)"
+var UserAgent = "DiscordBot (https://github.com/diamondburned/arikawa/v2)"
 
 type Client struct {
 	*httputil.Client
diff --git a/discord/channel.go b/discord/channel.go
index dd297ad..5a7aa72 100644
--- a/discord/channel.go
+++ b/discord/channel.go
@@ -128,7 +128,7 @@ const (
 	OverwriteMember
 )
 
-// UnmarshalJSON unmarshals both a string-quoteed number and a regular number
+// UnmarshalJSON unmarshalls both a string-quoted number and a regular number
 // into OverwriteType. We need to do this because Discord is so bad that they
 // can't even handle 1s and 0s properly.
 func (otype *OverwriteType) UnmarshalJSON(b []byte) error {
diff --git a/gateway/gateway.go b/gateway/gateway.go
index ea224be..7cad6cb 100644
--- a/gateway/gateway.go
+++ b/gateway/gateway.go
@@ -78,6 +78,13 @@ func BotURL(token string) (*BotData, error) {
 type Gateway struct {
 	WS        *wsutil.Websocket
 	WSTimeout time.Duration
+	// ReconnectTimeout is the timeout used during reconnection.
+	// If the a connection to the gateway can't be established before the
+	// duration passes, the Gateway will be closed and FatalErrorCallback will
+	// be called.
+	//
+	// Setting this to 0 is equivalent to no timeout.
+	ReconnectTimeout time.Duration
 
 	// All events sent over are pointers to Event structs (structs suffixed with
 	// "Event"). This shouldn't be accessed if the Gateway is created with a
@@ -94,6 +101,16 @@ type Gateway struct {
 	PacerLoop wsutil.PacemakerLoop
 
 	ErrorLog func(err error) // default to log.Println
+	// FatalErrorCallback is called, if the Gateway exits fatally. At the point
+	// of calling, the gateway will be already closed.
+	//
+	// Currently this will only be called, if the ReconnectTimeout was changed
+	// to a definite timeout, and connection could not be established during
+	// that time.
+	// err will be ErrWSMaxTries in that case.
+	//
+	// Defaults to noop.
+	FatalErrorCallback func(err error)
 
 	// AfterClose is called after each close. Error can be non-nil, as this is
 	// called even when the Gateway is gracefully closed. It's used mainly for
@@ -182,20 +199,25 @@ func (g *Gateway) Close() error {
 	return err
 }
 
-// Reconnect tries to reconnect forever. It will resume the connection if
-// possible. If an Invalid Session is received, it will start a fresh one.
+// Reconnect tries to reconnect until the ReconnectTimeout is reached, or if
+// set to 0 reconnects indefinitely.
 func (g *Gateway) Reconnect() {
-	for {
-		if err := g.ReconnectCtx(context.Background()); err != nil {
-			g.ErrorLog(err)
-		} else {
-			return
-		}
+	ctx := context.Background()
+
+	if g.ReconnectTimeout > 0 {
+		var cancel func()
+		ctx, cancel = context.WithTimeout(context.Background(), g.WSTimeout)
+
+		defer cancel()
 	}
+
+	// ignore the error, it is already logged and FatalErrorCallback was called
+	g.ReconnectCtx(ctx)
 }
 
-// ReconnectCtx attempts to reconnect until context expires. If context cannot
-// expire, then the gateway will try to reconnect forever.
+// ReconnectCtx attempts to reconnect until context expires.
+// If the context expires FatalErrorCallback will be called with ErrWSMaxTries,
+// and the last error returned by Open will be returned.
 func (g *Gateway) ReconnectCtx(ctx context.Context) (err error) {
 	wsutil.WSDebug("Reconnecting...")
 
@@ -206,6 +228,7 @@ func (g *Gateway) ReconnectCtx(ctx context.Context) (err error) {
 	for i := 1; ; i++ {
 		select {
 		case <-ctx.Done():
+			g.FatalErrorCallback(ErrWSMaxTries)
 			return err
 		default:
 		}