From bc7f3e6c65a6f5f3280c44bd3d5e28879256951f Mon Sep 17 00:00:00 2001
From: Damjan Cvetko <damjan.cvetko@monotek.net>
Date: Wed, 23 Oct 2019 00:28:15 +0200
Subject: [PATCH] ExtendedAgent support

---
 src/daemon/ssh_agent.go | 34 ++++++++++++++++++++++++----------
 1 file changed, 24 insertions(+), 10 deletions(-)

diff --git a/src/daemon/ssh_agent.go b/src/daemon/ssh_agent.go
index 406ecd92..1de2546c 100644
--- a/src/daemon/ssh_agent.go
+++ b/src/daemon/ssh_agent.go
@@ -32,7 +32,7 @@ type sessionIDSig struct {
 
 type hostAuthCallback chan *HostAuth
 
-func (a *Agent) withOriginalAgent(do func(agent.Agent)) error {
+func (a *Agent) withOriginalAgent(do func(extendedAgent agent.ExtendedAgent)) error {
 	conn, err := getOriginalAgentConn()
 	if conn == nil {
 		return nil
@@ -79,7 +79,7 @@ func (a *Agent) List() (keys []*agent.Key, err error) {
 		a.notify("", Yellow("Krypton ▶ "+ErrNotPaired.Error()))
 	}
 
-	a.withOriginalAgent(func(fallbackAgent agent.Agent) {
+	a.withOriginalAgent(func(fallbackAgent agent.ExtendedAgent) {
 		fallbackKeys, err := fallbackAgent.List()
 		if err == nil {
 			keys = append(keys, fallbackKeys...)
@@ -92,14 +92,18 @@ func (a *Agent) List() (keys []*agent.Key, err error) {
 // Sign has the agent sign the data using a protocol 2 key as defined
 // in [PROTOCOL.agent] section 2.6.2.
 func (a *Agent) Sign(key ssh.PublicKey, data []byte) (sshSignature *ssh.Signature, err error) {
+	return a.SignWithFlags(key, data, 0)
+}
+
+func (a *Agent) SignWithFlags(key ssh.PublicKey, data []byte, flags agent.SignatureFlags) (sshSignature *ssh.Signature, err error) {
 	keyFingerprint := sha256.Sum256(key.Marshal())
 
-	a.withOriginalAgent(func(fallbackAgent agent.Agent) {
+	a.withOriginalAgent(func(fallbackAgent agent.ExtendedAgent) {
 		fallbackKeys, fallbackErr := fallbackAgent.List()
 		if fallbackErr == nil {
 			for _, fallbackKey := range fallbackKeys {
 				if bytes.Equal(fallbackKey.Marshal(), key.Marshal()) {
-					sshSignature, err = fallbackAgent.Sign(key, data)
+					sshSignature, err = fallbackAgent.SignWithFlags(key, data, flags)
 					return
 				}
 			}
@@ -219,7 +223,7 @@ func (a *Agent) Sign(key ssh.PublicKey, data []byte) (sshSignature *ssh.Signatur
 
 // Add adds a private key to the agent.
 func (a *Agent) Add(key agent.AddedKey) (err error) {
-	connErr := a.withOriginalAgent(func(fallbackAgent agent.Agent) {
+	connErr := a.withOriginalAgent(func(fallbackAgent agent.ExtendedAgent) {
 		err = fallbackAgent.Add(key)
 	})
 	if connErr != nil {
@@ -230,7 +234,7 @@ func (a *Agent) Add(key agent.AddedKey) (err error) {
 
 // Remove removes all identities with the given public key.
 func (a *Agent) Remove(key ssh.PublicKey) (err error) {
-	connErr := a.withOriginalAgent(func(fallbackAgent agent.Agent) {
+	connErr := a.withOriginalAgent(func(fallbackAgent agent.ExtendedAgent) {
 		err = fallbackAgent.Remove(key)
 	})
 	if connErr != nil {
@@ -241,7 +245,7 @@ func (a *Agent) Remove(key ssh.PublicKey) (err error) {
 
 // RemoveAll removes all identities.
 func (a *Agent) RemoveAll() (err error) {
-	connErr := a.withOriginalAgent(func(fallbackAgent agent.Agent) {
+	connErr := a.withOriginalAgent(func(fallbackAgent agent.ExtendedAgent) {
 		err = fallbackAgent.RemoveAll()
 	})
 	if connErr != nil {
@@ -252,7 +256,7 @@ func (a *Agent) RemoveAll() (err error) {
 
 // Lock locks the agent. Sign and Remove will fail, and List will empty an empty list.
 func (a *Agent) Lock(passphrase []byte) (err error) {
-	connErr := a.withOriginalAgent(func(fallbackAgent agent.Agent) {
+	connErr := a.withOriginalAgent(func(fallbackAgent agent.ExtendedAgent) {
 		err = fallbackAgent.Lock(passphrase)
 	})
 	if connErr != nil {
@@ -263,7 +267,7 @@ func (a *Agent) Lock(passphrase []byte) (err error) {
 
 // Unlock undoes the effect of Lock
 func (a *Agent) Unlock(passphrase []byte) (err error) {
-	connErr := a.withOriginalAgent(func(fallbackAgent agent.Agent) {
+	connErr := a.withOriginalAgent(func(fallbackAgent agent.ExtendedAgent) {
 		err = fallbackAgent.Unlock(passphrase)
 	})
 	if connErr != nil {
@@ -274,7 +278,7 @@ func (a *Agent) Unlock(passphrase []byte) (err error) {
 
 // Signers returns signers for all the known keys.
 func (a *Agent) Signers() (signers []ssh.Signer, err error) {
-	connErr := a.withOriginalAgent(func(fallbackAgent agent.Agent) {
+	connErr := a.withOriginalAgent(func(fallbackAgent agent.ExtendedAgent) {
 		signers, err = fallbackAgent.Signers()
 	})
 	if connErr != nil {
@@ -283,6 +287,16 @@ func (a *Agent) Signers() (signers []ssh.Signer, err error) {
 	return
 }
 
+func (a *Agent) Extension(extensionType string, contents []byte) (ret []byte, err error) {
+	connErr := a.withOriginalAgent(func(fallbackAgent agent.ExtendedAgent) {
+		ret, err = fallbackAgent.Extension(extensionType, contents)
+	})
+	if connErr != nil {
+		err = agent.ErrExtensionUnsupported
+	}
+	return
+}
+
 func (a *Agent) notify(prefix, body string) {
 	n, err := socket.OpenNotifier(prefix)
 	if err != nil {