From 24016c35c7f011ec8aa31376175b202a8ebf7343 Mon Sep 17 00:00:00 2001 From: Elle Mouton Date: Wed, 29 Mar 2023 12:15:17 +0200 Subject: [PATCH] watchtower: add forceNext to NextSessionKeyIndex This commit adds a forceNext boolean parameter to NextSessionKeyIndex. Setting this param to true will force the index to cycle over 1000 key indices before returning the new key. --- watchtower/wtclient/interface.go | 4 +-- watchtower/wtclient/session_negotiator.go | 2 +- watchtower/wtdb/client_db.go | 44 +++++++++++++++-------- watchtower/wtdb/client_db_test.go | 42 +++++++++++++++------- watchtower/wtmock/client_db.go | 28 ++++++++++----- 5 files changed, 81 insertions(+), 39 deletions(-) diff --git a/watchtower/wtclient/interface.go b/watchtower/wtclient/interface.go index 8a6d9b320..06126ae21 100644 --- a/watchtower/wtclient/interface.go +++ b/watchtower/wtclient/interface.go @@ -49,8 +49,8 @@ type DB interface { // (tower, blob type) pair until CreateClientSession is invoked for that // tower and index, at which point a new index for that tower can be // reserved. Multiple calls to this method before CreateClientSession is - // invoked should return the same index. - NextSessionKeyIndex(wtdb.TowerID, blob.Type) (uint32, error) + // invoked should return the same index unless forceNext is true. + NextSessionKeyIndex(wtdb.TowerID, blob.Type, bool) (uint32, error) // CreateClientSession saves a newly negotiated client session to the // client's database. This enables the session to be used across diff --git a/watchtower/wtclient/session_negotiator.go b/watchtower/wtclient/session_negotiator.go index db0f543a1..aafc44b2a 100644 --- a/watchtower/wtclient/session_negotiator.go +++ b/watchtower/wtclient/session_negotiator.go @@ -306,7 +306,7 @@ retryWithBackoff: // with this specific tower. If one is already reserved, the // existing index will be returned. keyIndex, err := n.cfg.DB.NextSessionKeyIndex( - tower.ID, n.cfg.Policy.BlobType, + tower.ID, n.cfg.Policy.BlobType, false, ) if err != nil { n.log.Debugf("Unable to reserve session key index "+ diff --git a/watchtower/wtdb/client_db.go b/watchtower/wtdb/client_db.go index 5a18095b8..a949021d1 100644 --- a/watchtower/wtdb/client_db.go +++ b/watchtower/wtdb/client_db.go @@ -643,9 +643,10 @@ func (c *ClientDB) ListTowers() ([]*Tower, error) { // particular tower id. The index is reserved for that tower until // CreateClientSession is invoked for that tower and index, at which point a new // index for that tower can be reserved. Multiple calls to this method before -// CreateClientSession is invoked should return the same index. +// CreateClientSession is invoked should return the same index unless forceNext +// is true. func (c *ClientDB) NextSessionKeyIndex(towerID TowerID, - blobType blob.Type) (uint32, error) { + blobType blob.Type, forceNext bool) (uint32, error) { var index uint32 err := kvdb.Update(c.db, func(tx kvdb.RwTx) error { @@ -654,33 +655,48 @@ func (c *ClientDB) NextSessionKeyIndex(towerID TowerID, return ErrUninitializedDB } - // Check the session key index to see if a key has already been - // reserved for this tower. If so, we'll deserialize and return - // the index directly. var err error - index, err = getSessionKeyIndex(keyIndex, towerID, blobType) - if err == nil { - return nil + if !forceNext { + // Check the session key index to see if a key has + // already been reserved for this tower. If so, we'll + // deserialize and return the index directly. + index, err = getSessionKeyIndex( + keyIndex, towerID, blobType, + ) + if err == nil { + return nil + } } - // Otherwise, generate a new session key index since the node - // doesn't already have reserved index. The error is ignored - // since NextSequence can't fail inside Update. - index64, _ := keyIndex.NextSequence() + // By default, we use the next available bucket sequence as the + // key index. But if forceNext is true, then it is assumed that + // some data loss occurred and so the sequence is incremented a + // by a jump of 1000 so that we can arrive at a brand new key + // index quicker. + currentSequence := keyIndex.Sequence() + nextIndex := currentSequence + 1 + if forceNext { + nextIndex = currentSequence + 1000 + } + + if err = keyIndex.SetSequence(nextIndex); err != nil { + return fmt.Errorf("could not set next bucket "+ + "sequence: %w", err) + } // As a sanity check, assert that the index is still in the // valid range of unhardened pubkeys. In the future, we should // move to only using hardened keys, and this will prevent any // overlap from occurring until then. This also prevents us from // overflowing uint32s. - if index64 > math.MaxInt32 { + if nextIndex > math.MaxInt32 { return fmt.Errorf("exhausted session key indexes") } // Create the key that will used to be store the reserved index. keyBytes := createSessionKeyIndexKey(towerID, blobType) - index = uint32(index64) + index = uint32(nextIndex) var indexBuf [4]byte byteOrder.PutUint32(indexBuf[:], index) diff --git a/watchtower/wtdb/client_db_test.go b/watchtower/wtdb/client_db_test.go index 26d02bcf3..80557f1f8 100644 --- a/watchtower/wtdb/client_db_test.go +++ b/watchtower/wtdb/client_db_test.go @@ -62,12 +62,12 @@ func (h *clientDBHarness) listSessions(id *wtdb.TowerID, return sessions } -func (h *clientDBHarness) nextKeyIndex(id wtdb.TowerID, - blobType blob.Type) uint32 { +func (h *clientDBHarness) nextKeyIndex(id wtdb.TowerID, blobType blob.Type, + forceNext bool) uint32 { h.t.Helper() - index, err := h.db.NextSessionKeyIndex(id, blobType) + index, err := h.db.NextSessionKeyIndex(id, blobType, forceNext) require.NoError(h.t, err, "unable to create next session key index") require.NotZero(h.t, index, "next key index should never be 0") @@ -309,7 +309,7 @@ func testCreateClientSession(h *clientDBHarness) { h.insertSession(session, wtdb.ErrNoReservedKeyIndex) // Now, reserve a session key for this tower. - keyIndex := h.nextKeyIndex(session.TowerID, blobType) + keyIndex := h.nextKeyIndex(session.TowerID, blobType, false) // The client session hasn't been updated with the reserved key index // (since it's still zero). Inserting should fail due to the mismatch. @@ -318,7 +318,7 @@ func testCreateClientSession(h *clientDBHarness) { // Reserve another key for the same index. Since no session has been // successfully created, it should return the same index to maintain // idempotency across restarts. - keyIndex2 := h.nextKeyIndex(session.TowerID, blobType) + keyIndex2 := h.nextKeyIndex(session.TowerID, blobType, false) require.Equalf(h.t, keyIndex, keyIndex2, "next key index should "+ "be idempotent: want: %v, got %v", keyIndex, keyIndex2) @@ -335,11 +335,23 @@ func testCreateClientSession(h *clientDBHarness) { // session already existing. h.insertSession(session, wtdb.ErrClientSessionAlreadyExists) - // Finally, assert that reserving another key index succeeds with a - // different key index, now that the first one has been finalized. - keyIndex3 := h.nextKeyIndex(session.TowerID, blobType) + // Assert that reserving another key index succeeds with a different key + // index, now that the first one has been finalized. + keyIndex3 := h.nextKeyIndex(session.TowerID, blobType, false) require.NotEqualf(h.t, keyIndex, keyIndex3, "key index still "+ "reserved after creating session") + + // Show that calling NextSessionKeyIndex again now will result in the + // same key being returned as long as forceNext remains false. + keyIndex4 := h.nextKeyIndex(session.TowerID, blobType, false) + require.Equal(h.t, keyIndex3, keyIndex4) + + // Finally, assert that if the forceNext param of the + // NextSessionKeyIndex method is true, then the key index returned will + // be different. + keyIndex5 := h.nextKeyIndex(session.TowerID, blobType, true) + require.NotEqual(h.t, keyIndex5, keyIndex4) + require.Equal(h.t, keyIndex3+1000, keyIndex5) } // testFilterClientSessions asserts that we can correctly filter client sessions @@ -352,7 +364,7 @@ func testFilterClientSessions(h *clientDBHarness) { towerSessions := make(map[wtdb.TowerID][]wtdb.SessionID) for i := 0; i < numSessions; i++ { tower := h.newTower() - keyIndex := h.nextKeyIndex(tower.ID, blobType) + keyIndex := h.nextKeyIndex(tower.ID, blobType, false) sessionID := wtdb.SessionID([33]byte{byte(i)}) h.insertSession(&wtdb.ClientSession{ ClientSessionBody: wtdb.ClientSessionBody{ @@ -499,7 +511,9 @@ func testRemoveTower(h *clientDBHarness) { MaxUpdates: 100, }, RewardPkScript: []byte{0x01, 0x02, 0x03}, - KeyIndex: h.nextKeyIndex(tower.ID, blobType), + KeyIndex: h.nextKeyIndex( + tower.ID, blobType, false, + ), }, ID: wtdb.SessionID([33]byte{0x01}), } @@ -581,7 +595,7 @@ func testCommitUpdate(h *clientDBHarness) { ) // Reserve a session key index and insert the session. - session.KeyIndex = h.nextKeyIndex(session.TowerID, blobType) + session.KeyIndex = h.nextKeyIndex(session.TowerID, blobType, false) h.insertSession(session, nil) // Now, try to commit the update that failed initially which should @@ -774,7 +788,7 @@ func testAckUpdate(h *clientDBHarness) { h.ackUpdate(&session.ID, 1, 0, wtdb.ErrClientSessionNotFound) // Reserve a session key and insert the client session. - session.KeyIndex = h.nextKeyIndex(session.TowerID, blobType) + session.KeyIndex = h.nextKeyIndex(session.TowerID, blobType, false) h.insertSession(session, nil) // Now, try to ack update 1. This should fail since update 1 was never @@ -1051,7 +1065,9 @@ func (h *clientDBHarness) randSession(t *testing.T, MaxUpdates: maxUpdates, }, RewardPkScript: []byte{0x01, 0x02, 0x03}, - KeyIndex: h.nextKeyIndex(towerID, blobType), + KeyIndex: h.nextKeyIndex( + towerID, blobType, false, + ), }, ID: id, } diff --git a/watchtower/wtmock/client_db.go b/watchtower/wtmock/client_db.go index 8d43a0e45..5b77da744 100644 --- a/watchtower/wtmock/client_db.go +++ b/watchtower/wtmock/client_db.go @@ -401,9 +401,10 @@ func (m *ClientDB) CreateClientSession(session *wtdb.ClientSession) error { // particular tower id. The index is reserved for that tower until // CreateClientSession is invoked for that tower and index, at which point a new // index for that tower can be reserved. Multiple calls to this method before -// CreateClientSession is invoked should return the same index. -func (m *ClientDB) NextSessionKeyIndex(towerID wtdb.TowerID, - blobType blob.Type) (uint32, error) { +// CreateClientSession is invoked should return the same index unless forceNext +// is set to true. +func (m *ClientDB) NextSessionKeyIndex(towerID wtdb.TowerID, blobType blob.Type, + forceNext bool) (uint32, error) { m.mu.Lock() defer m.mu.Unlock() @@ -413,15 +414,24 @@ func (m *ClientDB) NextSessionKeyIndex(towerID wtdb.TowerID, blobType: blobType, } - if index, err := m.getSessionKeyIndex(key); err == nil { - return index, nil + if !forceNext { + if index, err := m.getSessionKeyIndex(key); err == nil { + return index, nil + } } - m.nextIndex++ - index := m.nextIndex - m.indexes[key] = index + // By default, we use the next available bucket sequence as the key + // index. But if forceNext is true, then it is assumed that some data + // loss occurred and so the sequence is incremented a by a jump of 1000 + // so that we can arrive at a brand new key index quicker. + nextIndex := m.nextIndex + 1 + if forceNext { + nextIndex = m.nextIndex + 1000 + } + m.nextIndex = nextIndex + m.indexes[key] = nextIndex - return index, nil + return nextIndex, nil } func (m *ClientDB) getSessionKeyIndex(key keyIndexKey) (uint32, error) {