diff --git a/watchtower/wtdb/mock.go b/watchtower/wtdb/mock.go index c47fe2b1..fa53d6cf 100644 --- a/watchtower/wtdb/mock.go +++ b/watchtower/wtdb/mock.go @@ -70,6 +70,33 @@ func (db *MockDB) InsertSessionInfo(info *SessionInfo) error { 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() diff --git a/watchtower/wtserver/interface.go b/watchtower/wtserver/interface.go index 69d4ec9c..d2364543 100644 --- a/watchtower/wtserver/interface.go +++ b/watchtower/wtserver/interface.go @@ -63,4 +63,8 @@ type DB interface { // validates the update against the current SessionInfo stored under the // update's session id.. InsertStateUpdate(*wtdb.SessionStateUpdate) (uint16, error) + + // DeleteSession removes all data associated with a particular session + // id from the tower's database. + DeleteSession(wtdb.SessionID) error }