From eb23d43cee3dd929f2a7bcaf11f3a9943175733b Mon Sep 17 00:00:00 2001
From: diamondburned <diamond@arikawa-hi.me>
Date: Fri, 18 Nov 2022 02:09:12 -0800
Subject: [PATCH] session: Connect to retry on non-fatal errors

---
 session/session.go | 40 +++++++++++++++++++++++++++++++++-------
 1 file changed, 33 insertions(+), 7 deletions(-)

diff --git a/session/session.go b/session/session.go
index 9f98f7f..b52bd63 100644
--- a/session/session.go
+++ b/session/session.go
@@ -14,6 +14,7 @@ import (
 	"github.com/diamondburned/arikawa/v3/gateway"
 	"github.com/diamondburned/arikawa/v3/utils/handler"
 	"github.com/diamondburned/arikawa/v3/utils/json/option"
+	"github.com/diamondburned/arikawa/v3/utils/ws"
 	"github.com/diamondburned/arikawa/v3/utils/ws/ophandler"
 )
 
@@ -160,6 +161,21 @@ func (s *Session) Gateway() *gateway.Gateway {
 	return s.state.gateway
 }
 
+// GatewayOpts returns a copy of the current session's gateway options. If Open
+// has never been called or Session was never constructed with a gateway, then
+// the default gateway options are returned.
+func (s *Session) GatewayOpts() *ws.GatewayOpts {
+	s.state.Lock()
+	defer s.state.Unlock()
+
+	opts := &gateway.DefaultGatewayOpts
+	if s.state.gateway != nil {
+		opts = s.state.gateway.Opts()
+	}
+
+	return opts
+}
+
 // GatewayError returns the gateway's error if the gateway is dead. If it's not
 // dead, then nil is always returned. The check is done with GatewayIsAlive().
 // If the gateway has never been started, nil will be returned (even though
@@ -211,15 +227,25 @@ func (s *Session) gatewayIsAlive() bool {
 // error). This is contrary to the common behavior of a ctx function returning
 // ctx.Err().
 func (s *Session) Connect(ctx context.Context) error {
-	if err := s.Open(ctx); err != nil {
-		return err
-	}
+	opts := s.GatewayOpts()
 
-	if err := s.Wait(ctx); err != nil && ctx.Err() == nil {
-		return err
-	}
+	for {
+		if err := s.Open(ctx); err != nil {
+			if opts.ErrorIsFatalClose(err) {
+				return err
+			}
+			continue
+		}
 
-	return nil
+		if err := s.Wait(ctx); err != nil {
+			if opts.ErrorIsFatalClose(err) {
+				return err
+			}
+			if ctx.Err() == nil {
+				return err
+			}
+		}
+	}
 }
 
 // Open opens the Discord gateway and its handler, then waits until either the