From 8b3e5e7465dcb9742dfa850be10ff0b5ac1cfb18 Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Mon, 8 Nov 2021 18:11:24 +0100
Subject: [PATCH] Plumb context in downstreamConn.handleMessageRegistered

References: https://todo.sr.ht/~emersion/soju/141
---
 downstream.go | 22 +++++++++++-----------
 server.go     |  3 ++-
 2 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/downstream.go b/downstream.go
index 0f1a244..016f84f 100644
--- a/downstream.go
+++ b/downstream.go
@@ -1335,7 +1335,7 @@ func (dc *downstreamConn) sendTargetBacklog(net *network, target, msgID string)
 
 	ch := net.channels.Value(target)
 
-	ctx, cancel := context.WithTimeout(context.TODO(), messageStoreTimeout)
+	ctx, cancel := context.WithTimeout(context.TODO(), backlogTimeout)
 	defer cancel()
 
 	targetCM := net.casemap(target)
@@ -1393,6 +1393,9 @@ func (dc *downstreamConn) runUntilRegistered() error {
 }
 
 func (dc *downstreamConn) handleMessageRegistered(msg *irc.Message) error {
+	ctx, cancel := context.WithTimeout(context.TODO(), handleDownstreamMessageTimeout)
+	defer cancel()
+
 	switch msg.Command {
 	case "CAP":
 		var subCmd string
@@ -1468,7 +1471,7 @@ func (dc *downstreamConn) handleMessageRegistered(msg *irc.Message) error {
 				return
 			}
 			n.Nick = nick
-			err = dc.srv.db.StoreNetwork(context.TODO(), dc.user.ID, &n.Network)
+			err = dc.srv.db.StoreNetwork(ctx, dc.user.ID, &n.Network)
 		})
 		if err != nil {
 			return err
@@ -1518,7 +1521,7 @@ func (dc *downstreamConn) handleMessageRegistered(msg *irc.Message) error {
 				})
 
 				n.Realname = storeRealname
-				if err := dc.srv.db.StoreNetwork(context.TODO(), dc.user.ID, &n.Network); err != nil {
+				if err := dc.srv.db.StoreNetwork(ctx, dc.user.ID, &n.Network); err != nil {
 					dc.logger.Printf("failed to store network realname: %v", err)
 					storeErr = err
 				}
@@ -1607,7 +1610,7 @@ func (dc *downstreamConn) handleMessageRegistered(msg *irc.Message) error {
 				}
 				uc.network.channels.SetValue(upstreamName, ch)
 			}
-			if err := dc.srv.db.StoreChannel(context.TODO(), uc.network.ID, ch); err != nil {
+			if err := dc.srv.db.StoreChannel(ctx, uc.network.ID, ch); err != nil {
 				dc.logger.Printf("failed to create or update channel %q: %v", upstreamName, err)
 			}
 		}
@@ -1639,7 +1642,7 @@ func (dc *downstreamConn) handleMessageRegistered(msg *irc.Message) error {
 					}
 					uc.network.channels.SetValue(upstreamName, ch)
 				}
-				if err := dc.srv.db.StoreChannel(context.TODO(), uc.network.ID, ch); err != nil {
+				if err := dc.srv.db.StoreChannel(ctx, uc.network.ID, ch); err != nil {
 					dc.logger.Printf("failed to create or update channel %q: %v", upstreamName, err)
 				}
 			} else {
@@ -2168,7 +2171,7 @@ func (dc *downstreamConn) handleMessageRegistered(msg *irc.Message) error {
 			}
 
 			if msg.Command == "PRIVMSG" && uc.network.casemap(upstreamName) == "nickserv" {
-				dc.handleNickServPRIVMSG(uc, text)
+				dc.handleNickServPRIVMSG(ctx, uc, text)
 			}
 
 			unmarshaledText := text
@@ -2343,9 +2346,6 @@ func (dc *downstreamConn) handleMessageRegistered(msg *irc.Message) error {
 
 		eventPlayback := dc.caps["draft/event-playback"]
 
-		ctx, cancel := context.WithTimeout(context.TODO(), messageStoreTimeout)
-		defer cancel()
-
 		var history []*irc.Message
 		switch subcommand {
 		case "BEFORE":
@@ -2545,7 +2545,7 @@ func (dc *downstreamConn) handleMessageRegistered(msg *irc.Message) error {
 	return nil
 }
 
-func (dc *downstreamConn) handleNickServPRIVMSG(uc *upstreamConn, text string) {
+func (dc *downstreamConn) handleNickServPRIVMSG(ctx context.Context, uc *upstreamConn, text string) {
 	username, password, ok := parseNickServCredentials(text, uc.nick)
 	if !ok {
 		return
@@ -2562,7 +2562,7 @@ func (dc *downstreamConn) handleNickServPRIVMSG(uc *upstreamConn, text string) {
 	n.SASL.Mechanism = "PLAIN"
 	n.SASL.Plain.Username = username
 	n.SASL.Plain.Password = password
-	if err := dc.srv.db.StoreNetwork(context.TODO(), dc.user.ID, &n.Network); err != nil {
+	if err := dc.srv.db.StoreNetwork(ctx, dc.user.ID, &n.Network); err != nil {
 		dc.logger.Printf("failed to save NickServ credentials: %v", err)
 	}
 }
diff --git a/server.go b/server.go
index ea83b95..116a3d4 100644
--- a/server.go
+++ b/server.go
@@ -25,7 +25,8 @@ var connectTimeout = 15 * time.Second
 var writeTimeout = 10 * time.Second
 var upstreamMessageDelay = 2 * time.Second
 var upstreamMessageBurst = 10
-var messageStoreTimeout = 10 * time.Second
+var backlogTimeout = 10 * time.Second
+var handleDownstreamMessageTimeout = 10 * time.Second
 var chatHistoryLimit = 1000
 var backlogLimit = 4000