feat: support of key's metadata for localkms

Signed-off-by: Volodymyr Kit <[email protected]>

kms/localkms/localkms.go

@@ -106,7 +106,7 @@ func (l *LocalKMS) Create(kt kmsapi.KeyType, opts ...kmsapi.KeyOpts) (string, in
 		return "", nil, fmt.Errorf("create: failed to create new keyset handle: %w", err)
 	}
 
-	keyID, err := l.storeKeySet(kh, kt)
+	keyID, err := l.storeKeySet(kh, kt, opts...)
 	if err != nil {
 		return "", nil, fmt.Errorf("create: failed to store keyset: %w", err)
 	}
@@ -122,6 +122,15 @@ func (l *LocalKMS) Get(keyID string) (interface{}, error) {
 	return l.getKeySet(keyID)
 }
 
+// GetWithMetadata key handle for the given keyID
+// Returns:
+//   - handle instance (to private key)
+//   - metadata if any saved
+//   - error if failure
+func (l *LocalKMS) GetWithMetadata(keyID string) (any, map[string]any, error) {
+	return l.getKeySetWithMetadata(keyID)
+}
+
 // Rotate a key referenced by keyID and return a new handle of a keyset including old key and
 // new key with type kt. It also returns the updated keyID as the first return value
 // Returns:
@@ -164,7 +173,7 @@ func (l *LocalKMS) Rotate(kt kmsapi.KeyType, keyID string, opts ...kmsapi.KeyOpt
 	return newID, updatedKH, nil
 }
 
