fix/conn: prevent re-entrant logger deadlock

Protocol errors invoke the configured logger while the connection is shutting down. Calling user code under the connection mutex meant a logger that sends over the same connection could never observe closure and instead deadlocked. Complete teardown before logging so re-entrant calls return ErrClosed.

Amp-Thread-ID: https://ampcode.com/threads/T-019fb774-aef7-75df-8eb3-8888d6835754
Co-authored-by: Amp <amp@ampcode.com>

Author: keegancsmith

conn.go

@@ -180,24 +180,29 @@ func (c *Conn) SendResponse(ctx context.Context, resp *Response) error {
 
 func (c *Conn) close(cause error) error {
 	c.mu.Lock()
-	defer c.mu.Unlock()
-
 	if c.closed {
+		c.mu.Unlock()
 		return ErrClosed
 	}
+	c.closed = true
 
 	for _, call := range c.pending {
 		close(call.done)
 	}
 
+	close(c.disconnect)
+	c.mu.Unlock()
+
+	c.cancelCtx()
+	err := c.stream.Close()
+
+	// The logger may call back into c, so invoke it only after shutdown is
+	// complete and c.mu is unlocked.
 	if cause != nil && cause != io.EOF && cause != io.ErrUnexpectedEOF {
 		c.logger.Printf("jsonrpc2: protocol error: %v\n", cause)
 	}
 
-	close(c.disconnect)
-	c.cancelCtx()
-	c.closed = true
-	return c.stream.Close()
+	return err
 }
 
 func (c *Conn) readMessages(ctx context.Context) {

conn_test.go

@@ -167,6 +167,38 @@ func TestConn_DisconnectNotify(t *testing.T) {
 		connA.Write([]byte("invalid json"))
 		assertDisconnect(t, c, connB)
 	})
+
+	t.Run("protocol error logger uses connection", func(t *testing.T) {
+		connA, connB := net.Pipe()
+		logResult := make(chan error, 1)
+		var c *jsonrpc2.Conn
+		c = jsonrpc2.NewConn(
+			context.Background(),
+			jsonrpc2.NewPlainObjectStream(connB),
+			noopHandler{},
+			jsonrpc2.SetLogger(connNotifyLogger{conn: &c, result: logResult}),
+		)
+		connA.Write([]byte("invalid json"))
+		assertDisconnect(t, c, connB)
+
+		select {
+		case err := <-logResult:
+			if err != jsonrpc2.ErrClosed {
+				t.Errorf("logger Notify: got %v, want %v", err, jsonrpc2.ErrClosed)
+			}
+		case <-time.After(200 * time.Millisecond):
+			t.Error("logger blocked using connection")
+		}
+	})
+}
+
+type connNotifyLogger struct {
+	conn   **jsonrpc2.Conn
+	result chan<- error
+}
+
+func (l connNotifyLogger) Printf(string, ...interface{}) {
+	l.result <- (*l.conn).Notify(context.Background(), "log", nil)
 }
 
 func TestConn_Close(t *testing.T) {