From 375ac53f7465fb67748649ef9f34feaece355278 Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Tue, 8 Feb 2022 16:37:31 +0100
Subject: [PATCH] Add context to upstreamConn.runUntilRegistered

---
 upstream.go | 6 +++---
 user.go     | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/upstream.go b/upstream.go
index 7d64e55..5899ce6 100644
--- a/upstream.go
+++ b/upstream.go
@@ -1955,14 +1955,14 @@ func (uc *upstreamConn) ReadMessage() (*irc.Message, error) {
 	return msg, nil
 }
 
-func (uc *upstreamConn) runUntilRegistered() error {
+func (uc *upstreamConn) runUntilRegistered(ctx context.Context) error {
 	for !uc.registered {
 		msg, err := uc.ReadMessage()
 		if err != nil {
 			return fmt.Errorf("failed to read message: %v", err)
 		}
 
-		if err := uc.handleMessage(context.TODO(), msg); err != nil {
+		if err := uc.handleMessage(ctx, msg); err != nil {
 			if _, ok := err.(registrationError); ok {
 				return err
 			} else {
@@ -1977,7 +1977,7 @@ func (uc *upstreamConn) runUntilRegistered() error {
 		if err != nil {
 			uc.logger.Printf("failed to parse connect command %q: %v", command, err)
 		} else {
-			uc.SendMessage(context.TODO(), m)
+			uc.SendMessage(ctx, m)
 		}
 	}
 
diff --git a/user.go b/user.go
index a1ff2d9..9530e0f 100644
--- a/user.go
+++ b/user.go
@@ -221,7 +221,7 @@ func (net *network) run() {
 		}
 
 		uc.register()
-		if err := uc.runUntilRegistered(); err != nil {
+		if err := uc.runUntilRegistered(context.TODO()); err != nil {
 			text := err.Error()
 			temp := true
 			if regErr, ok := err.(registrationError); ok {