From 43120fd08cec2b63f67c4b1e13b730f7e5cc0cd4 Mon Sep 17 00:00:00 2001 From: Conner Fromknecht Date: Thu, 31 Jan 2019 18:16:02 -0800 Subject: [PATCH] watchtower/multi: move MockSigner+MockPeer to wtmock --- watchtower/lookout/justice_descriptor_test.go | 56 +------------ .../{wtserver/mock.go => wtmock/peer.go} | 25 +++++- watchtower/wtmock/signer.go | 81 +++++++++++++++++++ watchtower/wtserver/server_test.go | 29 +++---- 4 files changed, 122 insertions(+), 69 deletions(-) rename watchtower/{wtserver/mock.go => wtmock/peer.go} (58%) create mode 100644 watchtower/wtmock/signer.go diff --git a/watchtower/lookout/justice_descriptor_test.go b/watchtower/lookout/justice_descriptor_test.go index 9502193d..63c2c29f 100644 --- a/watchtower/lookout/justice_descriptor_test.go +++ b/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) - toRemoteKeyLoc = signer.addPrivKey(toRemoteSK) + revKeyLoc = signer.AddPrivKey(revSK) + toRemoteKeyLoc = signer.AddPrivKey(toRemoteSK) ) // Construct the to-local witness script. diff --git a/watchtower/wtserver/mock.go b/watchtower/wtmock/peer.go similarity index 58% rename from watchtower/wtserver/mock.go rename to watchtower/wtmock/peer.go index 758e4ae2..9103c499 100644 --- a/watchtower/wtserver/mock.go +++ b/watchtower/wtmock/peer.go @@ -1,6 +1,4 @@ -// +build dev - -package wtserver +package wtmock 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) diff --git a/watchtower/wtmock/signer.go b/watchtower/wtmock/signer.go new file mode 100644 index 00000000..c41e4f2f --- /dev/null +++ b/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) diff --git a/watchtower/wtserver/server_test.go b/watchtower/wtserver/server_test.go index d8a8911c..61984d6c 100644 --- a/watchtower/wtserver/server_test.go +++ b/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) - peer2 := wtserver.NewMockPeer(peerPub, nil, 0) + peer1 := wtmock.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 - acceptedPeer *wtserver.MockPeer + rejectedPeer *wtmock.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 {