// Copyright 2012 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 implements the ssh-agent protocol, and provides both
// a client and a server. The client can talk to a standard ssh-agent
// that uses UNIX sockets, and one could implement an alternative
// ssh-agent process using the sample server.
//
// References:
//
//	[PROTOCOL.agent]: https://tools.ietf.org/html/draft-miller-ssh-agent-00
package agent

import (
	"bytes"
	"crypto/dsa"
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	"crypto/rsa"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"math/big"
	"sync"
	"sync/atomic"

	"golang.org/x/crypto/ssh"
)

// SignatureFlags represent additional flags that can be passed to the signature
// requests an defined in [PROTOCOL.agent] section 4.5.1.
type SignatureFlags uint32

// SignatureFlag values as defined in [PROTOCOL.agent] section 5.3.
const (
	SignatureFlagReserved SignatureFlags = 1 << iota
	SignatureFlagRsaSha256
	SignatureFlagRsaSha512
)

// Agent represents the capabilities of an ssh-agent.
type Agent interface {
	// List returns the identities known to the agent.
	List() ([]*Key, error)

	// Sign has the agent sign the data using a protocol 2 key as defined
	// in [PROTOCOL.agent] section 2.6.2.
	Sign(key ssh.PublicKey, data []byte) (*ssh.Signature, error)

	// Add adds a private key to the agent.
	Add(key AddedKey) error

	// Remove removes all identities with the given public key.
	Remove(key ssh.PublicKey) error

	// RemoveAll removes all identities.
	RemoveAll() error

	// Lock locks the agent. Sign and Remove will fail, and List will empty an empty list.
	Lock(passphrase []byte) error

	// Unlock undoes the effect of Lock
	Unlock(passphrase []byte) error

	// Signers returns signers for all the known keys.
	Signers() ([]ssh.Signer, error)
}

type ExtendedAgent interface {
	Agent

	// SignWithFlags signs like Sign, but allows for additional flags to be sent/received
	SignWithFlags(key ssh.PublicKey, data []byte, flags SignatureFlags) (*ssh.Signature, error)

	// Extension processes a custom extension request. Standard-compliant agents are not
	// required to support any extensions, but this method allows agents to implement
	// vendor-specific methods or add experimental features. See [PROTOCOL.agent] section 4.7.
	// If agent extensions are unsupported entirely this method MUST return an
	// ErrExtensionUnsupported error. Similarly, if just the specific extensionType in
	// the request is unsupported by the agent then ErrExtensionUnsupported MUST be
	// returned.
	//
	// In the case of success, since [PROTOCOL.agent] section 4.7 specifies that the contents
	// of the response are unspecified (including the type of the message), the complete
	// response will be returned as a []byte slice, including the "type" byte of the message.
	Extension(extensionType string, contents []byte) ([]byte, error)
}

// ConstraintExtension describes an optional constraint defined by users.
type ConstraintExtension struct {
	// ExtensionName consist of a UTF-8 string suffixed by the
	// implementation domain following the naming scheme defined
	// in Section 4.2 of RFC 4251, e.g.  "foo@example.com".
	ExtensionName string
	// ExtensionDetails contains the actual content of the extended
	// constraint.
	ExtensionDetails []byte
}

// AddedKey describes an SSH key to be added to an Agent.
type AddedKey struct {
	// PrivateKey must be a *rsa.PrivateKey, *dsa.PrivateKey,
	// ed25519.PrivateKey or *ecdsa.PrivateKey, which will be inserted into the
	// agent.
	PrivateKey interface{}
	// Certificate, if not nil, is communicated to the agent and will be
	// stored with the key.
	Certificate *ssh.Certificate
	// Comment is an optional, free-form string.
	Comment string
	// LifetimeSecs, if not zero, is the number of seconds that the
	// agent will store the key for.
	LifetimeSecs uint32
	// ConfirmBeforeUse, if true, requests that the agent confirm with the
	// user before each use of this key.
	ConfirmBeforeUse bool
	// ConstraintExtensions are the experimental or private-use constraints
	// defined by users.
	ConstraintExtensions []ConstraintExtension
}

