lnd.xprv/watchtower/wtmock/tower_db.go
2019-06-13 17:40:04 -07:00

164 lines
4.2 KiB
Go

package wtmock
import (
"sync"
"github.com/lightningnetwork/lnd/chainntnfs"
"github.com/lightningnetwork/lnd/watchtower/blob"
"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[blob.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[blob.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 []blob.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
}