eiac/crypto
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

ssh/agent: support parallel signing using request pipelining

Browse Source 
Make NewClient automatically pipeline concurrent requests over its
connection when the supplied io.ReadWriter also implements io.Closer.
In this mode the client writes requests to the wire as soon as the
write path is available and dispatches responses back to callers in
FIFO order via a dedicated reader goroutine, instead of fully
serializing each call. Up to 32 requests may be in flight on a single
connection.

This lets an agent that load-balances signing across multiple backend
devices keep several of them busy concurrently without forcing callers
to maintain their own connection pool. Aggregate throughput scales
roughly linearly with the number of backend devices up to the
in-flight cap; the protocol's in-order response requirement means
slow requests still delay subsequent ones on the same connection,
which is unchanged.

The pipelined path requires io.Closer because, on a Write error, the
background reader goroutine must be unblocked by closing the
underlying connection; otherwise it would remain parked forever
waiting for a response that will never arrive, leaking the goroutine
and desynchronising the FIFO routing of responses for subsequent
successful writes. When the supplied transport does not implement
io.Closer, NewClient falls back to the previous fully-serialized
behavior: a single in-flight call at a time, with no background
goroutine.

Fixes golang/go#78473

Change-Id: Icf14e5e8ca897506d68fb32c14fc72c774cd97b2
Reviewed-on: https://go-review.googlesource.com/c/crypto/+/768483
Reviewed-by: Cherry Mui <cherryyz@google.com>
Reviewed-by: Dmitri Shuralyov <dmitshur@google.com>
Reviewed-by: Filippo Valsorda <filippo@golang.org>
LUCI-TryBot-Result: golang-scoped@luci-project-accounts.iam.gserviceaccount.com <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
This commit is contained in:
Nicola Murino
2026-05-21 19:39:50 -07:00
parent d2fe592df1
commit e5306b2f3e
2 changed files with 863 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ssh/agent/client.go
+214 -3
View File
@@ -26,6 +26,7 @@ import (
"io"
"math/big"
"sync"
"sync/atomic"
"golang.org/x/crypto/ssh"
)
@@ -307,17 +308,50 @@ func parseKey(in []byte) (out *Key, rest []byte, err error) {
}, record.Rest, nil
}
// pipelineMaxInFlight is the maximum number of outstanding requests the
// client will pipeline to the agent before applying backpressure.
const pipelineMaxInFlight = 32
// client is a client for an ssh-agent process.
//
// Exactly one of pipeline / (mu, conn) is set, chosen by NewClient
// based on whether the underlying transport implements io.Closer.
type client struct {
// conn is typically a *net.UnixConn
// pipeline, if non-nil, dispatches requests over a pipelined
// connection: requests are written as soon as the wire is
// available and responses are routed back to per-call reply
// channels in FIFO order by a background reader goroutine.
pipeline *pipeline
// mu and conn are used in fully-serialized mode, when the
// transport does not implement io.Closer. Each call takes mu,
// writes its request, reads the matching response, and releases
// mu before returning. There is no background goroutine.
mu sync.Mutex
conn io.ReadWriter
// mu is used to prevent concurrent access to the agent
mu sync.Mutex
}
// NewClient returns an Agent that talks to an ssh-agent process over
// the given connection.
//
// If rw also implements io.Closer (like *net.UnixConn and ssh.Channel
// do), the returned client pipelines concurrent requests over the
// connection: callers can issue Sign and other operations from
// multiple goroutines and they will be written to the agent as soon
// as the wire is available, rather than waiting for the previous
// responses. The ssh-agent protocol still requires responses to be
// returned in request order, so a slow request delays subsequent
// responses on the same connection (head-of-line blocking).
//
// Pipelining requires io.Closer because, on a Write error, the
// background reader goroutine must be unblocked by closing the
// underlying connection. When rw does not implement io.Closer
// this is not possible, so NewClient falls back to fully
// serializing each request: a single in-flight call at a time.
func NewClient(rw io.ReadWriter) ExtendedAgent {
if rwc, ok := rw.(io.ReadWriteCloser); ok {
return &client{pipeline: newPipeline(rwc)}
}
return &client{conn: rw}
}
@@ -340,6 +374,16 @@ func (c *client) call(req []byte) (reply interface{}, err error) {
// bytes of the response are returned; no unmarshalling is
// performed on the response.
func (c *client) callRaw(req []byte) (reply []byte, err error) {
if c.pipeline != nil {
return c.pipeline.call(req)
}
return c.serialCall(req)
}
// serialCall implements the fully-serialized request/response path
// used when the transport is not an io.Closer. It writes req under mu
// and reads the matching response before returning.
func (c *client) serialCall(req []byte) (reply []byte, err error) {
c.mu.Lock()
defer c.mu.Unlock()
@@ -861,3 +905,170 @@ func (c *client) Extension(extensionType string, contents []byte) ([]byte, error
return buf, nil
}
// pipelineResult carries either a raw agent reply or an error back to a
// caller waiting on the response channel.
type pipelineResult struct {
reply []byte
err error
}
// pipeline implements request pipelining over a single agent connection.
//
// Writers serialize on writeMu to both register a reply channel in the
// pending FIFO queue and write the request bytes on the wire; the two
// must be atomic so the queue order matches the wire order. A single
// reader goroutine decodes responses from the connection and dispatches
// each one to the channel at the head of the queue.
//
// pending is a chan-of-chan acting as a FIFO queue with a fixed
// capacity of pipelineMaxInFlight. The outer channel provides ordering
// (reads happen in send order) and natural backpressure (a full queue
// blocks new writers). Each inner channel is buffered with capacity
// one and is sent to exactly once: either by the reader goroutine
// with the agent reply, or by shutdown with the terminal error during
// drain. The cap-one buffer makes the producer's send non-blocking,
// so the reader and shutdown never have to wait for the caller to be
// scheduled on the receive.
//
// When the reader goroutine exits (on read error or protocol
// violation), it closes exitCh to wake any writer blocked on the
// pending queue, then serializes with any in-flight writer to close
// the pending channel, and finally drains the remaining entries
// delivering the terminal error to each waiting caller. The
// pipeline relies on conn implementing io.Closer so a writer that
// hits a Write error can close the connection to unblock the reader
// goroutine; NewClient is responsible for only constructing a
// pipeline when this guarantee holds.
type pipeline struct {
conn io.ReadWriteCloser
writeMu sync.Mutex
// pending is the FIFO queue of reply channels with capacity
// pipelineMaxInFlight. See type-level documentation.
pending chan chan pipelineResult
exitCh chan struct{}
// err carries the terminal error to callers blocked on a closed
// pipeline. It is stored exactly once by the reader goroutine
// before exitCh is closed; every read happens after observing
// exitCh closed, so the load synchronises through the close and
// is guaranteed to return the stored value (never nil).
err atomic.Pointer[error]
}
func newPipeline(conn io.ReadWriteCloser) *pipeline {
p := &pipeline{
conn: conn,
pending: make(chan chan pipelineResult, pipelineMaxInFlight),
exitCh: make(chan struct{}),
}
go p.readLoop()
return p
}
// readLoop decodes responses from conn and dispatches them in FIFO order
// to reply channels in pending. On any failure it invokes shutdown.
func (p *pipeline) readLoop() {
var finalErr error
for {
var sizeBuf [4]byte
if _, err := io.ReadFull(p.conn, sizeBuf[:]); err != nil {
finalErr = err
break
}
respSize := binary.BigEndian.Uint32(sizeBuf[:])
if respSize > maxAgentResponseBytes {
finalErr = errors.New("response too large")
break
}
buf := make([]byte, respSize)
if _, err := io.ReadFull(p.conn, buf); err != nil {
finalErr = err
break
}
// Successful writes always enqueue before sending bytes, so
// pending has a waiting channel for this response.
ch := <-p.pending
// The reply channel is buffered with capacity 1 and is only
// ever written to once, so this send cannot block.
ch <- pipelineResult{reply: buf}
}
p.shutdown(clientErr(finalErr))
}
// shutdown is called exactly once, from readLoop, when the reader is
// terminating. It unblocks pending writers and fails all in-flight
// requests with finalErr.
func (p *pipeline) shutdown(finalErr error) {
// Publish the terminal error before closing exitCh so any
// writer that subsequently observes exitCh closed sees err.
p.err.Store(&finalErr)
// Wake any writer blocked waiting for a slot in the pending queue.
close(p.exitCh)
// Wait for any writer currently inside its critical section to
// complete. After this lock, no new writer can reach the send on
// pending: they will observe exitCh closed in the select and bail
// out before attempting the send.
p.writeMu.Lock()
close(p.pending)
p.writeMu.Unlock()
// Drain entries that were enqueued but never answered, delivering
// the terminal error to their waiting callers. The reply channels
// are buffered (cap 1) and written to exactly once, so these sends
// cannot block.
for ch := range p.pending {
ch <- pipelineResult{err: finalErr}
}
}
// call sends req to the agent and returns the matching raw response.
func (p *pipeline) call(req []byte) ([]byte, error) {
replyCh := make(chan pipelineResult, 1)
p.writeMu.Lock()
// Priority check: if the reader has already finished shutdown,
// pending is closed and sending to it would panic. Bail out now.
// Once we pass this check while holding writeMu, shutdown cannot
// complete close(pending) until we release writeMu, so the send
// below is safe against concurrent closure.
select {
case <-p.exitCh:
p.writeMu.Unlock()
return nil, *p.err.Load()
default:
}
// Enqueue the reply channel before writing the request, so FIFO
// order on the wire matches FIFO order in the pending queue. The
// exitCh arm handles the case where the reader errors while we
// block on a full queue.
select {
case p.pending <- replyCh:
case <-p.exitCh:
p.writeMu.Unlock()
return nil, *p.err.Load()
}
msg := make([]byte, 4+len(req))
binary.BigEndian.PutUint32(msg, uint32(len(req)))
copy(msg[4:], req)
_, werr := p.conn.Write(msg)
p.writeMu.Unlock()
if werr != nil {
// The connection is in an undefined state. Close it so the
// reader unblocks promptly and triggers shutdown for every
// other in-flight caller. NewClient guarantees conn is a
// real io.Closer when the pipeline is in use.
p.conn.Close()
return nil, clientErr(werr)
}
res := <-replyCh
return res.reply, res.err
}