// See [PROTOCOL.agent], section 3.
const (
	agentRequestV1Identities   = 1
	agentRemoveAllV1Identities = 9

	// 3.2 Requests from client to agent for protocol 2 key operations
	agentAddIdentity         = 17
	agentRemoveIdentity      = 18
	agentRemoveAllIdentities = 19
	agentAddIDConstrained    = 25

	// 3.3 Key-type independent requests from client to agent
	agentAddSmartcardKey            = 20
	agentRemoveSmartcardKey         = 21
	agentLock                       = 22
	agentUnlock                     = 23
	agentAddSmartcardKeyConstrained = 26

	// 3.7 Key constraint identifiers
	agentConstrainLifetime = 1
	agentConstrainConfirm  = 2
	// Constraint extension identifier up to version 2 of the protocol. A
	// backward incompatible change will be required if we want to add support
	// for SSH_AGENT_CONSTRAIN_MAXSIGN which uses the same ID.
	agentConstrainExtensionV00 = 3
	// Constraint extension identifier in version 3 and later of the protocol.
	agentConstrainExtension = 255
)

// maxAgentResponseBytes is the maximum agent reply size that is accepted. This
// is a sanity check, not a limit in the spec.
const maxAgentResponseBytes = 16 << 20

// Agent messages:
// These structures mirror the wire format of the corresponding ssh agent
// messages found in [PROTOCOL.agent].

// 3.4 Generic replies from agent to client
const agentFailure = 5

type failureAgentMsg struct{}

const agentSuccess = 6

type successAgentMsg struct{}

// See [PROTOCOL.agent], section 2.5.2.
const agentRequestIdentities = 11

type requestIdentitiesAgentMsg struct{}

// See [PROTOCOL.agent], section 2.5.2.
const agentIdentitiesAnswer = 12

type identitiesAnswerAgentMsg struct {
	NumKeys uint32 `sshtype:"12"`
	Keys    []byte `ssh:"rest"`
}

// See [PROTOCOL.agent], section 2.6.2.
const agentSignRequest = 13

type signRequestAgentMsg struct {
	KeyBlob []byte `sshtype:"13"`
	Data    []byte
	Flags   uint32
}

// See [PROTOCOL.agent], section 2.6.2.

// 3.6 Replies from agent to client for protocol 2 key operations
const agentSignResponse = 14

type signResponseAgentMsg struct {
	SigBlob []byte `sshtype:"14"`
}

type publicKey struct {
	Format string
	Rest   []byte `ssh:"rest"`
}

// 3.7 Key constraint identifiers
type constrainLifetimeAgentMsg struct {
	LifetimeSecs uint32 `sshtype:"1"`
}

type constrainExtensionAgentMsg struct {
	ExtensionName    string `sshtype:"255|3"`
	ExtensionDetails []byte

	// Rest is a field used for parsing, not part of message
	Rest []byte `ssh:"rest"`
}

// See [PROTOCOL.agent], section 4.7
const agentExtension = 27
const agentExtensionFailure = 28

// ErrExtensionUnsupported indicates that an extension defined in
// [PROTOCOL.agent] section 4.7 is unsupported by the agent. Specifically this
// error indicates that the agent returned a standard SSH_AGENT_FAILURE message
// as the result of a SSH_AGENTC_EXTENSION request. Note that the protocol
// specification (and therefore this error) does not distinguish between a
// specific extension being unsupported and extensions being unsupported entirely.
var ErrExtensionUnsupported = errors.New("agent: extension unsupported")

type extensionAgentMsg struct {
	ExtensionType string `sshtype:"27"`
	// NOTE: this matches OpenSSH's PROTOCOL.agent, not the IETF draft [PROTOCOL.agent],
	// so that it matches what OpenSSH actually implements in the wild.
	Contents []byte `ssh:"rest"`
}

