feat: use associated data along with the key (#22)

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

kms/localkms/localkms.go

@@ -16,12 +16,11 @@ import (
 
 	"github.com/google/tink/go/aead"
 	"github.com/google/tink/go/keyset"
+	"github.com/google/tink/go/tink"
 	"github.com/trustbloc/bbs-signature-go/bbs12381g2pub"
 
 	kmsapi "github.com/trustbloc/kms-go/spi/kms"
 
-	"github.com/trustbloc/kms-go/spi/secretlock"
-
 	cryptoapi "github.com/trustbloc/kms-go/spi/crypto"
 
 	"github.com/trustbloc/kms-go/doc/util/jwkkid"
@@ -50,28 +49,43 @@ var errInvalidKeyType = errors.New("key type is not supported")
 // It uses an underlying secret lock service (default local secretLock) to wrap (encrypt) keys
 // prior to storing them.
 type LocalKMS struct {
-	secretLock        secretlock.Service
-	primaryKeyURI     string
 	store             kmsapi.Store
 	primaryKeyEnvAEAD *aead.KMSEnvelopeAEAD
 }
 
 // New will create a new (local) KMS service.
 func New(primaryKeyURI string, p kmsapi.Provider) (*LocalKMS, error) {
-	secretLock := p.SecretLock()
+	return NewWithOpts(
+		WithPrimaryKeyURI(primaryKeyURI),
+		WithStore(p.StorageProvider()),
+		WithSecretLock(p.SecretLock()))
+}
 
-	kw, err := keywrapper.New(secretLock, primaryKeyURI)
-	if err != nil {
-		return nil, fmt.Errorf("new: failed to create new keywrapper: %w", err)
+// NewWithOpts will create a new KMS service with options.
+func NewWithOpts(opts ...KMSOpts) (*LocalKMS, error) {
+	options := NewKMSOpt()
+
+	for _, opt := range opts {
+		opt(options)
+	}
+
+	var aeadService tink.AEAD
+
+	if options.AEADService() != nil {
+		aeadService = options.AEADService()
+	} else {
+		kw, err := keywrapper.New(options.SecretLock(), options.PrimaryKeyURI())
+		if err != nil {
+			return nil, fmt.Errorf("new: failed to create new keywrapper: %w", err)
+		}
+
+		aeadService = kw
 	}
 
-	// create a KMSEnvelopeAEAD instance to wrap/unwrap keys managed by LocalKMS
-	keyEnvelopeAEAD := aead.NewKMSEnvelopeAEAD2(aead.AES256GCMKeyTemplate(), kw)
+	keyEnvelopeAEAD := aead.NewKMSEnvelopeAEAD2(aead.AES256GCMKeyTemplate(), aeadService)
 
 	return &LocalKMS{
-			store:             p.StorageProvider(),
-			secretLock:        secretLock,
-			primaryKeyURI:     primaryKeyURI,
+			store:             options.Store(),
 			primaryKeyEnvAEAD: keyEnvelopeAEAD,
 		},
 		nil
@@ -138,7 +152,13 @@ func (l *LocalKMS) GetWithOpts(keyID string, opts ...kmsapi.ExportKeyOpts) (any,
 //   - handle instance (to private key)
 //   - error if failure
 func (l *LocalKMS) Rotate(kt kmsapi.KeyType, keyID string, opts ...kmsapi.KeyOpts) (string, interface{}, error) {
-	kh, err := l.getKeySet(keyID)
+	keyOpts := kmsapi.NewKeyOpt()
+
+	for _, opt := range opts {
+		opt(keyOpts)
+	}
+
+	kh, _, err := l.getKeySetWithOpts(keyID, kmsapi.ExportAssociatedData(keyOpts.AssociatedData()))
 	if err != nil {
 		return "", nil, fmt.Errorf("rotate: failed to getKeySet: %w", err)
 	}
@@ -165,7 +185,7 @@ func (l *LocalKMS) Rotate(kt kmsapi.KeyType, keyID string, opts ...kmsapi.KeyOpt
 		return "", nil, fmt.Errorf("rotate: failed to delete entry for kid '%s': %w", keyID, err)
 	}
 
-	newID, err := l.storeKeySet(updatedKH, kt)
+	newID, err := l.storeKeySet(updatedKH, kt, opts...)
 	if err != nil {
 		return "", nil, fmt.Errorf("rotate: failed to store keySet: %w", err)
 	}
@@ -193,20 +213,20 @@ func (l *LocalKMS) storeKeySet(kh *keyset.Handle, kt kmsapi.KeyType, opts ...kms
 		}
 	}
 
+	keyOpts := kmsapi.NewKeyOpt()
+
+	for _, opt := range opts {
+		opt(keyOpts)
+	}
+
 	buf := new(bytes.Buffer)
 	jsonKeysetWriter := keyset.NewJSONWriter(buf)
 
-	err = kh.Write(jsonKeysetWriter, l.primaryKeyEnvAEAD)
+	err = kh.WriteWithAssociatedData(jsonKeysetWriter, l.primaryKeyEnvAEAD, keyOpts.AssociatedData())
 	if err != nil {
 		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 != "" {
@@ -229,18 +249,8 @@ func writeToStore(store kmsapi.Store, buf *bytes.Buffer, opts ...kmsapi.PrivateK
 }
 
 func (l *LocalKMS) getKeySet(id string) (*keyset.Handle, error) {
-	localDBReader := newReader(l.store, id)
-
-	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, fmt.Errorf("getKeySet: failed to read json keyset from reader: %w", err)
-	}
-
-	return kh, nil
+	ks, _, err := l.getKeySetWithOpts(id)
+	return ks, err
 }
 
 func (l *LocalKMS) getKeySetWithOpts(id string, opts ...kmsapi.ExportKeyOpts) (*keyset.Handle, map[string]any, error) {
@@ -250,7 +260,7 @@ func (l *LocalKMS) getKeySetWithOpts(id string, opts ...kmsapi.ExportKeyOpts) (*
 
 	// Read reads the encrypted keyset handle back from the io.reader implementation
 	// and decrypts it using primaryKeyEnvAEAD.
-	kh, err := keyset.Read(jsonKeysetReader, l.primaryKeyEnvAEAD)
+	kh, err := keyset.ReadWithAssociatedData(jsonKeysetReader, l.primaryKeyEnvAEAD, localDBReader.associatedData)
 	if err != nil {
 		return nil, nil, fmt.Errorf("getKeySet: failed to read json keyset from reader: %w", err)
 	}

kms/localkms/localkms_reader.go

@@ -24,19 +24,21 @@ func newReader(store kms.Store, keysetID string, opts ...kmsapi.ExportKeyOpts) *
 	}
 
 	return &storeReader{
-		storage:     store,
-		keysetID:    keysetID,
-		getMetadata: pOpts.GetMetadata(),
+		storage:        store,
+		keysetID:       keysetID,
+		getMetadata:    pOpts.GetMetadata(),
+		associatedData: pOpts.AssociatedData(),
 	}
 }
 
 // storeReader struct to load a keyset from a local storage.
 type storeReader struct {
-	buf         *bytes.Buffer
-	storage     kms.Store
-	keysetID    string
-	getMetadata bool
-	metadata    map[string]any
+	buf            *bytes.Buffer
+	storage        kms.Store
+	keysetID       string
+	getMetadata    bool
+	metadata       map[string]any
+	associatedData []byte
 }
 
 // Read the keyset from local storage into p.

kms/localkms/opts.go

@@ -0,0 +1,73 @@
+/*
+ Copyright SecureKey Technologies Inc. All Rights Reserved.
+
+ SPDX-License-Identifier: Apache-2.0
+*/
+
+package localkms
+
+import (
+	"github.com/google/tink/go/tink"
+	kmsapi "github.com/trustbloc/kms-go/spi/kms"
+	"github.com/trustbloc/kms-go/spi/secretlock"
+)
+
+type kmsOpts struct {
+	store         kmsapi.Store
+	lock          secretlock.Service
+	aeadService   tink.AEAD
+	primaryKeyURI string
+}
+
+// NewKMSOpt creates a new empty KMS options.
+func NewKMSOpt() *kmsOpts { // nolint
+	return &kmsOpts{}
+}
+
+func (k *kmsOpts) Store() kmsapi.Store {
+	return k.store
+}
+
+func (k *kmsOpts) SecretLock() secretlock.Service {
+	return k.lock
+}
+
+func (k *kmsOpts) AEADService() tink.AEAD {
+	return k.aeadService
+}
+
+func (k *kmsOpts) PrimaryKeyURI() string {
+	return k.primaryKeyURI
+}
+
+// KMSOpts are the create KMS option.
+type KMSOpts func(opts *kmsOpts)
+
+// WithStore option is for setting store for KMS.
+func WithStore(store kmsapi.Store) KMSOpts {
+	return func(opts *kmsOpts) {
+		opts.store = store
+	}
+}
+
+// WithSecretLock option is for setting secret-lock for KMS.
+func WithSecretLock(secretLock secretlock.Service) KMSOpts {
+	return func(opts *kmsOpts) {
+		opts.lock = secretLock
+	}
+}
+
+// WithPrimaryKeyURI option is for setting secret-lock for KMS.
+func WithPrimaryKeyURI(primaryKeyURI string) KMSOpts {
+	return func(opts *kmsOpts) {
+		opts.primaryKeyURI = primaryKeyURI
+	}
+}
+
+// WithAEAD option is for setting AEAD service directly for KMS.
+// If not set, secretLock and primaryKeyURI will be used.
+func WithAEAD(aeadService tink.AEAD) KMSOpts {
+	return func(opts *kmsOpts) {
+		opts.aeadService = aeadService
+	}
+}

spi/kms/key_opts.go

@@ -8,15 +8,18 @@ package kms
 
 // keyOpts holds options for Create, Rotate and CreateAndExportPubKeyBytes.
 type keyOpts struct {
-	attrs    []string
-	metadata map[string]any
+	attrs          []string
+	metadata       map[string]any
+	associatedData []byte
 }
 
 // NewKeyOpt creates a new empty key option.
 // Not to be used directly. It's intended for implementations of KeyManager interface
 // Use WithAttrs() option function below instead.
 func NewKeyOpt() *keyOpts { // nolint
-	return &keyOpts{}
+	return &keyOpts{
+		associatedData: []byte{},
+	}
 }
 
 // Attrs gets the additional attributes to be used for a key creation.
@@ -31,6 +34,11 @@ func (pk *keyOpts) Metadata() map[string]any {
 	return pk.metadata
 }
 
+// AssociatedData gets the associated data to be stored along with the key.
+func (pk *keyOpts) AssociatedData() []byte {
+	return pk.associatedData
+}
+
 // KeyOpts are the create key option.
 type KeyOpts func(opts *keyOpts)
 
@@ -47,3 +55,10 @@ func WithMetadata(metadata map[string]any) KeyOpts {
 		opts.metadata = metadata
 	}
 }
+
+// WithAssociatedData option is for creating a key that can have associated data.
+func WithAssociatedData(associatedData []byte) KeyOpts {
+	return func(opts *keyOpts) {
+		opts.associatedData = associatedData
+	}
+}

spi/kms/privKey_opts.go

@@ -50,19 +50,27 @@ func ImportWithMetadata(metadata map[string]any) PrivateKeyOpts {
 
 // exportKeyOpts holds options for ExportPubKey.
 type exportKeyOpts struct {
-	getMetadata bool
+	getMetadata    bool
+	associatedData []byte
 }
 
 // NewExportOpt creates a new empty export pub key option.
 func NewExportOpt() *exportKeyOpts { // nolint
-	return &exportKeyOpts{}
+	return &exportKeyOpts{
+		associatedData: []byte{},
+	}
 }
 
 // GetMetadata indicates that metadata have to be exported along with the key.
 func (pk *exportKeyOpts) GetMetadata() bool {
 	return pk.getMetadata
 }
 
+// AssociatedData returns associated data for the key.
+func (pk *exportKeyOpts) AssociatedData() []byte {
+	return pk.associatedData
+}
+
 // ExportKeyOpts are the export public key option.
 type ExportKeyOpts func(opts *exportKeyOpts)
 
@@ -72,3 +80,10 @@ func ExportWithMetadata(getMetadata bool) ExportKeyOpts {
 		opts.getMetadata = getMetadata
 	}
 }
+
+// ExportAssociatedData option is for exporting key saved using associated data.
+func ExportAssociatedData(associatedData []byte) ExportKeyOpts {
+	return func(opts *exportKeyOpts) {
+		opts.associatedData = associatedData
+	}
+}