diff --git a/slasher/db/BUILD.bazel b/slasher/db/BUILD.bazel index 4a919ac80..2063e01af 100644 --- a/slasher/db/BUILD.bazel +++ b/slasher/db/BUILD.bazel @@ -8,6 +8,7 @@ go_library( "indexed_attestations.go", "schema.go", "setup_db.go", + "validator_id_pubkey.go", ], importpath = "github.com/prysmaticlabs/prysm/slasher/db", visibility = ["//slasher:__subpackages__"], @@ -29,6 +30,7 @@ go_test( "block_header_test.go", "indexed_attestations_test.go", "setup_db_test.go", + "validator_id_pubkey_test.go", ], embed = [":go_default_library"], deps = ["//proto/eth/v1alpha1:go_default_library"], diff --git a/slasher/db/db.go b/slasher/db/db.go index e9fd22b71..36f497f75 100644 --- a/slasher/db/db.go +++ b/slasher/db/db.go @@ -85,6 +85,7 @@ func NewKVStore(dirPath string) (*Store, error) { historicIndexedAttestationsBucket, historicBlockHeadersBucket, indexedAttestationsIndicesBucket, + validatorsPublicKeysBucket, ) }); err != nil { return nil, err diff --git a/slasher/db/schema.go b/slasher/db/schema.go index e16a9307d..cede6407b 100644 --- a/slasher/db/schema.go +++ b/slasher/db/schema.go @@ -9,6 +9,7 @@ var ( historicIndexedAttestationsBucket = []byte("historic-indexed-attestations-bucket") historicBlockHeadersBucket = []byte("historic-block-headers-bucket") indexedAttestationsIndicesBucket = []byte("indexed-attestations-indices-bucket") + validatorsPublicKeysBucket = []byte("validators-public-keys-bucket") ) func encodeEpochValidatorID(epoch uint64, validatorID uint64) []byte { diff --git a/slasher/db/validator_id_pubkey.go b/slasher/db/validator_id_pubkey.go new file mode 100644 index 000000000..93b7c152f --- /dev/null +++ b/slasher/db/validator_id_pubkey.go @@ -0,0 +1,44 @@ +package db + +import ( + "github.com/boltdb/bolt" + "github.com/pkg/errors" + "github.com/prysmaticlabs/prysm/shared/bytesutil" +) + +// ValidatorPubKey accepts validator id and returns the corresponding pubkey. +// Returns nil if the pubkey for this validator id does not exist. +func (db *Store) ValidatorPubKey(validatorID uint64) ([]byte, error) { + var pk []byte + err := db.view(func(tx *bolt.Tx) error { + b := tx.Bucket(validatorsPublicKeysBucket) + pk = b.Get(bytesutil.Bytes4(validatorID)) + return nil + }) + return pk, err +} + +// SavePubKey accepts a validator id and its public key and writes it to disk. +func (db *Store) SavePubKey(validatorID uint64, pubKey []byte) error { + err := db.update(func(tx *bolt.Tx) error { + bucket := tx.Bucket(validatorsPublicKeysBucket) + key := bytesutil.Bytes4(validatorID) + if err := bucket.Put(key, pubKey); err != nil { + return errors.Wrap(err, "failed to add validator public key to slasher db.") + } + return nil + }) + return err +} + +// DeletePubKey deletes a public key of a validator id. +func (db *Store) DeletePubKey(validatorID uint64) error { + return db.update(func(tx *bolt.Tx) error { + bucket := tx.Bucket(validatorsPublicKeysBucket) + key := bytesutil.Bytes4(validatorID) + if err := bucket.Delete(key); err != nil { + return errors.Wrap(err, "failed to delete public key from validators public key bucket") + } + return bucket.Delete(key) + }) +} diff --git a/slasher/db/validator_id_pubkey_test.go b/slasher/db/validator_id_pubkey_test.go new file mode 100644 index 000000000..311110260 --- /dev/null +++ b/slasher/db/validator_id_pubkey_test.go @@ -0,0 +1,105 @@ +package db + +import ( + "bytes" + "testing" +) + +type publicKeyTestStruct struct { + validatorID uint64 + pk []byte +} + +var pkTests []publicKeyTestStruct + +func init() { + pkTests = []publicKeyTestStruct{ + { + validatorID: 1, + pk: []byte{1, 2, 3}, + }, + { + validatorID: 2, + pk: []byte{4, 5, 6}, + }, + { + validatorID: 3, + pk: []byte{7, 8, 9}, + }, + } +} + +func TestNilDBValidatorPublicKey(t *testing.T) { + db := SetupSlasherDB(t) + defer TeardownSlasherDB(t, db) + + validatorID := uint64(1) + + pk, err := db.ValidatorPubKey(validatorID) + if err != nil { + t.Fatal("nil ValidatorPubKey should not return error") + } + if pk != nil { + t.Fatal("ValidatorPubKey should return nil") + } + +} + +func TestSavePubKey(t *testing.T) { + db := SetupSlasherDB(t) + defer TeardownSlasherDB(t, db) + + for _, tt := range pkTests { + err := db.SavePubKey(tt.validatorID, tt.pk) + if err != nil { + t.Fatalf("save validator public key failed: %v", err) + } + + pk, err := db.ValidatorPubKey(tt.validatorID) + if err != nil { + t.Fatalf("failed to get validator public key: %v", err) + } + + if pk == nil || !bytes.Equal(pk, tt.pk) { + t.Fatalf("get should return validator public key: %v", tt.pk) + } + } + +} + +func TestDeletePublicKey(t *testing.T) { + db := SetupSlasherDB(t) + defer TeardownSlasherDB(t, db) + + for _, tt := range pkTests { + + err := db.SavePubKey(tt.validatorID, tt.pk) + if err != nil { + t.Fatalf("save validator public key failed: %v", err) + } + } + + for _, tt := range pkTests { + pk, err := db.ValidatorPubKey(tt.validatorID) + if err != nil { + t.Fatalf("failed to get validator public key: %v", err) + } + + if pk == nil || !bytes.Equal(pk, tt.pk) { + t.Fatalf("get should return validator public key: %v", pk) + } + err = db.DeletePubKey(tt.validatorID) + if err != nil { + t.Fatalf("delete validator public key: %v", err) + } + pk, err = db.ValidatorPubKey(tt.validatorID) + if err != nil { + t.Fatal(err) + } + if pk != nil { + t.Errorf("Expected validator public key to be deleted, received: %v", pk) + } + + } + +}