prysm-pulse/validator/keymanager/v2/direct/migrate_test.go

104 lines
3.8 KiB
Go

package direct
import (
"context"
"encoding/json"
"fmt"
"strings"
"testing"
"github.com/google/uuid"
"github.com/prysmaticlabs/prysm/shared/bls"
"github.com/prysmaticlabs/prysm/shared/petnames"
"github.com/prysmaticlabs/prysm/shared/testutil"
"github.com/prysmaticlabs/prysm/shared/testutil/assert"
"github.com/prysmaticlabs/prysm/shared/testutil/require"
mock "github.com/prysmaticlabs/prysm/validator/accounts/v2/testing"
v2keymanager "github.com/prysmaticlabs/prysm/validator/keymanager/v2"
logTest "github.com/sirupsen/logrus/hooks/test"
keystorev4 "github.com/wealdtech/go-eth2-wallet-encryptor-keystorev4"
)
func TestDirectKeymanager_MigrateToSingleKeystoreFormat(t *testing.T) {
hook := logTest.NewGlobal()
password := "secretPassw0rd$1999"
wallet := &mock.Wallet{
Files: make(map[string]map[string][]byte),
AccountPasswords: make(map[string]string),
WalletPassword: password,
}
dr := &Keymanager{
keysCache: make(map[[48]byte]bls.SecretKey),
wallet: wallet,
accountsStore: &AccountStore{},
}
ctx := context.Background()
// Generate several old account keystore and save to the wallet path.
numAccounts := 5
wallet.Directories = make([]string, numAccounts)
wantedValidatingKeys := make([]bls.SecretKey, numAccounts)
for i := 0; i < numAccounts; i++ {
validatingKey := bls.RandKey()
accountName, keystore := generateOldAccountKeystore(t, dr, validatingKey, password)
wallet.Directories[i] = accountName
wantedValidatingKeys[i] = validatingKey
encodedKeystore, err := json.MarshalIndent(keystore, "", "\t")
require.NoError(t, err)
require.NoError(t, dr.wallet.WritePasswordToDisk(ctx, accountName+".pass", password))
require.NoError(t, dr.wallet.WriteFileAtPath(ctx, accountName, KeystoreFileName, encodedKeystore))
}
// Now, we run the migration strategy.
require.NoError(t, dr.migrateToSingleKeystore(ctx))
// We retrieve the new accounts keystore format containing all keys in a single file.
var encodedAccountsFile []byte
for k, v := range wallet.Files[AccountsPath] {
if strings.Contains(k, "keystore") {
encodedAccountsFile = v
}
}
require.NotNil(t, encodedAccountsFile, "could not find keystore file")
accountsKeystore := &v2keymanager.Keystore{}
require.NoError(t, json.Unmarshal(encodedAccountsFile, accountsKeystore))
// We extract the accounts from the keystore.
decryptor := keystorev4.New()
encodedAccounts, err := decryptor.Decrypt(accountsKeystore.Crypto, password)
require.NoError(t, err, "Could not decrypt validator accounts")
store := &AccountStore{}
require.NoError(t, json.Unmarshal(encodedAccounts, store))
// We expect the migration strategy to have succeeded, with all accounts from earlier
// now being stored in a single keystore file.
require.Equal(t, numAccounts, len(store.PublicKeys))
require.Equal(t, numAccounts, len(store.PrivateKeys))
for i := 0; i < numAccounts; i++ {
privKey, err := bls.SecretKeyFromBytes(store.PrivateKeys[i])
require.NoError(t, err)
assert.DeepEqual(t, privKey.Marshal(), wantedValidatingKeys[i].Marshal())
}
testutil.AssertLogsContain(t, hook, "Now migrating accounts to a more efficient format")
}
func generateOldAccountKeystore(
t testing.TB, dr *Keymanager, validatingKey bls.SecretKey, password string,
) (string, *v2keymanager.Keystore) {
accountName := petnames.DeterministicName(validatingKey.PublicKey().Marshal(), "-")
// Generates a new EIP-2335 compliant keystore file
// from a BLS private key and marshals it as JSON.
encryptor := keystorev4.New()
cryptoFields, err := encryptor.Encrypt(validatingKey.Marshal(), password)
require.NoError(t, err)
id, err := uuid.NewRandom()
require.NoError(t, err)
return accountName, &v2keymanager.Keystore{
Crypto: cryptoFields,
ID: id.String(),
Pubkey: fmt.Sprintf("%x", validatingKey.PublicKey().Marshal()),
Version: encryptor.Version(),
Name: encryptor.Name(),
}
}