Update cluster pk manager to assign multiple keys to validators (#2112)

* checkpoint

* move lock

* checkpoint

* checkpoint

* chkpt

* update readme so i know what im supposed to do

* ckpt

* checkpoint

* checkpoint

* chkpt

* Fix image build

* handle errors

* add port to pod

* lint

* Update beacon-config.config.yaml

* lint

tools/cluster-pk-manager/README.md

@@ -13,13 +13,13 @@ pods.
 
 Workflow for bootstraping a validator pod
 
-1. Request a private key from the pk manager.
-1. If an unallocated private key exists (from previously terminated pod), assign
+1. Request `n` private keys from the pk manager.
+1. If unallocated private keys exist (from previously terminated pods), assign
    to the requesting pod.
-1. If all available private keys are in use, generate a new private key, and
-   make the deposit on behalf of this newly generated private key.
-1. Write the assignment to some persistent datastore and fulfill the request.
-1. The validator uses this private key to act as a deposited validator in the
+1. If there are not at least `n` keys not in use, generate new private keys, 
+   and make the deposits on behalf of these newly generated private keys.
+1. Write the key allocations to a persistent datastore and fulfill the request.
+1. The client uses these private keys to act as deposited validators in the
    system. 
 
 ## Server
@@ -27,6 +27,40 @@ Workflow for bootstraping a validator pod
 The server manages the private key database, allocates new private keys, makes
 validator deposits, and fulfills requests from pods for private key allocation.
 
+### Database structure
+
+There are two buckets for the server, unallocated keys and allocated keys.
+
+Unallocated keys bucket:
+
+| key         | value |
+|-------------|-------|
+| private key | nil   |
+
+Allocated keys bucket:
+
+| key      | value                |
+|----------|----------------------|
+| pod name | list of private keys |
+
+### Key management design
+
+There are two types of operations with regards to private keys: 
+
+- Allocate(podName, keys)
+- UnallocateAllKeys(podName)
+
+Allocating keys will first check and attempt to recycle existing, unused keys.
+If there are no unused keys available (or not enough), new keys are deposited.
+
+Unallocating keys happens when a pod is destroyed. This should return all of 
+that's pods' keys to the unallocated keys bucket.
+
+### Assignments HTTP Page `/assignments`
+
+The server exposes an HTTP page which maps pod names to public keys.
+This may be useful for determining which logs to follow for a given validator.
+
 ## Client
 
 The client makes the private key request with a given pod name and generates a 

tools/cluster-pk-manager/client/main.go

@@ -17,6 +17,7 @@ var (
 	podName     = flag.String("pod-name", "", "The name of the pod running this tool")
 	keystoreDir = flag.String("keystore-dir", "", "The directory to generate keystore with received validator key")
 	password    = flag.String("keystore-password", "", "The password to unlock the new keystore")
+	numKeys     = flag.Uint64("keys", 1, "The number of keys to request")
 )
 
 func main() {
@@ -34,25 +35,30 @@ func main() {
 
 	client := pb.NewPrivateKeyServiceClient(conn)
 
-	resp, err := client.Request(ctx, &pb.PrivateKeyRequest{PodName: *podName})
+	resp, err := client.Request(ctx, &pb.PrivateKeyRequest{
+		PodName:      *podName,
+		NumberOfKeys: *numKeys,
+	})
 	if err != nil {
 		panic(err)
 	}
 
-	pk, err := bls.SecretKeyFromBytes(resp.PrivateKey)
-	if err != nil {
-		panic(err)
-	}
+	for i, privateKey := range resp.PrivateKeys.PrivateKeys {
+		pk, err := bls.SecretKeyFromBytes(privateKey)
+		if err != nil {
+			panic(err)
+		}
 
-	k := &keystore.Key{
-		PublicKey: pk.PublicKey(),
-		SecretKey: pk,
-	}
+		k := &keystore.Key{
+			PublicKey: pk.PublicKey(),
+			SecretKey: pk,
+		}
 
-	validatorKeyFile := *keystoreDir + params.BeaconConfig().ValidatorPrivkeyFileName
-	if err := store.StoreKey(validatorKeyFile, k, *password); err != nil {
-		panic(err)
-	}
+		validatorKeyFile := *keystoreDir + params.BeaconConfig().ValidatorPrivkeyFileName + "-" + fmt.Sprintf("%d", i)
+		if err := store.StoreKey(validatorKeyFile, k, *password); err != nil {
+			panic(err)
+		}
 
-	fmt.Printf("New key written to %s\n", validatorKeyFile)
+		fmt.Printf("New key written to %s\n", validatorKeyFile)
+	}
 }

tools/cluster-pk-manager/server/BUILD.bazel

@@ -5,6 +5,7 @@ load("@io_bazel_rules_docker//container:container.bzl", "container_push")
 go_library(
     name = "go_default_library",
     srcs = [
+        "allocations.go",
         "db.go",
         "logger.go",
         "main.go",
@@ -16,6 +17,7 @@ go_library(
     deps = [
         "//contracts/deposit-contract:go_default_library",
         "//proto/cluster:go_default_library",
+        "//shared/bls:go_default_library",
         "//shared/keystore:go_default_library",
         "//shared/prometheus:go_default_library",
         "//shared/ssz:go_default_library",
@@ -25,6 +27,7 @@ go_library(
         "@com_github_ethereum_go_ethereum//crypto:go_default_library",
         "@com_github_ethereum_go_ethereum//ethclient:go_default_library",
         "@com_github_ethereum_go_ethereum//rpc:go_default_library",
+        "@com_github_gogo_protobuf//proto:go_default_library",
         "@com_github_prometheus_client_golang//prometheus:go_default_library",
         "@com_github_prometheus_client_golang//prometheus/promauto:go_default_library",
         "@com_github_sirupsen_logrus//:go_default_library",
@@ -52,6 +55,7 @@ go_test(
 go_image(
     name = "image",
     srcs = [
+        "allocations.go",
         "db.go",
         "logger.go",
         "main.go",
@@ -63,12 +67,14 @@ go_image(
     deps = [
         "//contracts/deposit-contract:go_default_library",
         "//proto/cluster:go_default_library",
+        "//shared/bls:go_default_library",
         "//shared/keystore:go_default_library",
         "//shared/prometheus:go_default_library",
         "//shared/ssz:go_default_library",
         "@com_github_boltdb_bolt//:go_default_library",
         "@com_github_ethereum_go_ethereum//accounts/abi/bind:go_default_library",
         "@com_github_ethereum_go_ethereum//common:go_default_library",
+        "@com_github_gogo_protobuf//proto:go_default_library",
         "@com_github_ethereum_go_ethereum//crypto:go_default_library",
         "@com_github_ethereum_go_ethereum//ethclient:go_default_library",
         "@com_github_ethereum_go_ethereum//rpc:go_default_library",

tools/cluster-pk-manager/server/allocations.go

@@ -0,0 +1,37 @@
+package main
+
+import (
+	"fmt"
+	"io"
+	"net/http"
+)
+
+func (s *server) serveAllocationsHTTPPage() {
+	mux := http.NewServeMux()
+
+	mux.HandleFunc("/allocations", func(w http.ResponseWriter, _ *http.Request) {
+		res := ""
+		a, err := s.db.Allocations()
+		if err != nil {
+			w.WriteHeader(http.StatusInternalServerError)
+			if _, err := io.WriteString(w, err.Error()); err != nil {
+				log.Error(err)
+			}
+			return
+		}
+		for podName, pubkeys := range a {
+			for _, pk := range pubkeys {
+				res += fmt.Sprintf("%s=%#x\n", podName, pk)
+			}
+		}
+		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
+		w.WriteHeader(http.StatusOK)
+		if _, err := io.WriteString(w, res); err != nil {
+			log.Error(err)
+		}
+	})
+
+	srv := &http.Server{Addr: ":8080", Handler: mux}
+	go srv.ListenAndServe()
+	log.Info("Serving allocations page at :8080/allocations")
+}

tools/cluster-pk-manager/server/db.go

@@ -3,13 +3,15 @@ package main
 import (
 	"bytes"
 	"context"
-	"errors"
 	"path"
 	"time"
 
 	"github.com/boltdb/bolt"
+	"github.com/gogo/protobuf/proto"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
+	pb "github.com/prysmaticlabs/prysm/proto/cluster"
+	"github.com/prysmaticlabs/prysm/shared/bls"
 	"github.com/prysmaticlabs/prysm/shared/keystore"
 )
 
@@ -55,8 +57,20 @@ func newDB(dbPath string) *db {
 	}
 
 	if err := boltdb.View(func(tx *bolt.Tx) error {
-		keys := tx.Bucket(assignedPkBucket).Stats().KeyN
+		keys := 0
+
+		// Iterate over all of the pod assigned keys (one to many).
+		c := tx.Bucket(assignedPkBucket).Cursor()
+		for k, v := c.First(); k != nil; k, v = c.Next() {
+			pks := &pb.PrivateKeys{}
+			if err := proto.Unmarshal(v, pks); err != nil {
+				return err
+			}
+			keys += len(pks.PrivateKeys)
+		}
 		assignedPkCount.Set(float64(keys))
+
+		// Add the unassigned keys count (one to one).
 		keys += tx.Bucket(unassignedPkBucket).Stats().KeyN
 		allocatedPkCount.Add(float64(keys))
 		return nil
@@ -67,33 +81,36 @@ func newDB(dbPath string) *db {
 	return &db{db: boltdb}
 }
 
-// UnallocatedPK returns the first unassigned private key, if any are
-// available.
-func (d *db) UnallocatedPK(_ context.Context) ([]byte, error) {
-	var pk []byte
+// UnallocatedPKs returns unassigned private keys, if any are available.
+func (d *db) UnallocatedPKs(_ context.Context, numKeys uint64) (*pb.PrivateKeys, error) {
+	pks := &pb.PrivateKeys{}
 	if err := d.db.View(func(tx *bolt.Tx) error {
 		c := tx.Bucket(unassignedPkBucket).Cursor()
-		k, _ := c.First()
+		i := uint64(0)
+		for k, _ := c.First(); k != nil && i < numKeys; k, _ = c.Next() {
+			pks.PrivateKeys = append(pks.PrivateKeys, k)
+			i++
+		}
 
-		pk = k
 		return nil
 	}); err != nil {
 		return nil, err
 	}
-	return pk, nil
+	return pks, nil
 }
 
 // PodPK returns an assigned private key to the given pod name, if one exists.
-func (d *db) PodPK(_ context.Context, podName string) ([]byte, error) {
-	var pk []byte
+func (d *db) PodPKs(_ context.Context, podName string) (*pb.PrivateKeys, error) {
+	pks := &pb.PrivateKeys{}
 	if err := d.db.View(func(tx *bolt.Tx) error {
-		pk = tx.Bucket(assignedPkBucket).Get([]byte(podName))
-		return nil
+		b := tx.Bucket(assignedPkBucket).Get([]byte(podName))
+
+		return proto.Unmarshal(b, pks)
 	}); err != nil {
 		return nil, err
 	}
 
-	return pk, nil
+	return pks, nil
 }
 
 // AllocateNewPkToPod records new private key assignment in DB.
@@ -129,17 +146,31 @@ func (d *db) RemovePKAssignment(_ context.Context, podName string) error {
 	})
 }
 
-// AssignExistingPK assigns a PK from the unassigned bucket to a given pod.
-func (d *db) AssignExistingPK(_ context.Context, pk []byte, podName string) error {
+// AssignExistingPKs assigns a PK from the unassigned bucket to a given pod.
+func (d *db) AssignExistingPKs(_ context.Context, pks *pb.PrivateKeys, podName string) error {
 	return d.db.Update(func(tx *bolt.Tx) error {
-		if !bytes.Equal(tx.Bucket(unassignedPkBucket).Get(pk), dummyVal) {
-			return errors.New("private key not in unassigned bucket")
+		for _, pk := range pks.PrivateKeys {
+			if bytes.Equal(tx.Bucket(unassignedPkBucket).Get(pk), dummyVal) {
+				if err := tx.Bucket(unassignedPkBucket).Delete(pk); err != nil {
+					return err
+				}
+			}
 		}
-		if err := tx.Bucket(unassignedPkBucket).Delete(pk); err != nil {
+		assignedPkCount.Add(float64(len(pks.PrivateKeys)))
+
+		// If pod assignment exists, append to it.
+		if existing := tx.Bucket(assignedPkBucket).Get([]byte(podName)); existing != nil {
+			existingKeys := &pb.PrivateKeys{}
+			if err := proto.Unmarshal(existing, existingKeys); err != nil {
+				pks.PrivateKeys = append(pks.PrivateKeys, existingKeys.PrivateKeys...)
+			}
+		}
+
+		data, err := proto.Marshal(pks)
+		if err != nil {
 			return err
 		}
-		assignedPkCount.Inc()
-		return tx.Bucket(assignedPkBucket).Put([]byte(podName), pk)
+		return tx.Bucket(assignedPkBucket).Put([]byte(podName), data)
 	})
 
 	return nil
@@ -159,3 +190,32 @@ func (d *db) AllocatedPodNames(_ context.Context) ([]string, error) {
 	}
 	return podNames, nil
 }
+
+func (d *db) Allocations() (map[string][][]byte, error) {
+	m := make(map[string][][]byte)
+	if err := d.db.View(func(tx *bolt.Tx) error {
+		return tx.Bucket(assignedPkBucket).ForEach(func(k, v []byte) error {
+			pks := &pb.PrivateKeys{}
+			if err := proto.Unmarshal(v, pks); err != nil {
+				return err
+			}
+			pubkeys := make([][]byte, len(pks.PrivateKeys))
+			for i, pk := range pks.PrivateKeys {
+				k, err := bls.SecretKeyFromBytes(pk)
+				if err != nil {
+					return err
+				}
+
+				pubkeys[i] = k.PublicKey().Marshal()
+			}
+			m[string(k)] = pubkeys
+
+			return nil
+		})
+	}); err != nil {
+		// do something
+		return nil, err
+	}
+
+	return m, nil
+}

tools/cluster-pk-manager/server/main.go

@@ -15,7 +15,7 @@ var (
 	port                = flag.Int("port", 8000, "The port to server gRPC")
 	metricsPort         = flag.Int("metrics-port", 9090, "The port to serve /metrics")
 	privateKey          = flag.String("private-key", "", "The private key of funder")
-	rpcPath             = flag.String("rpc", "", "RPC address of a running ETH1 node")
+	rpcPath             = flag.String("rpc", "https://goerli.prylabs.net", "RPC address of a running ETH1 node")
 	depositContractAddr = flag.String("deposit-contract", "", "Address of the deposit contract")
 	depositAmount       = flag.Int64("deposit-amount", 0, "The amount of wei to deposit into the contract")
 	dbPath              = flag.String("db-path", "", "The file path for database storage")
@@ -40,6 +40,7 @@ func main() {
 	pb.RegisterPrivateKeyServiceServer(s, srv)
 
 	go prometheus.RunSimpleServerOrDie(fmt.Sprintf(":%d", *metricsPort))
+	srv.serveAllocationsHTTPPage()
 
 	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", *port))
 	if err != nil {

tools/cluster-pk-manager/server/server.go

@@ -67,9 +67,6 @@ func newServer(
 }
 
 func (s *server) makeDeposit(data []byte) error {
-	s.clientLock.Lock()
-	defer s.clientLock.Unlock()
-
 	txOps := bind.NewKeyedTransactor(s.txPk)
 	txOps.Value = s.depositAmount
 	txOps.GasLimit = gasLimit
@@ -83,57 +80,89 @@ func (s *server) makeDeposit(data []byte) error {
 }
 
 func (s *server) Request(ctx context.Context, req *pb.PrivateKeyRequest) (*pb.PrivateKeyResponse, error) {
-	pk, err := s.db.PodPK(ctx, req.PodName)
-	if err != nil {
-		return nil, err
-	}
-	if pk != nil {
-		log.WithField("pod", req.PodName).Debug("Returning existing assignment")
-		return &pb.PrivateKeyResponse{PrivateKey: pk}, nil
+	s.clientLock.Lock()
+	defer s.clientLock.Unlock()
+
+	if req.NumberOfKeys == 0 {
+		req.NumberOfKeys = 1
 	}
 
-	pk, err = s.db.UnallocatedPK(ctx)
+	// build the list of PKs in the following order, until the requested
+	// amount is ready to return.
+	// - PKs already assigned to the pod
+	// - PKs that have not yet been allocated
+	// - PKs that are newly initialized with deposits
+
+	pks, err := s.db.PodPKs(ctx, req.PodName)
 	if err != nil {
 		return nil, err
 	}
-	if pk != nil {
-		log.WithField("pod", req.PodName).Debug("Recycling existing private key")
-		if err := s.db.AssignExistingPK(ctx, pk, req.PodName); err != nil {
+	if pks != nil && len(pks.PrivateKeys) > 0 {
+		log.WithField("pod", req.PodName).Debug("Returning existing assignment(s)")
+		return &pb.PrivateKeyResponse{
+			PrivateKeys: pks,
+		}, nil
+	}
+
+	unallocated, err := s.db.UnallocatedPKs(ctx, req.NumberOfKeys)
+	if err != nil {
+		return nil, err
+	}
+	log.WithField(
+		"pod", req.PodName,
+	).WithField(
+		"keys", len(unallocated.PrivateKeys),
+	).Debug("Recycling existing private key(s)")
+
+	pks.PrivateKeys = append(pks.PrivateKeys, unallocated.PrivateKeys...)
+
+	if len(pks.PrivateKeys) < int(req.NumberOfKeys) {
+		c := int(req.NumberOfKeys) - len(pks.PrivateKeys)
+		newKeys, err := s.allocateNewKeys(ctx, req.PodName, c)
+		if err != nil {
 			return nil, err
 		}
-		return &pb.PrivateKeyResponse{PrivateKey: pk}, nil
+		pks.PrivateKeys = append(pks.PrivateKeys, newKeys.PrivateKeys...)
 	}
 
-	log.WithField("pod", req.PodName).Debug("Allocating a new private key")
-	return s.allocateNewKey(ctx, req.PodName)
+	if err := s.db.AssignExistingPKs(ctx, pks, req.PodName); err != nil {
+		return nil, err
+	}
+
+	return &pb.PrivateKeyResponse{PrivateKeys: pks}, nil
 }
 
-func (s *server) allocateNewKey(ctx context.Context, podName string) (*pb.PrivateKeyResponse, error) {
-	key, err := keystore.NewKey(rand.Reader)
-	if err != nil {
-		return nil, err
+func (s *server) allocateNewKeys(ctx context.Context, podName string, numKeys int) (*pb.PrivateKeys, error) {
+	pks := make([][]byte, numKeys)
+
+	for i := 0; i < numKeys; i++ {
+		key, err := keystore.NewKey(rand.Reader)
+		if err != nil {
+			return nil, err
+		}
+
+		// Make the validator deposit
+		// NOTE: This uses the validator key as the withdrawal key
+		di, err := keystore.DepositInput(key /*depositKey*/, key /*withdrawalKey*/)
+		if err != nil {
+			return nil, err
+		}
+		serializedData := new(bytes.Buffer)
+		if err := ssz.Encode(serializedData, di); err != nil {
+			return nil, fmt.Errorf("could not serialize deposit data: %v", err)
+		}
+
+		// Do the actual deposit
+		if err := s.makeDeposit(serializedData.Bytes()); err != nil {
+			return nil, err
+		}
+		// Store in database
+		if err := s.db.AllocateNewPkToPod(ctx, key, podName); err != nil {
+			return nil, err
+		}
+		secret := key.SecretKey.Marshal()
+		pks[i] = secret
 	}
 
-	// Make the validator deposit
-	// NOTE: This uses the validator key as the withdrawal key
-	di, err := keystore.DepositInput(key /*depositKey*/, key /*withdrawalKey*/)
-	if err != nil {
-		return nil, err
-	}
-	serializedData := new(bytes.Buffer)
-	if err := ssz.Encode(serializedData, di); err != nil {
-		return nil, fmt.Errorf("could not serialize deposit data: %v", err)
-	}
-
-	// Do the actual deposit
-	if err := s.makeDeposit(serializedData.Bytes()); err != nil {
-		return nil, err
-	}
-	// Store in database
-	if err := s.db.AllocateNewPkToPod(ctx, key, podName); err != nil {
-		return nil, err
-	}
-
-	return &pb.PrivateKeyResponse{PrivateKey: key.SecretKey.Marshal()}, nil
-
+	return &pb.PrivateKeys{PrivateKeys: pks}, nil
 }