// Key represents a protocol 2 public key as defined in
// [PROTOCOL.agent], section 2.5.2.
type Key struct {
	Format  string
	Blob    []byte
	Comment string
}

func clientErr(err error) error {
	return fmt.Errorf("agent: client error: %v", err)
}

// String returns the storage form of an agent key with the format, base64
// encoded serialized key, and the comment if it is not empty.
func (k *Key) String() string {
	s := string(k.Format) + " " + base64.StdEncoding.EncodeToString(k.Blob)

	if k.Comment != "" {
		s += " " + k.Comment
	}

	return s
}

// Type returns the public key type.
func (k *Key) Type() string {
	return k.Format
}

// Marshal returns key blob to satisfy the ssh.PublicKey interface.
func (k *Key) Marshal() []byte {
	return k.Blob
}

// Verify satisfies the ssh.PublicKey interface.
func (k *Key) Verify(data []byte, sig *ssh.Signature) error {
	pubKey, err := ssh.ParsePublicKey(k.Blob)
	if err != nil {
		return fmt.Errorf("agent: bad public key: %v", err)
	}
	return pubKey.Verify(data, sig)
}

type wireKey struct {
	Format string
	Rest   []byte `ssh:"rest"`
}

func parseKey(in []byte) (out *Key, rest []byte, err error) {
	var record struct {
		Blob    []byte
		Comment string
		Rest    []byte `ssh:"rest"`
	}

	if err := ssh.Unmarshal(in, &record); err != nil {
		return nil, nil, err
	}

	var wk wireKey
	if err := ssh.Unmarshal(record.Blob, &wk); err != nil {
		return nil, nil, err
	}

	return &Key{
		Format:  wk.Format,
		Blob:    record.Blob,
		Comment: record.Comment,
	}, 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 {
	// 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
}

// 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}
}

// call sends an RPC to the agent. On success, the reply is
// unmarshaled into reply and replyType is set to the first byte of
// the reply, which contains the type of the message.
func (c *client) call(req []byte) (reply interface{}, err error) {
	buf, err := c.callRaw(req)
	if err != nil {
		return nil, err
	}
	reply, err = unmarshal(buf)
	if err != nil {
		return nil, clientErr(err)
	}
	return reply, nil
}

// callRaw sends an RPC to the agent. On success, the raw
// 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()

	msg := make([]byte, 4+len(req))
	binary.BigEndian.PutUint32(msg, uint32(len(req)))
	copy(msg[4:], req)
	if _, err = c.conn.Write(msg); err != nil {
		return nil, clientErr(err)
	}

	var respSizeBuf [4]byte
	if _, err = io.ReadFull(c.conn, respSizeBuf[:]); err != nil {
		return nil, clientErr(err)
	}
	respSize := binary.BigEndian.Uint32(respSizeBuf[:])
	if respSize > maxAgentResponseBytes {
		return nil, clientErr(errors.New("response too large"))
	}

	buf := make([]byte, respSize)
	if _, err = io.ReadFull(c.conn, buf); err != nil {
		return nil, clientErr(err)
	}
	return buf, nil
}

func (c *client) simpleCall(req []byte) error {
	resp, err := c.call(req)
	if err != nil {
		return err
	}
	if _, ok := resp.(*successAgentMsg); ok {
		return nil
	}
	return errors.New("agent: failure")
}

func (c *client) RemoveAll() error {
	return c.simpleCall([]byte{agentRemoveAllIdentities})
}

func (c *client) Remove(key ssh.PublicKey) error {
	req := ssh.Marshal(&agentRemoveIdentityMsg{
		KeyBlob: key.Marshal(),
	})
	return c.simpleCall(req)
}

func (c *client) Lock(passphrase []byte) error {
	req := ssh.Marshal(&agentLockMsg{
		Passphrase: passphrase,
	})
	return c.simpleCall(req)
}

func (c *client) Unlock(passphrase []byte) error {
	req := ssh.Marshal(&agentUnlockMsg{
		Passphrase: passphrase,
	})
	return c.simpleCall(req)
}

