lnwallet: Move mocks to separate file and augment mockSigner.

To implement the BOLT 03 test vectors, a more powerful mockSigner is
required. The new version of mockSigner stores multiple keys and signs
the transaction outputs with the appropriate one.

lnwallet/channel_test.go

@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"crypto/sha256"
 	"errors"
-	"fmt"
 	"io/ioutil"
 	"math/big"
 	"math/rand"
@@ -60,92 +59,6 @@ var (
 	numReqConfs = uint16(1)
 )
 
-type mockSigner struct {
-	key *btcec.PrivateKey
-}
-
-func (m *mockSigner) SignOutputRaw(tx *wire.MsgTx, signDesc *SignDescriptor) ([]byte, error) {
-	amt := signDesc.Output.Value
-	witnessScript := signDesc.WitnessScript
-	privKey := m.key
-
-	if !privKey.PubKey().IsEqual(signDesc.PubKey) {
-		return nil, fmt.Errorf("incorrect key passed")
-	}
-
-	switch {
-	case signDesc.SingleTweak != nil:
-		privKey = TweakPrivKey(privKey,
-			signDesc.SingleTweak)
-	case signDesc.DoubleTweak != nil:
-		privKey = DeriveRevocationPrivKey(privKey,
-			signDesc.DoubleTweak)
-	}
-
-	sig, err := txscript.RawTxInWitnessSignature(tx, signDesc.SigHashes,
-		signDesc.InputIndex, amt, witnessScript, signDesc.HashType,
-		privKey)
-	if err != nil {
-		return nil, err
-	}
-
-	return sig[:len(sig)-1], nil
-}
-func (m *mockSigner) ComputeInputScript(tx *wire.MsgTx, signDesc *SignDescriptor) (*InputScript, error) {
-
-	// TODO(roasbeef): expose tweaked signer from lnwallet so don't need to
-	// duplicate this code?
-
-	privKey := m.key
-
-	switch {
-	case signDesc.SingleTweak != nil:
-		privKey = TweakPrivKey(privKey,
-			signDesc.SingleTweak)
-	case signDesc.DoubleTweak != nil:
-		privKey = DeriveRevocationPrivKey(privKey,
-			signDesc.DoubleTweak)
-	}
-
-	witnessScript, err := txscript.WitnessSignature(tx, signDesc.SigHashes,
-		signDesc.InputIndex, signDesc.Output.Value, signDesc.Output.PkScript,
-		signDesc.HashType, privKey, true)
-	if err != nil {
-		return nil, err
-	}
-
-	return &InputScript{
-		Witness: witnessScript,
-	}, nil
-}
-
-type mockNotfier struct {
-}
-
-func (m *mockNotfier) RegisterConfirmationsNtfn(txid *chainhash.Hash, numConfs, heightHint uint32) (*chainntnfs.ConfirmationEvent, error) {
-	return nil, nil
-}
-
-func (m *mockNotfier) RegisterBlockEpochNtfn() (*chainntnfs.BlockEpochEvent, error) {
-	return nil, nil
-}
-
-func (m *mockNotfier) Start() error {
-	return nil
-}
-
-func (m *mockNotfier) Stop() error {
-	return nil
-}
-
-func (m *mockNotfier) RegisterSpendNtfn(outpoint *wire.OutPoint, heightHint uint32) (*chainntnfs.SpendEvent, error) {
-	return &chainntnfs.SpendEvent{
-		Spend: make(chan *chainntnfs.SpendDetail),
-		Cancel: func() {
-		},
-	}, nil
-}
-
 // mockSpendNotifier extends the mockNotifier so that spend notifications can be
 // triggered and delivered to subscribers.
 type mockSpendNotifier struct {
@@ -428,8 +341,8 @@ func createTestChannelsWithNotifier(revocationWindow int,
 		Db:                      dbBob,
 	}
 
-	aliceSigner := &mockSigner{aliceKeyPriv}
+	aliceSigner := &mockSigner{privkeys: []*btcec.PrivateKey{aliceKeyPriv}}
-	bobSigner := &mockSigner{bobKeyPriv}
+	bobSigner := &mockSigner{privkeys: []*btcec.PrivateKey{bobKeyPriv}}
 
 	channelAlice, err := NewLightningChannel(aliceSigner, notifier,
 		estimator, aliceChannelState)

lnwallet/common_test.go

@@ -0,0 +1,148 @@
+package lnwallet
+
+import (
+	"bytes"
+	"fmt"
+
+	"github.com/lightningnetwork/lnd/chainntnfs"
+	"github.com/roasbeef/btcd/btcec"
+	"github.com/roasbeef/btcd/chaincfg"
+	"github.com/roasbeef/btcd/chaincfg/chainhash"
+	"github.com/roasbeef/btcd/txscript"
+	"github.com/roasbeef/btcd/wire"
+	"github.com/roasbeef/btcutil"
+)
+
+// mockSigner is a simple implementation of the Signer interface. Each one has
+// a set of private keys in a slice and can sign messages using the appropriate
+// one.
+type mockSigner struct {
+	privkeys  []*btcec.PrivateKey
+	netParams *chaincfg.Params
+}
+
+func (m *mockSigner) SignOutputRaw(tx *wire.MsgTx, signDesc *SignDescriptor) ([]byte, error) {
+	pubkey := signDesc.PubKey
+	switch {
+	case signDesc.SingleTweak != nil:
+		pubkey = TweakPubKeyWithTweak(pubkey, signDesc.SingleTweak)
+	case signDesc.DoubleTweak != nil:
+		pubkey = DeriveRevocationPubkey(pubkey, signDesc.DoubleTweak.PubKey())
+	}
+
+	hash160 := btcutil.Hash160(pubkey.SerializeCompressed())
+	privKey := m.findKey(hash160, signDesc.SingleTweak, signDesc.DoubleTweak)
+	if privKey == nil {
+		return nil, fmt.Errorf("Mock signer does not have key")
+	}
+
+	sig, err := txscript.RawTxInWitnessSignature(tx, signDesc.SigHashes,
+		signDesc.InputIndex, signDesc.Output.Value, signDesc.WitnessScript,
+		txscript.SigHashAll, privKey)
+	if err != nil {
+		return nil, err
+	}
+
+	return sig[:len(sig)-1], nil
+}
+
+func (m *mockSigner) ComputeInputScript(tx *wire.MsgTx, signDesc *SignDescriptor) (*InputScript, error) {
+	scriptType, addresses, _, err := txscript.ExtractPkScriptAddrs(
+		signDesc.Output.PkScript, m.netParams)
+	if err != nil {
+		return nil, err
+	}
+
+	switch scriptType {
+	case txscript.PubKeyHashTy:
+		privKey := m.findKey(addresses[0].ScriptAddress(), signDesc.SingleTweak,
+			signDesc.DoubleTweak)
+		if privKey == nil {
+			return nil, fmt.Errorf("Mock signer does not have key for "+
+				"address %v", addresses[0])
+		}
+
+		scriptSig, err := txscript.SignatureScript(tx, signDesc.InputIndex,
+			signDesc.Output.PkScript, txscript.SigHashAll, privKey, true)
+		if err != nil {
+			return nil, err
+		}
+
+		return &InputScript{ScriptSig: scriptSig}, nil
+
+	case txscript.WitnessV0PubKeyHashTy:
+		privKey := m.findKey(addresses[0].ScriptAddress(), signDesc.SingleTweak,
+			signDesc.DoubleTweak)
+		if privKey == nil {
+			return nil, fmt.Errorf("Mock signer does not have key for "+
+				"address %v", addresses[0])
+		}
+
+		witnessScript, err := txscript.WitnessSignature(tx, signDesc.SigHashes,
+			signDesc.InputIndex, signDesc.Output.Value,
+			signDesc.Output.PkScript, txscript.SigHashAll, privKey, true)
+		if err != nil {
+			return nil, err
+		}
+
+		return &InputScript{Witness: witnessScript}, nil
+
+	default:
+		return nil, fmt.Errorf("Unexpected script type: %v", scriptType)
+	}
+}
+
+// findKey searches through all stored private keys and returns one
+// corresponding to the hashed pubkey if it can be found. The public key may
+// either correspond directly to the private key or to the private key with a
+// tweak applied.
+func (m *mockSigner) findKey(needleHash160 []byte, singleTweak []byte,
+	doubleTweak *btcec.PrivateKey) *btcec.PrivateKey {
+
+	for _, privkey := range m.privkeys {
+		// First check whether public key is directly derived from private key.
+		hash160 := btcutil.Hash160(privkey.PubKey().SerializeCompressed())
+		if bytes.Equal(hash160, needleHash160) {
+			return privkey
+		}
+
+		// Otherwise check if public key is derived from tweaked private key.
+		switch {
+		case singleTweak != nil:
+			privkey = TweakPrivKey(privkey, singleTweak)
+		case doubleTweak != nil:
+			privkey = DeriveRevocationPrivKey(privkey, doubleTweak)
+		default:
+			continue
+		}
+		hash160 = btcutil.Hash160(privkey.PubKey().SerializeCompressed())
+		if bytes.Equal(hash160, needleHash160) {
+			return privkey
+		}
+	}
+	return nil
+}
+
+type mockNotfier struct {
+}
+
+func (m *mockNotfier) RegisterConfirmationsNtfn(txid *chainhash.Hash, numConfs, heightHint uint32) (*chainntnfs.ConfirmationEvent, error) {
+	return nil, nil
+}
+func (m *mockNotfier) RegisterBlockEpochNtfn() (*chainntnfs.BlockEpochEvent, error) {
+	return nil, nil
+}
+
+func (m *mockNotfier) Start() error {
+	return nil
+}
+
+func (m *mockNotfier) Stop() error {
+	return nil
+}
+func (m *mockNotfier) RegisterSpendNtfn(outpoint *wire.OutPoint, heightHint uint32) (*chainntnfs.SpendEvent, error) {
+	return &chainntnfs.SpendEvent{
+		Spend:  make(chan *chainntnfs.SpendDetail),
+		Cancel: func() {},
+	}, nil
+}

lnwallet/script_utils_test.go

@@ -62,7 +62,9 @@ func TestCommitmentSpendValidation(t *testing.T) {
 	aliceCommitTweak := SingleTweakBytes(commitPoint, aliceKeyPub)
 	bobCommitTweak := SingleTweakBytes(commitPoint, bobKeyPub)
 
-	aliceSelfOutputSigner := &mockSigner{aliceKeyPriv}
+	aliceSelfOutputSigner := &mockSigner{
+		privkeys: []*btcec.PrivateKey{aliceKeyPriv},
+	}
 
 	// With all the test data set up, we create the commitment transaction.
 	// We only focus on a single party's transactions, as the scripts are
@@ -135,7 +137,7 @@ func TestCommitmentSpendValidation(t *testing.T) {
 		t.Fatalf("spend from delay output is invalid: %v", err)
 	}
 
-	bobSigner := &mockSigner{bobKeyPriv}
+	bobSigner := &mockSigner{privkeys: []*btcec.PrivateKey{bobKeyPriv}}
 
 	// Next, we'll test bob spending with the derived revocation key to
 	// simulate the scenario when Alice broadcasts this commitment
@@ -385,8 +387,8 @@ func TestHTLCSenderSpendValidation(t *testing.T) {
 	// Finally, we'll create mock signers for both of them based on their
 	// private keys. This test simplifies a bit and uses the same key as
 	// the base point for all scripts and derivations.
-	bobSigner := &mockSigner{bobKeyPriv}
+	bobSigner := &mockSigner{privkeys: []*btcec.PrivateKey{bobKeyPriv}}
-	aliceSigner := &mockSigner{aliceKeyPriv}
+	aliceSigner := &mockSigner{privkeys: []*btcec.PrivateKey{aliceKeyPriv}}
 
 	// We'll also generate a signature on the sweep transaction above
 	// that'll act as Bob's signature to Alice for the second level HTLC
@@ -630,8 +632,8 @@ func TestHTLCReceiverSpendValidation(t *testing.T) {
 	// Finally, we'll create mock signers for both of them based on their
 	// private keys. This test simplifies a bit and uses the same key as
 	// the base point for all scripts and derivations.
-	bobSigner := &mockSigner{bobKeyPriv}
+	bobSigner := &mockSigner{privkeys: []*btcec.PrivateKey{bobKeyPriv}}
-	aliceSigner := &mockSigner{aliceKeyPriv}
+	aliceSigner := &mockSigner{privkeys: []*btcec.PrivateKey{aliceKeyPriv}}
 
 	// We'll also generate a signature on the sweep transaction above
 	// that'll act as Alice's signature to Bob for the second level HTLC
@@ -866,8 +868,8 @@ func TestSecondLevelHtlcSpends(t *testing.T) {
 	// Finally, we'll create mock signers for both of them based on their
 	// private keys. This test simplifies a bit and uses the same key as
 	// the base point for all scripts and derivations.
-	bobSigner := &mockSigner{bobKeyPriv}
+	bobSigner := &mockSigner{privkeys: []*btcec.PrivateKey{bobKeyPriv}}
-	aliceSigner := &mockSigner{aliceKeyPriv}
+	aliceSigner := &mockSigner{privkeys: []*btcec.PrivateKey{aliceKeyPriv}}
 
 	testCases := []struct {
 		witness func() wire.TxWitness