go/hybrid/hpke_private_key_manager_test.go - third_party/tink - Git at Google

 // Copyright 2022 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //      http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
 //
 ////////////////////////////////////////////////////////////////////////////////

 package hybrid

 import (
 	"bytes"
 	"testing"

 	"github.com/google/go-cmp/cmp"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/testing/protocmp"
 	"github.com/google/tink/go/core/registry"
 	"github.com/google/tink/go/hybrid/internal/hpke"
 	"github.com/google/tink/go/subtle/random"
 	hpkepb "github.com/google/tink/go/proto/hpke_go_proto"
 	tinkpb "github.com/google/tink/go/proto/tink_go_proto"
 )

 func TestPrivateKeyManagerPrimitiveRejectsInvalidKeyVersion(t *testing.T) {
 	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
 	if err != nil {
 		t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
 	}
 	params := &hpkepb.HpkeParams{
 		Kem:  hpkepb.HpkeKem_DHKEM_X25519_HKDF_SHA256,
 		Kdf:  hpkepb.HpkeKdf_HKDF_SHA256,
 		Aead: hpkepb.HpkeAead_AES_256_GCM,
 	}
 	_, privKey := pubPrivKeys(t, params)
 	privKey.Version = 1
 	serializedPrivKey, err := proto.Marshal(privKey)
 	if err != nil {
 		t.Fatal(err)
 	}
 	if _, err := km.Primitive(serializedPrivKey); err == nil {
 		t.Error("Primitive() err = nil, want error")
 	}
 }

 func TestPrivateKeyManagerPrimitiveRejectsInvalidParams(t *testing.T) {
 	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
 	if err != nil {
 		t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
 	}

 	tests := []struct {
 		name   string
 		params *hpkepb.HpkeParams
 	}{
 		{"kem", &hpkepb.HpkeParams{
 			Kem:  hpkepb.HpkeKem_KEM_UNKNOWN,
 			Kdf:  hpkepb.HpkeKdf_HKDF_SHA256,
 			Aead: hpkepb.HpkeAead_AES_256_GCM,
 		}},
 		{"kdf", &hpkepb.HpkeParams{
 			Kem:  hpkepb.HpkeKem_DHKEM_X25519_HKDF_SHA256,
 			Kdf:  hpkepb.HpkeKdf_KDF_UNKNOWN,
 			Aead: hpkepb.HpkeAead_AES_256_GCM,
 		}},
 		{"aead", &hpkepb.HpkeParams{
 			Kem:  hpkepb.HpkeKem_DHKEM_X25519_HKDF_SHA256,
 			Kdf:  hpkepb.HpkeKdf_HKDF_SHA256,
 			Aead: hpkepb.HpkeAead_AEAD_UNKNOWN,
 		}},
 	}

 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			_, serializedPrivKey := serializedPubPrivKeys(t, test.params)
 			if _, err := km.Primitive(serializedPrivKey); err == nil {
 				t.Error("Primitive() err = nil, want error")
 			}
 		})
 	}
 }

 func TestPrivateKeyManagerPrimitiveRejectsMissingParams(t *testing.T) {
 	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
 	if err != nil {
 		t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
 	}
 	_, serializedPrivKey := serializedPubPrivKeys(t, nil)
 	if _, err := km.Primitive(serializedPrivKey); err == nil {
 		t.Error("Primitive() err = nil, want error")
 	}
 }

 func TestPrivateKeyManagerPrimitiveRejectsNilKey(t *testing.T) {
 	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
 	if err != nil {
 		t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
 	}
 	if _, err := km.Primitive(nil); err == nil {
 		t.Error("Primitive() err = nil, want error")
 	}
 }

 func TestPrivateKeyManagerPrimitiveEncryptDecrypt(t *testing.T) {
 	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
 	if err != nil {
 		t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
 	}
 	pt := random.GetRandomBytes(200)
 	ctxInfo := random.GetRandomBytes(100)

 	for _, aeadID := range hpkeAEADs {
 		params := &hpkepb.HpkeParams{
 			Kem:  hpkepb.HpkeKem_DHKEM_X25519_HKDF_SHA256,
 			Kdf:  hpkepb.HpkeKdf_HKDF_SHA256,
 			Aead: aeadID,
 		}
 		pubKey, privKey := pubPrivKeys(t, params)
 		serializedPrivKey, err := proto.Marshal(privKey)
 		if err != nil {
 			t.Fatal(err)
 		}

 		enc, err := hpke.NewEncrypt(pubKey)
 		if err != nil {
 			t.Fatalf("hpke.NewEncrypt() err = %v, want nil", err)
 		}
 		d, err := km.Primitive(serializedPrivKey)
 		if err != nil {
 			t.Fatalf("Primitive() err = %v, want nil", err)
 		}
 		dec, ok := d.(*hpke.Decrypt)
 		if !ok {
 			t.Fatal("primitive is not Decrypt")
 		}

 		ct, err := enc.Encrypt(pt, ctxInfo)
 		if err != nil {
 			t.Fatalf("Encrypt() err = %v, want nil", err)
 		}
 		gotPT, err := dec.Decrypt(ct, ctxInfo)
 		if err != nil {
 			t.Fatalf("Decrypt() err = %v, want nil", err)
 		}
 		if want := pt; !bytes.Equal(gotPT, want) {
 			t.Errorf("Decrypt() = %x, want %x", gotPT, want)
 		}
 	}
 }

 func TestPrivateKeyManagerNewKeyRejectsNilKeyFormat(t *testing.T) {
 	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
 	if err != nil {
 		t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
 	}
 	if _, err := km.NewKey(nil); err == nil {
 		t.Error("NewKey() err = nil, want error")
 	}
 }

 func TestPrivateKeyManagerNewKeyRejectsInvalidKeyFormat(t *testing.T) {
 	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
 	if err != nil {
 		t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
 	}
 	serializedKeyFormatUnknownKEM, err := proto.Marshal(
 		&hpkepb.HpkeParams{
 			Kem:  hpkepb.HpkeKem_KEM_UNKNOWN,
 			Kdf:  hpkepb.HpkeKdf_HKDF_SHA256,
 			Aead: hpkepb.HpkeAead_AES_256_GCM,
 		})
 	if err != nil {
 		t.Fatal(err)
 	}
 	serializedKeyFormatUnknownKDF, err := proto.Marshal(
 		&hpkepb.HpkeParams{
 			Kem:  hpkepb.HpkeKem_DHKEM_X25519_HKDF_SHA256,
 			Kdf:  hpkepb.HpkeKdf_KDF_UNKNOWN,
 			Aead: hpkepb.HpkeAead_AES_256_GCM,
 		})
 	if err != nil {
 		t.Fatal(err)
 	}
 	serializedKeyFormatUnknownAEAD, err := proto.Marshal(
 		&hpkepb.HpkeParams{
 			Kem:  hpkepb.HpkeKem_DHKEM_X25519_HKDF_SHA256,
 			Kdf:  hpkepb.HpkeKdf_HKDF_SHA256,
 			Aead: hpkepb.HpkeAead_AEAD_UNKNOWN,
 		})
 	if err != nil {
 		t.Fatal(err)
 	}

 	tests := []struct {
 		name      string
 		keyFormat []byte
 	}{
 		{"kem", serializedKeyFormatUnknownKEM},
 		{"kdf", serializedKeyFormatUnknownKDF},
 		{"aead", serializedKeyFormatUnknownAEAD},
 	}

 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			if _, err := km.NewKey(test.keyFormat); err == nil {
 				t.Error("NewKey() err = nil, want error")
 			}
 		})
 	}
 }

 func TestPrivateKeyManagerNewKeyEncryptDecrypt(t *testing.T) {
 	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
 	if err != nil {
 		t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
 	}

 	wantPT := random.GetRandomBytes(200)
 	ctxInfo := random.GetRandomBytes(100)

 	for _, aeadID := range hpkeAEADs {
 		keyFormat := &hpkepb.HpkeKeyFormat{
 			Params: &hpkepb.HpkeParams{
 				Kem:  hpkepb.HpkeKem_DHKEM_X25519_HKDF_SHA256,
 				Kdf:  hpkepb.HpkeKdf_HKDF_SHA256,
 				Aead: aeadID,
 			},
 		}
 		serializedKeyFormat, err := proto.Marshal(keyFormat)
 		if err != nil {
 			t.Fatal(err)
 		}

 		privKeyProto, err := km.NewKey(serializedKeyFormat)
 		if err != nil {
 			t.Fatalf("NewKey() err = %v, want nil", err)
 		}
 		privKey, ok := privKeyProto.(*hpkepb.HpkePrivateKey)
 		if !ok {
 			t.Fatal("primitive is not HpkePrivateKey")
 		}
 		if privKey.GetVersion() != 0 {
 			t.Errorf("private key version = %d, want %d", privKey.GetVersion(), 0)
 		}
 		if len(privKey.GetPrivateKey()) == 0 {
 			t.Error("private key is missing")
 		}

 		pubKey := privKey.GetPublicKey()
 		if pubKey.GetVersion() != 0 {
 			t.Errorf("public key version = %d, want %d", pubKey.GetVersion(), 0)
 		}
 		if !cmp.Equal(pubKey.GetParams(), keyFormat.GetParams(), protocmp.Transform()) {
 			t.Errorf("key params = %v, want %v", pubKey.GetParams(), keyFormat.GetParams())
 		}
 		if len(pubKey.GetPublicKey()) == 0 {
 			t.Error("public key is missing")
 		}

 		enc, err := hpke.NewEncrypt(pubKey)
 		if err != nil {
 			t.Fatalf("hpke.NewEncrypt() err = %v, want nil", err)
 		}
 		serializedPrivKey, err := proto.Marshal(privKeyProto)
 		if err != nil {
 			t.Fatal(err)
 		}
 		d, err := km.Primitive(serializedPrivKey)
 		if err != nil {
 			t.Fatalf("Primitive() err = %v, want nil", err)
 		}
 		dec, ok := d.(*hpke.Decrypt)
 		if !ok {
 			t.Fatal("primitive is not Decrypt")
 		}

 		ct, err := enc.Encrypt(wantPT, ctxInfo)
 		if err != nil {
 			t.Fatalf("Encrypt() err = %v, want nil", err)
 		}
 		gotPT, err := dec.Decrypt(ct, ctxInfo)
 		if err != nil {
 			t.Fatalf("Decrypt() err = %v, want nil", err)
 		}
 		if !bytes.Equal(gotPT, wantPT) {
 			t.Errorf("Decrypt() = %x, want %x", gotPT, wantPT)
 		}
 	}
 }

 func TestPrivateKeyManagerNewKeyDataRejectsNilKeyFormat(t *testing.T) {
 	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
 	if err != nil {
 		t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
 	}
 	if _, err := km.NewKeyData(nil); err == nil {
 		t.Error("NewKey() err = nil, want error")
 	}
 }

 func TestPrivateKeyManagerNewKeyData(t *testing.T) {
 	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
 	if err != nil {
 		t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
 	}

 	for _, aeadID := range hpkeAEADs {
 		keyFormat := &hpkepb.HpkeKeyFormat{
 			Params: &hpkepb.HpkeParams{
 				Kem:  hpkepb.HpkeKem_DHKEM_X25519_HKDF_SHA256,
 				Kdf:  hpkepb.HpkeKdf_HKDF_SHA256,
 				Aead: aeadID,
 			},
 		}
 		serializedKeyFormat, err := proto.Marshal(keyFormat)
 		if err != nil {
 			t.Fatal(err)
 		}

 		keyData, err := km.NewKeyData(serializedKeyFormat)
 		if err != nil {
 			t.Fatalf("NewKeyData() err = %v, want nil", err)
 		}
 		if got, want := keyData.GetTypeUrl(), hpkePrivateKeyTypeURL; got != want {
 			t.Errorf("type URL = %q, want %q", got, want)
 		}
 		if got, want := keyData.GetKeyMaterialType(), tinkpb.KeyData_ASYMMETRIC_PRIVATE; got != want {
 			t.Errorf("key material type = %d, want %d", got, want)
 		}

 		privKey := new(hpkepb.HpkePrivateKey)
 		if err := proto.Unmarshal(keyData.GetValue(), privKey); err != nil {
 			t.Fatalf("Unmarshal err = %v, want nil", err)
 		}
 		if privKey.GetVersion() != 0 {
 			t.Errorf("private key version = %d, want %d", privKey.GetVersion(), 0)
 		}
 		if len(privKey.GetPrivateKey()) == 0 {
 			t.Error("private key is missing")
 		}

 		pubKey := privKey.GetPublicKey()
 		if pubKey.GetVersion() != 0 {
 			t.Errorf("public key version = %d, want %d", pubKey.GetVersion(), 0)
 		}
 		if !cmp.Equal(pubKey.GetParams(), keyFormat.GetParams(), protocmp.Transform()) {
 			t.Errorf("key params = %v, want %v", pubKey.GetParams(), keyFormat.GetParams())
 		}
 		if len(pubKey.GetPublicKey()) == 0 {
 			t.Error("public key is missing")
 		}
 	}
 }

 func TestPrivateKeyManagerPublicKeyDataAcceptsNilKey(t *testing.T) {
 	k, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
 	if err != nil {
 		t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
 	}
 	km, ok := k.(registry.PrivateKeyManager)
 	if !ok {
 		t.Errorf("primitive is not PrivateKeyManager")
 	}
 	if _, err := km.PublicKeyData(nil); err != nil {
 		t.Errorf("PublicKeyData() err = %v, want nil", err)
 	}
 }

 func TestPrivateKeyManagerPublicKeyData(t *testing.T) {
 	k, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
 	if err != nil {
 		t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
 	}
 	km, ok := k.(registry.PrivateKeyManager)
 	if !ok {
 		t.Errorf("primitive is not PrivateKeyManager")
 	}

 	params := &hpkepb.HpkeParams{
 		Kem:  hpkepb.HpkeKem_DHKEM_X25519_HKDF_SHA256,
 		Kdf:  hpkepb.HpkeKdf_HKDF_SHA256,
 		Aead: hpkepb.HpkeAead_AES_256_GCM,
 	}
 	_, privKey := pubPrivKeys(t, params)
 	serializedPrivKey, err := proto.Marshal(privKey)
 	if err != nil {
 		t.Fatal(err)
 	}
 	wantPubKey := privKey.GetPublicKey()
 	serializedPubKey, err := proto.Marshal(wantPubKey)
 	if err != nil {
 		t.Fatal(err)
 	}

 	pubKey, err := km.PublicKeyData(serializedPrivKey)
 	if err != nil {
 		t.Fatalf("PublicKeyData() err = %v, want nil", err)
 	}
 	if got, want := pubKey.GetTypeUrl(), hpkePublicKeyTypeURL; got != want {
 		t.Errorf("type URL = %q, want %q", got, want)
 	}
 	if !bytes.Equal(pubKey.GetValue(), serializedPubKey) {
 		t.Errorf("value = %v, want %v", pubKey.GetValue(), serializedPubKey)
 	}
 	if got, want := pubKey.GetKeyMaterialType(), tinkpb.KeyData_ASYMMETRIC_PUBLIC; got != want {
 		t.Errorf("Key material type = %d, want %d", got, want)
 	}
 }

 func TestPrivateKeyManagerDoesSupport(t *testing.T) {
 	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
 	if err != nil {
 		t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
 	}
 	if !km.DoesSupport(hpkePrivateKeyTypeURL) {
 		t.Errorf("DoesSupport(%q) = false, want true", hpkePrivateKeyTypeURL)
 	}
 	unsupportedKeyTypeURL := "unsupported.key.type"
 	if km.DoesSupport(unsupportedKeyTypeURL) {
 		t.Errorf("DoesSupport(%q) = true, want false", unsupportedKeyTypeURL)
 	}
 }

 func TestPrivateKeyManagerTypeURL(t *testing.T) {
 	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
 	if err != nil {
 		t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
 	}
 	if km.TypeURL() != hpkePrivateKeyTypeURL {
 		t.Errorf("TypeURL = %q, want %q", km.TypeURL(), hpkePrivateKeyTypeURL)
 	}
 }
	// Copyright 2022 Google LLC
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.
	//
	////////////////////////////////////////////////////////////////////////////////

	package hybrid

	import (
	"bytes"
	"testing"

	"github.com/google/go-cmp/cmp"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/testing/protocmp"
	"github.com/google/tink/go/core/registry"
	"github.com/google/tink/go/hybrid/internal/hpke"
	"github.com/google/tink/go/subtle/random"
	hpkepb "github.com/google/tink/go/proto/hpke_go_proto"
	tinkpb "github.com/google/tink/go/proto/tink_go_proto"
	)

	func TestPrivateKeyManagerPrimitiveRejectsInvalidKeyVersion(t *testing.T) {
	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
	if err != nil {
	t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
	}
	params := &hpkepb.HpkeParams{
	Kem: hpkepb.HpkeKem_DHKEM_X25519_HKDF_SHA256,
	Kdf: hpkepb.HpkeKdf_HKDF_SHA256,
	Aead: hpkepb.HpkeAead_AES_256_GCM,
	}
	_, privKey := pubPrivKeys(t, params)
	privKey.Version = 1
	serializedPrivKey, err := proto.Marshal(privKey)
	if err != nil {
	t.Fatal(err)
	}
	if _, err := km.Primitive(serializedPrivKey); err == nil {
	t.Error("Primitive() err = nil, want error")
	}
	}

	func TestPrivateKeyManagerPrimitiveRejectsInvalidParams(t *testing.T) {
	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
	if err != nil {
	t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
	}

	tests := []struct {
	name string
	params *hpkepb.HpkeParams
	}{
	{"kem", &hpkepb.HpkeParams{
	Kem: hpkepb.HpkeKem_KEM_UNKNOWN,
	Kdf: hpkepb.HpkeKdf_HKDF_SHA256,
	Aead: hpkepb.HpkeAead_AES_256_GCM,
	}},
	{"kdf", &hpkepb.HpkeParams{
	Kem: hpkepb.HpkeKem_DHKEM_X25519_HKDF_SHA256,
	Kdf: hpkepb.HpkeKdf_KDF_UNKNOWN,
	Aead: hpkepb.HpkeAead_AES_256_GCM,
	}},
	{"aead", &hpkepb.HpkeParams{
	Kem: hpkepb.HpkeKem_DHKEM_X25519_HKDF_SHA256,
	Kdf: hpkepb.HpkeKdf_HKDF_SHA256,
	Aead: hpkepb.HpkeAead_AEAD_UNKNOWN,
	}},
	}

	for _, test := range tests {
	t.Run(test.name, func(t *testing.T) {
	_, serializedPrivKey := serializedPubPrivKeys(t, test.params)
	if _, err := km.Primitive(serializedPrivKey); err == nil {
	t.Error("Primitive() err = nil, want error")
	}
	})
	}
	}

	func TestPrivateKeyManagerPrimitiveRejectsMissingParams(t *testing.T) {
	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
	if err != nil {
	t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
	}
	_, serializedPrivKey := serializedPubPrivKeys(t, nil)
	if _, err := km.Primitive(serializedPrivKey); err == nil {
	t.Error("Primitive() err = nil, want error")
	}
	}

	func TestPrivateKeyManagerPrimitiveRejectsNilKey(t *testing.T) {
	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
	if err != nil {
	t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
	}
	if _, err := km.Primitive(nil); err == nil {
	t.Error("Primitive() err = nil, want error")
	}
	}

	func TestPrivateKeyManagerPrimitiveEncryptDecrypt(t *testing.T) {
	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
	if err != nil {
	t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
	}
	pt := random.GetRandomBytes(200)
	ctxInfo := random.GetRandomBytes(100)

	for _, aeadID := range hpkeAEADs {
	params := &hpkepb.HpkeParams{
	Kem: hpkepb.HpkeKem_DHKEM_X25519_HKDF_SHA256,
	Kdf: hpkepb.HpkeKdf_HKDF_SHA256,
	Aead: aeadID,
	}
	pubKey, privKey := pubPrivKeys(t, params)
	serializedPrivKey, err := proto.Marshal(privKey)
	if err != nil {
	t.Fatal(err)
	}

	enc, err := hpke.NewEncrypt(pubKey)
	if err != nil {
	t.Fatalf("hpke.NewEncrypt() err = %v, want nil", err)
	}
	d, err := km.Primitive(serializedPrivKey)
	if err != nil {
	t.Fatalf("Primitive() err = %v, want nil", err)
	}
	dec, ok := d.(*hpke.Decrypt)
	if !ok {
	t.Fatal("primitive is not Decrypt")
	}

	ct, err := enc.Encrypt(pt, ctxInfo)
	if err != nil {
	t.Fatalf("Encrypt() err = %v, want nil", err)
	}
	gotPT, err := dec.Decrypt(ct, ctxInfo)
	if err != nil {
	t.Fatalf("Decrypt() err = %v, want nil", err)
	}
	if want := pt; !bytes.Equal(gotPT, want) {
	t.Errorf("Decrypt() = %x, want %x", gotPT, want)
	}
	}
	}

	func TestPrivateKeyManagerNewKeyRejectsNilKeyFormat(t *testing.T) {
	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
	if err != nil {
	t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
	}
	if _, err := km.NewKey(nil); err == nil {
	t.Error("NewKey() err = nil, want error")
	}
	}

	func TestPrivateKeyManagerNewKeyRejectsInvalidKeyFormat(t *testing.T) {
	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
	if err != nil {
	t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
	}
	serializedKeyFormatUnknownKEM, err := proto.Marshal(
	&hpkepb.HpkeParams{
	Kem: hpkepb.HpkeKem_KEM_UNKNOWN,
	Kdf: hpkepb.HpkeKdf_HKDF_SHA256,
	Aead: hpkepb.HpkeAead_AES_256_GCM,
	})
	if err != nil {
	t.Fatal(err)
	}
	serializedKeyFormatUnknownKDF, err := proto.Marshal(
	&hpkepb.HpkeParams{
	Kem: hpkepb.HpkeKem_DHKEM_X25519_HKDF_SHA256,
	Kdf: hpkepb.HpkeKdf_KDF_UNKNOWN,
	Aead: hpkepb.HpkeAead_AES_256_GCM,
	})
	if err != nil {
	t.Fatal(err)
	}
	serializedKeyFormatUnknownAEAD, err := proto.Marshal(
	&hpkepb.HpkeParams{
	Kem: hpkepb.HpkeKem_DHKEM_X25519_HKDF_SHA256,
	Kdf: hpkepb.HpkeKdf_HKDF_SHA256,
	Aead: hpkepb.HpkeAead_AEAD_UNKNOWN,
	})
	if err != nil {
	t.Fatal(err)
	}

	tests := []struct {
	name string
	keyFormat []byte
	}{
	{"kem", serializedKeyFormatUnknownKEM},
	{"kdf", serializedKeyFormatUnknownKDF},
	{"aead", serializedKeyFormatUnknownAEAD},
	}

	for _, test := range tests {
	t.Run(test.name, func(t *testing.T) {
	if _, err := km.NewKey(test.keyFormat); err == nil {
	t.Error("NewKey() err = nil, want error")
	}
	})
	}
	}

	func TestPrivateKeyManagerNewKeyEncryptDecrypt(t *testing.T) {
	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
	if err != nil {
	t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
	}

	wantPT := random.GetRandomBytes(200)
	ctxInfo := random.GetRandomBytes(100)

	for _, aeadID := range hpkeAEADs {
	keyFormat := &hpkepb.HpkeKeyFormat{
	Params: &hpkepb.HpkeParams{
	Kem: hpkepb.HpkeKem_DHKEM_X25519_HKDF_SHA256,
	Kdf: hpkepb.HpkeKdf_HKDF_SHA256,
	Aead: aeadID,
	},
	}
	serializedKeyFormat, err := proto.Marshal(keyFormat)
	if err != nil {
	t.Fatal(err)
	}

	privKeyProto, err := km.NewKey(serializedKeyFormat)
	if err != nil {
	t.Fatalf("NewKey() err = %v, want nil", err)
	}
	privKey, ok := privKeyProto.(*hpkepb.HpkePrivateKey)
	if !ok {
	t.Fatal("primitive is not HpkePrivateKey")
	}
	if privKey.GetVersion() != 0 {
	t.Errorf("private key version = %d, want %d", privKey.GetVersion(), 0)
	}
	if len(privKey.GetPrivateKey()) == 0 {
	t.Error("private key is missing")
	}

	pubKey := privKey.GetPublicKey()
	if pubKey.GetVersion() != 0 {
	t.Errorf("public key version = %d, want %d", pubKey.GetVersion(), 0)
	}
	if !cmp.Equal(pubKey.GetParams(), keyFormat.GetParams(), protocmp.Transform()) {
	t.Errorf("key params = %v, want %v", pubKey.GetParams(), keyFormat.GetParams())
	}
	if len(pubKey.GetPublicKey()) == 0 {
	t.Error("public key is missing")
	}

	enc, err := hpke.NewEncrypt(pubKey)
	if err != nil {
	t.Fatalf("hpke.NewEncrypt() err = %v, want nil", err)
	}
	serializedPrivKey, err := proto.Marshal(privKeyProto)
	if err != nil {
	t.Fatal(err)
	}
	d, err := km.Primitive(serializedPrivKey)
	if err != nil {
	t.Fatalf("Primitive() err = %v, want nil", err)
	}
	dec, ok := d.(*hpke.Decrypt)
	if !ok {
	t.Fatal("primitive is not Decrypt")
	}

	ct, err := enc.Encrypt(wantPT, ctxInfo)
	if err != nil {
	t.Fatalf("Encrypt() err = %v, want nil", err)
	}
	gotPT, err := dec.Decrypt(ct, ctxInfo)
	if err != nil {
	t.Fatalf("Decrypt() err = %v, want nil", err)
	}
	if !bytes.Equal(gotPT, wantPT) {
	t.Errorf("Decrypt() = %x, want %x", gotPT, wantPT)
	}
	}
	}

	func TestPrivateKeyManagerNewKeyDataRejectsNilKeyFormat(t *testing.T) {
	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
	if err != nil {
	t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
	}
	if _, err := km.NewKeyData(nil); err == nil {
	t.Error("NewKey() err = nil, want error")
	}
	}

	func TestPrivateKeyManagerNewKeyData(t *testing.T) {
	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
	if err != nil {
	t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
	}

	for _, aeadID := range hpkeAEADs {
	keyFormat := &hpkepb.HpkeKeyFormat{
	Params: &hpkepb.HpkeParams{
	Kem: hpkepb.HpkeKem_DHKEM_X25519_HKDF_SHA256,
	Kdf: hpkepb.HpkeKdf_HKDF_SHA256,
	Aead: aeadID,
	},
	}
	serializedKeyFormat, err := proto.Marshal(keyFormat)
	if err != nil {
	t.Fatal(err)
	}

	keyData, err := km.NewKeyData(serializedKeyFormat)
	if err != nil {
	t.Fatalf("NewKeyData() err = %v, want nil", err)
	}
	if got, want := keyData.GetTypeUrl(), hpkePrivateKeyTypeURL; got != want {
	t.Errorf("type URL = %q, want %q", got, want)
	}
	if got, want := keyData.GetKeyMaterialType(), tinkpb.KeyData_ASYMMETRIC_PRIVATE; got != want {
	t.Errorf("key material type = %d, want %d", got, want)
	}

	privKey := new(hpkepb.HpkePrivateKey)
	if err := proto.Unmarshal(keyData.GetValue(), privKey); err != nil {
	t.Fatalf("Unmarshal err = %v, want nil", err)
	}
	if privKey.GetVersion() != 0 {
	t.Errorf("private key version = %d, want %d", privKey.GetVersion(), 0)
	}
	if len(privKey.GetPrivateKey()) == 0 {
	t.Error("private key is missing")
	}

	pubKey := privKey.GetPublicKey()
	if pubKey.GetVersion() != 0 {
	t.Errorf("public key version = %d, want %d", pubKey.GetVersion(), 0)
	}
	if !cmp.Equal(pubKey.GetParams(), keyFormat.GetParams(), protocmp.Transform()) {
	t.Errorf("key params = %v, want %v", pubKey.GetParams(), keyFormat.GetParams())
	}
	if len(pubKey.GetPublicKey()) == 0 {
	t.Error("public key is missing")
	}
	}
	}

	func TestPrivateKeyManagerPublicKeyDataAcceptsNilKey(t *testing.T) {
	k, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
	if err != nil {
	t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
	}
	km, ok := k.(registry.PrivateKeyManager)
	if !ok {
	t.Errorf("primitive is not PrivateKeyManager")
	}
	if _, err := km.PublicKeyData(nil); err != nil {
	t.Errorf("PublicKeyData() err = %v, want nil", err)
	}
	}

	func TestPrivateKeyManagerPublicKeyData(t *testing.T) {
	k, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
	if err != nil {
	t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
	}
	km, ok := k.(registry.PrivateKeyManager)
	if !ok {
	t.Errorf("primitive is not PrivateKeyManager")
	}

	params := &hpkepb.HpkeParams{
	Kem: hpkepb.HpkeKem_DHKEM_X25519_HKDF_SHA256,
	Kdf: hpkepb.HpkeKdf_HKDF_SHA256,
	Aead: hpkepb.HpkeAead_AES_256_GCM,
	}
	_, privKey := pubPrivKeys(t, params)
	serializedPrivKey, err := proto.Marshal(privKey)
	if err != nil {
	t.Fatal(err)
	}
	wantPubKey := privKey.GetPublicKey()
	serializedPubKey, err := proto.Marshal(wantPubKey)
	if err != nil {
	t.Fatal(err)
	}

	pubKey, err := km.PublicKeyData(serializedPrivKey)
	if err != nil {
	t.Fatalf("PublicKeyData() err = %v, want nil", err)
	}
	if got, want := pubKey.GetTypeUrl(), hpkePublicKeyTypeURL; got != want {
	t.Errorf("type URL = %q, want %q", got, want)
	}
	if !bytes.Equal(pubKey.GetValue(), serializedPubKey) {
	t.Errorf("value = %v, want %v", pubKey.GetValue(), serializedPubKey)
	}
	if got, want := pubKey.GetKeyMaterialType(), tinkpb.KeyData_ASYMMETRIC_PUBLIC; got != want {
	t.Errorf("Key material type = %d, want %d", got, want)
	}
	}

	func TestPrivateKeyManagerDoesSupport(t *testing.T) {
	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
	if err != nil {
	t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
	}
	if !km.DoesSupport(hpkePrivateKeyTypeURL) {
	t.Errorf("DoesSupport(%q) = false, want true", hpkePrivateKeyTypeURL)
	}
	unsupportedKeyTypeURL := "unsupported.key.type"
	if km.DoesSupport(unsupportedKeyTypeURL) {
	t.Errorf("DoesSupport(%q) = true, want false", unsupportedKeyTypeURL)
	}
	}

	func TestPrivateKeyManagerTypeURL(t *testing.T) {
	km, err := registry.GetKeyManager(hpkePrivateKeyTypeURL)
	if err != nil {
	t.Fatalf("GetKeyManager(%q) err = %v, want nil", hpkePrivateKeyTypeURL, err)
	}
	if km.TypeURL() != hpkePrivateKeyTypeURL {
	t.Errorf("TypeURL = %q, want %q", km.TypeURL(), hpkePrivateKeyTypeURL)
	}
	}