From 0f9c2e1d17b1e6a3d96524bc2167568db6c78e10 Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Fri, 14 Apr 2017 11:08:07 -0700
Subject: [PATCH] discovery: modify SignAnnouncement to take a MessageSigner

This commit modifies the SignAnnouncement helper function to instead
take a MessageSigner interface rather than the deprecated MessageSigner
struct.  We also modify the function to explicitly take the identity
public key to use when signing.
---
 discovery/service_test.go | 52 ++++++++++++++++++++++++++++-----------
 discovery/utils.go        |  4 +--
 2 files changed, 39 insertions(+), 17 deletions(-)

diff --git a/discovery/service_test.go b/discovery/service_test.go
index 767314a84..2c2f8f488 100644
--- a/discovery/service_test.go
+++ b/discovery/service_test.go
@@ -16,7 +16,6 @@ import (
 	"github.com/go-errors/errors"
 	"github.com/lightningnetwork/lnd/chainntnfs"
 	"github.com/lightningnetwork/lnd/channeldb"
-	"github.com/lightningnetwork/lnd/lnwallet"
 	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/lightningnetwork/lnd/routing"
 	"github.com/roasbeef/btcd/btcec"
@@ -56,6 +55,26 @@ var (
 	proofMatureDelta uint32
 )
 
+type mockSigner struct {
+	privKey *btcec.PrivateKey
+}
+
+func (n *mockSigner) SignMessage(pubKey *btcec.PublicKey,
+	msg []byte) (*btcec.Signature, error) {
+
+	if !pubKey.IsEqual(n.privKey.PubKey()) {
+		return nil, fmt.Errorf("unknown public key")
+	}
+
+	digest := chainhash.DoubleHashB(msg)
+	sign, err := n.privKey.Sign(digest)
+	if err != nil {
+		return nil, fmt.Errorf("can't sign the message: %v", err)
+	}
+
+	return sign, nil
+}
+
 type mockGraphSource struct {
 	nodes      []*channeldb.LightningNode
 	infos      map[uint64]*channeldb.ChannelEdgeInfo
@@ -278,8 +297,8 @@ func createNodeAnnouncement(priv *btcec.PrivateKey) (*lnwire.NodeAnnouncement,
 		Features:  testFeatures,
 	}
 
-	signer := lnwallet.NewMessageSigner(nodeKeyPriv1)
-	if a.Signature, err = SignAnnouncement(signer, a); err != nil {
+	signer := mockSigner{priv}
+	if a.Signature, err = SignAnnouncement(&signer, priv.PubKey(), a); err != nil {
 		return nil, err
 	}
 
@@ -301,8 +320,9 @@ func createUpdateAnnouncement(blockHeight uint32) (*lnwire.ChannelUpdateAnnounce
 		FeeProportionalMillionths: uint32(prand.Int31()),
 	}
 
-	signer := lnwallet.NewMessageSigner(nodeKeyPriv1)
-	if a.Signature, err = SignAnnouncement(signer, a); err != nil {
+	pub := nodeKeyPriv1.PubKey()
+	signer := mockSigner{nodeKeyPriv1}
+	if a.Signature, err = SignAnnouncement(&signer, pub, a); err != nil {
 		return nil, err
 	}
 
@@ -325,13 +345,15 @@ func createRemoteChannelAnnouncement(blockHeight uint32) (*lnwire.ChannelAnnounc
 		BitcoinKey2: bitcoinKeyPub2,
 	}
 
-	signer := lnwallet.NewMessageSigner(nodeKeyPriv1)
-	if a.NodeSig1, err = SignAnnouncement(signer, a); err != nil {
+	pub := nodeKeyPriv1.PubKey()
+	signer := mockSigner{nodeKeyPriv1}
+	if a.NodeSig1, err = SignAnnouncement(&signer, pub, a); err != nil {
 		return nil, err
 	}
 
-	signer = lnwallet.NewMessageSigner(nodeKeyPriv2)
-	if a.NodeSig2, err = SignAnnouncement(signer, a); err != nil {
+	pub = nodeKeyPriv2.PubKey()
+	signer = mockSigner{nodeKeyPriv2}
+	if a.NodeSig2, err = SignAnnouncement(&signer, pub, a); err != nil {
 		return nil, err
 	}
 
@@ -584,7 +606,7 @@ func TestSignatureAnnouncement(t *testing.T) {
 	}
 	select {
 	case <-ctx.broadcastedMessage:
-		t.Fatal("channel announcement was broadcasted")
+		t.Fatal("channel announcement was broadcast")
 	case <-time.After(2 * trickleDelay):
 	}
 
@@ -594,7 +616,7 @@ func TestSignatureAnnouncement(t *testing.T) {
 	}
 	select {
 	case <-ctx.broadcastedMessage:
-		t.Fatal("channel update announcement was broadcasted")
+		t.Fatal("channel update announcement was broadcast")
 	case <-time.After(2 * trickleDelay):
 	}
 
@@ -604,7 +626,7 @@ func TestSignatureAnnouncement(t *testing.T) {
 	}
 	select {
 	case <-ctx.broadcastedMessage:
-		t.Fatal("channel update announcement was broadcasted")
+		t.Fatal("channel update announcement was broadcast")
 	case <-time.After(2 * trickleDelay):
 	}
 
@@ -617,12 +639,12 @@ func TestSignatureAnnouncement(t *testing.T) {
 
 	select {
 	case <-ctx.broadcastedMessage:
-		t.Fatal("announcements were broadcasted")
+		t.Fatal("announcements were broadcast")
 	case <-time.After(2 * trickleDelay):
 	}
 
 	if len(ctx.discovery.waitingProofs) != 1 {
-		t.Fatal("local proof annoucement should be stored")
+		t.Fatal("local proof announcement should be stored")
 	}
 
 	err = <-ctx.discovery.ProcessRemoteAnnouncement(batch.remoteProofAnn, remoteKey)
@@ -634,7 +656,7 @@ func TestSignatureAnnouncement(t *testing.T) {
 		select {
 		case <-ctx.broadcastedMessage:
 		case <-time.After(time.Second):
-			t.Fatal("announcement wasn't broadcasted")
+			t.Fatal("announcement wasn't broadcast")
 		}
 	}
 }
diff --git a/discovery/utils.go b/discovery/utils.go
index e4e1be541..a2be11412 100644
--- a/discovery/utils.go
+++ b/discovery/utils.go
@@ -109,7 +109,7 @@ func copyPubKey(pub *btcec.PublicKey) *btcec.PublicKey {
 
 // SignAnnouncement is a helper function which is used to sign any outgoing
 // channel node node announcement messages.
-func SignAnnouncement(signer *lnwallet.MessageSigner,
+func SignAnnouncement(signer lnwallet.MessageSigner, pubKey *btcec.PublicKey,
 	msg lnwire.Message) (*btcec.Signature, error) {
 
 	var (
@@ -132,5 +132,5 @@ func SignAnnouncement(signer *lnwallet.MessageSigner,
 		return nil, errors.Errorf("unable to get data to sign: %v", err)
 	}
 
-	return signer.SignData(data)
+	return signer.SignMessage(pubKey, data)
 }