// List returns the identities known to the agent.
func (c *client) List() ([]*Key, error) {
	// see [PROTOCOL.agent] section 2.5.2.
	req := []byte{agentRequestIdentities}

	msg, err := c.call(req)
	if err != nil {
		return nil, err
	}

	switch msg := msg.(type) {
	case *identitiesAnswerAgentMsg:
		if msg.NumKeys > maxAgentResponseBytes/8 {
			return nil, errors.New("agent: too many keys in agent reply")
		}
		keys := make([]*Key, msg.NumKeys)
		data := msg.Keys
		for i := uint32(0); i < msg.NumKeys; i++ {
			var key *Key
			var err error
			if key, data, err = parseKey(data); err != nil {
				return nil, err
			}
			keys[i] = key
		}
		return keys, nil
	case *failureAgentMsg:
		return nil, errors.New("agent: failed to list keys")
	default:
		return nil, fmt.Errorf("agent: failed to list keys, unexpected message type %T", msg)
	}
}

// Sign has the agent sign the data using a protocol 2 key as defined
// in [PROTOCOL.agent] section 2.6.2.
func (c *client) Sign(key ssh.PublicKey, data []byte) (*ssh.Signature, error) {
	return c.SignWithFlags(key, data, 0)
}

func (c *client) SignWithFlags(key ssh.PublicKey, data []byte, flags SignatureFlags) (*ssh.Signature, error) {
	req := ssh.Marshal(signRequestAgentMsg{
		KeyBlob: key.Marshal(),
		Data:    data,
		Flags:   uint32(flags),
	})

	msg, err := c.call(req)
	if err != nil {
		return nil, err
	}

	switch msg := msg.(type) {
	case *signResponseAgentMsg:
		var sig ssh.Signature
		if err := ssh.Unmarshal(msg.SigBlob, &sig); err != nil {
			return nil, err
		}

		return &sig, nil
	case *failureAgentMsg:
		return nil, errors.New("agent: failed to sign challenge")
	default:
		return nil, fmt.Errorf("agent: failed to sign challenge, unexpected message type %T", msg)
	}
}

// unmarshal parses an agent message in packet, returning the parsed
// form and the message type of packet.
func unmarshal(packet []byte) (interface{}, error) {
	if len(packet) < 1 {
		return nil, errors.New("agent: empty packet")
	}
	var msg interface{}
	switch packet[0] {
	case agentFailure:
		return new(failureAgentMsg), nil
	case agentSuccess:
		return new(successAgentMsg), nil
	case agentIdentitiesAnswer:
		msg = new(identitiesAnswerAgentMsg)
	case agentSignResponse:
		msg = new(signResponseAgentMsg)
	case agentV1IdentitiesAnswer:
		msg = new(agentV1IdentityMsg)
	default:
		return nil, fmt.Errorf("agent: unknown type tag %d", packet[0])
	}
	if err := ssh.Unmarshal(packet, msg); err != nil {
		return nil, err
	}
	return msg, nil
}

type rsaKeyMsg struct {
	Type        string `sshtype:"17|25"`
	N           *big.Int
	E           *big.Int
	D           *big.Int
	Iqmp        *big.Int // IQMP = Inverse Q Mod P
	P           *big.Int
	Q           *big.Int
	Comments    string
	Constraints []byte `ssh:"rest"`
}

type dsaKeyMsg struct {
	Type        string `sshtype:"17|25"`
	P           *big.Int
	Q           *big.Int
	G           *big.Int
	Y           *big.Int
	X           *big.Int
	Comments    string
	Constraints []byte `ssh:"rest"`
}

type ecdsaKeyMsg struct {
	Type        string `sshtype:"17|25"`
	Curve       string
	KeyBytes    []byte
	D           *big.Int
	Comments    string
	Constraints []byte `ssh:"rest"`
}

type ed25519KeyMsg struct {
	Type        string `sshtype:"17|25"`
	Pub         []byte
	Priv        []byte
	Comments    string
	Constraints []byte `ssh:"rest"`
}

