ssh/agent/client_pipeline_test.go - crypto - Git at Google

 // Copyright 2026 The Go Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.

 package agent

 import (
 	"encoding/binary"
 	"errors"
 	"fmt"
 	"io"
 	"net"
 	"strings"
 	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
 )

 func TestPipelineParallelSigns(t *testing.T) {
 	c1, c2, err := netPipe()
 	if err != nil {
 		t.Fatalf("netPipe: %v", err)
 	}
 	defer c1.Close()
 	defer c2.Close()

 	keyring := NewKeyring()
 	if err := keyring.Add(AddedKey{PrivateKey: testPrivateKeys["rsa"]}); err != nil {
 		t.Fatal(err)
 	}
 	go ServeAgent(keyring, c2)

 	client := NewClient(c1)
 	pubKey := testPublicKeys["rsa"]

 	const N = 64
 	var wg sync.WaitGroup
 	errCh := make(chan error, N)
 	for i := range N {
 		wg.Go(func() {
 			data := []byte{byte(i), byte(i >> 8)}
 			sig, err := client.Sign(pubKey, data)
 			if err != nil {
 				errCh <- err
 				return
 			}
 			if err := pubKey.Verify(data, sig); err != nil {
 				errCh <- err
 			}
 		})
 	}
 	wg.Wait()
 	close(errCh)
 	for err := range errCh {
 		t.Error(err)
 	}
 }

 func TestPipelineBackpressure(t *testing.T) {
 	c1, c2, err := netPipe()
 	if err != nil {
 		t.Fatalf("netPipe: %v", err)
 	}

 	const total = pipelineMaxInFlight + 8

 	// Fake agent: the reader drains requests from the wire independently
 	// of response production, so the write path is never self-blocked.
 	// The responder only writes a canned success reply once unblocked
 	// via gate.
 	var reqsRead int64
 	gate := make(chan struct{}, total)
 	pending := make(chan struct{}, total+pipelineMaxInFlight)
 	readerDone := make(chan struct{})
 	responderDone := make(chan struct{})
 	go func() {
 		defer close(readerDone)
 		defer close(pending)
 		for {
 			var sizeBuf [4]byte
 			if _, err := io.ReadFull(c2, sizeBuf[:]); err != nil {
 				return
 			}
 			size := binary.BigEndian.Uint32(sizeBuf[:])
 			req := make([]byte, size)
 			if _, err := io.ReadFull(c2, req); err != nil {
 				return
 			}
 			atomic.AddInt64(&reqsRead, 1)
 			pending <- struct{}{}
 		}
 	}()
 	go func() {
 		defer close(responderDone)
 		for range pending {
 			<-gate
 			var out [5]byte
 			binary.BigEndian.PutUint32(out[:4], 1)
 			out[4] = agentSuccess
 			if _, err := c2.Write(out[:]); err != nil {
 				return
 			}
 		}
 	}()
 	defer func() {
 		c1.Close()
 		c2.Close()
 		<-readerDone
 		<-responderDone
 	}()

 	client := NewClient(c1)

 	var wg sync.WaitGroup
 	errs := make(chan error, total)
 	for range total {
 		wg.Go(func() {
 			// RemoveAll maps to a simpleCall expecting agentSuccess.
 			if err := client.RemoveAll(); err != nil {
 				errs <- err
 			}
 		})
 	}

 	// Wait for the client to fill the pipeline.
 	deadline := time.Now().Add(2 * time.Second)
 	for time.Now().Before(deadline) {
 		if atomic.LoadInt64(&reqsRead) >= pipelineMaxInFlight {
 			break
 		}
 		time.Sleep(5 * time.Millisecond)
 	}
 	if got := atomic.LoadInt64(&reqsRead); got != pipelineMaxInFlight {
 		t.Errorf("requests in flight before any response: got %d, want %d", got, pipelineMaxInFlight)
 	}

 	// Release all responses; every caller should eventually complete.
 	for range total {
 		gate <- struct{}{}
 	}
 	wg.Wait()
 	close(errs)
 	for err := range errs {
 		t.Error(err)
 	}

 	if got := atomic.LoadInt64(&reqsRead); got != total {
 		t.Errorf("total requests received: got %d, want %d", got, total)
 	}
 }

 func TestPipelineConnCloseFailsInFlight(t *testing.T) {
 	c1, c2, err := netPipe()
 	if err != nil {
 		t.Fatalf("netPipe: %v", err)
 	}

 	// Server that reads requests but never replies, so all calls are
 	// parked waiting for responses.
 	serverExit := make(chan struct{})
 	go func() {
 		defer close(serverExit)
 		buf := make([]byte, 4096)
 		for {
 			if _, err := c2.Read(buf); err != nil {
 				return
 			}
 		}
 	}()

 	client := NewClient(c1)

 	const N = 8
 	var wg sync.WaitGroup
 	errs := make(chan error, N)
 	for range N {
 		wg.Go(func() {
 			errs <- client.RemoveAll()
 		})
 	}

 	// Give callers time to enter the pipeline, then tear down the conn.
 	time.Sleep(50 * time.Millisecond)
 	c1.Close()
 	c2.Close()

 	// All callers must return promptly with an error.
 	doneCh := make(chan struct{})
 	go func() {
 		wg.Wait()
 		close(doneCh)
 	}()
 	select {
 	case <-doneCh:
 	case <-time.After(5 * time.Second):
 		t.Fatal("pipelined callers did not complete after conn close")
 	}

 	close(errs)
 	for err := range errs {
 		if err == nil {
 			t.Error("expected error after conn close, got nil")
 		}
 	}

 	// Subsequent calls also fail fast rather than hanging.
 	errCh := make(chan error, 1)
 	go func() { errCh <- client.RemoveAll() }()
 	select {
 	case err := <-errCh:
 		if err == nil {
 			t.Error("expected error on post-close call, got nil")
 		}
 	case <-time.After(2 * time.Second):
 		t.Fatal("post-close call hung")
 	}

 	<-serverExit
 }

 // readWriter exposes only io.Reader and io.Writer, hiding any Close
 // method on the wrapped value so NewClient cannot take the pipelined
 // path.
 type readWriter struct {
 	io.Reader
 	io.Writer
 }

 func TestNewClientSerialFallback(t *testing.T) {
 	c1, c2, err := netPipe()
 	if err != nil {
 		t.Fatalf("netPipe: %v", err)
 	}
 	defer c1.Close()
 	defer c2.Close()

 	keyring := NewKeyring()
 	if err := keyring.Add(AddedKey{PrivateKey: testPrivateKeys["rsa"]}); err != nil {
 		t.Fatal(err)
 	}
 	go ServeAgent(keyring, c2)

 	rw := readWriter{Reader: c1, Writer: c1}
 	ag := NewClient(rw)

 	ci, ok := ag.(*client)
 	if !ok {
 		t.Fatalf("NewClient: got %T, want *client", ag)
 	}
 	if ci.pipeline != nil {
 		t.Error("NewClient with non-Closer rw: pipeline must be nil")
 	}
 	if ci.conn == nil {
 		t.Error("NewClient with non-Closer rw: conn must be set")
 	}

 	// The serialized path correctly handles a request/response cycle.
 	keys, err := ag.List()
 	if err != nil {
 		t.Fatalf("List: %v", err)
 	}
 	if len(keys) != 1 {
 		t.Errorf("List: got %d keys, want 1", len(keys))
 	}
 }

 func TestSerialFallbackParallelSigns(t *testing.T) {
 	c1, c2, err := netPipe()
 	if err != nil {
 		t.Fatalf("netPipe: %v", err)
 	}
 	defer c1.Close()
 	defer c2.Close()

 	keyring := NewKeyring()
 	if err := keyring.Add(AddedKey{PrivateKey: testPrivateKeys["rsa"]}); err != nil {
 		t.Fatal(err)
 	}
 	go ServeAgent(keyring, c2)

 	rw := readWriter{Reader: c1, Writer: c1}
 	ag := NewClient(rw)

 	if ci, ok := ag.(*client); !ok {
 		t.Fatalf("NewClient: got %T, want *client", ag)
 	} else if ci.pipeline != nil {
 		t.Fatal("pipeline must be nil for non-Closer rw")
 	}

 	pubKey := testPublicKeys["rsa"]

 	const N = 64
 	var wg sync.WaitGroup
 	errCh := make(chan error, N)
 	for i := range N {
 		wg.Go(func() {
 			data := []byte{byte(i), byte(i >> 8)}
 			sig, err := ag.Sign(pubKey, data)
 			if err != nil {
 				errCh <- err
 				return
 			}
 			if err := pubKey.Verify(data, sig); err != nil {
 				errCh <- err
 			}
 		})
 	}
 	wg.Wait()
 	close(errCh)
 	for err := range errCh {
 		t.Error(err)
 	}
 }

 // writeFailConn wraps a net.Conn and forces Write to return an error when
 // the fail flag is set. Close delegates to the underlying conn so the test
 // can observe whether the pipeline closed it.
 type writeFailConn struct {
 	net.Conn
 	fail atomic.Bool
 }

 func (c *writeFailConn) Write(p []byte) (int, error) {
 	if c.fail.Load() {
 		return 0, errors.New("synthetic write failure")
 	}
 	return c.Conn.Write(p)
 }

 func TestPipelineWriteErrorClosesConn(t *testing.T) {
 	c1, c2, err := netPipe()
 	if err != nil {
 		t.Fatalf("netPipe: %v", err)
 	}
 	defer c2.Close()

 	// Peer that reads and discards, exiting when the conn closes.
 	serverExit := make(chan struct{})
 	go func() {
 		defer close(serverExit)
 		io.Copy(io.Discard, c2)
 	}()

 	wrapper := &writeFailConn{Conn: c1}
 	wrapper.fail.Store(true)

 	client := NewClient(wrapper)

 	errCh := make(chan error, 1)
 	go func() { errCh <- client.RemoveAll() }()

 	select {
 	case err := <-errCh:
 		if err == nil {
 			t.Fatal("expected error on forced write failure, got nil")
 		}
 		// The caller whose Write failed must receive the actual
 		// write error, not the reader's "closed connection" error
 		// from the post-shutdown drain.
 		if !strings.Contains(err.Error(), "synthetic write failure") {
 			t.Errorf("got %q, want error containing %q", err, "synthetic write failure")
 		}
 	case <-time.After(2 * time.Second):
 		t.Fatal("RemoveAll hung after write failure")
 	}

 	// The pipeline must have closed the underlying conn, so the peer
 	// goroutine returns from its Read.
 	select {
 	case <-serverExit:
 	case <-time.After(2 * time.Second):
 		t.Fatal("connection was not closed after write failure")
 	}

 	// After shutdown, subsequent calls must also fail promptly.
 	errCh2 := make(chan error, 1)
 	go func() { errCh2 <- client.RemoveAll() }()
 	select {
 	case err := <-errCh2:
 		if err == nil {
 			t.Error("expected error on post-shutdown call, got nil")
 		}
 	case <-time.After(2 * time.Second):
 		t.Fatal("post-shutdown call hung")
 	}
 }

 func TestPipelineResponseTooLarge(t *testing.T) {
 	c1, c2, err := netPipe()
 	if err != nil {
 		t.Fatalf("netPipe: %v", err)
 	}
 	defer c1.Close()
 	defer c2.Close()

 	// Server: read the request, then reply with a size header that
 	// exceeds maxAgentResponseBytes. The size check fires before the
 	// body is read, so no body needs to follow.
 	serverDone := make(chan struct{})
 	go func() {
 		defer close(serverDone)
 		var sizeBuf [4]byte
 		if _, err := io.ReadFull(c2, sizeBuf[:]); err != nil {
 			return
 		}
 		size := binary.BigEndian.Uint32(sizeBuf[:])
 		if _, err := io.CopyN(io.Discard, c2, int64(size)); err != nil {
 			return
 		}
 		var out [4]byte
 		binary.BigEndian.PutUint32(out[:], maxAgentResponseBytes+1)
 		c2.Write(out[:])
 	}()

 	client := NewClient(c1)

 	errCh := make(chan error, 1)
 	go func() { errCh <- client.RemoveAll() }()
 	select {
 	case err := <-errCh:
 		if err == nil {
 			t.Fatal("expected error, got nil")
 		}
 		if !strings.Contains(err.Error(), "response too large") {
 			t.Errorf("got %q, want error containing %q", err, "response too large")
 		}
 	case <-time.After(2 * time.Second):
 		t.Fatal("call hung")
 	}

 	// Subsequent calls must also fail: the pipeline is shut down.
 	errCh2 := make(chan error, 1)
 	go func() { errCh2 <- client.RemoveAll() }()
 	select {
 	case err := <-errCh2:
 		if err == nil {
 			t.Error("expected error on post-shutdown call, got nil")
 		}
 	case <-time.After(2 * time.Second):
 		t.Fatal("post-shutdown call hung")
 	}

 	<-serverDone
 }

 func TestPipelineShutdownWithBlockedWriter(t *testing.T) {
 	c1, c2, err := netPipe()
 	if err != nil {
 		t.Fatalf("netPipe: %v", err)
 	}
 	defer c1.Close()

 	// Server drains requests off the wire so writes never block on
 	// backpressure from the transport, but never replies. The reader
 	// goroutine therefore stays parked on Read, never dequeues from
 	// pending, and the queue fills up.
 	var reqsRead int64
 	serverExit := make(chan struct{})
 	go func() {
 		defer close(serverExit)
 		for {
 			var sizeBuf [4]byte
 			if _, err := io.ReadFull(c2, sizeBuf[:]); err != nil {
 				return
 			}
 			size := binary.BigEndian.Uint32(sizeBuf[:])
 			if _, err := io.CopyN(io.Discard, c2, int64(size)); err != nil {
 				return
 			}
 			atomic.AddInt64(&reqsRead, 1)
 		}
 	}()

 	client := NewClient(c1)

 	const total = pipelineMaxInFlight + 4
 	var wg sync.WaitGroup
 	errs := make(chan error, total)
 	for range total {
 		wg.Go(func() {
 			errs <- client.RemoveAll()
 		})
 	}

 	// Wait for the queue to fill: once pipelineMaxInFlight requests
 	// are on the wire, pending is at capacity and the +1 writer is
 	// blocked in the second-select arm; the rest are queued behind
 	// writeMu.
 	deadline := time.Now().Add(2 * time.Second)
 	for time.Now().Before(deadline) {
 		if atomic.LoadInt64(&reqsRead) >= pipelineMaxInFlight {
 			break
 		}
 		time.Sleep(5 * time.Millisecond)
 	}
 	if got := atomic.LoadInt64(&reqsRead); got < pipelineMaxInFlight {
 		t.Fatalf("requests on wire: got %d, want >= %d", got, pipelineMaxInFlight)
 	}
 	// Give the +1 writer a moment to actually park in the second
 	// select before we trigger shutdown.
 	time.Sleep(50 * time.Millisecond)

 	// Trigger shutdown: closing the server side makes the client's
 	// reader hit EOF, which calls shutdown(...).
 	c2.Close()

 	doneCh := make(chan struct{})
 	go func() {
 		wg.Wait()
 		close(doneCh)
 	}()
 	select {
 	case <-doneCh:
 	case <-time.After(5 * time.Second):
 		t.Fatal("blocked callers did not complete after shutdown")
 	}

 	close(errs)
 	for err := range errs {
 		if err == nil {
 			t.Error("expected error after shutdown, got nil")
 		}
 	}

 	<-serverExit
 }

 // multiDeviceAgent simulates an ssh-agent backed by numDevices signing
 // devices, each of which takes latency to process a request. Incoming
 // requests are round-robin dispatched to devices, and responses are
 // serialized back to the wire in request order as the protocol requires.
 //
 // The simulation does not emit a semantically meaningful reply; it just
 // writes a 1-byte SSH_AGENT_SUCCESS after the simulated latency has
 // elapsed. This is enough to exercise callRaw/simpleCall (RemoveAll).
 func multiDeviceAgent(t testing.TB, conn io.ReadWriteCloser, numDevices int, latency time.Duration) (done <-chan struct{}) {
 	t.Helper()

 	type job struct {
 		ready chan struct{}
 	}

 	devices := make([]chan job, numDevices)
 	var devicesWG sync.WaitGroup
 	for i := range numDevices {
 		devices[i] = make(chan job, 1024)
 		ch := devices[i]
 		devicesWG.Go(func() {
 			for j := range ch {
 				time.Sleep(latency)
 				close(j.ready)
 			}
 		})
 	}

 	// Order-preserving responder: responses are written in the same
 	// order the requests were read, regardless of which device finished
 	// first.
 	orderCh := make(chan chan struct{}, 1024)
 	responderDone := make(chan struct{})
 	go func() {
 		defer close(responderDone)
 		for ready := range orderCh {
 			<-ready
 			var out [5]byte
 			binary.BigEndian.PutUint32(out[:4], 1)
 			out[4] = agentSuccess
 			if _, err := conn.Write(out[:]); err != nil {
 				return
 			}
 		}
 	}()

 	finished := make(chan struct{})
 	go func() {
 		defer close(finished)
 		defer func() {
 			close(orderCh)
 			<-responderDone
 			for _, d := range devices {
 				close(d)
 			}
 			devicesWG.Wait()
 		}()
 		next := 0
 		for {
 			var sizeBuf [4]byte
 			if _, err := io.ReadFull(conn, sizeBuf[:]); err != nil {
 				return
 			}
 			size := binary.BigEndian.Uint32(sizeBuf[:])
 			if _, err := io.CopyN(io.Discard, conn, int64(size)); err != nil {
 				return
 			}
 			ready := make(chan struct{})
 			devices[next%numDevices] <- job{ready: ready}
 			next++
 			orderCh <- ready
 		}
 	}()

 	return finished
 }

 // BenchmarkPipelineMultiDevice measures throughput of the pipelined
 // agent client against a simulated multi-device agent, varying the
 // number of backend devices.
 //
 // Setup: each simulated signing device takes 2ms per request. The
 // protocol serializes responses in request order, but pipelining lets
 // the client keep multiple devices busy concurrently; throughput
 // should scale roughly linearly with numDevices up to the pipeline
 // depth (pipelineMaxInFlight).
 //
 // Run with: go test -benchmem -bench BenchmarkPipelineMultiDevice -benchtime=2s ./ssh/agent/
 func BenchmarkPipelineMultiDevice(b *testing.B) {
 	const (
 		deviceLatency = 2 * time.Millisecond
 		parallelism   = 32
 	)

 	for _, numDevices := range []int{1, 2, 4, 8, 16} {
 		b.Run(fmt.Sprintf("Devices=%d", numDevices), func(b *testing.B) {
 			c1, c2, err := netPipe()
 			if err != nil {
 				b.Fatalf("netPipe: %v", err)
 			}
 			serverDone := multiDeviceAgent(b, c2, numDevices, deviceLatency)

 			client := NewClient(c1)

 			b.ResetTimer()
 			b.SetParallelism(parallelism)
 			b.RunParallel(func(pb *testing.PB) {
 				for pb.Next() {
 					if err := client.RemoveAll(); err != nil {
 						b.Fatal(err)
 					}
 				}
 			})
 			b.StopTimer()

 			c1.Close()
 			c2.Close()
 			<-serverDone
 		})
 	}
 }
	// Copyright 2026 The Go Authors. All rights reserved.
	// Use of this source code is governed by a BSD-style
	// license that can be found in the LICENSE file.

	package agent

	import (
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
	)

	func TestPipelineParallelSigns(t *testing.T) {
	c1, c2, err := netPipe()
	if err != nil {
	t.Fatalf("netPipe: %v", err)
	}
	defer c1.Close()
	defer c2.Close()

	keyring := NewKeyring()
	if err := keyring.Add(AddedKey{PrivateKey: testPrivateKeys["rsa"]}); err != nil {
	t.Fatal(err)
	}
	go ServeAgent(keyring, c2)

	client := NewClient(c1)
	pubKey := testPublicKeys["rsa"]

	const N = 64
	var wg sync.WaitGroup
	errCh := make(chan error, N)
	for i := range N {
	wg.Go(func() {
	data := []byte{byte(i), byte(i >> 8)}
	sig, err := client.Sign(pubKey, data)
	if err != nil {
	errCh <- err
	return
	}
	if err := pubKey.Verify(data, sig); err != nil {
	errCh <- err
	}
	})
	}
	wg.Wait()
	close(errCh)
	for err := range errCh {
	t.Error(err)
	}
	}

	func TestPipelineBackpressure(t *testing.T) {
	c1, c2, err := netPipe()
	if err != nil {
	t.Fatalf("netPipe: %v", err)
	}

	const total = pipelineMaxInFlight + 8

	// Fake agent: the reader drains requests from the wire independently
	// of response production, so the write path is never self-blocked.
	// The responder only writes a canned success reply once unblocked
	// via gate.
	var reqsRead int64
	gate := make(chan struct{}, total)
	pending := make(chan struct{}, total+pipelineMaxInFlight)
	readerDone := make(chan struct{})
	responderDone := make(chan struct{})
	go func() {
	defer close(readerDone)
	defer close(pending)
	for {
	var sizeBuf [4]byte
	if _, err := io.ReadFull(c2, sizeBuf[:]); err != nil {
	return
	}
	size := binary.BigEndian.Uint32(sizeBuf[:])
	req := make([]byte, size)
	if _, err := io.ReadFull(c2, req); err != nil {
	return
	}
	atomic.AddInt64(&reqsRead, 1)
	pending <- struct{}{}
	}
	}()
	go func() {
	defer close(responderDone)
	for range pending {
	<-gate
	var out [5]byte
	binary.BigEndian.PutUint32(out[:4], 1)
	out[4] = agentSuccess
	if _, err := c2.Write(out[:]); err != nil {
	return
	}
	}
	}()
	defer func() {
	c1.Close()
	c2.Close()
	<-readerDone
	<-responderDone
	}()

	client := NewClient(c1)

	var wg sync.WaitGroup
	errs := make(chan error, total)
	for range total {
	wg.Go(func() {
	// RemoveAll maps to a simpleCall expecting agentSuccess.
	if err := client.RemoveAll(); err != nil {
	errs <- err
	}
	})
	}

	// Wait for the client to fill the pipeline.
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
	if atomic.LoadInt64(&reqsRead) >= pipelineMaxInFlight {
	break
	}
	time.Sleep(5 * time.Millisecond)
	}
	if got := atomic.LoadInt64(&reqsRead); got != pipelineMaxInFlight {
	t.Errorf("requests in flight before any response: got %d, want %d", got, pipelineMaxInFlight)
	}

	// Release all responses; every caller should eventually complete.
	for range total {
	gate <- struct{}{}
	}
	wg.Wait()
	close(errs)
	for err := range errs {
	t.Error(err)
	}

	if got := atomic.LoadInt64(&reqsRead); got != total {
	t.Errorf("total requests received: got %d, want %d", got, total)
	}
	}

	func TestPipelineConnCloseFailsInFlight(t *testing.T) {
	c1, c2, err := netPipe()
	if err != nil {
	t.Fatalf("netPipe: %v", err)
	}

	// Server that reads requests but never replies, so all calls are
	// parked waiting for responses.
	serverExit := make(chan struct{})
	go func() {
	defer close(serverExit)
	buf := make([]byte, 4096)
	for {
	if _, err := c2.Read(buf); err != nil {
	return
	}
	}
	}()

	client := NewClient(c1)

	const N = 8
	var wg sync.WaitGroup
	errs := make(chan error, N)
	for range N {
	wg.Go(func() {
	errs <- client.RemoveAll()
	})
	}

	// Give callers time to enter the pipeline, then tear down the conn.
	time.Sleep(50 * time.Millisecond)
	c1.Close()
	c2.Close()

	// All callers must return promptly with an error.
	doneCh := make(chan struct{})
	go func() {
	wg.Wait()
	close(doneCh)
	}()
	select {
	case <-doneCh:
	case <-time.After(5 * time.Second):
	t.Fatal("pipelined callers did not complete after conn close")
	}

	close(errs)
	for err := range errs {
	if err == nil {
	t.Error("expected error after conn close, got nil")
	}
	}

	// Subsequent calls also fail fast rather than hanging.
	errCh := make(chan error, 1)
	go func() { errCh <- client.RemoveAll() }()
	select {
	case err := <-errCh:
	if err == nil {
	t.Error("expected error on post-close call, got nil")
	}
	case <-time.After(2 * time.Second):
	t.Fatal("post-close call hung")
	}

	<-serverExit
	}

	// readWriter exposes only io.Reader and io.Writer, hiding any Close
	// method on the wrapped value so NewClient cannot take the pipelined
	// path.
	type readWriter struct {
	io.Reader
	io.Writer
	}

	func TestNewClientSerialFallback(t *testing.T) {
	c1, c2, err := netPipe()
	if err != nil {
	t.Fatalf("netPipe: %v", err)
	}
	defer c1.Close()
	defer c2.Close()

	keyring := NewKeyring()
	if err := keyring.Add(AddedKey{PrivateKey: testPrivateKeys["rsa"]}); err != nil {
	t.Fatal(err)
	}
	go ServeAgent(keyring, c2)

	rw := readWriter{Reader: c1, Writer: c1}
	ag := NewClient(rw)

	ci, ok := ag.(*client)
	if !ok {
	t.Fatalf("NewClient: got %T, want *client", ag)
	}
	if ci.pipeline != nil {
	t.Error("NewClient with non-Closer rw: pipeline must be nil")
	}
	if ci.conn == nil {
	t.Error("NewClient with non-Closer rw: conn must be set")
	}

	// The serialized path correctly handles a request/response cycle.
	keys, err := ag.List()
	if err != nil {
	t.Fatalf("List: %v", err)
	}
	if len(keys) != 1 {
	t.Errorf("List: got %d keys, want 1", len(keys))
	}
	}

	func TestSerialFallbackParallelSigns(t *testing.T) {
	c1, c2, err := netPipe()
	if err != nil {
	t.Fatalf("netPipe: %v", err)
	}
	defer c1.Close()
	defer c2.Close()

	keyring := NewKeyring()
	if err := keyring.Add(AddedKey{PrivateKey: testPrivateKeys["rsa"]}); err != nil {
	t.Fatal(err)
	}
	go ServeAgent(keyring, c2)

	rw := readWriter{Reader: c1, Writer: c1}
	ag := NewClient(rw)

	if ci, ok := ag.(*client); !ok {
	t.Fatalf("NewClient: got %T, want *client", ag)
	} else if ci.pipeline != nil {
	t.Fatal("pipeline must be nil for non-Closer rw")
	}

	pubKey := testPublicKeys["rsa"]

	const N = 64
	var wg sync.WaitGroup
	errCh := make(chan error, N)
	for i := range N {
	wg.Go(func() {
	data := []byte{byte(i), byte(i >> 8)}
	sig, err := ag.Sign(pubKey, data)
	if err != nil {
	errCh <- err
	return
	}
	if err := pubKey.Verify(data, sig); err != nil {
	errCh <- err
	}
	})
	}
	wg.Wait()
	close(errCh)
	for err := range errCh {
	t.Error(err)
	}
	}

	// writeFailConn wraps a net.Conn and forces Write to return an error when
	// the fail flag is set. Close delegates to the underlying conn so the test
	// can observe whether the pipeline closed it.
	type writeFailConn struct {
	net.Conn
	fail atomic.Bool
	}

	func (c *writeFailConn) Write(p []byte) (int, error) {
	if c.fail.Load() {
	return 0, errors.New("synthetic write failure")
	}
	return c.Conn.Write(p)
	}

	func TestPipelineWriteErrorClosesConn(t *testing.T) {
	c1, c2, err := netPipe()
	if err != nil {
	t.Fatalf("netPipe: %v", err)
	}
	defer c2.Close()

	// Peer that reads and discards, exiting when the conn closes.
	serverExit := make(chan struct{})
	go func() {
	defer close(serverExit)
	io.Copy(io.Discard, c2)
	}()

	wrapper := &writeFailConn{Conn: c1}
	wrapper.fail.Store(true)

	client := NewClient(wrapper)

	errCh := make(chan error, 1)
	go func() { errCh <- client.RemoveAll() }()

	select {
	case err := <-errCh:
	if err == nil {
	t.Fatal("expected error on forced write failure, got nil")
	}
	// The caller whose Write failed must receive the actual
	// write error, not the reader's "closed connection" error
	// from the post-shutdown drain.
	if !strings.Contains(err.Error(), "synthetic write failure") {
	t.Errorf("got %q, want error containing %q", err, "synthetic write failure")
	}
	case <-time.After(2 * time.Second):
	t.Fatal("RemoveAll hung after write failure")
	}

	// The pipeline must have closed the underlying conn, so the peer
	// goroutine returns from its Read.
	select {
	case <-serverExit:
	case <-time.After(2 * time.Second):
	t.Fatal("connection was not closed after write failure")
	}

	// After shutdown, subsequent calls must also fail promptly.
	errCh2 := make(chan error, 1)
	go func() { errCh2 <- client.RemoveAll() }()
	select {
	case err := <-errCh2:
	if err == nil {
	t.Error("expected error on post-shutdown call, got nil")
	}
	case <-time.After(2 * time.Second):
	t.Fatal("post-shutdown call hung")
	}
	}

	func TestPipelineResponseTooLarge(t *testing.T) {
	c1, c2, err := netPipe()
	if err != nil {
	t.Fatalf("netPipe: %v", err)
	}
	defer c1.Close()
	defer c2.Close()

	// Server: read the request, then reply with a size header that
	// exceeds maxAgentResponseBytes. The size check fires before the
	// body is read, so no body needs to follow.
	serverDone := make(chan struct{})
	go func() {
	defer close(serverDone)
	var sizeBuf [4]byte
	if _, err := io.ReadFull(c2, sizeBuf[:]); err != nil {
	return
	}
	size := binary.BigEndian.Uint32(sizeBuf[:])
	if _, err := io.CopyN(io.Discard, c2, int64(size)); err != nil {
	return
	}
	var out [4]byte
	binary.BigEndian.PutUint32(out[:], maxAgentResponseBytes+1)
	c2.Write(out[:])
	}()

	client := NewClient(c1)

	errCh := make(chan error, 1)
	go func() { errCh <- client.RemoveAll() }()
	select {
	case err := <-errCh:
	if err == nil {
	t.Fatal("expected error, got nil")
	}
	if !strings.Contains(err.Error(), "response too large") {
	t.Errorf("got %q, want error containing %q", err, "response too large")
	}
	case <-time.After(2 * time.Second):
	t.Fatal("call hung")
	}

	// Subsequent calls must also fail: the pipeline is shut down.
	errCh2 := make(chan error, 1)
	go func() { errCh2 <- client.RemoveAll() }()
	select {
	case err := <-errCh2:
	if err == nil {
	t.Error("expected error on post-shutdown call, got nil")
	}
	case <-time.After(2 * time.Second):
	t.Fatal("post-shutdown call hung")
	}

	<-serverDone
	}

	func TestPipelineShutdownWithBlockedWriter(t *testing.T) {
	c1, c2, err := netPipe()
	if err != nil {
	t.Fatalf("netPipe: %v", err)
	}
	defer c1.Close()

	// Server drains requests off the wire so writes never block on
	// backpressure from the transport, but never replies. The reader
	// goroutine therefore stays parked on Read, never dequeues from
	// pending, and the queue fills up.
	var reqsRead int64
	serverExit := make(chan struct{})
	go func() {
	defer close(serverExit)
	for {
	var sizeBuf [4]byte
	if _, err := io.ReadFull(c2, sizeBuf[:]); err != nil {
	return
	}
	size := binary.BigEndian.Uint32(sizeBuf[:])
	if _, err := io.CopyN(io.Discard, c2, int64(size)); err != nil {
	return
	}
	atomic.AddInt64(&reqsRead, 1)
	}
	}()

	client := NewClient(c1)

	const total = pipelineMaxInFlight + 4
	var wg sync.WaitGroup
	errs := make(chan error, total)
	for range total {
	wg.Go(func() {
	errs <- client.RemoveAll()
	})
	}

	// Wait for the queue to fill: once pipelineMaxInFlight requests
	// are on the wire, pending is at capacity and the +1 writer is
	// blocked in the second-select arm; the rest are queued behind
	// writeMu.
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
	if atomic.LoadInt64(&reqsRead) >= pipelineMaxInFlight {
	break
	}
	time.Sleep(5 * time.Millisecond)
	}
	if got := atomic.LoadInt64(&reqsRead); got < pipelineMaxInFlight {
	t.Fatalf("requests on wire: got %d, want >= %d", got, pipelineMaxInFlight)
	}
	// Give the +1 writer a moment to actually park in the second
	// select before we trigger shutdown.
	time.Sleep(50 * time.Millisecond)

	// Trigger shutdown: closing the server side makes the client's
	// reader hit EOF, which calls shutdown(...).
	c2.Close()

	doneCh := make(chan struct{})
	go func() {
	wg.Wait()
	close(doneCh)
	}()
	select {
	case <-doneCh:
	case <-time.After(5 * time.Second):
	t.Fatal("blocked callers did not complete after shutdown")
	}

	close(errs)
	for err := range errs {
	if err == nil {
	t.Error("expected error after shutdown, got nil")
	}
	}

	<-serverExit
	}

	// multiDeviceAgent simulates an ssh-agent backed by numDevices signing
	// devices, each of which takes latency to process a request. Incoming
	// requests are round-robin dispatched to devices, and responses are
	// serialized back to the wire in request order as the protocol requires.
	//
	// The simulation does not emit a semantically meaningful reply; it just
	// writes a 1-byte SSH_AGENT_SUCCESS after the simulated latency has
	// elapsed. This is enough to exercise callRaw/simpleCall (RemoveAll).
	func multiDeviceAgent(t testing.TB, conn io.ReadWriteCloser, numDevices int, latency time.Duration) (done <-chan struct{}) {
	t.Helper()

	type job struct {
	ready chan struct{}
	}

	devices := make([]chan job, numDevices)
	var devicesWG sync.WaitGroup
	for i := range numDevices {
	devices[i] = make(chan job, 1024)
	ch := devices[i]
	devicesWG.Go(func() {
	for j := range ch {
	time.Sleep(latency)
	close(j.ready)
	}
	})
	}

	// Order-preserving responder: responses are written in the same
	// order the requests were read, regardless of which device finished
	// first.
	orderCh := make(chan chan struct{}, 1024)
	responderDone := make(chan struct{})
	go func() {
	defer close(responderDone)
	for ready := range orderCh {
	<-ready
	var out [5]byte
	binary.BigEndian.PutUint32(out[:4], 1)
	out[4] = agentSuccess
	if _, err := conn.Write(out[:]); err != nil {
	return
	}
	}
	}()

	finished := make(chan struct{})
	go func() {
	defer close(finished)
	defer func() {
	close(orderCh)
	<-responderDone
	for _, d := range devices {
	close(d)
	}
	devicesWG.Wait()
	}()
	next := 0
	for {
	var sizeBuf [4]byte
	if _, err := io.ReadFull(conn, sizeBuf[:]); err != nil {
	return
	}
	size := binary.BigEndian.Uint32(sizeBuf[:])
	if _, err := io.CopyN(io.Discard, conn, int64(size)); err != nil {
	return
	}
	ready := make(chan struct{})
	devices[next%numDevices] <- job{ready: ready}
	next++
	orderCh <- ready
	}
	}()

	return finished
	}

	// BenchmarkPipelineMultiDevice measures throughput of the pipelined
	// agent client against a simulated multi-device agent, varying the
	// number of backend devices.
	//
	// Setup: each simulated signing device takes 2ms per request. The
	// protocol serializes responses in request order, but pipelining lets
	// the client keep multiple devices busy concurrently; throughput
	// should scale roughly linearly with numDevices up to the pipeline
	// depth (pipelineMaxInFlight).
	//
	// Run with: go test -benchmem -bench BenchmarkPipelineMultiDevice -benchtime=2s ./ssh/agent/
	func BenchmarkPipelineMultiDevice(b *testing.B) {
	const (
	deviceLatency = 2 * time.Millisecond
	parallelism = 32
	)

	for _, numDevices := range []int{1, 2, 4, 8, 16} {
	b.Run(fmt.Sprintf("Devices=%d", numDevices), func(b *testing.B) {
	c1, c2, err := netPipe()
	if err != nil {
	b.Fatalf("netPipe: %v", err)
	}
	serverDone := multiDeviceAgent(b, c2, numDevices, deviceLatency)

	client := NewClient(c1)

	b.ResetTimer()
	b.SetParallelism(parallelism)
	b.RunParallel(func(pb *testing.PB) {
	for pb.Next() {
	if err := client.RemoveAll(); err != nil {
	b.Fatal(err)
	}
	}
	})
	b.StopTimer()

	c1.Close()
	c2.Close()
	<-serverDone
	})
	}
	}