From 6cf1844b0eaeda4919e39ad04973ba24b4a50a4f Mon Sep 17 00:00:00 2001
From: Oliver Gugger <gugger@gmail.com>
Date: Tue, 28 Apr 2020 10:06:19 +0200
Subject: [PATCH] keychain: add new interface methods to wallet

---
 keychain/btcwallet.go  | 60 +++++++++++++++++++++++++++++++++++++++++-
 keychain/derivation.go |  4 +++
 mock.go                | 20 ++++++++++++++
 3 files changed, 83 insertions(+), 1 deletion(-)

diff --git a/keychain/btcwallet.go b/keychain/btcwallet.go
index b8a47b12..42737daf 100644
--- a/keychain/btcwallet.go
+++ b/keychain/btcwallet.go
@@ -247,7 +247,9 @@ func (b *BtcWalletKeyRing) DeriveKey(keyLoc KeyLocator) (KeyDescriptor, error) {
 // passed key descriptor.
 //
 // NOTE: This is part of the keychain.SecretKeyRing interface.
-func (b *BtcWalletKeyRing) DerivePrivKey(keyDesc KeyDescriptor) (*btcec.PrivateKey, error) {
+func (b *BtcWalletKeyRing) DerivePrivKey(keyDesc KeyDescriptor) (
+	*btcec.PrivateKey, error) {
+
 	var key *btcec.PrivateKey
 
 	db := b.wallet.Database()
@@ -371,3 +373,59 @@ func (b *BtcWalletKeyRing) ScalarMult(keyDesc KeyDescriptor,
 
 	return h[:], nil
 }
+
+// ECDH performs a scalar multiplication (ECDH-like operation) between the
+// target key descriptor and remote public key. The output returned will be
+// the sha256 of the resulting shared point serialized in compressed format. If
+// k is our private key, and P is the public key, we perform the following
+// operation:
+//
+//  sx := k*P s := sha256(sx.SerializeCompressed())
+//
+// NOTE: This is part of the keychain.ECDHRing interface.
+func (b *BtcWalletKeyRing) ECDH(keyDesc KeyDescriptor,
+	pub *btcec.PublicKey) ([32]byte, error) {
+
+	privKey, err := b.DerivePrivKey(keyDesc)
+	if err != nil {
+		return [32]byte{}, err
+	}
+
+	s := &btcec.PublicKey{}
+	x, y := btcec.S256().ScalarMult(pub.X, pub.Y, privKey.D.Bytes())
+	s.X = x
+	s.Y = y
+
+	h := sha256.Sum256(s.SerializeCompressed())
+
+	return h, nil
+}
+
+// SignDigest signs the given SHA256 message digest with the private key
+// described in the key descriptor.
+//
+// NOTE: This is part of the keychain.DigestSignerRing interface.
+func (b *BtcWalletKeyRing) SignDigest(keyDesc KeyDescriptor,
+	digest [32]byte) (*btcec.Signature, error) {
+
+	privKey, err := b.DerivePrivKey(keyDesc)
+	if err != nil {
+		return nil, err
+	}
+	return privKey.Sign(digest[:])
+}
+
+// SignDigestCompact signs the given SHA256 message digest with the private key
+// described in the key descriptor and returns the signature in the compact,
+// public key recoverable format.
+//
+// NOTE: This is part of the keychain.DigestSignerRing interface.
+func (b *BtcWalletKeyRing) SignDigestCompact(keyDesc KeyDescriptor,
+	digest [32]byte) ([]byte, error) {
+
+	privKey, err := b.DerivePrivKey(keyDesc)
+	if err != nil {
+		return nil, err
+	}
+	return btcec.SignCompact(btcec.S256(), privKey, digest[:], true)
+}
diff --git a/keychain/derivation.go b/keychain/derivation.go
index 18a1b345..67c901d0 100644
--- a/keychain/derivation.go
+++ b/keychain/derivation.go
@@ -176,6 +176,10 @@ type KeyRing interface {
 type SecretKeyRing interface {
 	KeyRing
 
+	ECDHRing
+
+	DigestSignerRing
+
 	// DerivePrivKey attempts to derive the private key that corresponds to
 	// the passed key descriptor.  If the public key is set, then this
 	// method will perform an in-order scan over the key set, with a max of
diff --git a/mock.go b/mock.go
index 2b13e75f..7729bb76 100644
--- a/mock.go
+++ b/mock.go
@@ -363,3 +363,23 @@ func (m *mockSecretKeyRing) ScalarMult(keyDesc keychain.KeyDescriptor,
 	pubKey *btcec.PublicKey) ([]byte, error) {
 	return nil, nil
 }
+
+func (m *mockSecretKeyRing) ECDH(_ keychain.KeyDescriptor,
+	pubKey *btcec.PublicKey) ([32]byte, error) {
+
+	return [32]byte{}, nil
+}
+
+func (m *mockSecretKeyRing) SignDigest(_ keychain.KeyDescriptor,
+	digest [32]byte) (*btcec.Signature, error) {
+
+	return m.rootKey.Sign(digest[:])
+}
+
+func (m *mockSecretKeyRing) SignDigestCompact(_ keychain.KeyDescriptor,
+	digest [32]byte) ([]byte, error) {
+
+	return btcec.SignCompact(btcec.S256(), m.rootKey, digest[:], true)
+}
+
+var _ keychain.SecretKeyRing = (*mockSecretKeyRing)(nil)