Extend aead_factory to support Configs for primitive creation.

PiperOrigin-RevId: 603678235
Change-Id: I41eef01b893f817b5457649ad0354959226bc275

aead/BUILD.bazel

@@ -25,6 +25,7 @@ go_library(
         "//core/cryptofmt",
         "//core/primitiveset",
         "//core/registry",
+        "//internal/internalapi",
         "//internal/internalregistry",
         "//internal/monitoringutil",
         "//internal/tinkerror",
@@ -70,6 +71,7 @@ go_test(
         "//core/cryptofmt",
         "//core/registry",
         "//insecurecleartextkeyset",
+        "//internal/internalapi",
         "//internal/internalregistry",
         "//internal/testing/stubkeymanager",
         "//internal/tinkerror/tinkerrortest",

aead/aead.go

@@ -23,9 +23,23 @@ import (
 	"fmt"
 
 	"github.com/tink-crypto/tink-go/v2/core/registry"
+	"github.com/tink-crypto/tink-go/v2/internal/internalapi"
 	"github.com/tink-crypto/tink-go/v2/internal/internalregistry"
+	tinkpb "github.com/tink-crypto/tink-go/v2/proto/tink_go_proto"
 )
 
+// Config is an interface used by key managers and primitive factories to
+// represent the relevant Config functionality.
+//
+// Note that the interface is not public and meant for package-internal use
+// only, despite it being exported.
+//
+// TODO: b/286235179 -- link the implementation.
+type Config interface {
+	RegisterKeyManager(km registry.KeyManager, t internalapi.Token) error
+	PrimitiveFromKeyData(keyData *tinkpb.KeyData, _ internalapi.Token) (any, error)
+}
+
 func init() {
 	if err := registry.RegisterKeyManager(new(aesCTRHMACAEADKeyManager)); err != nil {
 		panic(fmt.Sprintf("aead.init() failed: %v", err))

aead/aead_factory.go

@@ -28,12 +28,41 @@ import (
 	"github.com/tink-crypto/tink-go/v2/tink"
 )
 
+type newOptions struct {
+	Config Config
+}
+
+// NewOption configures [New].
+type NewOption func(*newOptions) error
+
+// WithConfig sets the Config used by [New] to create primitives.
+func WithConfig(c Config) NewOption {
+	return func(options *newOptions) error {
+		if options.Config != nil {
+			return fmt.Errorf("aead_factory: more than one config provided")
+		}
+		options.Config = c
+		return nil
+	}
+}
+
 // New returns an AEAD primitive from the given keyset handle.
-func New(handle *keyset.Handle) (tink.AEAD, error) {
-	ps, err := handle.Primitives()
+//
+// It will use a Config when provided via [WithConfig]. If no Config is
+// provided, New uses the global registry instead.
+func New(handle *keyset.Handle, opts ...NewOption) (tink.AEAD, error) {
+	var options newOptions
+	for _, opt := range opts {
+		if err := opt(&options); err != nil {
+			return nil, err
+		}
+	}
+
+	ps, err := handle.Primitives(keyset.WithConfig(options.Config))
 	if err != nil {
 		return nil, fmt.Errorf("aead_factory: cannot obtain primitive set: %s", err)
 	}
+
 	return newWrappedAead(ps)
 }
 

aead/aead_factory_test.go

@@ -29,6 +29,7 @@ import (
 	"github.com/tink-crypto/tink-go/v2/core/cryptofmt"
 	"github.com/tink-crypto/tink-go/v2/core/registry"
 	"github.com/tink-crypto/tink-go/v2/insecurecleartextkeyset"
+	"github.com/tink-crypto/tink-go/v2/internal/internalapi"
 	"github.com/tink-crypto/tink-go/v2/internal/internalregistry"
 	"github.com/tink-crypto/tink-go/v2/internal/testing/stubkeymanager"
 	"github.com/tink-crypto/tink-go/v2/keyset"
@@ -107,6 +108,59 @@ func TestFactoryMultipleKeys(t *testing.T) {
 	}
 }
 
+type fakeAEAD struct{}
+
+func (a *fakeAEAD) Encrypt(_, _ []byte) ([]byte, error) {
+	return []byte("fakeAEAD Encrypt"), nil
+}
+func (a *fakeAEAD) Decrypt(_, _ []byte) ([]byte, error) {
+	return []byte("fakeAEAD Decrypt"), nil
+}
+
+type fakeConfig struct{}
+
+func (tc *fakeConfig) RegisterKeyManager(_ registry.KeyManager, _ internalapi.Token) error {
+	return nil // not needed in the test
+}
+func (tc *fakeConfig) PrimitiveFromKeyData(_ *tinkpb.KeyData, _ internalapi.Token) (any, error) {
+	return new(fakeAEAD), nil
+}
+
+func TestNewWithConfig(t *testing.T) {
+	ks := testutil.NewTestAESGCMKeyset(tinkpb.OutputPrefixType_RAW)
+	keysetHandle, err := testkeyset.NewHandle(ks)
+	if err != nil {
+		t.Fatalf("testkeyset.NewHandle(keyset) err = %v, want nil", err)
+	}
+
+	config := &fakeConfig{}
+	a, err := aead.New(keysetHandle, aead.WithConfig(config))
+	if err != nil {
+		t.Fatalf("aead.New(keysetHandle, aead.WithConfig(config)) err = %v, want nil", err)
+	}
+	ct, err := a.Encrypt([]byte("plaintext"), []byte("aad"))
+	if err != nil {
+		t.Fatalf("aead.Encrypt() err = %v, want nil", err)
+	}
+
+	if wantCT := "fakeAEAD Encrypt"; !bytes.Equal(ct, []byte(wantCT)) {
+		t.Errorf("aead.Encrypt() = %q, want %q", ct, wantCT)
+	}
+}
+
+func TestNewFailsWithMultipleConfigs(t *testing.T) {
+	ks := testutil.NewTestAESGCMKeyset(tinkpb.OutputPrefixType_RAW)
+	keysetHandle, err := testkeyset.NewHandle(ks)
+	if err != nil {
+		t.Fatalf("testkeyset.NewHandle(keyset) err = %v, want nil", err)
+	}
+
+	config := &fakeConfig{}
+	if _, err := aead.New(keysetHandle, aead.WithConfig(config), aead.WithConfig(config)); err == nil { // if NO error
+		t.Errorf("aead.New() with multiple configs err = nil, want error")
+	}
+}
+
 func TestFactoryRawKeyAsPrimary(t *testing.T) {
 	keyset := testutil.NewTestAESGCMKeyset(tinkpb.OutputPrefixType_RAW)
 	if keyset.Key[0].OutputPrefixType != tinkpb.OutputPrefixType_RAW {