Changeset 704 in code for trunk/downstream.go

Timestamp:

Nov 17, 2021, 1:54:03 PM (4 years ago)

Author:

contact

Message:

Add timeout for downstream connection registration

File:

• trunk/downstream.go (modified) (3 diffs)

trunk/downstream.go

@@ -615 +615 @@
 }
 
-func (dc *downstreamConn) handleMessage(msg *irc.Message) error {
-        ctx, cancel := dc.conn.NewContext(context.TODO())
+func (dc *downstreamConn) handleMessage(ctx context.Context, msg *irc.Message) error {
+        ctx, cancel := dc.conn.NewContext(ctx)
         defer cancel()
 
@@ -1402 +1402 @@
 
 func (dc *downstreamConn) runUntilRegistered() error {
+        ctx, cancel := context.WithTimeout(context.TODO(), downstreamRegisterTimeout)
+        defer cancel()
+
+        // Close the connection with an error if the deadline is exceeded
+        go func() {
+                <-ctx.Done()
+                if err := ctx.Err(); err == context.DeadlineExceeded {
+                        dc.SendMessage(&irc.Message{
+                                Prefix:  dc.srv.prefix(),
+                                Command: "ERROR",
+                                Params:  []string{"Connection registration timed out"},
+                        })
+                        dc.Close()
+                }
+        }()
+
         for !dc.registered {
                 msg, err := dc.ReadMessage()
@@ -1408 +1424 @@
                 }
 
-                err = dc.handleMessage(msg)
+                err = dc.handleMessage(ctx, msg)
                 if ircErr, ok := err.(ircError); ok {
                         ircErr.Message.Prefix = dc.srv.prefix()