Constrain NewHybridEncrypt/Decrypt to exclude AEAD primitives.

PiperOrigin-RevId: 547982249

6 files changed, 72 insertions, 87 deletions

diff --git a/go/hybrid/BUILD.bazel b/go/hybrid/BUILD.bazel
index 7bca0d0df..0e8c13ef1 100644
--- a/go/hybrid/BUILD.bazel
+++ b/go/hybrid/BUILD.bazel
@@ -63,7 +63,6 @@ go_test(
         "//hybrid/subtle",
         "//insecurecleartextkeyset",
         "//internal/internalregistry",
-        "//internal/testing/stubkeymanager",
         "//keyset",
         "//mac",
         "//monitoring",
diff --git a/go/hybrid/hybrid_b243759652_test.go b/go/hybrid/hybrid_b243759652_test.go
index a2aa98e21..329836cec 100644
--- a/go/hybrid/hybrid_b243759652_test.go
+++ b/go/hybrid/hybrid_b243759652_test.go
@@ -24,7 +24,7 @@ import (
 	"github.com/google/tink/go/keyset"
 )
 
-/* This test demonstrates buggy behavior of Tink. Please do not rely on it. */
+// TODO(b/243759652): Move to hybrid_factory_test.go.
 func TestBHybridDecrypt(t *testing.T) {
 	manager := keyset.NewManager()
 	id, err := manager.Add(aead.AES256GCMKeyTemplate())
@@ -44,14 +44,11 @@ func TestBHybridDecrypt(t *testing.T) {
 		t.Fatalf("manager.Handle gives err = '%v', want nil", err)
 	}
 
-	// TODO(b/243759652): This should fail.
-	_, err = hybrid.NewHybridDecrypt(handle)
-	if err != nil {
-		t.Fatalf("hybrid.NewHybridDecrypt gives err = '%v', you fixed b/243759652, please change the test.", err)
+	if _, err := hybrid.NewHybridDecrypt(handle); err == nil {
+		t.Error("hybrid.NewHybridDecrypt err = nil, want err")
 	}
 }
 
-/* This test demonstrates buggy behavior of Tink. Please do not rely on it. */
 func TestBHybridEncrypt(t *testing.T) {
 	handle, err := keyset.NewHandle(hybrid.DHKEM_X25519_HKDF_SHA256_HKDF_SHA256_AES_128_GCM_Key_Template())
 	if err != nil {
@@ -70,9 +67,7 @@ func TestBHybridEncrypt(t *testing.T) {
 	if err != nil {
 		t.Fatalf("manager.Handle gives err = '%v', want nil", err)
 	}
-	// TODO(b/243759652): We now have a HybridEncrypt and a Aead key, so this should fail
-	_, err = hybrid.NewHybridEncrypt(mixedHandle)
-	if err != nil {
-		t.Fatalf("hybrid.NewHybridEncrypt gives err = '%v', you fixed b/243759652, please change the test.", err)
+	if _, err := hybrid.NewHybridEncrypt(mixedHandle); err == nil {
+		t.Error("hybrid.NewHybridDecrypt err = nil, want err")
 	}
 }
diff --git a/go/hybrid/hybrid_decrypt_factory.go b/go/hybrid/hybrid_decrypt_factory.go
index 614da1727..23a48d4d0 100644
--- a/go/hybrid/hybrid_decrypt_factory.go
+++ b/go/hybrid/hybrid_decrypt_factory.go
@@ -48,14 +48,14 @@ type wrappedHybridDecrypt struct {
 var _ tink.HybridDecrypt = (*wrappedHybridDecrypt)(nil)
 
 func newWrappedHybridDecrypt(ps *primitiveset.PrimitiveSet) (*wrappedHybridDecrypt, error) {
-	if _, ok := (ps.Primary.Primitive).(tink.HybridDecrypt); !ok {
-		return nil, fmt.Errorf("hybrid_factory: not a HybridDecrypt primitive")
+	if err := isHybridDecrypt(ps.Primary.Primitive); err != nil {
+		return nil, err
 	}
 
 	for _, primitives := range ps.Entries {
 		for _, p := range primitives {
-			if _, ok := (p.Primitive).(tink.HybridDecrypt); !ok {
-				return nil, fmt.Errorf("hybrid_factory: not a HybridDecrypt primitive")
+			if err := isHybridDecrypt(p.Primitive); err != nil {
+				return nil, err
 			}
 		}
 	}
@@ -95,11 +95,7 @@ func (a *wrappedHybridDecrypt) Decrypt(ciphertext, contextInfo []byte) ([]byte,
 		entries, err := a.ps.EntriesForPrefix(string(prefix))
 		if err == nil {
 			for i := 0; i < len(entries); i++ {
-				p, ok := (entries[i].Primitive).(tink.HybridDecrypt)
-				if !ok {
-					return nil, fmt.Errorf("hybrid_factory: not a HybridDecrypt primitive")
-				}
-
+				p := entries[i].Primitive.(tink.HybridDecrypt) // verified in newWrappedHybridDecrypt
 				pt, err := p.Decrypt(ctNoPrefix, contextInfo)
 				if err == nil {
 					a.logger.Log(entries[i].KeyID, len(ctNoPrefix))
@@ -113,11 +109,7 @@ func (a *wrappedHybridDecrypt) Decrypt(ciphertext, contextInfo []byte) ([]byte,
 	entries, err := a.ps.RawEntries()
 	if err == nil {
 		for i := 0; i < len(entries); i++ {
-			p, ok := (entries[i].Primitive).(tink.HybridDecrypt)
-			if !ok {
-				return nil, fmt.Errorf("hybrid_factory: not a HybridDecrypt primitive")
-			}
-
+			p := entries[i].Primitive.(tink.HybridDecrypt) // verified in newWrappedHybridDecrypt
 			pt, err := p.Decrypt(ciphertext, contextInfo)
 			if err == nil {
 				a.logger.Log(entries[i].KeyID, len(ciphertext))
@@ -130,3 +122,15 @@ func (a *wrappedHybridDecrypt) Decrypt(ciphertext, contextInfo []byte) ([]byte,
 	a.logger.LogFailure()
 	return nil, fmt.Errorf("hybrid_factory: decryption failed")
 }
+
+// Asserts `p` implements tink.HybridDecrypt and not tink.AEAD. The latter check
+// is required as implementations of tink.AEAD also satisfy tink.HybridDecrypt.
+func isHybridDecrypt(p any) error {
+	if _, ok := p.(tink.AEAD); ok {
+		return fmt.Errorf("hybrid_factory: tink.AEAD is not tink.HybridDecrypt")
+	}
+	if _, ok := p.(tink.HybridDecrypt); !ok {
+		return fmt.Errorf("hybrid_factory: not tink.HybridDecrypt")
+	}
+	return nil
+}
diff --git a/go/hybrid/hybrid_encrypt_factory.go b/go/hybrid/hybrid_encrypt_factory.go
index 04db9c235..7725a76ab 100644
--- a/go/hybrid/hybrid_encrypt_factory.go
+++ b/go/hybrid/hybrid_encrypt_factory.go
@@ -46,14 +46,14 @@ type wrappedHybridEncrypt struct {
 var _ tink.HybridEncrypt = (*wrappedHybridEncrypt)(nil)
 
 func newEncryptPrimitiveSet(ps *primitiveset.PrimitiveSet) (*wrappedHybridEncrypt, error) {
-	if _, ok := (ps.Primary.Primitive).(tink.HybridEncrypt); !ok {
-		return nil, fmt.Errorf("hybrid_factory: not a HybridEncrypt primitive")
+	if err := isHybridEncrypt(ps.Primary.Primitive); err != nil {
+		return nil, err
 	}
 
 	for _, primitives := range ps.Entries {
 		for _, p := range primitives {
-			if _, ok := (p.Primitive).(tink.HybridEncrypt); !ok {
-				return nil, fmt.Errorf("hybrid_factory: not a HybridEncrypt primitive")
+			if err := isHybridEncrypt(p.Primitive); err != nil {
+				return nil, err
 			}
 		}
 	}
@@ -86,10 +86,7 @@ func createEncryptLogger(ps *primitiveset.PrimitiveSet) (monitoring.Logger, erro
 // It returns the concatenation of the primary's identifier and the ciphertext.
 func (a *wrappedHybridEncrypt) Encrypt(plaintext, contextInfo []byte) ([]byte, error) {
 	primary := a.ps.Primary
-	p, ok := (primary.Primitive).(tink.HybridEncrypt)
-	if !ok {
-		return nil, fmt.Errorf("hybrid_factory: not a HybridEncrypt primitive")
-	}
+	p := primary.Primitive.(tink.HybridEncrypt) // verified in newEncryptPrimitiveSet
 
 	ct, err := p.Encrypt(plaintext, contextInfo)
 	if err != nil {
@@ -105,3 +102,15 @@ func (a *wrappedHybridEncrypt) Encrypt(plaintext, contextInfo []byte) ([]byte, e
 	output = append(output, ct...)
 	return output, nil
 }
+
+// Asserts `p` implements tink.HybridEncrypt and not tink.AEAD. The latter check
+// is required as implementations of tink.AEAD also satisfy tink.HybridEncrypt.
+func isHybridEncrypt(p any) error {
+	if _, ok := p.(tink.AEAD); ok {
+		return fmt.Errorf("hybrid_factory: tink.AEAD is not tink.HybridEncrypt")
+	}
+	if _, ok := p.(tink.HybridEncrypt); !ok {
+		return fmt.Errorf("hybrid_factory: not tink.HybridEncrypt")
+	}
+	return nil
+}
diff --git a/go/hybrid/hybrid_factory_test.go b/go/hybrid/hybrid_factory_test.go
index 1d6b3c3a3..941b5182e 100644
--- a/go/hybrid/hybrid_factory_test.go
+++ b/go/hybrid/hybrid_factory_test.go
@@ -18,7 +18,6 @@ package hybrid_test
 
 import (
 	"bytes"
-	"fmt"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -26,11 +25,9 @@ import (
 	"google.golang.org/protobuf/proto"
 	"github.com/google/tink/go/aead"
 	"github.com/google/tink/go/core/cryptofmt"
-	"github.com/google/tink/go/core/registry"
 	"github.com/google/tink/go/hybrid"
 	"github.com/google/tink/go/insecurecleartextkeyset"
 	"github.com/google/tink/go/internal/internalregistry"
-	"github.com/google/tink/go/internal/testing/stubkeymanager"
 	"github.com/google/tink/go/keyset"
 	"github.com/google/tink/go/monitoring"
 	"github.com/google/tink/go/signature"
@@ -524,41 +521,8 @@ func TestPrimitiveFactoryMonitoringWithAnnotationsEncryptFailureIsLogged(t *test
 	if err := internalregistry.RegisterMonitoringClient(client); err != nil {
 		t.Fatalf("internalregistry.RegisterMonitoringClient() err = %v, want nil", err)
 	}
-	// Since this key type will be registered in the registry,
-	// we create a very unique typeURL to avoid colliding with other tests.
-	privKeyTypeURL := "TestPrimitiveFactoryWithMonitoringEncryptionFailureIsLogged" + "PrivateKeyManager"
-	pubKeyTypeURL := "TestPrimitiveFactoryWithMonitoringEncryptionFailureIsLogged" + "PublicKeyManager"
-	template := &tinkpb.KeyTemplate{
-		TypeUrl:          privKeyTypeURL,
-		OutputPrefixType: tinkpb.OutputPrefixType_LEGACY,
-	}
-	privKeyManager := &stubkeymanager.StubPrivateKeyManager{
-		StubKeyManager: stubkeymanager.StubKeyManager{
-			URL: privKeyTypeURL,
-			KeyData: &tinkpb.KeyData{
-				TypeUrl:         template.TypeUrl,
-				Value:           []byte("some_data"),
-				KeyMaterialType: tinkpb.KeyData_ASYMMETRIC_PRIVATE,
-			},
-		},
-		PubKeyData: &tinkpb.KeyData{
-			TypeUrl:         pubKeyTypeURL,
-			Value:           []byte("some_data"),
-			KeyMaterialType: tinkpb.KeyData_ASYMMETRIC_PUBLIC,
-		},
-	}
-	pubKeyManager := &stubkeymanager.StubKeyManager{
-		URL: pubKeyTypeURL,
-		// TODO(b/243759652): implement an always failing HybridEncrypt
-		Prim: &testutil.AlwaysFailingAead{Error: fmt.Errorf("panic at the kernel")},
-	}
-	if err := registry.RegisterKeyManager(privKeyManager); err != nil {
-		t.Fatalf("registry.RegisterKeyManager() err = %v, want nil", err)
-	}
-	if err := registry.RegisterKeyManager(pubKeyManager); err != nil {
-		t.Fatalf("registry.RegisterKeyManager() err = %v, want nil", err)
-	}
-	handle, err := keyset.NewHandle(template)
+
+	handle, err := keyset.NewHandle(hybrid.DHKEM_X25519_HKDF_SHA256_HKDF_SHA256_AES_128_GCM_Key_Template())
 	if err != nil {
 		t.Fatalf("keyset.NewHandle() err = %v, want nil", err)
 	}
@@ -571,11 +535,12 @@ func TestPrimitiveFactoryMonitoringWithAnnotationsEncryptFailureIsLogged(t *test
 	if err != nil {
 		t.Fatalf("insecurecleartextkeyset.Read() err = %v, want nil", err)
 	}
+	buff.Reset()
+
 	pubHandle, err := privHandle.Public()
 	if err != nil {
 		t.Fatalf("privHandle.Public() err = %v, want nil", err)
 	}
-	buff.Reset()
 	if err := insecurecleartextkeyset.Write(pubHandle, keyset.NewBinaryWriter(buff)); err != nil {
 		t.Fatalf("insecurecleartextkeyset.Write() err = %v, want nil", err)
 	}
@@ -583,28 +548,40 @@ func TestPrimitiveFactoryMonitoringWithAnnotationsEncryptFailureIsLogged(t *test
 	if err != nil {
 		t.Fatalf("insecurecleartextkeyset.Read() err = %v, want nil", err)
 	}
+
 	e, err := hybrid.NewHybridEncrypt(pubHandle)
 	if err != nil {
-		t.Fatalf("hybrid.NewHybridEncrypt() err = %v, want nil", err)
+		t.Fatalf("NewHybridEncrypt() err = %v, want nil", err)
+	}
+	d, err := hybrid.NewHybridDecrypt(privHandle)
+	if err != nil {
+		t.Fatalf("NewHybridDecrypt() err = %v, want nil", err)
+	}
+
+	ct, err := e.Encrypt([]byte("plaintext"), []byte("info"))
+	if err != nil {
+		t.Fatalf("Encrypt() err = nil, want non-nil")
 	}
-	if _, err := e.Encrypt(nil, nil); err == nil {
-		t.Fatalf("e.Encrypt() err = nil, want non-nil error")
+	if _, err := d.Decrypt(ct, []byte("wrong info")); err == nil {
+		t.Fatalf("Decrypt() err = nil, want non-nil")
 	}
+
 	got := client.Failures()
+	primaryKeyID := privHandle.KeysetInfo().GetPrimaryKeyId()
 	want := []*fakemonitoring.LogFailure{
 		{
 			Context: monitoring.NewContext(
-				"hybrid_encrypt",
-				"encrypt",
+				"hybrid_decrypt",
+				"decrypt",
 				monitoring.NewKeysetInfo(
 					annotations,
-					pubHandle.KeysetInfo().GetPrimaryKeyId(),
+					primaryKeyID,
 					[]*monitoring.Entry{
 						{
-							KeyID:     pubHandle.KeysetInfo().GetPrimaryKeyId(),
+							KeyID:     primaryKeyID,
 							Status:    monitoring.Enabled,
-							KeyType:   pubKeyTypeURL,
-							KeyPrefix: "LEGACY",
+							KeyType:   "tink.HpkePrivateKey",
+							KeyPrefix: "TINK",
 						},
 					},
 				),
diff --git a/testing/cross_language/primitive_creation_test.py b/testing/cross_language/primitive_creation_test.py
index f0e1f297e..45b9b974b 100644
--- a/testing/cross_language/primitive_creation_test.py
+++ b/testing/cross_language/primitive_creation_test.py
@@ -84,6 +84,7 @@ def named_testcases():
         case_num += 1
 
 
+# Delete.
 def _is_b243759652_test_case(lang: str, keyset: bytes, primitive: Any) -> bool:
   """Returns whether the test case falls under b/243759652.
 
@@ -150,8 +151,8 @@ class SupportedKeyTypesTest(parameterized.TestCase):
     keytype = keytypes[0]
 
     if _is_b243759652_test_case(lang, keyset, primitive):
-      # TODO(b/243759652): This should raise a TinkError, but doesn't
-      _ = testing_servers.remote_primitive(lang, keyset, primitive)
+      with self.assertRaises(tink.TinkError):
+        _ = testing_servers.remote_primitive(lang, keyset, primitive)
       return
 
     if (lang in tink_config.supported_languages_for_key_type(keytype) and
@@ -184,8 +185,8 @@ class SupportedKeyTypesTest(parameterized.TestCase):
     keytype = keytypes[0]
 
     if _is_b243759652_test_case(lang, public_keyset, primitive):
-      # TODO(b/243759652): This should raise a TinkError, but doesn't
-      _ = testing_servers.remote_primitive(lang, public_keyset, primitive)
+      with self.assertRaises(tink.TinkError):
+        _ = testing_servers.remote_primitive(lang, public_keyset, primitive)
       return
 
     if (lang in tink_config.supported_languages_for_key_type(keytype) and
@@ -215,8 +216,8 @@ class SupportedKeyTypesTest(parameterized.TestCase):
     keytype = keytypes[0]
 
     if _is_b243759652_test_case(lang, modified_keyset, primitive):
-      # TODO(b/243759652): This should raise a TinkError, but doesn't
-      _ = testing_servers.remote_primitive(lang, modified_keyset, primitive)
+      with self.assertRaises(tink.TinkError):
+        _ = testing_servers.remote_primitive(lang, modified_keyset, primitive)
       return
 
     if (lang in tink_config.supported_languages_for_key_type(keytype) and