watchtower/multi: move MockSigner+MockPeer to wtmock
This commit is contained in:
parent
948646b58b
commit
43120fd08c
@ -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.
|
||||
|
@ -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
81
watchtower/wtmock/signer.go
Normal 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)
|
@ -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 {
|
||||
|
Loading…
Reference in New Issue
Block a user