// Insert adds a private key to the agent.
func (c *client) insertKey(s interface{}, comment string, constraints []byte) error {
	var req []byte
	switch k := s.(type) {
	case *rsa.PrivateKey:
		if len(k.Primes) != 2 {
			return fmt.Errorf("agent: unsupported RSA key with %d primes", len(k.Primes))
		}
		k.Precompute()
		req = ssh.Marshal(rsaKeyMsg{
			Type:        ssh.KeyAlgoRSA,
			N:           k.N,
			E:           big.NewInt(int64(k.E)),
			D:           k.D,
			Iqmp:        k.Precomputed.Qinv,
			P:           k.Primes[0],
			Q:           k.Primes[1],
			Comments:    comment,
			Constraints: constraints,
		})
	case *dsa.PrivateKey:
		req = ssh.Marshal(dsaKeyMsg{
			Type:        ssh.InsecureKeyAlgoDSA,
			P:           k.P,
			Q:           k.Q,
			G:           k.G,
			Y:           k.Y,
			X:           k.X,
			Comments:    comment,
			Constraints: constraints,
		})
	case *ecdsa.PrivateKey:
		nistID := fmt.Sprintf("nistp%d", k.Params().BitSize)
		req = ssh.Marshal(ecdsaKeyMsg{
			Type:        "ecdsa-sha2-" + nistID,
			Curve:       nistID,
			KeyBytes:    elliptic.Marshal(k.Curve, k.X, k.Y),
			D:           k.D,
			Comments:    comment,
			Constraints: constraints,
		})
	case ed25519.PrivateKey:
		if len(k) != ed25519.PrivateKeySize {
			return fmt.Errorf("agent: bad ED25519 key size: %d", len(k))
		}
		req = ssh.Marshal(ed25519KeyMsg{
			Type:        ssh.KeyAlgoED25519,
			Pub:         []byte(k)[32:],
			Priv:        []byte(k),
			Comments:    comment,
			Constraints: constraints,
		})
	// This function originally supported only *ed25519.PrivateKey, however the
	// general idiom is to pass ed25519.PrivateKey by value, not by pointer.
	// We still support the pointer variant for backwards compatibility.
	case *ed25519.PrivateKey:
		if len(*k) != ed25519.PrivateKeySize {
			return fmt.Errorf("agent: bad ED25519 key size: %d", len(*k))
		}
		req = ssh.Marshal(ed25519KeyMsg{
			Type:        ssh.KeyAlgoED25519,
			Pub:         []byte(*k)[32:],
			Priv:        []byte(*k),
			Comments:    comment,
			Constraints: constraints,
		})
	default:
		return fmt.Errorf("agent: unsupported key type %T", s)
	}

	// if constraints are present then the message type needs to be changed.
	if len(constraints) != 0 {
		req[0] = agentAddIDConstrained
	}

	resp, err := c.call(req)
	if err != nil {
		return err
	}
	if _, ok := resp.(*successAgentMsg); ok {
		return nil
	}
	return errors.New("agent: failure")
}

type rsaCertMsg struct {
	Type        string `sshtype:"17|25"`
	CertBytes   []byte
	D           *big.Int
	Iqmp        *big.Int // IQMP = Inverse Q Mod P
	P           *big.Int
	Q           *big.Int
	Comments    string
	Constraints []byte `ssh:"rest"`
}

type dsaCertMsg struct {
	Type        string `sshtype:"17|25"`
	CertBytes   []byte
	X           *big.Int
	Comments    string
	Constraints []byte `ssh:"rest"`
}

type ecdsaCertMsg struct {
	Type        string `sshtype:"17|25"`
	CertBytes   []byte
	D           *big.Int
	Comments    string
	Constraints []byte `ssh:"rest"`
}

type ed25519CertMsg struct {
	Type        string `sshtype:"17|25"`
	CertBytes   []byte
	Pub         []byte
	Priv        []byte
	Comments    string
	Constraints []byte `ssh:"rest"`
}

