watchtower/multi: move MockSigner+MockPeer to wtmock

watchtower/wtmock/peer.go

@@ -0,0 +1,117 @@
+package wtmock
+
+import (
+	"fmt"
+	"net"
+	"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
+
+	IncomingMsgs chan []byte
+	OutgoingMsgs chan []byte
+
+	writeDeadline <-chan time.Time
+	readDeadline  <-chan time.Time
+
+	Quit chan struct{}
+}
+
+// NewMockPeer returns a fresh MockPeer.
+func NewMockPeer(pk *btcec.PublicKey, addr net.Addr, bufferSize int) *MockPeer {
+	return &MockPeer{
+		remotePub:    pk,
+		remoteAddr:   addr,
+		IncomingMsgs: make(chan []byte, bufferSize),
+		OutgoingMsgs: make(chan []byte, bufferSize),
+		Quit:         make(chan struct{}),
+	}
+}
+
+// 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:
+		return len(b), nil
+	case <-p.writeDeadline:
+		return 0, fmt.Errorf("write timeout expired")
+	case <-p.Quit:
+		return 0, fmt.Errorf("connection closed")
+	}
+}
+
+// Close tearsdown the connection, and fails any pending reads or writes.
+func (p *MockPeer) Close() error {
+	select {
+	case <-p.Quit:
+		return fmt.Errorf("connection already closed")
+	default:
+		close(p.Quit)
+		return nil
+	}
+}
+
+// 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:
+		return b, nil
+	case <-p.readDeadline:
+		return nil, fmt.Errorf("read timeout expired")
+	case <-p.Quit:
+		return nil, fmt.Errorf("connection closed")
+	}
+}
+
+// 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
+		return nil
+	}
+
+	duration := time.Until(t)
+	p.writeDeadline = time.After(duration)
+
+	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
+		return nil
+	}
+
+	duration := time.Until(t)
+	p.readDeadline = time.After(duration)
+
+	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)