From f79ca706452464763f7e7b135adbc06e282ac894 Mon Sep 17 00:00:00 2001
From: thattommyhall <thattommyhall@gmail.com>
Date: Sat, 24 Jul 2021 01:31:48 +0100
Subject: [PATCH] Allow deterministic key generation from seed

---
 README.md      |  4 ++++
 go.mod         |  1 +
 idgen/idgen.go | 23 ++++++++++++++++++++---
 main.go        | 18 ++++++++++++++++++
 4 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 769eed8..eca0be0 100644
--- a/README.md
+++ b/README.md
@@ -102,6 +102,8 @@ Usage of hydra-booster:
         Specify the DHT protocol prefix (default "/ipfs") (default "/ipfs")
   -pstore string
         Peerstore directory for LevelDB store (defaults to in-memory store)
+  -random-seed string
+        Seed to use to generate IDs (useful if you want to have persistent IDs). Should be Base64 encoded and 256bits
   -stagger duration
         Duration to stagger nodes starts by
   -ui-theme string
@@ -135,6 +137,8 @@ Alternatively, some flags can be set via environment variables. Note that flags
         Specify the number of Hydra heads to create. (default -1)
   HYDRA_PORT_BEGIN int
         If set, begin port allocation here (default -1)
+  HYDRA_RANDOM_SEED string
+        Seed to use to generate IDs (useful if you want to have persistent IDs). Should be Base64 encoded and 256bits   
 ```
 
 ### Best Practices
diff --git a/go.mod b/go.mod
index 564cffa..dd4a951 100644
--- a/go.mod
+++ b/go.mod
@@ -33,6 +33,7 @@ require (
 	github.com/prometheus/client_golang v1.7.1
 	github.com/whyrusleeping/timecache v0.0.0-20160911033111-cfcb2f1abfee
 	go.opencensus.io v0.22.5
+	golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9 // indirect
 )
 
 go 1.14
diff --git a/idgen/idgen.go b/idgen/idgen.go
index 843a12a..ec4e0e4 100644
--- a/idgen/idgen.go
+++ b/idgen/idgen.go
@@ -1,13 +1,17 @@
 package idgen
 
 import (
+	"crypto/rand"
+	"crypto/sha256"
 	"fmt"
+	"io"
 	"math/bits"
 	"sync"
 
 	"github.com/libp2p/go-libp2p-core/crypto"
 	"github.com/libp2p/go-libp2p-core/peer"
 	kbucket "github.com/libp2p/go-libp2p-kbucket"
+	"golang.org/x/crypto/hkdf"
 )
 
 // HydraIdentityGenerator is a shared balanced ID generator.
@@ -34,12 +38,25 @@ type BalancedIdentityGenerator struct {
 	sync.Mutex
 	xorTrie *XorTrie
 	count   int
+	reader  io.Reader
 }
 
 // NewBalancedIdentityGenerator creates a new balanced identity generator.
 func NewBalancedIdentityGenerator() *BalancedIdentityGenerator {
 	return &BalancedIdentityGenerator{
 		xorTrie: NewXorTrie(),
+		reader:  rand.Reader,
+	}
+}
+
+func NewBalancedIdentityGeneratorFromSeed(seed []byte) *BalancedIdentityGenerator {
+	hash := sha256.New
+	info := []byte("hydra keys")
+	hkdf := hkdf.New(hash, seed, nil, info)
+
+	return &BalancedIdentityGenerator{
+		xorTrie: NewXorTrie(),
+		reader:  hkdf,
 	}
 }
 
@@ -85,7 +102,7 @@ func (bg *BalancedIdentityGenerator) AddBalanced() (crypto.PrivKey, error) {
 
 func (bg *BalancedIdentityGenerator) genUniqueID() (privKey crypto.PrivKey, trieKey TrieKey, depth int, err error) {
 	for {
-		if privKey, trieKey, err = genID(); err != nil {
+		if privKey, trieKey, err = bg.genID(); err != nil {
 			return nil, nil, 0, err
 		}
 		if depth, ok := bg.xorTrie.Insert(trieKey); ok {
@@ -121,8 +138,8 @@ func (bg *BalancedIdentityGenerator) Depth() int {
 	return bg.xorTrie.Depth()
 }
 
-func genID() (crypto.PrivKey, TrieKey, error) {
-	privKey, _, err := crypto.GenerateKeyPair(crypto.Ed25519, 0)
+func (bg *BalancedIdentityGenerator) genID() (crypto.PrivKey, TrieKey, error) {
+	privKey, _, err := crypto.GenerateKeyPairWithReader(crypto.Ed25519, 0, bg.reader)
 	if err != nil {
 		return nil, nil, fmt.Errorf("generating private key for trie, %w", err)
 	}
diff --git a/main.go b/main.go
index 1fcda6d..c6e86bb 100644
--- a/main.go
+++ b/main.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"encoding/base64"
 	"flag"
 	"fmt"
 	"log"
@@ -35,6 +36,7 @@ const (
 func main() {
 	start := time.Now()
 	nheads := flag.Int("nheads", -1, "Specify the number of Hydra heads to create.")
+	randomSeed := flag.String("random-seed", "", "Seed to use to generate IDs (useful if you want to have persistent IDs). Should be Base64 encoded and 256bits")
 	dbpath := flag.String("db", "", "Datastore directory (for LevelDB store) or postgresql:// connection URI (for PostgreSQL store)")
 	pstorePath := flag.String("pstore", "", "Peerstore directory for LevelDB store (defaults to in-memory store)")
 	httpAPIAddr := flag.String("httpapi-addr", defaultHTTPAPIAddr, "Specify an IP and port to run the HTTP API server on")
@@ -71,6 +73,9 @@ func main() {
 	if *nheads == -1 {
 		*nheads = mustGetEnvInt("HYDRA_NHEADS", 1)
 	}
+	if *randomSeed == "" {
+		*randomSeed = os.Getenv("HYDRA_RANDOM_SEED")
+	}
 	if *portBegin == -1 {
 		*portBegin = mustGetEnvInt("HYDRA_PORT_BEGIN", 0)
 	}
@@ -110,6 +115,19 @@ func main() {
 	defer cancel()
 
 	var idGenerator idgen.IdentityGenerator
+	if *randomSeed != "" && *idgenAddr != "" {
+		log.Fatalln("error: Should not set both idgen-addr and random-seed")
+	}
+	if *randomSeed != "" {
+		seed, err := base64.StdEncoding.DecodeString(*randomSeed)
+		if err != nil {
+			log.Fatalln("error: Could not base64 decode seed")
+		}
+		if len(seed) != 32 {
+			log.Fatalln("error: Seed should be 256bit in base64")
+		}
+		idGenerator = idgen.NewBalancedIdentityGeneratorFromSeed(seed)
+	}
 	if *idgenAddr != "" {
 		dg := idgen.NewCleaningIDGenerator(idgen.NewDelegatedIDGenerator(*idgenAddr))
 		defer func() {