// Add adds a private key to the agent. If a certificate is given,
// that certificate is added instead as public key.
func (c *client) Add(key AddedKey) error {
	var constraints []byte

	if secs := key.LifetimeSecs; secs != 0 {
		constraints = append(constraints, ssh.Marshal(constrainLifetimeAgentMsg{secs})...)
	}

	if key.ConfirmBeforeUse {
		constraints = append(constraints, agentConstrainConfirm)
	}

	for _, ext := range key.ConstraintExtensions {
		constraints = append(constraints, ssh.Marshal(constrainExtensionAgentMsg{
			ExtensionName:    ext.ExtensionName,
			ExtensionDetails: ext.ExtensionDetails,
		})...)
	}

	cert := key.Certificate
	if cert == nil {
		return c.insertKey(key.PrivateKey, key.Comment, constraints)
	}
	return c.insertCert(key.PrivateKey, cert, key.Comment, constraints)
}

func (c *client) insertCert(s interface{}, cert *ssh.Certificate, comment string, constraints []byte) error {
	var req []byte
	switch k := s.(type) {
	case *rsa.PrivateKey:
		if len(k.Primes) != 2 {
			return fmt.Errorf("agent: unsupported RSA key with %d primes", len(k.Primes))
		}
		k.Precompute()
		req = ssh.Marshal(rsaCertMsg{
			Type:        cert.Type(),
			CertBytes:   cert.Marshal(),
			D:           k.D,
			Iqmp:        k.Precomputed.Qinv,
			P:           k.Primes[0],
			Q:           k.Primes[1],
			Comments:    comment,
			Constraints: constraints,
		})
	case *dsa.PrivateKey:
		req = ssh.Marshal(dsaCertMsg{
			Type:        cert.Type(),
			CertBytes:   cert.Marshal(),
			X:           k.X,
			Comments:    comment,
			Constraints: constraints,
		})
	case *ecdsa.PrivateKey:
		req = ssh.Marshal(ecdsaCertMsg{
			Type:        cert.Type(),
			CertBytes:   cert.Marshal(),
			D:           k.D,
			Comments:    comment,
			Constraints: constraints,
		})
	case ed25519.PrivateKey:
		if len(k) != ed25519.PrivateKeySize {
			return fmt.Errorf("agent: bad ED25519 key size: %d", len(k))
		}
		req = ssh.Marshal(ed25519CertMsg{
			Type:        cert.Type(),
			CertBytes:   cert.Marshal(),
			Pub:         []byte(k)[32:],
			Priv:        []byte(k),
			Comments:    comment,
			Constraints: constraints,
		})
	// This function originally supported only *ed25519.PrivateKey, however the
	// general idiom is to pass ed25519.PrivateKey by value, not by pointer.
	// We still support the pointer variant for backwards compatibility.
	case *ed25519.PrivateKey:
		if len(*k) != ed25519.PrivateKeySize {
			return fmt.Errorf("agent: bad ED25519 key size: %d", len(*k))
		}
		req = ssh.Marshal(ed25519CertMsg{
			Type:        cert.Type(),
			CertBytes:   cert.Marshal(),
			Pub:         []byte(*k)[32:],
			Priv:        []byte(*k),
			Comments:    comment,
			Constraints: constraints,
		})
	default:
		return fmt.Errorf("agent: unsupported key type %T", s)
	}

	// if constraints are present then the message type needs to be changed.
	if len(constraints) != 0 {
		req[0] = agentAddIDConstrained
	}

	signer, err := ssh.NewSignerFromKey(s)
	if err != nil {
		return err
	}
	if !bytes.Equal(cert.Key.Marshal(), signer.PublicKey().Marshal()) {
		return errors.New("agent: signer and cert have different public key")
	}

	resp, err := c.call(req)
	if err != nil {
		return err
	}
	if _, ok := resp.(*successAgentMsg); ok {
		return nil
	}
	return errors.New("agent: failure")
}

