From 789afaa88f70e56fa29989654cec3c38bcd8c2fd Mon Sep 17 00:00:00 2001
From: diamondburned <diamond@arikawa-hi.me>
Date: Sun, 20 Nov 2022 20:30:57 -0800
Subject: [PATCH] session: Fix Connect->Open not handling ctx

---
 session/session.go | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/session/session.go b/session/session.go
index 18bc0f6..19ffc12 100644
--- a/session/session.go
+++ b/session/session.go
@@ -231,19 +231,25 @@ func (s *Session) Connect(ctx context.Context) error {
 
 	for {
 		if err := s.Open(ctx); err != nil {
-			if opts.ErrorIsFatalClose(err) {
+			if opts.ErrorIsFatalClose(err) || ctx.Err() != nil {
+				// Fatal error or context is done, return.
 				return err
 			}
+			// Non-fatal error, retry.
 			continue
 		}
 
 		if err := s.Wait(ctx); err != nil {
 			if opts.ErrorIsFatalClose(err) {
+				// Gateway returned a fatal error, so we can't recover.
 				return err
 			}
-			if ctx.Err() == nil {
-				return err
+			if ctx.Err() != nil {
+				// Context was done, so we can't recover. Exit with no error,
+				// since we're just waiting.
+				return nil
 			}
+			// Non-fatal error, retry.
 		}
 	}
 }