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

pkg/keystorev3/pbkdf2.go

@@ -21,19 +21,19 @@ import (
 	"encoding/json"
 	"fmt"
 
-	"github.com/hyperledger/firefly-signer/pkg/secp256k1"
 	"golang.org/x/crypto/pbkdf2"
 )
 
 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)
 }
 
@@ -45,9 +45,6 @@ func (w *walletFilePbkdf2) decrypt(password []byte) (err error) {
 	derivedKey := pbkdf2.Key(password, w.Crypto.KDFParams.Salt, w.Crypto.KDFParams.C, w.Crypto.KDFParams.DKLen, sha256.New)
 
 	w.privateKey, err = w.Crypto.decryptCommon(derivedKey)
-	if err == nil {
-		w.keypair, err = secp256k1.NewSecp256k1KeyPair(w.privateKey)
-	}
 	return err
 
 }

pkg/keystorev3/pbkdf2_test.go

@@ -42,10 +42,15 @@ func TestPbkdf2Wallet(t *testing.T) {
 
 	w1 := &walletFilePbkdf2{
 		walletFileBase: walletFileBase{
-			Address: ethtypes.AddressPlainHex(keypair.Address),
-			ID:      fftypes.NewUUID(),
-			Version: version3,
-			keypair: keypair,
+			walletFileCoreFields: walletFileCoreFields{
+				ID:      fftypes.NewUUID(),
+				Version: version3,
+			},
+			walletFileMetadata: walletFileMetadata{
+				metadata: map[string]interface{}{
+					"address": ethtypes.AddressPlainHex(keypair.Address).String(),
+				},
+			},
 		},
 		Crypto: cryptoPbkdf2{
 			cryptoCommon: cryptoCommon{
@@ -78,14 +83,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)
 }
 
@@ -46,14 +47,14 @@ func mustGenerateDerivedScryptKey(password string, salt []byte, n, p int) []byte
 }
 
 // creates an ethereum address wallet file
-func newScryptWalletFile(password string, keypair *secp256k1.KeyPair, n int, p int) WalletFile {
-	wf := newScryptWalletFileBytes(password, keypair.PrivateKeyBytes(), ethtypes.AddressPlainHex(keypair.Address), n, p)
-	wf.keypair = keypair
+func newScryptWalletFileSecp256k1(password string, keypair *secp256k1.KeyPair, n int, p int) WalletFile {
+	wf := newScryptWalletFileBytes(password, keypair.PrivateKeyBytes(), n, p)
+	wf.Metadata()["address"] = ethtypes.AddressPlainHex(keypair.Address).String()
 	return wf
 }
 
 // this allows creation of any size/type of key in the store
-func newScryptWalletFileBytes(password string, privateKey []byte, addr ethtypes.AddressPlainHex, n int, p int) *walletFileScrypt {
+func newScryptWalletFileBytes(password string, privateKey []byte, n int, p int) *walletFileScrypt {
 
 	// Generate a sale for the scrypt
 	salt := mustReadBytes(32, rand.Reader)
@@ -75,9 +76,13 @@ 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{
+				metadata: map[string]interface{}{},
+			},
 			privateKey: privateKey,
 		},
 		Crypto: cryptoScrypt{
@@ -107,8 +112,5 @@ func (w *walletFileScrypt) decrypt(password []byte) error {
 		return fmt.Errorf("invalid scrypt keystore: %s", err)
 	}
 	w.privateKey, err = w.Crypto.decryptCommon(derivedKey)
-	if err == nil {
-		w.keypair, err = secp256k1.NewSecp256k1KeyPair(w.privateKey)
-	}
 	return err
 }

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

@@ -21,7 +21,6 @@ import (
 	"fmt"
 	"io"
 
-	"github.com/hyperledger/firefly-signer/pkg/ethtypes"
 	"github.com/hyperledger/firefly-signer/pkg/secp256k1"
 	"golang.org/x/crypto/sha3"
 )
@@ -33,32 +32,28 @@ const (
 )
 
 func NewWalletFileLight(password string, keypair *secp256k1.KeyPair) WalletFile {
-	return newScryptWalletFile(password, keypair, nLight, pDefault)
+	return newScryptWalletFileSecp256k1(password, keypair, nLight, pDefault)
 }
 
 func NewWalletFileStandard(password string, keypair *secp256k1.KeyPair) WalletFile {
-	return newScryptWalletFile(password, keypair, nStandard, pDefault)
-}
-
-func addressFirst32(privateKey []byte) ethtypes.AddressPlainHex {
-	if len(privateKey) > 32 {
-		privateKey = privateKey[0:32]
-	}
-	kp, _ := secp256k1.NewSecp256k1KeyPair(privateKey)
-	return ethtypes.AddressPlainHex(kp.Address)
+	return newScryptWalletFileSecp256k1(password, keypair, nStandard, pDefault)
 }
 
 func NewWalletFileCustomBytesLight(password string, privateKey []byte) WalletFile {
-	return newScryptWalletFileBytes(password, privateKey, addressFirst32(privateKey), nStandard, pDefault)
+	return newScryptWalletFileBytes(password, privateKey, nStandard, pDefault)
 }
 
 func NewWalletFileCustomBytesStandard(password string, privateKey []byte) WalletFile {
-	return newScryptWalletFileBytes(password, privateKey, addressFirst32(privateKey), nStandard, pDefault)
+	return newScryptWalletFileBytes(password, privateKey, nStandard, pDefault)
 }
 
 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 +64,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,13 +86,20 @@ type cryptoPbkdf2 struct {
 	KDFParams kdfParamsPbkdf2 `json:"kdfparams"`
 }
 
-type walletFileBase struct {
-	Address ethtypes.AddressPlainHex `json:"address"`
-	ID      *fftypes.UUID            `json:"id"`
-	Version int                      `json:"version"`
+type walletFileCoreFields struct {
+	ID      *fftypes.UUID `json:"id"`
+	Version int           `json:"version"`
+}
+
+type walletFileMetadata struct {
+	// 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
 }
 
 type walletFileCommon struct {
@@ -95,13 +112,51 @@ 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{}{}
+	for k, v := range wc.metadata {
+		if v != nil {
+			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
+	return secp256k1.KeyPairFromBytes(w.privateKey)
 }
 
 func (w *walletFileBase) PrivateKey() []byte {

pkg/secp256k1/keypair.go

@@ -1,4 +1,4 @@
-// Copyright © 2022 Kaleido, Inc.
+// Copyright © 2024 Kaleido, Inc.
 //
 // SPDX-License-Identifier: Apache-2.0
 //
@@ -43,11 +43,17 @@ func GenerateSecp256k1KeyPair() (*KeyPair, error) {
 	return wrapSecp256k1Key(key, key.PubKey()), nil
 }
 
+// Note there is no error condition returned by this function (use KeyPairFromBytes)
 func NewSecp256k1KeyPair(b []byte) (*KeyPair, error) {
 	key, pubKey := btcec.PrivKeyFromBytes(b)
 	return wrapSecp256k1Key(key, pubKey), nil
 }
 
+func KeyPairFromBytes(b []byte) *KeyPair {
+	key, pubKey := btcec.PrivKeyFromBytes(b)
+	return wrapSecp256k1Key(key, pubKey)
+}
+
 func wrapSecp256k1Key(key *btcec.PrivateKey, pubKey *btcec.PublicKey) *KeyPair {
 	return &KeyPair{
 		PrivateKey: key,