// Signers provides a callback for client authentication.
func (c *client) Signers() ([]ssh.Signer, error) {
	keys, err := c.List()
	if err != nil {
		return nil, err
	}

	var result []ssh.Signer
	for _, k := range keys {
		result = append(result, &agentKeyringSigner{c, k})
	}
	return result, nil
}

type agentKeyringSigner struct {
	agent *client
	pub   ssh.PublicKey
}

func (s *agentKeyringSigner) PublicKey() ssh.PublicKey {
	return s.pub
}

func (s *agentKeyringSigner) Sign(rand io.Reader, data []byte) (*ssh.Signature, error) {
	// The agent has its own entropy source, so the rand argument is ignored.
	return s.agent.Sign(s.pub, data)
}

func (s *agentKeyringSigner) SignWithAlgorithm(rand io.Reader, data []byte, algorithm string) (*ssh.Signature, error) {
	if algorithm == "" || algorithm == underlyingAlgo(s.pub.Type()) {
		return s.Sign(rand, data)
	}

	var flags SignatureFlags
	switch algorithm {
	case ssh.KeyAlgoRSASHA256:
		flags = SignatureFlagRsaSha256
	case ssh.KeyAlgoRSASHA512:
		flags = SignatureFlagRsaSha512
	default:
		return nil, fmt.Errorf("agent: unsupported algorithm %q", algorithm)
	}

	return s.agent.SignWithFlags(s.pub, data, flags)
}

var _ ssh.AlgorithmSigner = &agentKeyringSigner{}

// certKeyAlgoNames is a mapping from known certificate algorithm names to the
// corresponding public key signature algorithm.
//
// This map must be kept in sync with the one in certs.go.
var certKeyAlgoNames = map[string]string{
	ssh.CertAlgoRSAv01:         ssh.KeyAlgoRSA,
	ssh.CertAlgoRSASHA256v01:   ssh.KeyAlgoRSASHA256,
	ssh.CertAlgoRSASHA512v01:   ssh.KeyAlgoRSASHA512,
	ssh.InsecureCertAlgoDSAv01: ssh.InsecureKeyAlgoDSA,
	ssh.CertAlgoECDSA256v01:    ssh.KeyAlgoECDSA256,
	ssh.CertAlgoECDSA384v01:    ssh.KeyAlgoECDSA384,
	ssh.CertAlgoECDSA521v01:    ssh.KeyAlgoECDSA521,
	ssh.CertAlgoSKECDSA256v01:  ssh.KeyAlgoSKECDSA256,
	ssh.CertAlgoED25519v01:     ssh.KeyAlgoED25519,
	ssh.CertAlgoSKED25519v01:   ssh.KeyAlgoSKED25519,
}

// underlyingAlgo returns the signature algorithm associated with algo (which is
// an advertised or negotiated public key or host key algorithm). These are
// usually the same, except for certificate algorithms.
func underlyingAlgo(algo string) string {
	if a, ok := certKeyAlgoNames[algo]; ok {
		return a
	}
	return algo
}

// Calls an extension method. It is up to the agent implementation as to whether or not
// any particular extension is supported and may always return an error. Because the
// type of the response is up to the implementation, this returns the bytes of the
// response and does not attempt any type of unmarshalling.
func (c *client) Extension(extensionType string, contents []byte) ([]byte, error) {
	req := ssh.Marshal(extensionAgentMsg{
		ExtensionType: extensionType,
		Contents:      contents,
	})
	buf, err := c.callRaw(req)
	if err != nil {
		return nil, err
	}
	if len(buf) == 0 {
		return nil, errors.New("agent: failure; empty response")
	}
	// [PROTOCOL.agent] section 4.7 indicates that an SSH_AGENT_FAILURE message
	// represents an agent that does not support the extension
	if buf[0] == agentFailure {
		return nil, ErrExtensionUnsupported
	}
	if buf[0] == agentExtensionFailure {
		return nil, errors.New("agent: generic extension failure")
	}

	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
}
