akovalenko/lnd.xprv
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Merge pull request #2514 from cfromknecht/add-wtmock-pkg

Browse Source 
watchtower/multi: move MockSigner+MockPeer to wtmock
This commit is contained in:
Olaoluwa Osuntokun 2019-01-31 19:13:47 -08:00 committed by GitHub
commit 55ed7769a5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 122 additions and 69 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

56
watchtower/lookout/justice_descriptor_test.go
View File

@ -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.

25
watchtower/wtserver/mock.go → watchtower/wtmock/peer.go
View File

@ -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)

81
watchtower/wtmock/signer.go Normal file
View File

@ -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)

29
watchtower/wtserver/server_test.go
View File

@ -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:
@ -217,7 +218,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.
@ -245,7 +246,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.
@ -525,7 +526,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.
@ -545,7 +546,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.
@ -556,7 +557,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
@ -580,7 +581,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)
@ -588,9 +589,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()
@ -612,7 +613,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()
@ -656,7 +657,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 {