From e373e6ba1aec3eab6d21788fff5b81d3654eee6c Mon Sep 17 00:00:00 2001 From: guillaumemichel Date: Tue, 15 Oct 2024 17:40:14 +0200 Subject: [PATCH 1/3] key persistence --- keys.go | 16 +++++++++++++++- queen.go | 16 ++++++++++++++-- 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/keys.go b/keys.go index e6257da..d38d996 100644 --- a/keys.go +++ b/keys.go @@ -91,6 +91,17 @@ func (db *KeysDB) writeKeysToFile(keysTrie *trie.Trie[bit256.Key, crypto.PrivKey } } +func integrateKeysIntoTrie(keysTrie *trie.Trie[bit256.Key, crypto.PrivKey], keys []crypto.PrivKey) { + for _, key := range keys { + pid, err := peer.IDFromPrivateKey(key) + if err != nil { + logger.Warnf("Error getting peer ID: %v", err) + continue + } + keysTrie.Add(PeeridToKadid(pid), key) + } +} + func genKey() crypto.PrivKey { priv, _, err := crypto.GenerateEd25519Key(rand.Reader) if err != nil { @@ -140,10 +151,13 @@ func getMatchingKeys(prefixes []bitstr.Key, keysTrie *trie.Trie[bit256.Key, cryp return keys } -func (db *KeysDB) MatchingKeys(prefixes []bitstr.Key) []crypto.PrivKey { +func (db *KeysDB) MatchingKeys(prefixes []bitstr.Key, returned []crypto.PrivKey) []crypto.PrivKey { // read keys from disk keysTrie := db.readKeysFromFile() + // integrate returned keys into the trie, they are available again + integrateKeysIntoTrie(keysTrie, returned) + // pop any matching keys from the keysTrie, generate new keys if needed // and store them in keysTrie keys := getMatchingKeys(prefixes, keysTrie) diff --git a/queen.go b/queen.go index bc037f1..6a4bb0b 100644 --- a/queen.go +++ b/queen.go @@ -14,6 +14,7 @@ import ( "github.com/ipfs/go-log/v2" "github.com/libp2p/go-libp2p-kad-dht/antslog" kadpb "github.com/libp2p/go-libp2p-kad-dht/pb" + "github.com/libp2p/go-libp2p/core/crypto" "github.com/libp2p/go-libp2p/core/peer" "github.com/libp2p/go-libp2p/core/peerstore" "github.com/libp2p/go-libp2p/p2p/host/peerstore/pstoremem" @@ -185,6 +186,7 @@ func (q *Queen) Run(ctx context.Context) { case <-t.C: q.routine(ctx) case <-ctx.Done(): + q.persistLiveAntsKeys() return } } @@ -277,6 +279,14 @@ func (q *Queen) consumeAntsLogs(ctx context.Context) { } } +func (q *Queen) persistLiveAntsKeys() { + antsKeys := make([]crypto.PrivKey, 0, len(q.ants)) + for _, ant := range q.ants { + antsKeys = append(antsKeys, ant.Host.Peerstore().PrivKey(ant.Host.ID())) + } + q.keysDB.MatchingKeys(nil, antsKeys) +} + func (q *Queen) routine(ctx context.Context) { networkPeers, err := q.nebulaDB.GetLatestPeerIds(ctx) if err != nil { @@ -322,7 +332,9 @@ func (q *Queen) routine(ctx context.Context) { logger.Debugf("removing %d ants", len(excessAntsIndices)) // remove ants - for _, index := range excessAntsIndices { + returnedKeys := make([]crypto.PrivKey, len(excessAntsIndices)) + for i, index := range excessAntsIndices { + returnedKeys[i] = q.ants[index].Host.Peerstore().PrivKey(q.ants[index].Host.ID()) port := q.ants[index].port q.ants[index].Close() q.ants = append(q.ants[:index], q.ants[index+1:]...) @@ -330,7 +342,7 @@ func (q *Queen) routine(ctx context.Context) { } // add missing ants - privKeys := q.keysDB.MatchingKeys(missingKeys) + privKeys := q.keysDB.MatchingKeys(missingKeys, returnedKeys) for _, key := range privKeys { port, err := q.takeAvailablePort() if err != nil { From f5c6b8ce8c8465bf29a6b0ef6c438efda698b1fb Mon Sep 17 00:00:00 2001 From: guillaumemichel Date: Wed, 16 Oct 2024 10:18:19 +0200 Subject: [PATCH 2/3] adding test for persisted returned keys --- keys_test.go | 36 ++++++++++++++++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/keys_test.go b/keys_test.go index 14dd64a..55aef94 100644 --- a/keys_test.go +++ b/keys_test.go @@ -70,7 +70,7 @@ func TestKeysDB(t *testing.T) { db := NewKeysDB(filename) require.NotNil(t, db) - privKeys := db.MatchingKeys(prefixes) + privKeys := db.MatchingKeys(prefixes, nil) // check that the returned private keys match the prefixes require.Len(t, privKeys, len(prefixes)) @@ -81,9 +81,41 @@ func TestKeysDB(t *testing.T) { } // check that the keys are not reused - privKeys2 := db.MatchingKeys(prefixes) + privKeys2 := db.MatchingKeys(prefixes, nil) require.Len(t, privKeys2, len(prefixes)) for _, key := range privKeys { require.NotContains(t, privKeys2, key) } } + +func TestReturnKeys(t *testing.T) { + nReturnedKeys := 2 + filename := "test_keys.db" + prefixDepth := 4 // 2**(prefixDepth) prefixes + defer os.Remove(filename) + + prefixes := genPrefixesForDepth(prefixDepth) + db := NewKeysDB(filename) + require.NotNil(t, db) + privKeys := db.MatchingKeys(prefixes[:len(prefixes)/2], nil) + // check that the provided private keys match the prefixes + require.Len(t, privKeys, len(prefixes)/2) + for i, prefix := range prefixes[:len(prefixes)/2] { + pid, err := peer.IDFromPrivateKey(privKeys[i]) + require.NoError(t, err) + require.Equal(t, prefix.BitLen(), key.CommonPrefixLength(PeeridToKadid(pid), prefix)) + } + // check that the keys are not reused + privKeys2 := db.MatchingKeys(prefixes[:len(prefixes)/2], privKeys[:nReturnedKeys]) + require.Len(t, privKeys2, len(prefixes)/2) + + // test that returned keys are written to the file + keysTrie := db.readKeysFromFile() + for i := range nReturnedKeys { + pid, err := peer.IDFromPrivateKey(privKeys[i]) + require.NoError(t, err) + found, key := trie.Find(keysTrie, PeeridToKadid(pid)) + require.True(t, found) + require.Equal(t, privKeys[i], key) + } +} From 4006e199bc18cc568d571038a82e7479cd13311e Mon Sep 17 00:00:00 2001 From: guillaumemichel Date: Wed, 16 Oct 2024 14:45:01 +0200 Subject: [PATCH 3/3] updated test --- keys_test.go | 37 ++++++++++++------------------------- 1 file changed, 12 insertions(+), 25 deletions(-) diff --git a/keys_test.go b/keys_test.go index 55aef94..bc71514 100644 --- a/keys_test.go +++ b/keys_test.go @@ -88,34 +88,21 @@ func TestKeysDB(t *testing.T) { } } -func TestReturnKeys(t *testing.T) { - nReturnedKeys := 2 +func TestReturnKeysToEmptyTrie(t *testing.T) { filename := "test_keys.db" - prefixDepth := 4 // 2**(prefixDepth) prefixes + db := NewKeysDB(filename) defer os.Remove(filename) - prefixes := genPrefixesForDepth(prefixDepth) - db := NewKeysDB(filename) - require.NotNil(t, db) - privKeys := db.MatchingKeys(prefixes[:len(prefixes)/2], nil) - // check that the provided private keys match the prefixes - require.Len(t, privKeys, len(prefixes)/2) - for i, prefix := range prefixes[:len(prefixes)/2] { - pid, err := peer.IDFromPrivateKey(privKeys[i]) - require.NoError(t, err) - require.Equal(t, prefix.BitLen(), key.CommonPrefixLength(PeeridToKadid(pid), prefix)) - } - // check that the keys are not reused - privKeys2 := db.MatchingKeys(prefixes[:len(prefixes)/2], privKeys[:nReturnedKeys]) - require.Len(t, privKeys2, len(prefixes)/2) + key := genKey() + privKeys := db.MatchingKeys(nil, []crypto.PrivKey{key}) + require.Len(t, privKeys, 0) - // test that returned keys are written to the file keysTrie := db.readKeysFromFile() - for i := range nReturnedKeys { - pid, err := peer.IDFromPrivateKey(privKeys[i]) - require.NoError(t, err) - found, key := trie.Find(keysTrie, PeeridToKadid(pid)) - require.True(t, found) - require.Equal(t, privKeys[i], key) - } + require.Equal(t, 1, keysTrie.Size()) + + pid, err := peer.IDFromPrivateKey(key) + require.NoError(t, err) + found, foundKey := trie.Find(keysTrie, PeeridToKadid(pid)) + require.True(t, found) + require.Equal(t, key, foundKey) }