Skip to content

Commit

Permalink
shared/keystore: cleanup (#5689)
Browse files Browse the repository at this point in the history 
* fixes inconsistent naming

* removes unused methods

* proper temp keystore cleanup in tests

* more naming fixes

* one more fix of keystore teardown in tests

* more robust teardown
  • Loading branch information
@farazdagi
farazdagi authored Apr 30, 2020
1 parent 5f53f9f commit cc07494
Show file tree
Hide file tree
Showing 4 changed files with 71 additions and 95 deletions.
8 changes: 4 additions & 4 deletions shared/keystore/key.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,11 +62,11 @@ type Key struct {
}

type keyStore interface {
// Loads and decrypts the key from disk.
// GetKey loads and decrypts the key from disk.
GetKey(filename string, password string) (*Key, error)
// Writes and encrypts the key.
// StoreKey writes and encrypts the key.
StoreKey(filename string, k *Key, auth string) error
// Joins filename with the key directory unless it is already absolute.
// JoinPath joins filename with the key directory unless it is already absolute.
JoinPath(filename string) string
}

Expand Down Expand Up @@ -95,7 +95,7 @@ type cipherparamsJSON struct {
IV string `json:"iv"`
}

// MarshalJSON marshalls a key struct into a JSON blob.
// MarshalJSON marshals a key struct into a JSON blob.
func (k *Key) MarshalJSON() (j []byte, err error) {
jStruct := plainKeyJSON{
hex.EncodeToString(k.PublicKey.Marshal()),
Expand Down
24 changes: 7 additions & 17 deletions shared/keystore/key_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,11 @@ package keystore
import (
"bytes"
"io/ioutil"
"os"
"testing"

"github.com/pborman/uuid"
"github.com/prysmaticlabs/prysm/shared/bls"
"github.com/prysmaticlabs/prysm/shared/bytesutil"
"github.com/prysmaticlabs/prysm/shared/testutil"
)

func TestMarshalAndUnmarshal(t *testing.T) {
Expand Down Expand Up @@ -42,21 +40,17 @@ func TestMarshalAndUnmarshal(t *testing.T) {
}

func TestStoreRandomKey(t *testing.T) {
tmpdir := testutil.TempDir()
filedir := tmpdir + "/keystore"
tempDir, teardown := setupTempKeystoreDir(t)
defer teardown()
ks := &Store{
keysDirPath: filedir,
keysDirPath: tempDir,
scryptN: LightScryptN,
scryptP: LightScryptP,
}

if err := storeNewRandomKey(ks, "password"); err != nil {
t.Fatalf("storage of random key unsuccessful %v", err)
}

if err := os.RemoveAll(filedir); err != nil {
t.Errorf("unable to remove temporary files %v", err)
}
}

func TestNewKeyFromBLS(t *testing.T) {
Expand All @@ -83,26 +77,22 @@ func TestNewKeyFromBLS(t *testing.T) {
}

func TestWriteFile(t *testing.T) {
tmpdir := testutil.TempDir()
filedir := tmpdir + "/keystore"
tempDir, teardown := setupTempKeystoreDir(t)
defer teardown()

testKeystore := []byte{'t', 'e', 's', 't'}

err := writeKeyFile(filedir, testKeystore)
err := writeKeyFile(tempDir, testKeystore)
if err != nil {
t.Fatalf("unable to write file %v", err)
}

keystore, err := ioutil.ReadFile(filedir)
keystore, err := ioutil.ReadFile(tempDir)
if err != nil {
t.Fatalf("unable to retrieve file %v", err)
}

if !bytes.Equal(keystore, testKeystore) {
t.Fatalf("retrieved keystore is not the same %v", keystore)
}

if err := os.RemoveAll(filedir); err != nil {
t.Errorf("unable to remove temporary files %v", err)
}
}
43 changes: 13 additions & 30 deletions shared/keystore/keystore.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,17 +52,6 @@ type Store struct {
scryptP int
}

// RetrievePubKey retrieves the public key from the keystore.
func RetrievePubKey(directory string, password string) (*bls.PublicKey, error) {
ks := Store{
keysDirPath: directory,
scryptN: StandardScryptN,
scryptP: StandardScryptP,
}
key, err := ks.GetKey(ks.keysDirPath, password)
return key.PublicKey, err
}

// NewKeystore from a directory.
func NewKeystore(directory string) Store {
return Store{
Expand All @@ -76,16 +65,16 @@ func NewKeystore(directory string) Store {
func (ks Store) GetKey(filename, password string) (*Key, error) {
// Load the key from the keystore and decrypt its contents
// #nosec G304
keyjson, err := ioutil.ReadFile(filename)
keyJSON, err := ioutil.ReadFile(filename)
if err != nil {
return nil, err
}
return DecryptKey(keyjson, password)
return DecryptKey(keyJSON, password)
}

// GetKeys from directory using the prefix to filter relevant files
// and a decryption password.
func (ks Store) GetKeys(directory, fileprefix, password string, warnOnFail bool) (map[string]*Key, error) {
func (ks Store) GetKeys(directory, filePrefix, password string, warnOnFail bool) (map[string]*Key, error) {
// Load the key from the keystore and decrypt its contents
// #nosec G304
files, err := ioutil.ReadDir(directory)
Expand All @@ -106,17 +95,17 @@ func (ks Store) GetKeys(directory, fileprefix, password string, warnOnFail bool)
}
}
}
cp := strings.Contains(n, strings.TrimPrefix(fileprefix, "/"))
cp := strings.Contains(n, strings.TrimPrefix(filePrefix, "/"))
if f.Mode().IsRegular() && cp {
// #nosec G304
keyjson, err := ioutil.ReadFile(filePath)
keyJSON, err := ioutil.ReadFile(filePath)
if err != nil {
return nil, err
}
key, err := DecryptKey(keyjson, password)
key, err := DecryptKey(keyJSON, password)
if err != nil {
if warnOnFail {
log.WithError(err).WithField("keyfile", string(keyjson)).Warn("Failed to decrypt key")
log.WithError(err).WithField("keyfile", string(keyJSON)).Warn("Failed to decrypt key")
}
continue
}
Expand All @@ -128,11 +117,11 @@ func (ks Store) GetKeys(directory, fileprefix, password string, warnOnFail bool)

// StoreKey in filepath and encrypt it with a password.
func (ks Store) StoreKey(filename string, key *Key, auth string) error {
keyjson, err := EncryptKey(key, auth, ks.scryptN, ks.scryptP)
keyJSON, err := EncryptKey(key, auth, ks.scryptN, ks.scryptP)
if err != nil {
return err
}
return writeKeyFile(filename, keyjson)
return writeKeyFile(filename, keyJSON)
}

// JoinPath joins the filename with the keystore directory path.
Expand All @@ -143,13 +132,7 @@ func (ks Store) JoinPath(filename string) string {
return filepath.Join(ks.keysDirPath, filename)
}

// StoreRandomKey generates a key, encrypts with 'auth' and stores in the given directory
func StoreRandomKey(dir, password string, scryptN, scryptP int) error {
err := storeNewRandomKey(Store{dir, scryptN, scryptP}, password)
return err
}

// EncryptKey encrypts a key using the specified scrypt parameters into a json
// EncryptKey encrypts a key using the specified scrypt parameters into a JSON
// blob that can be decrypted later on.
func EncryptKey(key *Key, password string, scryptN, scryptP int) ([]byte, error) {
authArray := []byte(password)
Expand Down Expand Up @@ -205,13 +188,13 @@ func EncryptKey(key *Key, password string, scryptN, scryptP int) ([]byte, error)
return json.Marshal(encryptedJSON)
}

// DecryptKey decrypts a key from a json blob, returning the private key itself.
func DecryptKey(keyjson []byte, password string) (*Key, error) {
// DecryptKey decrypts a key from a JSON blob, returning the private key itself.
func DecryptKey(keyJSON []byte, password string) (*Key, error) {
var keyBytes, keyID []byte
var err error

k := new(encryptedKeyJSON)
if err := json.Unmarshal(keyjson, k); err != nil {
if err := json.Unmarshal(keyJSON, k); err != nil {
return nil, err
}

Expand Down
91 changes: 47 additions & 44 deletions shared/keystore/keystore_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,11 @@ package keystore

import (
"bytes"
"crypto/rand"
"fmt"
"math/big"
"os"
"path"
"testing"

"github.com/pborman/uuid"
Expand All @@ -12,10 +16,10 @@ import (
)

func TestStoreAndGetKey(t *testing.T) {
tmpdir := testutil.TempDir()
filedir := tmpdir + "/keystore"
tempDir, teardown := setupTempKeystoreDir(t)
defer teardown()
ks := &Store{
keysDirPath: filedir,
keysDirPath: tempDir,
scryptN: LightScryptN,
scryptP: LightScryptP,
}
Expand All @@ -25,29 +29,25 @@ func TestStoreAndGetKey(t *testing.T) {
t.Fatalf("key generation failed %v", err)
}

if err := ks.StoreKey(filedir, key, "password"); err != nil {
if err := ks.StoreKey(tempDir, key, "password"); err != nil {
t.Fatalf("unable to store key %v", err)
}

newkey, err := ks.GetKey(filedir, "password")
decryptedKey, err := ks.GetKey(tempDir, "password")
if err != nil {
t.Fatalf("unable to get key %v", err)
}

if !bytes.Equal(newkey.SecretKey.Marshal(), key.SecretKey.Marshal()) {
t.Fatalf("retrieved secret keys are not equal %v , %v", newkey.SecretKey.Marshal(), key.SecretKey.Marshal())
}

if err := os.RemoveAll(filedir); err != nil {
t.Errorf("unable to remove temporary files %v", err)
if !bytes.Equal(decryptedKey.SecretKey.Marshal(), key.SecretKey.Marshal()) {
t.Fatalf("retrieved secret keys are not equal %v , %v", decryptedKey.SecretKey.Marshal(), key.SecretKey.Marshal())
}
}

func TestStoreAndGetKeys(t *testing.T) {
tmpdir := testutil.TempDir()
filePrefix := "/keystore"
tempDir, teardown := setupTempKeystoreDir(t)
defer teardown()
ks := &Store{
keysDirPath: tmpdir,
keysDirPath: tempDir,
scryptN: LightScryptN,
scryptP: LightScryptP,
}
Expand All @@ -57,33 +57,26 @@ func TestStoreAndGetKeys(t *testing.T) {
t.Fatalf("key generation failed %v", err)
}

if err := ks.StoreKey(tmpdir+filePrefix+"/test-1", key, "password"); err != nil {
if err := ks.StoreKey(tempDir+"/test-1", key, "password"); err != nil {
t.Fatalf("unable to store key %v", err)
}
key2, err := NewKey()
if err != nil {
t.Fatalf("key generation failed %v", err)
}
if err := ks.StoreKey(tmpdir+filePrefix+"/test-2", key2, "password"); err != nil {
if err := ks.StoreKey(tempDir+"/test-2", key2, "password"); err != nil {
t.Fatalf("unable to store key %v", err)
}
newkeys, err := ks.GetKeys(tmpdir+filePrefix, "test", "password", false)
decryptedKeys, err := ks.GetKeys(tempDir, "test", "password", false)
if err != nil {
t.Fatalf("unable to get key %v", err)
}
for _, s := range newkeys {
for _, s := range decryptedKeys {
if !bytes.Equal(s.SecretKey.Marshal(), key.SecretKey.Marshal()) && !bytes.Equal(s.SecretKey.Marshal(), key2.SecretKey.Marshal()) {
t.Fatalf("retrieved secret keys are not equal %v ", s.SecretKey.Marshal())
}

}

if err := os.RemoveAll(tmpdir + filePrefix + "-2"); err != nil {
t.Errorf("unable to remove temporary files %v", err)
}
if err := os.RemoveAll(tmpdir + filePrefix + "-1"); err != nil {
t.Errorf("unable to remove temporary files %v", err)
}
}

func TestEncryptDecryptKey(t *testing.T) {
Expand All @@ -102,34 +95,29 @@ func TestEncryptDecryptKey(t *testing.T) {
PublicKey: pk.PublicKey(),
}

keyjson, err := EncryptKey(key, password, LightScryptN, LightScryptP)
keyJSON, err := EncryptKey(key, password, LightScryptN, LightScryptP)
if err != nil {
t.Fatalf("unable to encrypt key %v", err)
}

newkey, err := DecryptKey(keyjson, password)
decryptedKey, err := DecryptKey(keyJSON, password)
if err != nil {
t.Fatalf("unable to decrypt keystore %v", err)
}

if !bytes.Equal(newkey.ID, newID) {
t.Fatalf("decrypted key's uuid doesn't match %v", newkey.ID)
if !bytes.Equal(decryptedKey.ID, newID) {
t.Fatalf("decrypted key's uuid doesn't match %v", decryptedKey.ID)
}

expected := pk.Marshal()
if !bytes.Equal(newkey.SecretKey.Marshal(), expected) {
t.Fatalf("decrypted key's value is not equal %v", newkey.SecretKey.Marshal())
if !bytes.Equal(decryptedKey.SecretKey.Marshal(), expected) {
t.Fatalf("decrypted key's value is not equal %v", decryptedKey.SecretKey.Marshal())
}

}

func TestGetSymlinkedKeys(t *testing.T) {
tmpdir := testutil.TempDir() + "/symlinked-keystore"
defer func() {
if err := os.RemoveAll(tmpdir); err != nil {
t.Logf("unable to remove temporary files: %v", err)
}
}()
tempDir, teardown := setupTempKeystoreDir(t)
defer teardown()
ks := &Store{
scryptN: LightScryptN,
scryptP: LightScryptP,
Expand All @@ -140,24 +128,39 @@ func TestGetSymlinkedKeys(t *testing.T) {
t.Fatalf("key generation failed %v", err)
}

if err := ks.StoreKey(tmpdir+"/files/test-1", key, "password"); err != nil {
if err := ks.StoreKey(tempDir+"/files/test-1", key, "password"); err != nil {
t.Fatalf("unable to store key %v", err)
}

if err := os.Symlink(tmpdir+"/files/test-1", tmpdir+"/test-1"); err != nil {
if err := os.Symlink(tempDir+"/files/test-1", tempDir+"/test-1"); err != nil {
t.Fatalf("unable to create symlink: %v", err)
}

newkeys, err := ks.GetKeys(tmpdir, "test", "password", false)
decryptedKeys, err := ks.GetKeys(tempDir, "test", "password", false)
if err != nil {
t.Fatalf("unable to get key %v", err)
}
if len(newkeys) != 1 {
t.Errorf("unexpected number of keys returned, want: %d, got: %d", 1, len(newkeys))
if len(decryptedKeys) != 1 {
t.Errorf("unexpected number of keys returned, want: %d, got: %d", 1, len(decryptedKeys))
}
for _, s := range newkeys {
for _, s := range decryptedKeys {
if !bytes.Equal(s.SecretKey.Marshal(), key.SecretKey.Marshal()) {
t.Fatalf("retrieved secret keys are not equal %v ", s.SecretKey.Marshal())
}
}
}

// setupTempKeystoreDir creates temporary directory for storing keystore files.
func setupTempKeystoreDir(t *testing.T) (string, func()) {
randPath, err := rand.Int(rand.Reader, big.NewInt(1000000))
if err != nil {
t.Fatalf("could not generate random file path: %v", err)
}
tempDir := path.Join(testutil.TempDir(), fmt.Sprintf("%d", randPath), "keystore")

return tempDir, func() {
if err := os.RemoveAll(tempDir); err != nil {
t.Logf("unable to remove temporary files: %v", err)
}
}
}

0 comments on commit cc07494

Please sign in to comment.