Allow additional/override of non-spec fields, including "address"

Signed-off-by: Peter Broadhurst <peter.broadhurst@kaleido.io>

pkg/keystorev3/pbkdf2.go

@@ -29,11 +29,12 @@ const (
 	prfHmacSHA256 = "hmac-sha256"
 )
 
-func readPbkdf2WalletFile(jsonWallet []byte, password []byte) (WalletFile, error) {
+func readPbkdf2WalletFile(jsonWallet []byte, password []byte, metadata map[string]interface{}) (WalletFile, error) {
 	var w *walletFilePbkdf2
 	if err := json.Unmarshal(jsonWallet, &w); err != nil {
 		return nil, fmt.Errorf("invalid pbkdf2 keystore: %s", err)
 	}
+	w.metadata = metadata
 	return w, w.decrypt(password)
 }
 

pkg/keystorev3/pbkdf2_test.go

@@ -42,9 +42,13 @@ func TestPbkdf2Wallet(t *testing.T) {
 
 	w1 := &walletFilePbkdf2{
 		walletFileBase: walletFileBase{
-			Address: ethtypes.AddressPlainHex(keypair.Address),
-			ID:      fftypes.NewUUID(),
-			Version: version3,
+			walletFileCoreFields: walletFileCoreFields{
+				ID:      fftypes.NewUUID(),
+				Version: version3,
+			},
+			walletFileMetadata: walletFileMetadata{
+				Address: ethtypes.AddressPlainHex(keypair.Address),
+			},
 			keypair: keypair,
 		},
 		Crypto: cryptoPbkdf2{
@@ -78,14 +82,14 @@ func TestPbkdf2Wallet(t *testing.T) {
 
 func TestPbkdf2WalletFileDecryptInvalid(t *testing.T) {
 
-	_, err := readPbkdf2WalletFile([]byte(`!! not json`), []byte(""))
+	_, err := readPbkdf2WalletFile([]byte(`!! not json`), []byte(""), nil)
 	assert.Regexp(t, "invalid pbkdf2 keystore", err)
 
 }
 
 func TestPbkdf2WalletFileUnsupportedPRF(t *testing.T) {
 
-	_, err := readPbkdf2WalletFile([]byte(`{}`), []byte(""))
+	_, err := readPbkdf2WalletFile([]byte(`{}`), []byte(""), nil)
 	assert.Regexp(t, "invalid pbkdf2 wallet file: unsupported prf", err)
 
 }

pkg/keystorev3/scrypt.go

@@ -29,11 +29,12 @@ import (
 
 const defaultR = 8
 
-func readScryptWalletFile(jsonWallet []byte, password []byte) (WalletFile, error) {
+func readScryptWalletFile(jsonWallet []byte, password []byte, metadata map[string]interface{}) (WalletFile, error) {
 	var w *walletFileScrypt
 	if err := json.Unmarshal(jsonWallet, &w); err != nil {
 		return nil, fmt.Errorf("invalid scrypt wallet file: %s", err)
 	}
+	w.metadata = metadata
 	return w, w.decrypt(password)
 }
 
@@ -75,9 +76,14 @@ func newScryptWalletFileBytes(password string, privateKey []byte, addr ethtypes.
 
 	return &walletFileScrypt{
 		walletFileBase: walletFileBase{
-			ID:         fftypes.NewUUID(),
-			Address:    addr,
-			Version:    version3,
+			walletFileCoreFields: walletFileCoreFields{
+				ID:      fftypes.NewUUID(),
+				Version: version3,
+			},
+			walletFileMetadata: walletFileMetadata{
+				Address:  addr,
+				metadata: map[string]interface{}{},
+			},
 			privateKey: privateKey,
 		},
 		Crypto: cryptoScrypt{

pkg/keystorev3/scrypt_test.go

@@ -58,7 +58,7 @@ func TestScryptWalletRoundTripStandard(t *testing.T) {
 
 func TestScryptReadInvalidFile(t *testing.T) {
 
-	_, err := readScryptWalletFile([]byte(`!bad JSON`), []byte(""))
+	_, err := readScryptWalletFile([]byte(`!bad JSON`), []byte(""), nil)
 	assert.Error(t, err)
 
 }

pkg/keystorev3/wallet.go

@@ -58,7 +58,11 @@ func NewWalletFileCustomBytesStandard(password string, privateKey []byte) Wallet
 
 func ReadWalletFile(jsonWallet []byte, password []byte) (WalletFile, error) {
 	var w walletFileCommon
-	if err := json.Unmarshal(jsonWallet, &w); err != nil {
+	err := json.Unmarshal(jsonWallet, &w)
+	if err == nil {
+		err = json.Unmarshal(jsonWallet, &w.metadata)
+	}
+	if err != nil {
 		return nil, fmt.Errorf("invalid wallet file: %s", err)
 	}
 	if w.ID == nil {
@@ -69,9 +73,9 @@ func ReadWalletFile(jsonWallet []byte, password []byte) (WalletFile, error) {
 	}
 	switch w.Crypto.KDF {
 	case kdfTypeScrypt:
-		return readScryptWalletFile(jsonWallet, password)
+		return readScryptWalletFile(jsonWallet, password, w.metadata)
 	case kdfTypePbkdf2:
-		return readPbkdf2WalletFile(jsonWallet, password)
+		return readPbkdf2WalletFile(jsonWallet, password, w.metadata)
 	default:
 		return nil, fmt.Errorf("unsupported kdf: %s", w.Crypto.KDF)
 	}

pkg/keystorev3/wallet_test.go

@@ -18,6 +18,7 @@ package keystorev3
 
 import (
 	"encoding/hex"
+	"encoding/json"
 	"fmt"
 	"testing"
 	"testing/iotest"
@@ -164,3 +165,33 @@ func TestWalletFileCustomBytesLight(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, kp.Address, w2.KeyPair().Address)
 }
+
+func TestMarshalWalletJSONFail(t *testing.T) {
+	_, err := marshalWalletJSON(&walletFileBase{}, map[bool]bool{false: true})
+	assert.Error(t, err)
+}
+
+func TestWalletFileCustomBytesUnsetAddress(t *testing.T) {
+	customBytes := ([]byte)("something deterministic for testing")
+
+	w := NewWalletFileCustomBytesLight("correcthorsebatterystaple", customBytes)
+
+	w.Metadata()["address"] = nil
+	w.Metadata()["myKeyIdentifier"] = "something I know works for me"
+	w.Metadata()["id"] = "attempting to set this does not work"
+	w.Metadata()["version"] = 42
+
+	jsonBytes, err := json.Marshal(w)
+	assert.NoError(t, err)
+
+	var roundTripBackFromJSON map[string]interface{}
+	err = json.Unmarshal(jsonBytes, &roundTripBackFromJSON)
+	assert.NoError(t, err)
+
+	_, hasAddress := roundTripBackFromJSON["address"]
+	assert.False(t, hasAddress)
+	assert.Equal(t, "something I know works for me", roundTripBackFromJSON["myKeyIdentifier"])
+	assert.Equal(t, float64(w.GetVersion()), roundTripBackFromJSON["version"])
+	assert.Equal(t, w.GetID().String(), roundTripBackFromJSON["id"])
+
+}

pkg/keystorev3/walletfile.go

@@ -37,6 +37,16 @@ type WalletFile interface {
 	PrivateKey() []byte
 	KeyPair() *secp256k1.KeyPair
 	JSON() []byte
+	GetID() *fftypes.UUID
+	GetVersion() int
+
+	// Any fields set into this that do not conflict with the base fields (id/version/crypto) will
+	// be serialized into the JSON when it is marshalled.
+	// This includes setting the "address" field (which is not a core part of the V3 standard) to
+	// an arbitrary string, adding new fields for different key identifiers (like "bjj" or "btc" for
+	// different public key compression algos).
+	// If you want to remove the address field completely, simple set "address": nil in the map.
+	Metadata() map[string]interface{}
 }
 
 type kdfParamsScrypt struct {
@@ -76,11 +86,21 @@ type cryptoPbkdf2 struct {
 	KDFParams kdfParamsPbkdf2 `json:"kdfparams"`
 }
 
-type walletFileBase struct {
+type walletFileCoreFields struct {
+	ID      *fftypes.UUID `json:"id"`
+	Version int           `json:"version"`
+}
+
+type walletFileMetadata struct {
+	// address is not technically part of keystorev3 syntax, and note this can be overridden/removed by callers of the package
 	Address ethtypes.AddressPlainHex `json:"address"`
-	ID      *fftypes.UUID            `json:"id"`
-	Version int                      `json:"version"`
+	// arbitrary additional fields that can be stored in the JSON, including overriding/removing the "address" field (other core fields cannot be overridden)
+	metadata map[string]interface{}
+}
 
+type walletFileBase struct {
+	walletFileCoreFields
+	walletFileMetadata
 	privateKey []byte
 	keypair    *secp256k1.KeyPair
 }
@@ -95,11 +115,53 @@ type walletFilePbkdf2 struct {
 	Crypto cryptoPbkdf2 `json:"crypto"`
 }
 
+func (w *walletFilePbkdf2) MarshalJSON() ([]byte, error) {
+	return marshalWalletJSON(&w.walletFileBase, w.Crypto)
+}
+
 type walletFileScrypt struct {
 	walletFileBase
 	Crypto cryptoScrypt `json:"crypto"`
 }
 
+func (w *walletFileScrypt) MarshalJSON() ([]byte, error) {
+	return marshalWalletJSON(&w.walletFileBase, w.Crypto)
+}
+
+func (w *walletFileBase) GetVersion() int {
+	return w.Version
+}
+
+func (w *walletFileBase) GetID() *fftypes.UUID {
+	return w.ID
+}
+
+func (w *walletFileBase) Metadata() map[string]interface{} {
+	return w.metadata
+}
+
+func marshalWalletJSON(wc *walletFileBase, crypto interface{}) ([]byte, error) {
+	cryptoJSON, err := json.Marshal(crypto)
+	if err != nil {
+		return nil, err
+	}
+	jsonMap := map[string]interface{}{}
+	// note address can be set to "nil" to remove it entirely
+	jsonMap["address"] = wc.Address
+	for k, v := range wc.metadata {
+		if v == nil {
+			delete(jsonMap, k)
+		} else {
+			jsonMap[k] = v
+		}
+	}
+	// cannot override these fields
+	jsonMap["id"] = wc.ID
+	jsonMap["version"] = wc.Version
+	jsonMap["crypto"] = json.RawMessage(cryptoJSON)
+	return json.Marshal(jsonMap)
+}
+
 func (w *walletFileBase) KeyPair() *secp256k1.KeyPair {
 	return w.keypair
 }