lnd.xprv/watchtower/wtdb/mock.go
2018-10-31 20:42:06 -07:00

115 lines
2.2 KiB
Go

// +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()
if _, ok := db.sessions[info.ID]; ok {
return ErrSessionAlreadyExists
}
db.sessions[info.ID] = info
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
}