watchtower/multi: move MockSigner+MockPeer to wtmock

watchtower/lookout/justice_descriptor_test.go

@@ -19,6 +19,7 @@ import (
 	"github.com/lightningnetwork/lnd/watchtower/blob"
 	"github.com/lightningnetwork/lnd/watchtower/lookout"
 	"github.com/lightningnetwork/lnd/watchtower/wtdb"
+	"github.com/lightningnetwork/lnd/watchtower/wtmock"
 	"github.com/lightningnetwork/lnd/watchtower/wtpolicy"
 )
 
@@ -47,55 +48,6 @@ var (
 	}
 )
 
-type mockSigner struct {
-	index uint32
-	keys  map[keychain.KeyLocator]*btcec.PrivateKey
-}
-
-func newMockSigner() *mockSigner {
-	return &mockSigner{
-		keys: make(map[keychain.KeyLocator]*btcec.PrivateKey),
-	}
-}
-
-func (s *mockSigner) SignOutputRaw(tx *wire.MsgTx,
-	signDesc *input.SignDescriptor) ([]byte, error) {
-
-	witnessScript := signDesc.WitnessScript
-	amt := signDesc.Output.Value
-
-	privKey, ok := s.keys[signDesc.KeyDesc.KeyLocator]
-	if !ok {
-		panic("cannot sign w/ unknown key")
-	}
-
-	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 (s *mockSigner) ComputeInputScript(tx *wire.MsgTx,
-	signDesc *input.SignDescriptor) (*input.Script, error) {
-	return nil, nil
-}
-
-func (s *mockSigner) addPrivKey(privKey *btcec.PrivateKey) keychain.KeyLocator {
-	keyLoc := keychain.KeyLocator{
-		Index: s.index,
-	}
-	s.index++
-
-	s.keys[keyLoc] = privKey
-
-	return keyLoc
-}
-
 func TestJusticeDescriptor(t *testing.T) {
 	const (
 		localAmount  = btcutil.Amount(100000)
@@ -115,10 +67,10 @@ func TestJusticeDescriptor(t *testing.T) {
 	)
 
 	// Create the signer, and add the revocation and to-remote privkeys.
-	signer := newMockSigner()
+	signer := wtmock.NewMockSigner()
 	var (
-		revKeyLoc      = signer.addPrivKey(revSK)
+		revKeyLoc      = signer.AddPrivKey(revSK)
-		toRemoteKeyLoc = signer.addPrivKey(toRemoteSK)
+		toRemoteKeyLoc = signer.AddPrivKey(toRemoteSK)
 	)
 
 	// Construct the to-local witness script.

watchtower/wtmock/peer.go

@@ -1,6 +1,4 @@
-// +build dev
+package wtmock
-
-package wtserver
 
 import (
 	"fmt"
@@ -8,8 +6,10 @@ import (
 	"time"
 
 	"github.com/btcsuite/btcd/btcec"
+	"github.com/lightningnetwork/lnd/watchtower/wtserver"
 )
 
+// MockPeer emulates a single endpoint of brontide transport.
 type MockPeer struct {
 	remotePub  *btcec.PublicKey
 	remoteAddr net.Addr
@@ -23,6 +23,7 @@ type MockPeer struct {
 	Quit chan struct{}
 }
 
+// NewMockPeer returns a fresh MockPeer.
 func NewMockPeer(pk *btcec.PublicKey, addr net.Addr, bufferSize int) *MockPeer {
 	return &MockPeer{
 		remotePub:    pk,
@@ -33,6 +34,10 @@ func NewMockPeer(pk *btcec.PublicKey, addr net.Addr, bufferSize int) *MockPeer {
 	}
 }
 
+// Write sends the raw bytes as the next full message read to the remote peer.
+// The write will fail if either party closes the connection or the write
+// deadline expires. The passed bytes slice is copied before sending, thus the
+// bytes may be reused once the method returns.
 func (p *MockPeer) Write(b []byte) (n int, err error) {
 	select {
 	case p.OutgoingMsgs <- b:
@@ -44,6 +49,7 @@ func (p *MockPeer) Write(b []byte) (n int, err error) {
 	}
 }
 
+// Close tearsdown the connection, and fails any pending reads or writes.
 func (p *MockPeer) Close() error {
 	select {
 	case <-p.Quit:
@@ -54,6 +60,9 @@ func (p *MockPeer) Close() error {
 	}
 }
 
+// ReadNextMessage returns the raw bytes of the next full message read from the
+// remote peer. The read will fail if either party closes the connection or the
+// read deadline expires.
 func (p *MockPeer) ReadNextMessage() ([]byte, error) {
 	select {
 	case b := <-p.IncomingMsgs:
@@ -65,6 +74,8 @@ func (p *MockPeer) ReadNextMessage() ([]byte, error) {
 	}
 }
 
+// SetWriteDeadline initializes a timer that will cause any pending writes to
+// fail at time t. If t is zero, the deadline is infinite.
 func (p *MockPeer) SetWriteDeadline(t time.Time) error {
 	if t.IsZero() {
 		p.writeDeadline = nil
@@ -77,6 +88,8 @@ func (p *MockPeer) SetWriteDeadline(t time.Time) error {
 	return nil
 }
 
+// SetReadDeadline initializes a timer that will cause any pending reads to fail
+// at time t. If t is zero, the deadline is infinite.
 func (p *MockPeer) SetReadDeadline(t time.Time) error {
 	if t.IsZero() {
 		p.readDeadline = nil
@@ -89,10 +102,16 @@ func (p *MockPeer) SetReadDeadline(t time.Time) error {
 	return nil
 }
 
+// RemotePub returns the public key of the remote peer.
 func (p *MockPeer) RemotePub() *btcec.PublicKey {
 	return p.remotePub
 }
 
+// RemoteAddr returns the net address of the remote peer.
 func (p *MockPeer) RemoteAddr() net.Addr {
 	return p.remoteAddr
 }
+
+// Compile-time constraint ensuring the MockPeer implements the wserver.Peer
+// interface.
+var _ wtserver.Peer = (*MockPeer)(nil)

watchtower/wtmock/signer.go

@@ -0,0 +1,81 @@
+package wtmock
+
+import (
+	"sync"
+
+	"github.com/btcsuite/btcd/btcec"
+	"github.com/btcsuite/btcd/txscript"
+	"github.com/btcsuite/btcd/wire"
+	"github.com/lightningnetwork/lnd/input"
+	"github.com/lightningnetwork/lnd/keychain"
+)
+
+// MockSigner is an input.Signer that allows one to add arbitrary private keys
+// and sign messages by passing the assigned keychain.KeyLocator.
+type MockSigner struct {
+	mu sync.Mutex
+
+	index uint32
+	keys  map[keychain.KeyLocator]*btcec.PrivateKey
+}
+
+// NewMockSigner returns a fresh MockSigner.
+func NewMockSigner() *MockSigner {
+	return &MockSigner{
+		keys: make(map[keychain.KeyLocator]*btcec.PrivateKey),
+	}
+}
+
+// SignOutputRaw signs an input on the passed transaction using the input index
+// in the sign descriptor. The returned signature is the raw DER-encoded
+// signature without the signhash flag.
+func (s *MockSigner) SignOutputRaw(tx *wire.MsgTx,
+	signDesc *input.SignDescriptor) ([]byte, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	witnessScript := signDesc.WitnessScript
+	amt := signDesc.Output.Value
+
+	privKey, ok := s.keys[signDesc.KeyDesc.KeyLocator]
+	if !ok {
+		panic("cannot sign w/ unknown key")
+	}
+
+	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
+}
+
+// ComputeInputScript is not implemented.
+func (s *MockSigner) ComputeInputScript(tx *wire.MsgTx,
+	signDesc *input.SignDescriptor) (*input.Script, error) {
+	panic("not implemented")
+}
+
+// AddPrivKey records the passed privKey in the MockSigner's registry of keys it
+// can sign with in the future. A unique key locator is returned, allowing the
+// caller to sign with this key when presented via an input.SignDescriptor.
+func (s *MockSigner) AddPrivKey(privKey *btcec.PrivateKey) keychain.KeyLocator {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	keyLoc := keychain.KeyLocator{
+		Index: s.index,
+	}
+	s.index++
+
+	s.keys[keyLoc] = privKey
+
+	return keyLoc
+}
+
+// Compile-time constraint ensuring the MockSigner implements the input.Signer
+// interface.
+var _ input.Signer = (*MockSigner)(nil)

watchtower/wtserver/server_test.go

@@ -15,6 +15,7 @@ import (
 	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/lightningnetwork/lnd/watchtower/blob"
 	"github.com/lightningnetwork/lnd/watchtower/wtdb"
+	"github.com/lightningnetwork/lnd/watchtower/wtmock"
 	"github.com/lightningnetwork/lnd/watchtower/wtserver"
 	"github.com/lightningnetwork/lnd/watchtower/wtwire"
 )
@@ -85,8 +86,8 @@ func TestServerOnlyAcceptOnePeer(t *testing.T) {
 
 	// Create two peers using the same session id.
 	peerPub := randPubKey(t)
-	peer1 := wtserver.NewMockPeer(peerPub, nil, 0)
+	peer1 := wtmock.NewMockPeer(peerPub, nil, 0)
-	peer2 := wtserver.NewMockPeer(peerPub, nil, 0)
+	peer2 := wtmock.NewMockPeer(peerPub, nil, 0)
 
 	// Serialize a Init message to be sent by both peers.
 	init := wtwire.NewInitMessage(
@@ -112,8 +113,8 @@ func TestServerOnlyAcceptOnePeer(t *testing.T) {
 	// Try to send a message on either peer, and record the opposite peer as
 	// the one we assume to be rejected.
 	var (
-		rejectedPeer *wtserver.MockPeer
+		rejectedPeer *wtmock.MockPeer
-		acceptedPeer *wtserver.MockPeer
+		acceptedPeer *wtmock.MockPeer
 	)
 	select {
 	case peer1.IncomingMsgs <- msg:
@@ -215,7 +216,7 @@ func testServerCreateSession(t *testing.T, i int, test createSessionTestCase) {
 
 	// Create a new client and connect to server.
 	peerPub := randPubKey(t)
-	peer := wtserver.NewMockPeer(peerPub, nil, 0)
+	peer := wtmock.NewMockPeer(peerPub, nil, 0)
 	connect(t, i, s, peer, test.initMsg, timeoutDuration)
 
 	// Send the CreateSession message, and wait for a reply.
@@ -243,7 +244,7 @@ func testServerCreateSession(t *testing.T, i int, test createSessionTestCase) {
 
 	// Simulate a peer with the same session id connection to the server
 	// again.
-	peer = wtserver.NewMockPeer(peerPub, nil, 0)
+	peer = wtmock.NewMockPeer(peerPub, nil, 0)
 	connect(t, i, s, peer, test.initMsg, timeoutDuration)
 
 	// Send the _same_ CreateSession message as the first attempt.
@@ -515,7 +516,7 @@ func testServerStateUpdates(t *testing.T, i int, test stateUpdateTestCase) {
 
 	// Create a new client and connect to the server.
 	peerPub := randPubKey(t)
-	peer := wtserver.NewMockPeer(peerPub, nil, 0)
+	peer := wtmock.NewMockPeer(peerPub, nil, 0)
 	connect(t, i, s, peer, test.initMsg, timeoutDuration)
 
 	// Register a session for this client to use in the subsequent tests.
@@ -535,7 +536,7 @@ func testServerStateUpdates(t *testing.T, i int, test stateUpdateTestCase) {
 
 	// Now that the original connection has been closed, connect a new
 	// client with the same session id.
-	peer = wtserver.NewMockPeer(peerPub, nil, 0)
+	peer = wtmock.NewMockPeer(peerPub, nil, 0)
 	connect(t, i, s, peer, test.initMsg, timeoutDuration)
 
 	// Send the intended StateUpdate messages in series.
@@ -546,7 +547,7 @@ func testServerStateUpdates(t *testing.T, i int, test stateUpdateTestCase) {
 		if update == nil {
 			assertConnClosed(t, peer, 2*timeoutDuration)
 
-			peer = wtserver.NewMockPeer(peerPub, nil, 0)
+			peer = wtmock.NewMockPeer(peerPub, nil, 0)
 			connect(t, i, s, peer, test.initMsg, timeoutDuration)
 
 			continue
@@ -570,7 +571,7 @@ func testServerStateUpdates(t *testing.T, i int, test stateUpdateTestCase) {
 	assertConnClosed(t, peer, 2*timeoutDuration)
 }
 
-func connect(t *testing.T, i int, s wtserver.Interface, peer *wtserver.MockPeer,
+func connect(t *testing.T, i int, s wtserver.Interface, peer *wtmock.MockPeer,
 	initMsg *wtwire.Init, timeout time.Duration) {
 
 	s.InboundPeerConnected(peer)
@@ -578,9 +579,9 @@ func connect(t *testing.T, i int, s wtserver.Interface, peer *wtserver.MockPeer,
 	recvReply(t, i, "Init", peer, timeout)
 }
 
-// sendMsg sends a wtwire.Message message via a wtserver.MockPeer.
+// sendMsg sends a wtwire.Message message via a wtmock.MockPeer.
 func sendMsg(t *testing.T, i int, msg wtwire.Message,
-	peer *wtserver.MockPeer, timeout time.Duration) {
+	peer *wtmock.MockPeer, timeout time.Duration) {
 
 	t.Helper()
 
@@ -602,7 +603,7 @@ func sendMsg(t *testing.T, i int, msg wtwire.Message,
 // expected reply type. The supported replies are CreateSessionReply and
 // StateUpdateReply.
 func recvReply(t *testing.T, i int, name string,
-	peer *wtserver.MockPeer, timeout time.Duration) wtwire.Message {
+	peer *wtmock.MockPeer, timeout time.Duration) wtwire.Message {
 
 	t.Helper()
 
@@ -646,7 +647,7 @@ func recvReply(t *testing.T, i int, name string,
 
 // assertConnClosed checks that the peer's connection is closed before the
 // timeout expires.
-func assertConnClosed(t *testing.T, peer *wtserver.MockPeer, duration time.Duration) {
+func assertConnClosed(t *testing.T, peer *wtmock.MockPeer, duration time.Duration) {
 	t.Helper()
 
 	select {