From c7cc1d646e482341f53262562ebca365d5909a97 Mon Sep 17 00:00:00 2001
From: Chris Sutcliff <chris@sutcliff.me>
Date: Fri, 23 Feb 2024 11:09:38 +0000
Subject: [PATCH] Implement ExtendedAgent

Allows clients to specify the algorithm used via the `SignWithFlags` function. Currently this is `rsa-sha2-256` or `rsa-sha2-512` which are the two allowable options in the spec but it is extendable in the future as required.

Aiming to reduce code duplication, Existing `Sign` function retains it's default behaviour by calling `SignWithFlags` with `SignatureFlagReserved`.

The other `ExtendedAgent` function `Extension` returns `ErrExtensionUnsupported` to all calls as required by the spec.
---
 agent/ssh/ssh.go | 29 ++++++++++++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

diff --git a/agent/ssh/ssh.go b/agent/ssh/ssh.go
index e0e1d65..b319552 100644
--- a/agent/ssh/ssh.go
+++ b/agent/ssh/ssh.go
@@ -81,6 +81,10 @@ func Eq(a, b ssh.PublicKey) bool {
 }
 
 func (vaultAgent vaultAgent) Sign(key ssh.PublicKey, data []byte) (*ssh.Signature, error) {
+	return vaultAgent.SignWithFlags(key, data, agent.SignatureFlagReserved)
+}
+
+func (vaultAgent vaultAgent) SignWithFlags(key ssh.PublicKey, data []byte, flags agent.SignatureFlags) (*ssh.Signature, error) {
 	log.Info("Sign Request for key: %s", ssh.FingerprintSHA256(key))
 	if vaultAgent.vault.Keyring.IsLocked() {
 		if !vaultAgent.unlockRequestAction() {
@@ -156,7 +160,30 @@ func (vaultAgent vaultAgent) Sign(key ssh.PublicKey, data []byte) (*ssh.Signatur
 	} else {
 		notify.Notify("Goldwarden", fmt.Sprintf("SSH Signing Request Approved for %s", sshKey.Name), "", 10*time.Second, func() {})
 	}
-	return signer.Sign(rand, data)
+
+	algo := ""
+
+	switch flags {
+	case agent.SignatureFlagRsaSha256:
+		algo = ssh.KeyAlgoRSASHA256
+	case agent.SignatureFlagRsaSha512:
+		algo = ssh.KeyAlgoRSASHA512
+	default:
+		return signer.Sign(rand, data)
+	}
+
+	log.Info("%s algorithm requested", algo)
+
+	algoSigner, err := ssh.NewSignerWithAlgorithms(signer.(ssh.AlgorithmSigner), []string{algo})
+	if err != nil {
+		return nil, err
+	}
+
+	return algoSigner.SignWithAlgorithm(rand, data, algo)
+}
+
+func (vaultAgent) Extension(extensionType string, contents []byte) ([]byte, error) {
+	return nil, agent.ErrExtensionUnsupported
 }
 
 func (vaultAgent) Signers() ([]ssh.Signer, error) {