You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
237 lines
5.7 KiB
237 lines
5.7 KiB
package sweep |
|
|
|
import ( |
|
"fmt" |
|
"sync" |
|
"testing" |
|
"time" |
|
|
|
"github.com/btcsuite/btcd/chaincfg/chainhash" |
|
"github.com/btcsuite/btcd/wire" |
|
"github.com/lightningnetwork/lnd/chainntnfs" |
|
) |
|
|
|
var ( |
|
defaultTestTimeout = 5 * time.Second |
|
processingDelay = 1 * time.Second |
|
mockChainHash, _ = chainhash.NewHashFromStr("00aabbccddeeff") |
|
mockChainHeight = int32(100) |
|
) |
|
|
|
// MockNotifier simulates the chain notifier for test purposes. This type is |
|
// exported because it is used in nursery tests. |
|
type MockNotifier struct { |
|
confChannel map[chainhash.Hash]chan *chainntnfs.TxConfirmation |
|
epochChan map[chan *chainntnfs.BlockEpoch]int32 |
|
spendChan map[wire.OutPoint][]chan *chainntnfs.SpendDetail |
|
spends map[wire.OutPoint]*wire.MsgTx |
|
mutex sync.RWMutex |
|
t *testing.T |
|
} |
|
|
|
// NewMockNotifier instantiates a new mock notifier. |
|
func NewMockNotifier(t *testing.T) *MockNotifier { |
|
return &MockNotifier{ |
|
confChannel: make(map[chainhash.Hash]chan *chainntnfs.TxConfirmation), |
|
epochChan: make(map[chan *chainntnfs.BlockEpoch]int32), |
|
spendChan: make(map[wire.OutPoint][]chan *chainntnfs.SpendDetail), |
|
spends: make(map[wire.OutPoint]*wire.MsgTx), |
|
t: t, |
|
} |
|
} |
|
|
|
// NotifyEpoch simulates a new epoch arriving. |
|
func (m *MockNotifier) NotifyEpoch(height int32) { |
|
m.t.Helper() |
|
|
|
for epochChan, chanHeight := range m.epochChan { |
|
// Only send notifications if the height is greater than the |
|
// height the caller passed into the register call. |
|
if chanHeight >= height { |
|
continue |
|
} |
|
|
|
log.Debugf("Notifying height %v to listener", height) |
|
|
|
select { |
|
case epochChan <- &chainntnfs.BlockEpoch{ |
|
Height: height, |
|
}: |
|
case <-time.After(defaultTestTimeout): |
|
m.t.Fatal("epoch event not consumed") |
|
} |
|
} |
|
} |
|
|
|
// ConfirmTx simulates a tx confirming. |
|
func (m *MockNotifier) ConfirmTx(txid *chainhash.Hash, height uint32) error { |
|
confirm := &chainntnfs.TxConfirmation{ |
|
BlockHeight: height, |
|
} |
|
select { |
|
case m.getConfChannel(txid) <- confirm: |
|
case <-time.After(defaultTestTimeout): |
|
return fmt.Errorf("confirmation not consumed") |
|
} |
|
return nil |
|
} |
|
|
|
// SpendOutpoint simulates a utxo being spent. |
|
func (m *MockNotifier) SpendOutpoint(outpoint wire.OutPoint, |
|
spendingTx wire.MsgTx) { |
|
|
|
log.Debugf("Spending outpoint %v", outpoint) |
|
|
|
m.mutex.Lock() |
|
defer m.mutex.Unlock() |
|
|
|
channels, ok := m.spendChan[outpoint] |
|
if ok { |
|
for _, channel := range channels { |
|
m.sendSpend(channel, &outpoint, &spendingTx) |
|
} |
|
} |
|
|
|
m.spends[outpoint] = &spendingTx |
|
} |
|
|
|
func (m *MockNotifier) sendSpend(channel chan *chainntnfs.SpendDetail, |
|
outpoint *wire.OutPoint, |
|
spendingTx *wire.MsgTx) { |
|
|
|
spenderTxHash := spendingTx.TxHash() |
|
channel <- &chainntnfs.SpendDetail{ |
|
SpenderTxHash: &spenderTxHash, |
|
SpendingTx: spendingTx, |
|
SpentOutPoint: outpoint, |
|
} |
|
} |
|
|
|
// RegisterConfirmationsNtfn registers for tx confirm notifications. |
|
func (m *MockNotifier) RegisterConfirmationsNtfn(txid *chainhash.Hash, |
|
_ []byte, numConfs, heightHint uint32) (*chainntnfs.ConfirmationEvent, |
|
error) { |
|
|
|
return &chainntnfs.ConfirmationEvent{ |
|
Confirmed: m.getConfChannel(txid), |
|
}, nil |
|
} |
|
|
|
func (m *MockNotifier) getConfChannel( |
|
txid *chainhash.Hash) chan *chainntnfs.TxConfirmation { |
|
|
|
m.mutex.Lock() |
|
defer m.mutex.Unlock() |
|
|
|
channel, ok := m.confChannel[*txid] |
|
if ok { |
|
return channel |
|
} |
|
channel = make(chan *chainntnfs.TxConfirmation) |
|
m.confChannel[*txid] = channel |
|
|
|
return channel |
|
} |
|
|
|
// RegisterBlockEpochNtfn registers a block notification. |
|
func (m *MockNotifier) RegisterBlockEpochNtfn( |
|
bestBlock *chainntnfs.BlockEpoch) (*chainntnfs.BlockEpochEvent, error) { |
|
|
|
log.Tracef("Mock block ntfn registered") |
|
|
|
m.mutex.Lock() |
|
epochChan := make(chan *chainntnfs.BlockEpoch, 1) |
|
|
|
// The real notifier returns a notification with the current block hash |
|
// and height immediately if no best block hash or height is specified |
|
// in the request. We want to emulate this behaviour as well for the |
|
// mock. |
|
switch { |
|
case bestBlock == nil: |
|
epochChan <- &chainntnfs.BlockEpoch{ |
|
Hash: mockChainHash, |
|
Height: mockChainHeight, |
|
} |
|
m.epochChan[epochChan] = mockChainHeight |
|
default: |
|
m.epochChan[epochChan] = bestBlock.Height |
|
} |
|
m.mutex.Unlock() |
|
|
|
return &chainntnfs.BlockEpochEvent{ |
|
Epochs: epochChan, |
|
Cancel: func() { |
|
log.Tracef("Mock block ntfn canceled") |
|
m.mutex.Lock() |
|
delete(m.epochChan, epochChan) |
|
m.mutex.Unlock() |
|
}, |
|
}, nil |
|
} |
|
|
|
// Start the notifier. |
|
func (m *MockNotifier) Start() error { |
|
return nil |
|
} |
|
|
|
// Started checks if started |
|
func (m *MockNotifier) Started() bool { |
|
return true |
|
} |
|
|
|
// Stop the notifier. |
|
func (m *MockNotifier) Stop() error { |
|
return nil |
|
} |
|
|
|
// RegisterSpendNtfn registers for spend notifications. |
|
func (m *MockNotifier) RegisterSpendNtfn(outpoint *wire.OutPoint, |
|
_ []byte, heightHint uint32) (*chainntnfs.SpendEvent, error) { |
|
|
|
// Add channel to global spend ntfn map. |
|
m.mutex.Lock() |
|
|
|
channels, ok := m.spendChan[*outpoint] |
|
if !ok { |
|
channels = make([]chan *chainntnfs.SpendDetail, 0) |
|
} |
|
|
|
channel := make(chan *chainntnfs.SpendDetail, 1) |
|
channels = append(channels, channel) |
|
m.spendChan[*outpoint] = channels |
|
|
|
// Check if this output has already been spent. |
|
spendingTx, spent := m.spends[*outpoint] |
|
|
|
m.mutex.Unlock() |
|
|
|
// If output has been spent already, signal now. Do this outside the |
|
// lock to prevent a deadlock. |
|
if spent { |
|
m.sendSpend(channel, outpoint, spendingTx) |
|
} |
|
|
|
return &chainntnfs.SpendEvent{ |
|
Spend: channel, |
|
Cancel: func() { |
|
log.Infof("Cancelling RegisterSpendNtfn for %v", |
|
outpoint) |
|
|
|
m.mutex.Lock() |
|
defer m.mutex.Unlock() |
|
channels := m.spendChan[*outpoint] |
|
for i, c := range channels { |
|
if c == channel { |
|
channels[i] = channels[len(channels)-1] |
|
m.spendChan[*outpoint] = |
|
channels[:len(channels)-1] |
|
} |
|
} |
|
|
|
close(channel) |
|
|
|
log.Infof("Spend ntfn channel closed for %v", |
|
outpoint) |
|
}, |
|
}, nil |
|
}
|
|
|