fix/conn: prevent close-response race panic

Closing pending calls must also remove them from response routing; otherwise a response already read from the stream can later send into a closed waiter channel. Keep registration and shutdown mutually exclusive so the pending map retains that invariant.

Amp-Thread-ID: https://ampcode.com/threads/T-01a08a19-4a3a-754a-80cb-56329e64cae5
Co-authored-by: Amp <amp@ampcode.com>

conn.go

@@ -186,8 +186,9 @@ func (c *Conn) close(cause error) error {
 	}
 	c.closed = true
 
-	for _, call := range c.pending {
+	for id, call := range c.pending {
 		close(call.done)
+		delete(c.pending, id)
 	}
 
 	close(c.disconnect)
@@ -318,6 +319,10 @@ func (c *Conn) send(_ context.Context, m *anyMessage, wait bool) (cc *call, err
 	// responses.
 	if m.request != nil && wait {
 		c.mu.Lock()
+		if c.closed {
+			c.mu.Unlock()
+			return nil, ErrClosed
+		}
 		id = m.request.ID
 		c.pending[id] = cc
 		c.mu.Unlock()

conn_test.go

@@ -275,6 +275,78 @@ func TestConn_Close(t *testing.T) {
 	}
 }
 
+func TestConn_CloseAfterResponseRead(t *testing.T) {
+	client, server := net.Pipe()
+	responseRead := make(chan struct{})
+	processResponse := make(chan struct{})
+	responseIgnored := make(chan struct{})
+	var resumeResponse sync.Once
+	resume := func() { resumeResponse.Do(func() { close(processResponse) }) }
+	defer resume()
+
+	clientStream := &pauseAfterReadStream{
+		ObjectStream: jsonrpc2.NewPlainObjectStream(client),
+		read:         responseRead,
+		resume:       processResponse,
+	}
+	clientConn := jsonrpc2.NewConn(
+		context.Background(),
+		clientStream,
+		noopHandler{},
+		jsonrpc2.SetLogger(signalLogger{called: responseIgnored}),
+	)
+	serverConn := jsonrpc2.NewConn(
+		context.Background(),
+		jsonrpc2.NewPlainObjectStream(server),
+		handlerFunc(func(ctx context.Context, conn *jsonrpc2.Conn, req *jsonrpc2.Request) {
+			if err := conn.Reply(ctx, req.ID, nil); err != nil {
+				t.Error(err)
+			}
+		}),
+	)
+	defer serverConn.Close()
+
+	call, err := clientConn.DispatchCall(context.Background(), "m", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	<-responseRead
+
+	if err := clientConn.Close(); err != nil {
+		t.Fatal(err)
+	}
+	resume()
+	<-responseIgnored
+
+	if err := call.Wait(context.Background(), nil); err != jsonrpc2.ErrClosed {
+		t.Fatalf("got error %v, want %v", err, jsonrpc2.ErrClosed)
+	}
+}
+
+type pauseAfterReadStream struct {
+	jsonrpc2.ObjectStream
+	read   chan<- struct{}
+	resume <-chan struct{}
+	once   sync.Once
+}
+
+func (s *pauseAfterReadStream) ReadObject(v interface{}) error {
+	err := s.ObjectStream.ReadObject(v)
+	s.once.Do(func() {
+		close(s.read)
+		<-s.resume
+	})
+	return err
+}
+
+type signalLogger struct {
+	called chan<- struct{}
+}
+
+func (l signalLogger) Printf(string, ...interface{}) {
+	close(l.called)
+}
+
 func testParams(t *testing.T, want *json.RawMessage, fn func(c *jsonrpc2.Conn) error) {
 	wg := &sync.WaitGroup{}
 	handler := handlerFunc(func(ctx context.Context, conn *jsonrpc2.Conn, r *jsonrpc2.Request) {