watchtower/wtmock/tower_db: move mock tower db to wtmock pkg
This commit is contained in:
Conner Fromknecht
parent
a783f35ad4
commit
2ce6228021
@ -15,6 +15,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"
|
||||
)
|
||||
|
||||
@ -66,7 +67,7 @@ func makeAddrSlice(size int) []byte {
|
||||
}
|
||||
|
||||
func TestLookoutBreachMatching(t *testing.T) {
|
||||
db := wtdb.NewMockDB()
|
||||
db := wtmock.NewTowerDB()
|
||||
|
||||
// Initialize an mock backend to feed the lookout blocks.
|
||||
backend := lookout.NewMockBackend()
|
||||
|
||||
@ -369,7 +369,7 @@ type testHarness struct {
|
||||
clientDB *wtmock.ClientDB
|
||||
clientCfg *wtclient.Config
|
||||
client wtclient.Client
|
||||
serverDB *wtdb.MockDB
|
||||
serverDB *wtmock.TowerDB
|
||||
serverCfg *wtserver.Config
|
||||
server *wtserver.Server
|
||||
net *mockNet
|
||||
@ -406,7 +406,7 @@ func newHarness(t *testing.T, cfg harnessCfg) *testHarness {
|
||||
}
|
||||
|
||||
const timeout = 200 * time.Millisecond
|
||||
serverDB := wtdb.NewMockDB()
|
||||
serverDB := wtmock.NewTowerDB()
|
||||
|
||||
serverCfg := &wtserver.Config{
|
||||
DB: serverDB,
|
||||
|
||||
@ -1,142 +0,0 @@
|
||||
// +build dev
|
||||
|
||||
package wtdb
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
"github.com/lightningnetwork/lnd/chainntnfs"
|
||||
)
|
||||
|
||||
type MockDB struct {
|
||||
mu sync.Mutex
|
||||
lastEpoch *chainntnfs.BlockEpoch
|
||||
sessions map[SessionID]*SessionInfo
|
||||
blobs map[BreachHint]map[SessionID]*SessionStateUpdate
|
||||
}
|
||||
|
||||
func NewMockDB() *MockDB {
|
||||
return &MockDB{
|
||||
sessions: make(map[SessionID]*SessionInfo),
|
||||
blobs: make(map[BreachHint]map[SessionID]*SessionStateUpdate),
|
||||
}
|
||||
}
|
||||
|
||||
func (db *MockDB) InsertStateUpdate(update *SessionStateUpdate) (uint16, error) {
|
||||
db.mu.Lock()
|
||||
defer db.mu.Unlock()
|
||||
|
||||
info, ok := db.sessions[update.ID]
|
||||
if !ok {
|
||||
return 0, ErrSessionNotFound
|
||||
}
|
||||
|
||||
err := info.AcceptUpdateSequence(update.SeqNum, update.LastApplied)
|
||||
if err != nil {
|
||||
return info.LastApplied, err
|
||||
}
|
||||
|
||||
sessionsToUpdates, ok := db.blobs[update.Hint]
|
||||
if !ok {
|
||||
sessionsToUpdates = make(map[SessionID]*SessionStateUpdate)
|
||||
db.blobs[update.Hint] = sessionsToUpdates
|
||||
}
|
||||
sessionsToUpdates[update.ID] = update
|
||||
|
||||
return info.LastApplied, nil
|
||||
}
|
||||
|
||||
func (db *MockDB) GetSessionInfo(id *SessionID) (*SessionInfo, error) {
|
||||
db.mu.Lock()
|
||||
defer db.mu.Unlock()
|
||||
|
||||
if info, ok := db.sessions[*id]; ok {
|
||||
return info, nil
|
||||
}
|
||||
|
||||
return nil, ErrSessionNotFound
|
||||
}
|
||||
|
||||
func (db *MockDB) InsertSessionInfo(info *SessionInfo) error {
|
||||
db.mu.Lock()
|
||||
defer db.mu.Unlock()
|
||||
|
||||
dbInfo, ok := db.sessions[info.ID]
|
||||
if ok && dbInfo.LastApplied > 0 {
|
||||
return ErrSessionAlreadyExists
|
||||
}
|
||||
|
||||
db.sessions[info.ID] = info
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (db *MockDB) DeleteSession(target SessionID) error {
|
||||
db.mu.Lock()
|
||||
defer db.mu.Unlock()
|
||||
|
||||
// Fail if the session doesn't exit.
|
||||
if _, ok := db.sessions[target]; !ok {
|
||||
return ErrSessionNotFound
|
||||
}
|
||||
|
||||
// Remove the target session.
|
||||
delete(db.sessions, target)
|
||||
|
||||
// Remove the state updates for any blobs stored under the target
|
||||
// session identifier.
|
||||
for hint, sessionUpdates := range db.blobs {
|
||||
delete(sessionUpdates, target)
|
||||
|
||||
//If this was the last state update, we can also remove the hint
|
||||
//that would map to an empty set.
|
||||
if len(sessionUpdates) == 0 {
|
||||
delete(db.blobs, hint)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (db *MockDB) GetLookoutTip() (*chainntnfs.BlockEpoch, error) {
|
||||
db.mu.Lock()
|
||||
defer db.mu.Unlock()
|
||||
|
||||
return db.lastEpoch, nil
|
||||
}
|
||||
|
||||
func (db *MockDB) QueryMatches(breachHints []BreachHint) ([]Match, error) {
|
||||
db.mu.Lock()
|
||||
defer db.mu.Unlock()
|
||||
|
||||
var matches []Match
|
||||
for _, hint := range breachHints {
|
||||
sessionsToUpdates, ok := db.blobs[hint]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
for id, update := range sessionsToUpdates {
|
||||
info, ok := db.sessions[id]
|
||||
if !ok {
|
||||
panic("session not found")
|
||||
}
|
||||
|
||||
match := Match{
|
||||
ID: id,
|
||||
SeqNum: update.SeqNum,
|
||||
Hint: hint,
|
||||
EncryptedBlob: update.EncryptedBlob,
|
||||
SessionInfo: info,
|
||||
}
|
||||
matches = append(matches, match)
|
||||
}
|
||||
}
|
||||
|
||||
return matches, nil
|
||||
}
|
||||
|
||||
func (db *MockDB) SetLookoutTip(epoch *chainntnfs.BlockEpoch) error {
|
||||
db.lastEpoch = epoch
|
||||
return nil
|
||||
}
|
||||
162
watchtower/wtmock/tower_db.go
Normal file
162
watchtower/wtmock/tower_db.go
Normal file
@ -0,0 +1,162 @@
|
||||
package wtmock
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
"github.com/lightningnetwork/lnd/chainntnfs"
|
||||
"github.com/lightningnetwork/lnd/watchtower/wtdb"
|
||||
)
|
||||
|
||||
// TowerDB is a mock, in-memory implementation of a watchtower.DB.
|
||||
type TowerDB struct {
|
||||
mu sync.Mutex
|
||||
lastEpoch *chainntnfs.BlockEpoch
|
||||
sessions map[wtdb.SessionID]*wtdb.SessionInfo
|
||||
blobs map[wtdb.BreachHint]map[wtdb.SessionID]*wtdb.SessionStateUpdate
|
||||
}
|
||||
|
||||
// NewTowerDB initializes a fresh mock TowerDB.
|
||||
func NewTowerDB() *TowerDB {
|
||||
return &TowerDB{
|
||||
sessions: make(map[wtdb.SessionID]*wtdb.SessionInfo),
|
||||
blobs: make(map[wtdb.BreachHint]map[wtdb.SessionID]*wtdb.SessionStateUpdate),
|
||||
}
|
||||
}
|
||||
|
||||
// InsertStateUpdate stores an update sent by the client after validating that
|
||||
// the update is well-formed in the context of other updates sent for the same
|
||||
// session. This include verifying that the sequence number is incremented
|
||||
// properly and the last applied values echoed by the client are sane.
|
||||
func (db *TowerDB) InsertStateUpdate(update *wtdb.SessionStateUpdate) (uint16, error) {
|
||||
db.mu.Lock()
|
||||
defer db.mu.Unlock()
|
||||
|
||||
info, ok := db.sessions[update.ID]
|
||||
if !ok {
|
||||
return 0, wtdb.ErrSessionNotFound
|
||||
}
|
||||
|
||||
err := info.AcceptUpdateSequence(update.SeqNum, update.LastApplied)
|
||||
if err != nil {
|
||||
return info.LastApplied, err
|
||||
}
|
||||
|
||||
sessionsToUpdates, ok := db.blobs[update.Hint]
|
||||
if !ok {
|
||||
sessionsToUpdates = make(map[wtdb.SessionID]*wtdb.SessionStateUpdate)
|
||||
db.blobs[update.Hint] = sessionsToUpdates
|
||||
}
|
||||
sessionsToUpdates[update.ID] = update
|
||||
|
||||
return info.LastApplied, nil
|
||||
}
|
||||
|
||||
// GetSessionInfo retrieves the session for the passed session id. An error is
|
||||
// returned if the session could not be found.
|
||||
func (db *TowerDB) GetSessionInfo(id *wtdb.SessionID) (*wtdb.SessionInfo, error) {
|
||||
db.mu.Lock()
|
||||
defer db.mu.Unlock()
|
||||
|
||||
if info, ok := db.sessions[*id]; ok {
|
||||
return info, nil
|
||||
}
|
||||
|
||||
return nil, wtdb.ErrSessionNotFound
|
||||
}
|
||||
|
||||
// InsertSessionInfo records a negotiated session in the tower database. An
|
||||
// error is returned if the session already exists.
|
||||
func (db *TowerDB) InsertSessionInfo(info *wtdb.SessionInfo) error {
|
||||
db.mu.Lock()
|
||||
defer db.mu.Unlock()
|
||||
|
||||
dbInfo, ok := db.sessions[info.ID]
|
||||
if ok && dbInfo.LastApplied > 0 {
|
||||
return wtdb.ErrSessionAlreadyExists
|
||||
}
|
||||
|
||||
db.sessions[info.ID] = info
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteSession removes all data associated with a particular session id from
|
||||
// the tower's database.
|
||||
func (db *TowerDB) DeleteSession(target wtdb.SessionID) error {
|
||||
db.mu.Lock()
|
||||
defer db.mu.Unlock()
|
||||
|
||||
// Fail if the session doesn't exit.
|
||||
if _, ok := db.sessions[target]; !ok {
|
||||
return wtdb.ErrSessionNotFound
|
||||
}
|
||||
|
||||
// Remove the target session.
|
||||
delete(db.sessions, target)
|
||||
|
||||
// Remove the state updates for any blobs stored under the target
|
||||
// session identifier.
|
||||
for hint, sessionUpdates := range db.blobs {
|
||||
delete(sessionUpdates, target)
|
||||
|
||||
// If this was the last state update, we can also remove the
|
||||
// hint that would map to an empty set.
|
||||
if len(sessionUpdates) == 0 {
|
||||
delete(db.blobs, hint)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// QueryMatches searches against all known state updates for any that match the
|
||||
// passed breachHints. More than one Match will be returned for a given hint if
|
||||
// they exist in the database.
|
||||
func (db *TowerDB) QueryMatches(
|
||||
breachHints []wtdb.BreachHint) ([]wtdb.Match, error) {
|
||||
|
||||
db.mu.Lock()
|
||||
defer db.mu.Unlock()
|
||||
|
||||
var matches []wtdb.Match
|
||||
for _, hint := range breachHints {
|
||||
sessionsToUpdates, ok := db.blobs[hint]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
for id, update := range sessionsToUpdates {
|
||||
info, ok := db.sessions[id]
|
||||
if !ok {
|
||||
panic("session not found")
|
||||
}
|
||||
|
||||
match := wtdb.Match{
|
||||
ID: id,
|
||||
SeqNum: update.SeqNum,
|
||||
Hint: hint,
|
||||
EncryptedBlob: update.EncryptedBlob,
|
||||
SessionInfo: info,
|
||||
}
|
||||
matches = append(matches, match)
|
||||
}
|
||||
}
|
||||
|
||||
return matches, nil
|
||||
}
|
||||
|
||||
// SetLookoutTip stores the provided epoch as the latest lookout tip epoch in
|
||||
// the tower database.
|
||||
func (db *TowerDB) SetLookoutTip(epoch *chainntnfs.BlockEpoch) error {
|
||||
db.lastEpoch = epoch
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetLookoutTip retrieves the current lookout tip block epoch from the tower
|
||||
// database.
|
||||
func (db *TowerDB) GetLookoutTip() (*chainntnfs.BlockEpoch, error) {
|
||||
db.mu.Lock()
|
||||
defer db.mu.Unlock()
|
||||
|
||||
return db.lastEpoch, nil
|
||||
}
|
||||
@ -1,5 +1,3 @@
|
||||
// +build dev
|
||||
|
||||
package wtserver_test
|
||||
|
||||
import (
|
||||
@ -53,7 +51,7 @@ func initServer(t *testing.T, db wtserver.DB,
|
||||
t.Helper()
|
||||
|
||||
if db == nil {
|
||||
db = wtdb.NewMockDB()
|
||||
db = wtmock.NewTowerDB()
|
||||
}
|
||||
|
||||
s, err := wtserver.New(&wtserver.Config{
|
||||
@ -687,7 +685,7 @@ func testServerStateUpdates(t *testing.T, test stateUpdateTestCase) {
|
||||
// checking that the proper error is returned when the session doesn't exist and
|
||||
// that a successful deletion does not disrupt other sessions.
|
||||
func TestServerDeleteSession(t *testing.T) {
|
||||
db := wtdb.NewMockDB()
|
||||
db := wtmock.NewTowerDB()
|
||||
|
||||
localPub := randPubKey(t)
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user