Merge pull request #6 from probe-lab/key-persistence

feat: key persistence

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)

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,28 @@ 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 TestReturnKeysToEmptyTrie(t *testing.T) {
+	filename := "test_keys.db"
+	db := NewKeysDB(filename)
+	defer os.Remove(filename)
+
+	key := genKey()
+	privKeys := db.MatchingKeys(nil, []crypto.PrivKey{key})
+	require.Len(t, privKeys, 0)
+
+	keysTrie := db.readKeysFromFile()
+	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)
+}

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,15 +332,17 @@ 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:]...)
 		q.freePort(port)
 	}
 
 	// 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 {