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 {