-func (l *LocalKMS) storeKeySet(kh *keyset.Handle, kt kmsapi.KeyType) (string, error) {
+func (l *LocalKMS) storeKeySet(kh *keyset.Handle, kt kmsapi.KeyType, opts ...kmsapi.KeyOpts) (string, error) {
 	var (
 		kid string
 		err error
@@ -192,13 +201,19 @@ func (l *LocalKMS) storeKeySet(kh *keyset.Handle, kt kmsapi.KeyType) (string, er
 		return "", fmt.Errorf("storeKeySet: failed to write json key to buffer: %w", err)
 	}
 
+	keyOpts := kmsapi.NewKeyOpt()
+
+	for _, opt := range opts {
+		opt(keyOpts)
+	}
+
 	// asymmetric keys are JWK thumbprints of the public key, base64URL encoded stored in kid.
 	// symmetric keys will have a randomly generated key ID (where kid is empty)
 	if kid != "" {
-		return writeToStore(l.store, buf, kmsapi.WithKeyID(kid))
+		return writeToStore(l.store, buf, kmsapi.WithKeyID(kid), kmsapi.ImportWithMetadata(keyOpts.Metadata()))
 	}
 
-	return writeToStore(l.store, buf)
+	return writeToStore(l.store, buf, kmsapi.ImportWithMetadata(keyOpts.Metadata()))
 }
 
 func writeToStore(store kmsapi.Store, buf *bytes.Buffer, opts ...kmsapi.PrivateKeyOpts) (string, error) {
@@ -228,6 +243,21 @@ func (l *LocalKMS) getKeySet(id string) (*keyset.Handle, error) {
 	return kh, nil
 }
 
+func (l *LocalKMS) getKeySetWithMetadata(id string) (*keyset.Handle, map[string]any, error) {
+	localDBReader := newReader(l.store, id, kmsapi.ExportWithMetadata(true))
+
+	jsonKeysetReader := keyset.NewJSONReader(localDBReader)
+
+	// Read reads the encrypted keyset handle back from the io.reader implementation
+	// and decrypts it using primaryKeyEnvAEAD.
+	kh, err := keyset.Read(jsonKeysetReader, l.primaryKeyEnvAEAD)
+	if err != nil {
+		return nil, nil, fmt.Errorf("getKeySet: failed to read json keyset from reader: %w", err)
+	}
+
+	return kh, localDBReader.metadata, nil
+}
+
 // ExportPubKeyBytes will fetch a key referenced by id then gets its public key in raw bytes and returns it.
 // The key must be an asymmetric key.
 // Returns:

kms/localkms/localkms_reader.go

@@ -11,38 +11,72 @@ import (
 	"fmt"
 
 	"github.com/trustbloc/kms-go/spi/kms"
+	kmsapi "github.com/trustbloc/kms-go/spi/kms"
 )
 
 // newReader will create a new local storage storeReader of a keyset with ID value = keysetID
 // it is used internally by local kms.
-func newReader(store kms.Store, keysetID string) *storeReader {
+func newReader(store kms.Store, keysetID string, opts ...kmsapi.ExportKeyOpts) *storeReader {
+	pOpts := kmsapi.NewExportOpt()
+
+	for _, opt := range opts {
+		opt(pOpts)
+	}
+
 	return &storeReader{
-		storage:  store,
-		keysetID: keysetID,
+		storage:     store,
+		keysetID:    keysetID,
+		getMetadata: pOpts.GetMetadata(),
 	}
 }
 
 // storeReader struct to load a keyset from a local storage.
 type storeReader struct {
-	buf      *bytes.Buffer
-	storage  kms.Store
-	keysetID string
+	buf         *bytes.Buffer
+	storage     kms.Store
+	keysetID    string
+	getMetadata bool
+	metadata    map[string]any
 }
 
 // Read the keyset from local storage into p.
 func (l *storeReader) Read(p []byte) (int, error) {
-	if l.buf == nil {
-		if l.keysetID == "" {
-			return 0, fmt.Errorf("keysetID is not set")
-		}
+	if l.buf != nil {
+		return l.buf.Read(p)
+	}
+
+	if l.keysetID == "" {
+		return 0, fmt.Errorf("keysetID is not set")
+	}
+
+	var data []byte
 
-		data, err := l.storage.Get(l.keysetID)
-		if err != nil {
-			return 0, fmt.Errorf("cannot read data for keysetID %s: %w", l.keysetID, err)
+	var err error
+
+	var metadata map[string]any
+
+	if l.getMetadata {
+		metadataStorage, ok := l.storage.(kmsapi.StoreWithMetadata)
+		if !ok {
+			return 0, fmt.Errorf("requested to get 'metadata', but storage doesn't support it")
 		}
 
-		l.buf = bytes.NewBuffer(data)
+		data, metadata, err = metadataStorage.GetWithMetadata(l.keysetID)
+	} else {
+		data, err = l.storage.Get(l.keysetID)
+	}
+
+	metadataStorage, ok := l.storage.(kmsapi.StoreWithMetadata)
+	if ok {
+		data, metadata, err = metadataStorage.GetWithMetadata(l.keysetID)
 	}
 
+	if err != nil {
+		return 0, fmt.Errorf("cannot read data for keysetID %s: %w", l.keysetID, err)
+	}
+
+	l.metadata = metadata
+	l.buf = bytes.NewBuffer(data)
+
 	return l.buf.Read(p)
 }

kms/localkms/localkms_writer.go

@@ -31,6 +31,7 @@ func newWriter(kmsStore kmsapi.Store, opts ...kmsapi.PrivateKeyOpts) *storeWrite
 	return &storeWriter{
 		storage:           kmsStore,
 		requestedKeysetID: pOpts.KsID(),
+		metadata:          pOpts.Metadata(),
 	}
 }
 
@@ -39,6 +40,7 @@ type storeWriter struct {
 	storage kmsapi.Store
 	//
 	requestedKeysetID string
+	metadata          map[string]any
 	// KeysetID is set when Write() is called
 	KeysetID string
 }
@@ -61,7 +63,17 @@ func (l *storeWriter) Write(p []byte) (int, error) {
 		}
 	}
 
-	err = l.storage.Put(ksID, p)
+	if len(l.metadata) != 0 {
+		metadataStorage, ok := l.storage.(kmsapi.StoreWithMetadata)
+		if !ok {
+			return 0, fmt.Errorf("requested to save 'metadata', but storage doesn't support it")
+		}
+
+		err = metadataStorage.PutWithMetadata(ksID, p, l.metadata)
+	} else {
+		err = l.storage.Put(ksID, p)
+	}
+
 	if err != nil {
 		return 0, err
 	}

spi/kms/key_opts.go

@@ -8,7 +8,8 @@ package kms
 
 // keyOpts holds options for Create, Rotate and CreateAndExportPubKeyBytes.
 type keyOpts struct {
-	attrs []string
+	attrs    []string
+	metadata map[string]any
 }
 
 // NewKeyOpt creates a new empty key option.
@@ -25,6 +26,11 @@ func (pk *keyOpts) Attrs() []string {
 	return pk.attrs
 }
 
+// Metadata gets the additional data to be stored along with the key.
+func (pk *keyOpts) Metadata() map[string]any {
+	return pk.metadata
+}
+
 // KeyOpts are the create key option.
 type KeyOpts func(opts *keyOpts)
 
@@ -34,3 +40,10 @@ func WithAttrs(attrs []string) KeyOpts {
 		opts.attrs = attrs
 	}
 }
+
+// WithMetadata option is for creating a key that can have additional metadata.
+func WithMetadata(metadata map[string]any) KeyOpts {
+	return func(opts *keyOpts) {
+		opts.metadata = metadata
+	}
+}

spi/kms/kms.go

@@ -79,6 +79,16 @@ type Store interface {
 	Delete(keysetID string) error
 }
 
+// StoreWithMetadata defines extended storage capability to work with key's metadata.
+type StoreWithMetadata interface {
+	// PutWithMetadata stores the given key and metadata under the given keysetID.
+	PutWithMetadata(keysetID string, key []byte, metadata map[string]any) error
+	// GetWithMetadata retrieves the key and its' metadata stored under the given keysetID.
+	// If no key is found, the returned error is expected to wrap ErrKeyNotFound.
+	// KMS implementations may check to see if the error wraps that error type for certain operations.
+	GetWithMetadata(keysetID string) (key []byte, metadata map[string]any, err error)
+}
+
 // Provider for KeyManager builder/constructor.
 type Provider interface {
 	StorageProvider() Store

spi/kms/privKey_opts.go

@@ -8,7 +8,8 @@ package kms
 
 // privateKeyOpts holds options for ImportPrivateKey.
 type privateKeyOpts struct {
-	ksID string
+	ksID     string
+	metadata map[string]any
 }
 
 // NewOpt creates a new empty private key option.
@@ -25,6 +26,11 @@ func (pk *privateKeyOpts) KsID() string {
 	return pk.ksID
 }
 
+// Metadata gets the additional data to be stored along with the key.
+func (pk *privateKeyOpts) Metadata() map[string]any {
+	return pk.metadata
+}
+
 // PrivateKeyOpts are the import private key option.
 type PrivateKeyOpts func(opts *privateKeyOpts)
 
@@ -34,3 +40,35 @@ func WithKeyID(keyID string) PrivateKeyOpts {
 		opts.ksID = keyID
 	}
 }
+
+// ImportWithMetadata option is for importing a private key that can have additional metadata.
+func ImportWithMetadata(metadata map[string]any) PrivateKeyOpts {
+	return func(opts *privateKeyOpts) {
+		opts.metadata = metadata
+	}
+}
+
+// exportKeyOpts holds options for ExportPubKey.
+type exportKeyOpts struct {
+	getMetadata bool
+}
+
+// NewExportOpt creates a new empty export pub key option.
+func NewExportOpt() *exportKeyOpts { // nolint
+	return &exportKeyOpts{}
+}
+
+// GetMetadata indicates that metadata have to be exported along with the key.
+func (pk *exportKeyOpts) GetMetadata() bool {
+	return pk.getMetadata
+}
+
+// ExportKeyOpts are the export public key option.
+type ExportKeyOpts func(opts *exportKeyOpts)
+
+// ExportWithMetadata option is for exporting public key with metadata.
+func ExportWithMetadata(getMetadata bool) ExportKeyOpts {
+	return func(opts *exportKeyOpts) {
+		opts.getMetadata = getMetadata
+	}
+}