// Copyright 2019 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 keyset_test

import (
	"bytes"
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"testing"

	"github.com/google/go-cmp/cmp"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/testing/protocmp"
	"github.com/tink-crypto/tink-go/v2/aead"
	"github.com/tink-crypto/tink-go/v2/core/registry"
	"github.com/tink-crypto/tink-go/v2/insecuresecretdataaccess"
	"github.com/tink-crypto/tink-go/v2/internal"
	"github.com/tink-crypto/tink-go/v2/internal/internalregistry"
	"github.com/tink-crypto/tink-go/v2/internal/monitoringutil"
	"github.com/tink-crypto/tink-go/v2/internal/protoserialization"
	"github.com/tink-crypto/tink-go/v2/keyset"
	"github.com/tink-crypto/tink-go/v2/mac/aescmac"
	"github.com/tink-crypto/tink-go/v2/mac/hmac"
	"github.com/tink-crypto/tink-go/v2/mac"
	"github.com/tink-crypto/tink-go/v2/monitoring"
	"github.com/tink-crypto/tink-go/v2/secretdata"
	"github.com/tink-crypto/tink-go/v2/signature/rsassapss"
	"github.com/tink-crypto/tink-go/v2/signature"
	"github.com/tink-crypto/tink-go/v2/testing/fakekms"
	"github.com/tink-crypto/tink-go/v2/testing/fakemonitoring"
	"github.com/tink-crypto/tink-go/v2/testkeyset"
	"github.com/tink-crypto/tink-go/v2/testutil"
	"github.com/tink-crypto/tink-go/v2/tink"
	tinkpb "github.com/tink-crypto/tink-go/v2/proto/tink_go_proto"
)

func TestNewHandle(t *testing.T) {
	template := mac.HMACSHA256Tag128KeyTemplate()
	handle, err := keyset.NewHandle(template)
	if err != nil {
		t.Errorf("keyset.NewHandle(template) = %v, want nil", err)
	}
	ks := testkeyset.KeysetMaterial(handle)
	if len(ks.Key) != 1 {
		t.Errorf("len(ks.Key) = %d, want 1", len(ks.Key))
	}
	key := ks.Key[0]
	if ks.PrimaryKeyId != key.KeyId {
		t.Errorf("ks.PrimaryKeyId = %d, want %d", ks.PrimaryKeyId, key.KeyId)
	}
	if key.KeyData.TypeUrl != template.TypeUrl {
		t.Errorf("key.KeyData.TypeUrl = %v, want %v", key.KeyData.TypeUrl, template.TypeUrl)
	}
	if _, err = mac.New(handle); err != nil {
		t.Errorf("mac.New(handle) err = %v, want nil", err)
	}
}

func TestKeysetMaterialMakesACopy(t *testing.T) {
	wantProtoKeyset := testutil.NewKeyset(1, []*tinkpb.Keyset_Key{
		testutil.NewKey(testutil.NewKeyData("some type url", []byte{0}, tinkpb.KeyData_SYMMETRIC), tinkpb.KeyStatusType_ENABLED, 1, tinkpb.OutputPrefixType_TINK),
	})
	handle, err := testkeyset.NewHandle(wantProtoKeyset)
	if err != nil {
		t.Errorf("testkeyset.NewHandle(wantProtoKeyset) = %v, want nil", err)
	}
	gotProtoKeyset := testkeyset.KeysetMaterial(handle)
	if wantProtoKeyset == gotProtoKeyset {
		t.Errorf("testkeyset.KeysetMaterial(handle) = %v, want a copy of %v", gotProtoKeyset, wantProtoKeyset)
	}
	if !proto.Equal(gotProtoKeyset, wantProtoKeyset) {
		t.Errorf("testkeyset.NewHandle(wantProtoKeyset) = %v, want %v", gotProtoKeyset, wantProtoKeyset)
	}
}

func TestNewHandleExistingKeyset(t *testing.T) {
	testCases := []struct {
		name string
		ks   *tinkpb.Keyset
	}{
		{
			name: "one enabled key",
			ks: &tinkpb.Keyset{
				PrimaryKeyId: 1,
				Key: []*tinkpb.Keyset_Key{
					&tinkpb.Keyset_Key{
						KeyId:            1,
						Status:           tinkpb.KeyStatusType_ENABLED,
						OutputPrefixType: tinkpb.OutputPrefixType_TINK,
						KeyData:          testutil.NewKeyData("some type url", []byte{0}, tinkpb.KeyData_SYMMETRIC),
					},
				},
			},
		},
		{
			name: "multiple keys",
			ks: &tinkpb.Keyset{
				PrimaryKeyId: 1,
				Key: []*tinkpb.Keyset_Key{
					&tinkpb.Keyset_Key{
						KeyId:            1,
						Status:           tinkpb.KeyStatusType_ENABLED,
						OutputPrefixType: tinkpb.OutputPrefixType_TINK,
						KeyData:          testutil.NewKeyData("some type url", []byte{0}, tinkpb.KeyData_SYMMETRIC),
					},
					&tinkpb.Keyset_Key{
						KeyId:            2,
						Status:           tinkpb.KeyStatusType_DISABLED,
						OutputPrefixType: tinkpb.OutputPrefixType_TINK,
						KeyData:          testutil.NewKeyData("some type url", []byte{0}, tinkpb.KeyData_SYMMETRIC),
					},
				},
			},
		},
	}
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			wantProtoKeyset := tc.ks
			handle, err := testkeyset.NewHandle(wantProtoKeyset)
			if err != nil {
				t.Fatalf("testkeyset.NewHandle(wantProtoKeyset) = %v, want nil", err)
			}
			gotProtoKeyset := testkeyset.KeysetMaterial(handle)
			if !proto.Equal(gotProtoKeyset, wantProtoKeyset) {
				t.Errorf("testkeyset.NewHandle(wantProtoKeyset) = %v, want %v", gotProtoKeyset, wantProtoKeyset)
			}
		})
	}
}

func TestNewHandleWithInvalidTypeURLFails(t *testing.T) {
	// template with unknown TypeURL
	invalidTemplate := mac.HMACSHA256Tag128KeyTemplate()
	invalidTemplate.TypeUrl = "some unknown TypeURL"
	if _, err := keyset.NewHandle(invalidTemplate); err == nil {
		t.Errorf("keyset.NewHandle(invalidTemplate) err = nil, want error")
	}
}

func TestNewHandleWithNilTemplateFails(t *testing.T) {
	if _, err := keyset.NewHandle(nil); err == nil {
		t.Error("keyset.NewHandle(nil) err = nil, want error")
	}
}

func TestWriteAndReadInBinary(t *testing.T) {
	keysetEncryptionHandle, err := keyset.NewHandle(aead.AES128GCMKeyTemplate())
	if err != nil {
		t.Errorf("keyset.NewHandle(aead.AES128GCMKeyTemplate()) err = %v, want nil", err)
	}
	keysetEncryptionAead, err := aead.New(keysetEncryptionHandle)
	if err != nil {
		t.Errorf("aead.New(keysetEncryptionHandle) err = %v, want nil", err)
	}

	handle, err := keyset.NewHandle(mac.HMACSHA256Tag128KeyTemplate())
	if err != nil {
		t.Fatalf("keyset.NewHandle(mac.HMACSHA256Tag128KeyTemplate()) err = %v, want nil", err)
	}

	buff := &bytes.Buffer{}
	err = handle.Write(keyset.NewBinaryWriter(buff), keysetEncryptionAead)
	if err != nil {
		t.Fatalf("handle.Write(keyset.NewBinaryWriter(buff), keysetEncryptionAead) err = %v, want nil", err)
	}
	encrypted := buff.Bytes()

	gotHandle, err := keyset.Read(keyset.NewBinaryReader(bytes.NewBuffer(encrypted)), keysetEncryptionAead)
	if err != nil {
		t.Fatalf("keyset.Read() err = %v, want nil", err)
	}

	if !proto.Equal(testkeyset.KeysetMaterial(gotHandle), testkeyset.KeysetMaterial(handle)) {
		t.Fatalf("keyset.Read() = %v, want %v", gotHandle, handle)
	}
}

func TestWriteAndReadInJSON(t *testing.T) {
	keysetEncryptionHandle, err := keyset.NewHandle(aead.AES128GCMKeyTemplate())
	if err != nil {
		t.Errorf("keyset.NewHandle(aead.AES128GCMKeyTemplate()) err = %v, want nil", err)
	}
	keysetEncryptionAead, err := aead.New(keysetEncryptionHandle)
	if err != nil {
		t.Errorf("aead.New(keysetEncryptionHandle) err = %v, want nil", err)
	}

	handle, err := keyset.NewHandle(mac.HMACSHA256Tag128KeyTemplate())
	if err != nil {
		t.Fatalf("keyset.NewHandle(mac.HMACSHA256Tag128KeyTemplate()) err = %v, want nil", err)
	}

	buff := &bytes.Buffer{}
	err = handle.Write(keyset.NewJSONWriter(buff), keysetEncryptionAead)
	if err != nil {
		t.Fatalf("h.Write(keyset.NewJSONWriter(buff), keysetEncryptionAead) err = %v, want nil", err)
	}
	encrypted := buff.Bytes()

	gotHandle, err := keyset.Read(keyset.NewJSONReader(bytes.NewBuffer(encrypted)), keysetEncryptionAead)
	if err != nil {
		t.Fatalf("keyset.Read() err = %v, want nil", err)
	}

	if !proto.Equal(testkeyset.KeysetMaterial(gotHandle), testkeyset.KeysetMaterial(handle)) {
		t.Fatalf("keyset.Read() = %v, want %v", gotHandle, handle)
	}
}

const fakeKeyURI = "fake-kms://CM2b3_MDElQKSAowdHlwZS5nb29nbGVhcGlzLmNvbS9nb29nbGUuY3J5cHRvLnRpbmsuQWVzR2NtS2V5EhIaEIK75t5L-adlUwVhWvRuWUwYARABGM2b3_MDIAE"

func TestWriteAndReadWithAssociatedData(t *testing.T) {
	keysetEncryptionAead, err := fakekms.NewAEAD(fakeKeyURI)
	if err != nil {
		t.Fatalf("fakekms.NewAEAD(fakeKeyURI) err = %v, want nil", err)
	}

	handle, err := keyset.NewHandle(mac.HMACSHA256Tag128KeyTemplate())
	if err != nil {
		t.Fatalf("keyset.NewHandle(mac.HMACSHA256Tag128KeyTemplate()) err = %v, want nil", err)
	}
	associatedData := []byte{0x01, 0x02}

	buff := &bytes.Buffer{}
	err = handle.WriteWithAssociatedData(keyset.NewBinaryWriter(buff), keysetEncryptionAead, associatedData)
	if err != nil {
		t.Fatalf("handle.WriteWithAssociatedData() err = %v, want nil", err)
	}
	encrypted := buff.Bytes()

	handle2, err := keyset.ReadWithAssociatedData(keyset.NewBinaryReader(bytes.NewBuffer(encrypted)), keysetEncryptionAead, associatedData)
	if err != nil {
		t.Fatalf("keyset.ReadWithAssociatedData() err = %v, want nil", err)
	}

	if !proto.Equal(testkeyset.KeysetMaterial(handle), testkeyset.KeysetMaterial(handle2)) {
		t.Errorf("keyset.ReadWithAssociatedData() = %v, want %v", handle2, handle)
	}

	// Test that ReadWithContext is compatible with WriteWithAssociatedData
	kekAEADWithContext, err := fakekms.NewAEADWithContext(fakeKeyURI)
	if err != nil {
		t.Fatalf("fakekms.NewAEADWithContext(fakeKeyURI) err = %v, want nil", err)
	}
	ctx := context.Background()
	handle3, err := keyset.ReadWithContext(ctx, keyset.NewBinaryReader(bytes.NewBuffer(encrypted)), kekAEADWithContext, associatedData)
	if err != nil {
		t.Fatalf("keyset.ReadWithContext() err = %v, want nil", err)
	}
	if !proto.Equal(testkeyset.KeysetMaterial(handle), testkeyset.KeysetMaterial(handle3)) {
		t.Errorf("keyset.ReadWithContext() = %v, want %v", handle3, handle)
	}
}

func TestReadWithMismatchedAssociatedData(t *testing.T) {
	keysetEncryptionHandle, err := keyset.NewHandle(aead.AES128GCMKeyTemplate())
	if err != nil {
		t.Errorf("keyset.NewHandle(aead.AES128GCMKeyTemplate()) err = %v, want nil", err)
	}
	keysetEncryptionAead, err := aead.New(keysetEncryptionHandle)
	if err != nil {
		t.Errorf("aead.New(keysetEncryptionHandle) err = %v, want nil", err)
	}

	handle, err := keyset.NewHandle(mac.HMACSHA256Tag128KeyTemplate())
	if err != nil {
		t.Fatalf("keyset.NewHandle(mac.HMACSHA256Tag128KeyTemplate()) err = %v, want nil", err)
	}
	associatedData := []byte{0x01, 0x02}

	buff := &bytes.Buffer{}
	err = handle.WriteWithAssociatedData(keyset.NewBinaryWriter(buff), keysetEncryptionAead, associatedData)
	if err != nil {
		t.Fatalf("handle.WriteWithAssociatedData() err = %v, want nil", err)
	}
	encrypted := buff.Bytes()

	invalidAssociatedData := []byte{0x01, 0x03}
	_, err = keyset.ReadWithAssociatedData(keyset.NewBinaryReader(bytes.NewBuffer(encrypted)), keysetEncryptionAead, invalidAssociatedData)
	if err == nil {
		t.Errorf("keyset.ReadWithAssociatedData() err = nil, want err")
	}
}

func TestWriteAndReadWithContext(t *testing.T) {
	kekAEADWithContext, err := fakekms.NewAEADWithContext(fakeKeyURI)
	if err != nil {
		t.Fatalf("fakekms.NewAEADWithContext(fakeKeyURI) err = %v, want nil", err)
	}

	handle, err := keyset.NewHandle(mac.HMACSHA256Tag128KeyTemplate())
	if err != nil {
		t.Fatalf("keyset.NewHandle(mac.HMACSHA256Tag128KeyTemplate()) err = %v, want nil", err)
	}
	associatedData := []byte{0x01, 0x02}

	ctx := context.Background()
	buff := &bytes.Buffer{}
	err = handle.WriteWithContext(ctx, keyset.NewBinaryWriter(buff), kekAEADWithContext, associatedData)
	if err != nil {
		t.Fatalf("handle.WriteWithContext() err = %v, want nil", err)
	}
	encrypted := buff.Bytes()

	handle2, err := keyset.ReadWithContext(ctx, keyset.NewBinaryReader(bytes.NewBuffer(encrypted)), kekAEADWithContext, associatedData)
	if err != nil {
		t.Fatalf("keyset.ReadWithContext() err = %v, want nil", err)
	}

	if !proto.Equal(testkeyset.KeysetMaterial(handle), testkeyset.KeysetMaterial(handle2)) {
		t.Errorf("keyset.ReadWithContext() = %v, want %v", handle2, handle)
	}

	invalidAssociatedData := []byte{0x01, 0x03}
	_, err = keyset.ReadWithContext(ctx, keyset.NewBinaryReader(bytes.NewBuffer(encrypted)), kekAEADWithContext, invalidAssociatedData)
	if err == nil {
		t.Errorf("keyset.ReadWithContext() err = nil, want err")
	}

	// Test that ReadWithAssociatedData is compatible with WriteWithContext
	kekAEAD, err := fakekms.NewAEAD(fakeKeyURI)
	if err != nil {
		t.Fatalf("fakekms.NewAEAD(fakeKeyURI) err = %v, want nil", err)
	}
	handle3, err := keyset.ReadWithAssociatedData(keyset.NewBinaryReader(bytes.NewBuffer(encrypted)), kekAEAD, associatedData)
	if err != nil {
		t.Fatalf("keyset.ReadWithAssociatedData() err = %v, want nil", err)
	}
	if !proto.Equal(testkeyset.KeysetMaterial(handle), testkeyset.KeysetMaterial(handle3)) {
		t.Errorf("keyset.ReadWithAssociatedData() = %v, want %v", handle3, handle)
	}
}

func TestWriteWithContextDoesNotIgnoreContext(t *testing.T) {
	kekAEADWithContext, err := fakekms.NewAEADWithContext(fakeKeyURI)
	if err != nil {
		t.Fatalf("fakekms.NewAEADWithContext(fakeKeyURI) err = %v, want nil", err)
	}

	handle, err := keyset.NewHandle(mac.HMACSHA256Tag128KeyTemplate())
	if err != nil {
		t.Fatalf("keyset.NewHandle(mac.HMACSHA256Tag128KeyTemplate()) err = %v, want nil", err)
	}
	associatedData := []byte{0x01, 0x02}

	canceledCtx, cancel := context.WithCancelCause(context.Background())
	causeErr := errors.New("cause error message")
	cancel(causeErr)

	buff := &bytes.Buffer{}
	err = handle.WriteWithContext(canceledCtx, keyset.NewBinaryWriter(buff), kekAEADWithContext, associatedData)
	if err == nil {
		t.Errorf("handle.WriteWithContext() err = nil, want error")
	}
}

func TestReadWithContextDoesNotIgnoreContext(t *testing.T) {
	kekAEADWithContext, err := fakekms.NewAEADWithContext(fakeKeyURI)
	if err != nil {
		t.Fatalf("fakekms.NewAEADWithContext(fakeKeyURI) err = %v, want nil", err)
	}

	handle, err := keyset.NewHandle(mac.HMACSHA256Tag128KeyTemplate())
	if err != nil {
		t.Fatalf("keyset.NewHandle(mac.HMACSHA256Tag128KeyTemplate()) err = %v, want nil", err)
	}
	associatedData := []byte{0x01, 0x02}

	ctx := context.Background()
	buff := &bytes.Buffer{}
	err = handle.WriteWithContext(ctx, keyset.NewBinaryWriter(buff), kekAEADWithContext, associatedData)
	if err != nil {
		t.Fatalf("handle.WriteWithContext() err = %v, want nil", err)
	}
	encrypted := buff.Bytes()

	canceledCtx, cancel := context.WithCancelCause(ctx)
	causeErr := errors.New("cause error message")
	cancel(causeErr)

	_, err = keyset.ReadWithContext(canceledCtx, keyset.NewBinaryReader(bytes.NewBuffer(encrypted)), kekAEADWithContext, associatedData)
	if err == nil {
		t.Errorf("keyset.ReadWithContext() err = nil, want error")
	}
}

func TestPrimaryReturnsError(t *testing.T) {
	testCases := []struct {
		name   string
		handle *keyset.Handle
	}{
		{
			name:   "zero value handle",
			handle: &keyset.Handle{},
		},
		{
			name:   "nil handle",
			handle: nil,
		},
	}
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			_, err := tc.handle.Primary()
			if err == nil {
				t.Errorf("handle.Primary() err = nil, want err")
			}
		})
	}
}

func TestLenReturnsZero(t *testing.T) {
	testCases := []struct {
		name   string
		handle *keyset.Handle
	}{
		{
			name:   "zero value handle",
			handle: &keyset.Handle{},
		},
		{
			name:   "nil handle",
			handle: nil,
		},
	}
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			length := tc.handle.Len()
			if length != 0 {
				t.Errorf("handle.Len() = %v, want 0", length)
			}
		})
	}
}

func TestPublicReturnsError(t *testing.T) {
	testCases := []struct {
		name   string
		handle *keyset.Handle
	}{
		{
			name:   "zero value handle",
			handle: &keyset.Handle{},
		},
		{
			name:   "nil handle",
			handle: nil,
		},
	}
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			_, err := tc.handle.Public()
			if err == nil {
				t.Errorf("handle.Public() err = nil, want err")
			}
		})
	}
}

func TestEntryReturnsError(t *testing.T) {
	testCases := []struct {
		name   string
		handle *keyset.Handle
	}{
		{
			name:   "zero value handle",
			handle: &keyset.Handle{},
		},
		{
			name:   "nil handle",
			handle: nil,
		},
	}
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			_, err := tc.handle.Entry(0)
			if err == nil {
				t.Errorf("handle.Entry(0) err = nil, want err")
			}
		})
	}
}

func TestKeysetInfoPanics(t *testing.T) {
	testCases := []struct {
		name   string
		handle *keyset.Handle
	}{
		{
			name:   "zero value handle",
			handle: &keyset.Handle{},
		},
		{
			name:   "nil handle",
			handle: nil,
		},
	}
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			defer func() {
				if r := recover(); r == nil {
					t.Errorf("tc.handle.KeysetInfo() did not panic")
				}
			}()
			_ = tc.handle.KeysetInfo()
		})
	}
}

func TestStringPanics(t *testing.T) {
	testCases := []struct {
		name   string
		handle *keyset.Handle
	}{
		{
			name:   "zero value handle",
			handle: &keyset.Handle{},
		},
		{
			name:   "nil handle",
			handle: nil,
		},
	}
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			defer func() {
				if r := recover(); r == nil {
					t.Errorf("tc.handle.String() did not panic")
				}
			}()
			_ = tc.handle.String()
		})
	}
}

func TestWriteReturnsError(t *testing.T) {
	keysetEncryptionHandle, err := keyset.NewHandle(aead.AES128GCMKeyTemplate())
	if err != nil {
		t.Errorf("keyset.NewHandle(aead.AES128GCMKeyTemplate()) err = %v, want nil", err)
	}
	keysetEncryptionAEAD, err := aead.New(keysetEncryptionHandle)
	if err != nil {
		t.Errorf("aead.New(keysetEncryptionHandle) err = %v, want nil", err)
	}
	testCases := []struct {
		name   string
		handle *keyset.Handle
	}{
		{
			name:   "zero value handle",
			handle: &keyset.Handle{},
		},
		{
			name:   "nil handle",
			handle: nil,
		},
	}
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			buff := &bytes.Buffer{}
			if err := tc.handle.Write(keyset.NewBinaryWriter(buff), keysetEncryptionAEAD); err == nil {
				t.Error("handle.Write() err = nil, want err")
			}
		})
	}
}

func TestWriteWithAssociatedDataReturnsError(t *testing.T) {
	keysetEncryptionHandle, err := keyset.NewHandle(aead.AES128GCMKeyTemplate())
	if err != nil {
		t.Errorf("keyset.NewHandle(aead.AES128GCMKeyTemplate()) err = %v, want nil", err)
	}
	keysetEncryptionAEAD, err := aead.New(keysetEncryptionHandle)
	if err != nil {
		t.Errorf("aead.New(keysetEncryptionHandle) err = %v, want nil", err)
	}

	testCases := []struct {
		name   string
		handle *keyset.Handle
	}{
		{
			name:   "zero value handle",
			handle: &keyset.Handle{},
		},
		{
			name:   "nil handle",
			handle: nil,
		},
	}
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			buff := &bytes.Buffer{}
			if err := tc.handle.WriteWithAssociatedData(keyset.NewBinaryWriter(buff), keysetEncryptionAEAD, []byte("aad")); err == nil {
				t.Error("handle.WriteWithAssociatedData() err = nil, want err")
			}
		})
	}
}

func TestWriteWithNoSecretsReturnsError(t *testing.T) {
	testCases := []struct {
		name   string
		handle *keyset.Handle
	}{
		{
			name:   "zero value handle",
			handle: &keyset.Handle{},
		},
		{
			name:   "nil handle",
			handle: nil,
		},
	}
	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			buff := &bytes.Buffer{}
			if err := tc.handle.WriteWithNoSecrets(keyset.NewBinaryWriter(buff)); err == nil {
				t.Error("handle.WriteWithNoSecrets() err = nil, want err")
			}
		})
	}
}

func TestWriteAndReadWithNoSecrets(t *testing.T) {
	// Create a keyset that contains a public key.
	privateHandle, err := keyset.NewHandle(signature.ECDSAP256KeyTemplate())
	if err != nil {
		t.Fatalf("keyset.NewHandle(signature.ECDSAP256KeyTemplate()) err = %v, want nil", err)
	}
	handle, err := privateHandle.Public()
	if err != nil {
		t.Fatalf("privateHandle.Public() err = %v, want nil", err)
	}

	buff := &bytes.Buffer{}
	err = handle.WriteWithNoSecrets(keyset.NewBinaryWriter(buff))
	if err != nil {
		t.Fatalf("handle.WriteWithAssociatedData(keyset.NewBinaryWriter(buff), masterKey, associatedData) err = %v, want nil", err)
	}
	serialized := buff.Bytes()

	// Using ReadWithNoSecrets.
	handle2, err := keyset.ReadWithNoSecrets(keyset.NewBinaryReader(bytes.NewBuffer(serialized)))
	if err != nil {
		t.Fatalf("keyset.ReadWithNoSecrets() err = %v, want nil", err)
	}

	if !proto.Equal(testkeyset.KeysetMaterial(handle), testkeyset.KeysetMaterial(handle2)) {
		t.Fatalf("keyset.ReadWithNoSecrets() = %v, want %v", handle2, handle)
	}

	// Using Read() and then NewHandleWithNoSecrets.
	reader := keyset.NewBinaryReader(bytes.NewBuffer(serialized))
	protoPublicKeyset, err := reader.Read()
	if err != nil {
		t.Fatalf("reader.Read() err = %v, want nil", err)
	}
	handle3, err := keyset.NewHandleWithNoSecrets(protoPublicKeyset)
	if err != nil {
		t.Fatalf("keyset.NewHandleWithNoSecrets() err = %v, want nil", err)
	}

	if !proto.Equal(testkeyset.KeysetMaterial(handle), testkeyset.KeysetMaterial(handle3)) {
		t.Fatalf("keyset.NewHandleWithNoSecrets() = %v, want %v", handle3, handle)
	}
}

func TestNewHandleWithNoSecretsReturnsErrorIfInputIsNil(t *testing.T) {
	if _, err := keyset.NewHandleWithNoSecrets(nil); err == nil {
		t.Fatal("keyset.NewHandleWithNoSecrets(nil) err = nil, want error")
	}
}

func TestWriteWithNoSecretsFailsWithSymmetricSecretKey(t *testing.T) {
	// Create a keyset that contains a symmetric secret key.
	handle, err := keyset.NewHandle(mac.HMACSHA256Tag128KeyTemplate())
	if err != nil {
		t.Fatalf("keyset.NewHandle(aead.HMACSHA256Tag128KeyTemplate()) err = %v, want nil", err)
	}

	buff := &bytes.Buffer{}
	err = handle.WriteWithNoSecrets(keyset.NewBinaryWriter(buff))
	if err == nil {
		t.Error("handle.WriteWithNoSecrets() = nil, want error")
	}
}

func TestReadWithNoSecretsFailsWithSymmetricSecretKey(t *testing.T) {
	// Create a keyset that contains a symmetric secret key.
	handle, err := keyset.NewHandle(mac.HMACSHA256Tag128KeyTemplate())
	if err != nil {
		t.Fatalf("keyset.NewHandle(aead.HMACSHA256Tag128KeyTemplate()) err = %v, want nil", err)
	}
	buff := &bytes.Buffer{}
	err = testkeyset.Write(handle, keyset.NewBinaryWriter(buff))
	if err != nil {
		t.Fatalf("insecurecleartextkeyset.Write(handle, keyset.NewBinaryWriter(buff)) err = %v, want nil", err)
	}
	serialized := buff.Bytes()

	_, err = keyset.ReadWithNoSecrets(keyset.NewBinaryReader(bytes.NewBuffer(serialized)))
	if err == nil {
		t.Error("keyset.ReadWithNoSecrets() = nil, want error")
	}
}

func TestWriteWithNoSecretsFailsWithPrivateKey(t *testing.T) {
	// Create a keyset that contains a private key.
	handle, err := keyset.NewHandle(signature.ECDSAP256KeyTemplate())
	if err != nil {
		t.Fatalf("keyset.NewHandle(signature.ECDSAP256KeyTemplate()) err = %v, want nil", err)
	}

	buff := &bytes.Buffer{}
	if err := handle.WriteWithNoSecrets(keyset.NewBinaryWriter(buff)); err == nil {
		t.Error("handle.WriteWithNoSecrets() = nil, want error")
	}
}

func TestReadWithNoSecretsFailsWithPrivateKey(t *testing.T) {
	// Create a keyset that contains a private key.
	handle, err := keyset.NewHandle(signature.ECDSAP256KeyTemplate())
	if err != nil {
		t.Fatalf("keyset.NewHandle(signature.ECDSAP256KeyTemplate()) err = %v, want nil", err)
	}
	buff := &bytes.Buffer{}
	err = testkeyset.Write(handle, keyset.NewBinaryWriter(buff))
	if err != nil {
		t.Fatalf("insecurecleartextkeyset.Write(handle, keyset.NewBinaryWriter(buff)) err = %v, want nil", err)
	}
	serialized := buff.Bytes()

	_, err = keyset.ReadWithNoSecrets(keyset.NewBinaryReader(bytes.NewBuffer(serialized)))
	if err == nil {
		t.Error("keyset.ReadWithNoSecrets() = nil, want error")
	}
}

func TestWriteAndReadWithNoSecretsFailsWithUnknownKeyMaterial(t *testing.T) {
	// Create a keyset that contains unknown key material.
	keyData := testutil.NewKeyData("some type url", []byte{0}, tinkpb.KeyData_UNKNOWN_KEYMATERIAL)
	key := testutil.NewKey(keyData, tinkpb.KeyStatusType_ENABLED, 1, tinkpb.OutputPrefixType_TINK)
	ks := testutil.NewKeyset(1, []*tinkpb.Keyset_Key{key})
	handle, err := testkeyset.NewHandle(ks)
	if err != nil {
		t.Fatal(err)
	}
	serialized, err := proto.Marshal(ks)
	if err != nil {
		t.Fatal(err)
	}

	buff := &bytes.Buffer{}
	err = handle.WriteWithNoSecrets(keyset.NewBinaryWriter(buff))
	if err == nil {
		t.Error("handle.WriteWithNoSecrets() = nil, want error")
	}

	_, err = keyset.ReadWithNoSecrets(keyset.NewBinaryReader(bytes.NewBuffer(serialized)))
	if err == nil {
		t.Error("handle.ReadWithNoSecrets() = nil, want error")
	}
}

func TestKeysetInfo(t *testing.T) {
	params, err := hmac.NewParameters(hmac.ParametersOpts{
		KeySizeInBytes: 16,
		TagSizeInBytes: 16,
		HashType:       hmac.SHA256,
		Variant:        hmac.VariantTink,
	})
	if err != nil {
		t.Fatalf("NewParameters() err = %v, want nil", err)
	}
	hmacKey, err := hmac.NewKey(
		secretdata.NewBytesFromData([]byte("0123456789012345"), insecuresecretdataaccess.Token{}),
		params,
		1234)
	if err != nil {
		t.Errorf("hmac.NewKey() err = %v, want nil", err)
	}
	aescmacParams, err := aescmac.NewParameters(aescmac.ParametersOpts{
		KeySizeInBytes: 16,
		TagSizeInBytes: 16,
		Variant:        aescmac.VariantCrunchy,
	})
	if err != nil {
		t.Fatalf("aescmac.NewParameters() err = %v, want nil", err)
	}
	aescmacKey, err := aescmac.NewKey(
		secretdata.NewBytesFromData([]byte("0123456789012345"), insecuresecretdataaccess.Token{}),
		aescmacParams,
		4567)
	if err != nil {
		t.Errorf("aescmac.NewKey() err = nil, want error")
	}

	manager := keyset.NewManager()
	if err != nil {
		t.Fatalf("keyset.NewManager() err = %v, want nil", err)
	}
	id0, err := manager.AddKey(hmacKey)
	if err != nil {
		t.Fatalf("manager.Add(kt) err = %v, want nil", err)
	}
	err = manager.Disable(id0)
	if err != nil {
		t.Fatalf("manager.Disable(id0) err = %v, want nil", err)
	}
	id, err := manager.AddKey(aescmacKey)
	if err != nil {
		t.Fatalf("manager.Add(kt) err = %v, want nil", err)
	}
	if err := manager.SetPrimary(id); err != nil {
		t.Fatalf("manager.SetPrimary(id) err = %v, want nil", err)
	}
	handle, err := manager.Handle()

	if err != nil {
		t.Errorf("unexpected error: %s", err)
	}
	gotInfo := handle.KeysetInfo()

	wantInfo := &tinkpb.KeysetInfo{
		PrimaryKeyId: 4567,
		KeyInfo: []*tinkpb.KeysetInfo_KeyInfo{
			&tinkpb.KeysetInfo_KeyInfo{
				KeyId:            1234,
				Status:           tinkpb.KeyStatusType_DISABLED,
				TypeUrl:          "type.googleapis.com/google.crypto.tink.HmacKey",
				OutputPrefixType: tinkpb.OutputPrefixType_TINK,
			},
			&tinkpb.KeysetInfo_KeyInfo{
				KeyId:            4567,
				Status:           tinkpb.KeyStatusType_ENABLED,
				TypeUrl:          "type.googleapis.com/google.crypto.tink.AesCmacKey",
				OutputPrefixType: tinkpb.OutputPrefixType_CRUNCHY,
			},
		},
	}
	if diff := cmp.Diff(wantInfo, gotInfo, protocmp.Transform()); diff != "" {
		t.Errorf("handle.KeysetInfo() returned unexpected diff (-want +got):\n%s", diff)
	}
}

// RSA-SSA-PSS keys with salt length is 0 are different in tink-go than in other Tink
// implementations. For backwards compatibility, we continue to allow these keys in
// tink-go, but we don't want to allow them to be serialized anymore. This test makes sure
// that KeysetInfo() works with these special keys.
func TestKeysetInfo_withRSAPSSAndSaltLengthZero_works(t *testing.T) {
	n2048Base64 := "s1EKK81M5kTFtZSuUFnhKy8FS2WNXaWVmi_fGHG4CLw98-Yo0nkuUarVwSS0O9pFPcpc3kvPKOe9Tv-6DLS3Qru21aATy2PRqjqJ4CYn71OYtSwM_ZfSCKvrjXybzgu-sBmobdtYm-sppbdL-GEHXGd8gdQw8DDCZSR6-dPJFAzLZTCdB-Ctwe_RXPF-ewVdfaOGjkZIzDoYDw7n-OHnsYCYozkbTOcWHpjVevipR-IBpGPi1rvKgFnlcG6d_tj0hWRl_6cS7RqhjoiNEtxqoJzpXs_Kg8xbCxXbCchkf11STA8udiCjQWuWI8rcDwl69XMmHJjIQAqhKvOOQ8rYTQ"
	n2048, err := base64.URLEncoding.WithPadding(base64.NoPadding).DecodeString(n2048Base64)
	if err != nil {
		t.Fatalf("base64 decoding failed: %v", err)
	}
	paramsValues := rsassapss.ParametersValues{
		ModulusSizeBits: 2048,
		SigHashType:     rsassapss.SHA256,
		MGF1HashType:    rsassapss.SHA256,
		PublicExponent:  65537,
		SaltLengthBytes: 0,
	}
	params, err := rsassapss.NewParameters(paramsValues, rsassapss.VariantTink)
	if err != nil {
		t.Fatalf("NewParameters(%v, %v) = %v, want nil", paramsValues, rsassapss.VariantTink, err)
	}
	publicKey, err := rsassapss.NewPublicKey(n2048, 1234, params)
	if err != nil {
		t.Fatalf("rsassapss.NewPublicKey() err = %v, want nil", err)
	}
	km := keyset.NewManager()
	id, err := km.AddKey(publicKey)
	if err != nil {
		t.Fatalf("km.AddKey(publicKey) err = %v, want nil", err)
	}
	if err := km.SetPrimary(id); err != nil {
		t.Fatalf("km.SetPrimary(id) err = %v, want nil", err)
	}
	handle, err := km.Handle()
	if err != nil {
		t.Fatalf("km.Handle() err = %v, want nil", err)
	}

	gotInfo := handle.KeysetInfo()

	wantInfo := &tinkpb.KeysetInfo{
		PrimaryKeyId: 1234,
		KeyInfo: []*tinkpb.KeysetInfo_KeyInfo{
			&tinkpb.KeysetInfo_KeyInfo{
				KeyId:            1234,
				Status:           tinkpb.KeyStatusType_ENABLED,
				TypeUrl:          "type.googleapis.com/google.crypto.tink.RsaSsaPssPublicKey",
				OutputPrefixType: tinkpb.OutputPrefixType_TINK,
			},
		},
	}
	if diff := cmp.Diff(wantInfo, gotInfo, protocmp.Transform()); diff != "" {
		t.Errorf("handle.KeysetInfo() returned unexpected diff (-want +got):\n%s", diff)
	}
}

func TestPrimitivesWithRegistry(t *testing.T) {
	template := mac.HMACSHA256Tag128KeyTemplate()
	template.OutputPrefixType = tinkpb.OutputPrefixType_RAW
	handle, err := keyset.NewHandle(template)
	if err != nil {
		t.Fatalf("keyset.NewHandle(%v) err = %v, want nil", template, err)
	}
	handleMAC, err := mac.New(handle)
	if err != nil {
		t.Fatalf("mac.New(%v) err = %v, want nil", handle, err)
	}

	ks := testkeyset.KeysetMaterial(handle)
	if len(ks.Key) != 1 {
		t.Fatalf("len(ks.Key) = %d, want 1", len(ks.Key))
	}
	keyDataPrimitive, err := registry.PrimitiveFromKeyData(ks.Key[0].KeyData)
	if err != nil {
		t.Fatalf("registry.PrimitiveFromKeyData(%v) err = %v, want nil", ks.Key[0].KeyData, err)
	}
	keyDataMAC, ok := keyDataPrimitive.(tink.MAC)
	if !ok {
		t.Fatal("registry.PrimitiveFromKeyData(keyData) is not of type tink.MAC")
	}

	plaintext := []byte("plaintext")
	handleMACTag, err := handleMAC.ComputeMAC(plaintext)
	if err != nil {
		t.Fatalf("handleMAC.ComputeMAC(%v) err = %v, want nil", plaintext, err)
	}
	if err = keyDataMAC.VerifyMAC(handleMACTag, plaintext); err != nil {
		t.Errorf("keyDataMAC.VerifyMAC(%v, %v) err = %v, want nil", handleMACTag, plaintext, err)
	}
	keyDataMACTag, err := keyDataMAC.ComputeMAC(plaintext)
	if err != nil {
		t.Fatalf("keyDataMAC.ComputeMAC(%v) err = %v, want nil", plaintext, err)
	}
	if err = handleMAC.VerifyMAC(keyDataMACTag, plaintext); err != nil {
		t.Errorf("handleMAC.VerifyMAC(%v, %v) err = %v, want nil", keyDataMACTag, plaintext, err)
	}
}

func TestLenWithOneKey(t *testing.T) {
	template := mac.HMACSHA256Tag128KeyTemplate()
	handle, err := keyset.NewHandle(template)
	if err != nil {
		t.Fatalf("keyset.NewHandle(%v) err = %v, want nil", template, err)
	}
	if handle.Len() != 1 {
		t.Errorf("handle.Len() = %d, want 1", handle.Len())
	}
}

func TestLenWithMultipleKeys(t *testing.T) {
	ks := &tinkpb.Keyset{
		Key: []*tinkpb.Keyset_Key{
			testutil.NewDummyKey(1, tinkpb.KeyStatusType_ENABLED, tinkpb.OutputPrefixType_TINK),
			testutil.NewDummyKey(2, tinkpb.KeyStatusType_ENABLED, tinkpb.OutputPrefixType_TINK),
			testutil.NewDummyKey(3, tinkpb.KeyStatusType_ENABLED, tinkpb.OutputPrefixType_TINK),
			testutil.NewDummyKey(4, tinkpb.KeyStatusType_ENABLED, tinkpb.OutputPrefixType_TINK),
		},
		PrimaryKeyId: 1,
	}
	handle, err := testkeyset.NewHandle(ks)
	if err != nil {
		t.Fatalf("testkeyset.NewHandle(%v) err = %v, want nil", ks, err)
	}
	if handle.Len() != len(ks.Key) {
		t.Errorf("handle.Len() = %d, want %d", handle.Len(), len(ks.Key))
	}
}

func TestEntryReturnsCorrectKey(t *testing.T) {
	ks := &tinkpb.Keyset{
		Key: []*tinkpb.Keyset_Key{
			testutil.NewDummyKey(0, tinkpb.KeyStatusType_ENABLED, tinkpb.OutputPrefixType_TINK),
			testutil.NewDummyKey(1, tinkpb.KeyStatusType_ENABLED, tinkpb.OutputPrefixType_TINK),
			testutil.NewDummyKey(2, tinkpb.KeyStatusType_ENABLED, tinkpb.OutputPrefixType_TINK),
			testutil.NewDummyKey(3, tinkpb.KeyStatusType_ENABLED, tinkpb.OutputPrefixType_TINK),
			testutil.NewDummyKey(4, tinkpb.KeyStatusType_ENABLED, tinkpb.OutputPrefixType_TINK),
		},
		PrimaryKeyId: 2,
	}
	handle, err := testkeyset.NewHandle(ks)
	if err != nil {
		t.Fatalf("testkeyset.NewHandle(%v) err = %v, want nil", ks, err)
	}

	for i := 0; i < handle.Len(); i++ {
		entry, err := handle.Entry(i)
		if err != nil {
			t.Errorf("handle.Entry(%d) err = %v, want nil", i, err)
		}
		if int(entry.KeyID()) != i {
			t.Errorf("entry.KeyID() = %v, want %v", entry.KeyID(), i)
		}
		if wantIsPrimary := i == 2; entry.IsPrimary() != wantIsPrimary {
			t.Errorf("entry.IsPrimary() = %v, want %v", entry.IsPrimary(), wantIsPrimary)
		}
		if entry.KeyStatus() != keyset.Enabled {
			t.Errorf("entry.KeyStatus() = %v, want Enabled", entry.KeyStatus())
		}
	}
}

func TestEntryFailsIfIndexOutOfRange(t *testing.T) {
	ks := &tinkpb.Keyset{
		Key: []*tinkpb.Keyset_Key{
			testutil.NewDummyKey(1, tinkpb.KeyStatusType_ENABLED, tinkpb.OutputPrefixType_TINK),
			testutil.NewDummyKey(2, tinkpb.KeyStatusType_ENABLED, tinkpb.OutputPrefixType_TINK),
			testutil.NewDummyKey(3, tinkpb.KeyStatusType_ENABLED, tinkpb.OutputPrefixType_TINK),
			testutil.NewDummyKey(4, tinkpb.KeyStatusType_ENABLED, tinkpb.OutputPrefixType_TINK),
		},
		PrimaryKeyId: 1,
	}
	handle, err := testkeyset.NewHandle(ks)
	if err != nil {
		t.Fatalf("testkeyset.NewHandle(%v) err = %v, want nil", ks, err)
	}
	_, err = handle.Entry(-1)
	if err == nil {
		t.Error("handle.Entry(-1) err = nil, want error")
	}
	_, err = handle.Entry(handle.Len())
	if err == nil {
		t.Errorf("handle.Entry(%d) err = nil, want error", handle.Len())
	}
}

func TestPrimaryReturnsPrimaryKey(t *testing.T) {
	primaryKey := testutil.NewDummyKey(2, tinkpb.KeyStatusType_ENABLED, tinkpb.OutputPrefixType_TINK)
	ks := &tinkpb.Keyset{
		Key: []*tinkpb.Keyset_Key{
			testutil.NewDummyKey(1, tinkpb.KeyStatusType_ENABLED, tinkpb.OutputPrefixType_TINK),
			primaryKey,
			testutil.NewDummyKey(3, tinkpb.KeyStatusType_ENABLED, tinkpb.OutputPrefixType_TINK),
		},
		PrimaryKeyId: 2,
	}
	handle, err := testkeyset.NewHandle(ks)
	if err != nil {
		t.Fatalf("testkeyset.NewHandle(%v) err = %v, want nil", ks, err)
	}
	primaryEntry, err := handle.Primary()
	if err != nil {
		t.Fatalf("handle.Primary() err = %v, want nil", err)
	}
	if primaryEntry.KeyID() != 2 {
		t.Errorf("primaryEntry.KeyID() = %v, want 2", primaryEntry.KeyID())
	}
	if primaryEntry.KeyStatus() != keyset.Enabled {
		t.Errorf("primaryEntry.KeyStatus() = %v, want Enabled", primaryEntry.KeyStatus())
	}
	primaryProtoKey, ok := primaryEntry.Key().(*protoserialization.FallbackProtoKey)
	if !ok {
		t.Fatalf("type mismatch: got %T, want *protoserialization.FallbackProtoKey", primaryEntry.Key())
	}

	primaryKeySerialization := protoserialization.GetKeySerialization(primaryProtoKey)
	wantKeySerialization, err := protoserialization.NewKeySerialization(primaryKey.GetKeyData(), primaryKey.GetOutputPrefixType(), primaryKey.GetKeyId())
	if err != nil {
		t.Fatalf("protoserialization.NewKeySerialization(%v) err = %v, want nil", primaryKey, err)
	}
	if !primaryKeySerialization.Equal(wantKeySerialization) {
		t.Errorf("primaryKeySerialization = %v, want %v", primaryKeySerialization, wantKeySerialization)
	}
	// Check that is the same as Entry(1).
	entry, err := handle.Entry(1)
	if err != nil {
		t.Fatalf("handle.Entry(1) err = %v, want nil", err)
	}
	entryProtoKey, ok := entry.Key().(*protoserialization.FallbackProtoKey)
	if !ok {
		t.Fatalf("type mismatch: got %T, want *protoserialization.FallbackProtoKey", entry.Key())
	}
	entryKeySerialization := protoserialization.GetKeySerialization(entryProtoKey)
	if !entryKeySerialization.Equal(wantKeySerialization) {
		t.Errorf("entryKeySerialization = %v, want %v", entryKeySerialization, wantKeySerialization)
	}
}

func TestPrimaryIsThreadSafe(t *testing.T) {
	template := signature.ECDSAP256KeyTemplate()
	manager := keyset.NewManager()
	// Add 10 keys. Last one is the primary.
	for i := 0; i < 10; i++ {
		keyID, err := manager.Add(template)
		if err != nil {
			t.Fatalf("manager.Add(template) err = %v, want nil", err)
		}
		if err = manager.SetPrimary(keyID); err != nil {
			t.Fatalf("manager.SetPrimary(%v) err = %v, want nil", keyID, err)
		}
	}
	handle, err := manager.Handle()
	if err != nil {
		t.Fatalf("manager.Handle() err = %v, want nil", err)
	}
	for i := 0; i < 50; i++ {
		t.Run(fmt.Sprintf("entry %d", i), func(t *testing.T) {
			t.Parallel()
			_, err := handle.Primary()
			if err != nil {
				t.Fatalf("handle.Primary() err = %v, want nil", err)
			}
		})
	}
}

func TestEntryIsThreadSafe(t *testing.T) {
	template := signature.ECDSAP256KeyTemplate()
	manager := keyset.NewManager()
	// Add 10 keys. Last one is the primary.
	for i := 0; i < 10; i++ {
		keyID, err := manager.Add(template)
		if err != nil {
			t.Fatalf("manager.Add(template) err = %v, want nil", err)
		}
		if err = manager.SetPrimary(keyID); err != nil {
			t.Fatalf("manager.SetPrimary(%v) err = %v, want nil", keyID, err)
		}
	}
	handle, err := manager.Handle()
	if err != nil {
		t.Fatalf("manager.Handle() err = %v, want nil", err)
	}
	for i := 0; i < 50; i++ {
		t.Run(fmt.Sprintf("entry %d", i), func(t *testing.T) {
			t.Parallel()
			_, err := handle.Entry(0) // Index doesn't matter.
			if err != nil {
				t.Fatalf("handle.Entry() err = %v, want nil", err)
			}
		})
	}
}

func TestKeysetInfoIsThreadSafe(t *testing.T) {
	template := signature.ECDSAP256KeyTemplate()
	manager := keyset.NewManager()
	// Add 10 keys. Last one is the primary.
	for i := 0; i < 10; i++ {
		keyID, err := manager.Add(template)
		if err != nil {
			t.Fatalf("manager.Add(template) err = %v, want nil", err)
		}
		if err = manager.SetPrimary(keyID); err != nil {
			t.Fatalf("manager.SetPrimary(%v) err = %v, want nil", keyID, err)
		}
	}
	handle, err := manager.Handle()
	if err != nil {
		t.Fatalf("manager.Handle() err = %v, want nil", err)
	}
	for i := 0; i < 50; i++ {
		t.Run(fmt.Sprintf("entry %d", i), func(t *testing.T) {
			t.Parallel()
			if handle.KeysetInfo() == nil {
				t.Fatalf("handle.KeysetInfo() == nul, want non-nil")
			}
		})
	}
}

func TestPublicKeysetHasPrimaryKey(t *testing.T) {
	template := signature.ECDSAP256KeyTemplate()
	manager := keyset.NewManager()
	// Add 10 keys. Last one is the primary.
	for i := 0; i < 10; i++ {
		keyID, err := manager.Add(template)
		if err != nil {
			t.Fatalf("manager.Add(template) err = %v, want nil", err)
		}
		if err = manager.SetPrimary(keyID); err != nil {
			t.Fatalf("manager.SetPrimary(%v) err = %v, want nil", keyID, err)
		}
	}
	handle, err := manager.Handle()
	if err != nil {
		t.Fatalf("manager.Handle() err = %v, want nil", err)
	}
	publicHandle, err := handle.Public()
	if err != nil {
		t.Fatalf("handle.Public() err = %v, want nil", err)
	}
	if _, err := publicHandle.Primary(); err != nil {
		t.Errorf("publicHandle.Primary() err = %v, want nil", err)
	}
}

func TestPublicIsThreadSafe(t *testing.T) {
	template := signature.ECDSAP256KeyTemplate()
	manager := keyset.NewManager()
	// Add 10 keys. Last one is the primary.
	for i := 0; i < 10; i++ {
		keyID, err := manager.Add(template)
		if err != nil {
			t.Fatalf("manager.Add(template) err = %v, want nil", err)
		}
		if err = manager.SetPrimary(keyID); err != nil {
			t.Fatalf("manager.SetPrimary(%v) err = %v, want nil", keyID, err)
		}
	}
	handle, err := manager.Handle()
	if err != nil {
		t.Fatalf("manager.Handle() err = %v, want nil", err)
	}
	for i := 0; i < 50; i++ {
		t.Run(fmt.Sprintf("entry %d", i), func(t *testing.T) {
			t.Parallel()
			_, err := handle.Public()
			if err != nil {
				t.Fatalf("handle.Public() err = %v, want nil", err)
			}
		})
	}
}

func TestWriteIsThreadSafe(t *testing.T) {
	template := signature.ECDSAP256KeyTemplate()
	manager := keyset.NewManager()
	// Add 10 keys. Last one is the primary.
	for i := 0; i < 10; i++ {
		keyID, err := manager.Add(template)
		if err != nil {
			t.Fatalf("manager.Add(template) err = %v, want nil", err)
		}
		if err = manager.SetPrimary(keyID); err != nil {
			t.Fatalf("manager.SetPrimary(%v) err = %v, want nil", keyID, err)
		}
	}
	handle, err := manager.Handle()
	if err != nil {
		t.Fatalf("manager.Handle() err = %v, want nil", err)
	}
	keysetEncryptionHandle, err := keyset.NewHandle(aead.AES128GCMKeyTemplate())
	if err != nil {
		t.Errorf("keyset.NewHandle(aead.AES128GCMKeyTemplate()) err = %v, want nil", err)
	}
	for i := 0; i < 50; i++ {
		t.Run(fmt.Sprintf("entry %d", i), func(t *testing.T) {
			t.Parallel()
			keysetEncryptionAead, err := aead.New(keysetEncryptionHandle)
			if err != nil {
				t.Errorf("aead.New(keysetEncryptionHandle) err = %v, want nil", err)
			}
			buff := &bytes.Buffer{}
			err = handle.Write(keyset.NewBinaryWriter(buff), keysetEncryptionAead)
			if err != nil {
				t.Fatalf("handle.Write() err = %v, want nil", err)
			}
		})
	}
}

func TestWriteWithAssociatedDataIsThreadSafe(t *testing.T) {
	template := signature.ECDSAP256KeyTemplate()
	manager := keyset.NewManager()
	// Add 10 keys. Last one is the primary.
	for i := 0; i < 10; i++ {
		keyID, err := manager.Add(template)
		if err != nil {
			t.Fatalf("manager.Add(template) err = %v, want nil", err)
		}
		if err = manager.SetPrimary(keyID); err != nil {
			t.Fatalf("manager.SetPrimary(%v) err = %v, want nil", keyID, err)
		}
	}
	handle, err := manager.Handle()
	if err != nil {
		t.Fatalf("manager.Handle() err = %v, want nil", err)
	}

	keysetEncryptionHandle, err := keyset.NewHandle(aead.AES128GCMKeyTemplate())
	if err != nil {
		t.Errorf("keyset.NewHandle(aead.AES128GCMKeyTemplate()) err = %v, want nil", err)
	}
	associatedData := []byte{0x01, 0x02}
	for i := 0; i < 50; i++ {
		t.Run(fmt.Sprintf("entry %d", i), func(t *testing.T) {
			t.Parallel()
			buff := &bytes.Buffer{}
			keysetEncryptionAead, err := aead.New(keysetEncryptionHandle)
			if err != nil {
				t.Errorf("aead.New(keysetEncryptionHandle) err = %v, want nil", err)
			}
			if err := handle.WriteWithAssociatedData(keyset.NewBinaryWriter(buff), keysetEncryptionAead, associatedData); err != nil {
				t.Fatalf("handle.WriteWithAssociatedData(keyset.NewBinaryWriter(buff), keysetEncryptionAead, associatedData) err = %v, want nil", err)
			}
		})
	}
}

func TestWriteWithNoSecretsIsThreadSafe(t *testing.T) {
	template := signature.ECDSAP256KeyTemplate()
	manager := keyset.NewManager()
	// Add 10 keys. Last one is the primary.
	for i := 0; i < 10; i++ {
		keyID, err := manager.Add(template)
		if err != nil {
			t.Fatalf("manager.Add(template) err = %v, want nil", err)
		}
		if err = manager.SetPrimary(keyID); err != nil {
			t.Fatalf("manager.SetPrimary(%v) err = %v, want nil", keyID, err)
		}
	}
	handle, err := manager.Handle()
	if err != nil {
		t.Fatalf("manager.Handle() err = %v, want nil", err)
	}
	publicHandle, err := handle.Public()
	if err != nil {
		t.Fatalf("manager.Public() err = %v, want nil", err)
	}
	for i := 0; i < 50; i++ {
		t.Run(fmt.Sprintf("entry %d", i), func(t *testing.T) {
			t.Parallel()
			buff := &bytes.Buffer{}
			if err := publicHandle.WriteWithNoSecrets(keyset.NewBinaryWriter(buff)); err != nil {
				t.Fatalf("publicHandle.WriteWithNoSecrets(keyset.NewBinaryWriter(buff)) err = %v, want nil", err)
			}
		})
	}
}

func TestKeyExportMonitoring(t *testing.T) {
	defer internalregistry.ClearMonitoringClient()
	fakeClient := fakemonitoring.NewClient("fake_client")
	if err := internalregistry.RegisterMonitoringClient(fakeClient); err != nil {
		t.Fatalf("internalregistry.RegisterMonitoringClient() err = %v, want nil", err)
	}

	km := keyset.NewManager()
	keyID1, err := km.Add(signature.ECDSAP256KeyTemplate())
	if err != nil {
		t.Fatalf("km.Add(template) err = %v, want nil", err)
	}
	keyID2, err := km.Add(signature.ECDSAP256KeyTemplate())
	if err != nil {
		t.Fatalf("km.Add(template) err = %v, want nil", err)
	}
	if err = km.SetPrimary(keyID1); err != nil {
		t.Fatalf("km.SetPrimary(%v) err = %v, want nil", keyID1, err)
	}
	if err := km.SetAnnotations(map[string]string{"foo": "bar"}); err != nil {
		t.Fatalf("km.SetAnnotations() err = %v, want nil", err)
	}

	handle, err := km.Handle()
	if err != nil {
		t.Fatalf("km.Handle() err = %v, want nil", err)
	}

	// Export all keys.
	if got, want := len(fakeClient.Events()), 0; got != want {
		t.Errorf("len(fakeClient.Events()) = %d, want %d", got, want)
	}
	if got, want := len(fakeClient.Failures()), 0; got != want {
		t.Errorf("len(fakeClient.Failures()) = %d, want %d", got, want)
	}
	if got, want := len(fakeClient.KeyExportsLogs()), 0; got != want {
		t.Errorf("len(fakeClient.KeyExportsLogs()) = %d, want %d", got, want)
	}
	for i := 0; i < handle.Len(); i++ {
		entry, err := handle.Entry(i)
		if err != nil {
			t.Fatalf("handle.Entry(%d) err = %v, want nil", i, err)
		}
		_ = entry.Key()
		if got, want := len(fakeClient.KeyExportsLogs()), i+1; got != want {
			t.Errorf("len(fakeClient.KeyExportsLogs()) = %d, want %d", got, want)
		}
	}

	info, err := monitoringutil.MonitoringKeysetInfoFromKeysetInfo(handle.KeysetInfo(), map[string]string{"foo": "bar"})
	if err != nil {
		t.Fatalf("monitoringutil.MonitoringKeysetInfoFromKeysetInfo() err = %v, want nil", err)
	}

	want := []*fakemonitoring.LogKeyExport{
		&fakemonitoring.LogKeyExport{
			KeyID: keyID1,
			Context: &monitoring.Context{
				Primitive:   "keyset_handle",
				APIFunction: "get_key",
				KeysetInfo:  info,
			},
		},
		&fakemonitoring.LogKeyExport{
			KeyID: keyID2,
			Context: &monitoring.Context{
				Primitive:   "keyset_handle",
				APIFunction: "get_key",
				KeysetInfo:  info,
			},
		},
	}
	if diff := cmp.Diff(want, fakeClient.KeyExportsLogs(), protocmp.Transform()); diff != "" {
		t.Errorf("fakeClient.KeyExportsLogs() returned unexpected diff (-want +got):\n%s", diff)
	}
}

func TestKeyExportMonitoring_keysetMaterial(t *testing.T) {
	defer internalregistry.ClearMonitoringClient()
	fakeClient := fakemonitoring.NewClient("fake_client")
	if err := internalregistry.RegisterMonitoringClient(fakeClient); err != nil {
		t.Fatalf("internalregistry.RegisterMonitoringClient() err = %v, want nil", err)
	}

	km := keyset.NewManager()
	keyID1, err := km.Add(signature.ECDSAP256KeyTemplate())
	if err != nil {
		t.Fatalf("km.Add(template) err = %v, want nil", err)
	}
	keyID2, err := km.Add(signature.ECDSAP256KeyTemplate())
	if err != nil {
		t.Fatalf("km.Add(template) err = %v, want nil", err)
	}
	if err = km.SetPrimary(keyID1); err != nil {
		t.Fatalf("km.SetPrimary(%v) err = %v, want nil", keyID1, err)
	}
	if err := km.SetAnnotations(map[string]string{"foo": "bar"}); err != nil {
		t.Fatalf("km.SetAnnotations() err = %v, want nil", err)
	}

	handle, err := km.Handle()
	if err != nil {
		t.Fatalf("km.Handle() err = %v, want nil", err)
	}

	if got, want := len(fakeClient.Events()), 0; got != want {
		t.Errorf("len(fakeClient.Events()) = %d, want %d", got, want)
	}
	if got, want := len(fakeClient.Failures()), 0; got != want {
		t.Errorf("len(fakeClient.Failures()) = %d, want %d", got, want)
	}
	if got, want := len(fakeClient.KeyExportsLogs()), 0; got != want {
		t.Errorf("len(fakeClient.KeyExportsLogs()) = %d, want %d", got, want)
	}

	// Export all keys.
	_ = internal.KeysetMaterial.(func(h *keyset.Handle) *tinkpb.Keyset)(handle)

	info, err := monitoringutil.MonitoringKeysetInfoFromKeysetInfo(handle.KeysetInfo(), map[string]string{"foo": "bar"})
	if err != nil {
		t.Fatalf("monitoringutil.MonitoringKeysetInfoFromKeysetInfo() err = %v, want nil", err)
	}

	want := []*fakemonitoring.LogKeyExport{
		&fakemonitoring.LogKeyExport{
			KeyID: keyID1,
			Context: &monitoring.Context{
				Primitive:   "keyset_handle",
				APIFunction: "get_key",
				KeysetInfo:  info,
			},
		},
		&fakemonitoring.LogKeyExport{
			KeyID: keyID2,
			Context: &monitoring.Context{
				Primitive:   "keyset_handle",
				APIFunction: "get_key",
				KeysetInfo:  info,
			},
		},
	}
	if diff := cmp.Diff(want, fakeClient.KeyExportsLogs(), protocmp.Transform()); diff != "" {
		t.Errorf("fakeClient.KeyExportsLogs() returned unexpected diff (-want +got):\n%s", diff)
	}
}

// Not monitored if no annotations are set.
func TestKeyExportNotMonitored(t *testing.T) {
	defer internalregistry.ClearMonitoringClient()
	fakeClient := fakemonitoring.NewClient("fake_client")
	if err := internalregistry.RegisterMonitoringClient(fakeClient); err != nil {
		t.Fatalf("internalregistry.RegisterMonitoringClient() err = %v, want nil", err)
	}

	km := keyset.NewManager()
	keyID1, err := km.Add(signature.ECDSAP256KeyTemplate())
	if err != nil {
		t.Fatalf("km.Add(template) err = %v, want nil", err)
	}
	if _, err = km.Add(signature.ECDSAP256KeyTemplate()); err != nil {
		t.Fatalf("km.Add(template) err = %v, want nil", err)
	}
	if err = km.SetPrimary(keyID1); err != nil {
		t.Fatalf("km.SetPrimary(%v) err = %v, want nil", keyID1, err)
	}
	handle, err := km.Handle()
	if err != nil {
		t.Fatalf("km.Handle() err = %v, want nil", err)
	}

	// Export all keys.
	if got, want := len(fakeClient.KeyExportsLogs()), 0; got != want {
		t.Errorf("len(fakeClient.KeyExportsLogs()) = %d, want %d", got, want)
	}
	for i := 0; i < handle.Len(); i++ {
		entry, err := handle.Entry(i)
		if err != nil {
			t.Fatalf("handle.Entry(%d) err = %v, want nil", i, err)
		}
		_ = entry.Key()
		if got, want := len(fakeClient.KeyExportsLogs()), 0; got != want {
			t.Errorf("len(fakeClient.KeyExportsLogs()) = %d, want %d", got, want)
		}
	}

	if fakeClient.KeyExportsLogs() != nil {
		t.Errorf("fakeClient.KeyExportsLogs() = %v, want nil", fakeClient.KeyExportsLogs())
	}
}

func TestWriteGeneratesNoKeyExport(t *testing.T) {
	defer internalregistry.ClearMonitoringClient()
	fakeClient := fakemonitoring.NewClient("fake_client")
	if err := internalregistry.RegisterMonitoringClient(fakeClient); err != nil {
		t.Fatalf("internalregistry.RegisterMonitoringClient() err = %v, want nil", err)
	}

	km := keyset.NewManager()
	keyID1, err := km.Add(signature.ECDSAP256KeyTemplate())
	if err != nil {
		t.Fatalf("km.Add(template) err = %v, want nil", err)
	}
	if _, err = km.Add(signature.ECDSAP256KeyTemplate()); err != nil {
		t.Fatalf("km.Add(template) err = %v, want nil", err)
	}
	if err = km.SetPrimary(keyID1); err != nil {
		t.Fatalf("km.SetPrimary(%v) err = %v, want nil", keyID1, err)
	}
	if err := km.SetAnnotations(map[string]string{"foo": "bar"}); err != nil {
		t.Fatalf("km.SetAnnotations() err = %v, want nil", err)
	}
	handle, err := km.Handle()
	if err != nil {
		t.Fatalf("km.Handle() err = %v, want nil", err)
	}

	if fakeClient.KeyExportsLogs() != nil {
		t.Errorf("fakeClient.KeyExportsLogs() = %v, want nil", fakeClient.KeyExportsLogs())
	}
	keysetEncryptionAEAD, err := fakekms.NewAEAD(fakeKeyURI)
	if err != nil {
		t.Fatalf("fakekms.NewAEAD(fakeKeyURI) err = %v, want nil", err)
	}
	buf1 := &bytes.Buffer{}
	if err := handle.Write(keyset.NewBinaryWriter(buf1), keysetEncryptionAEAD); err != nil {
		t.Fatalf("handle.Write() err = %v, want nil", err)
	}
	buf2 := &bytes.Buffer{}
	if err := handle.WriteWithAssociatedData(keyset.NewBinaryWriter(buf2), keysetEncryptionAEAD, []byte("AssociatedData")); err != nil {
		t.Fatalf("handle.WriteWithAssociatedData() err = %v, want nil", err)
	}
	keysetEncryptionAEADWithContext, err := fakekms.NewAEADWithContext(fakeKeyURI)
	if err != nil {
		t.Fatalf("fakekms.NewAEADWithContext(fakeKeyURI) err = %v, want nil", err)
	}
	buf3 := &bytes.Buffer{}
	if err := handle.WriteWithContext(context.Background(), keyset.NewBinaryWriter(buf3), keysetEncryptionAEADWithContext, []byte("AssociatedData")); err != nil {
		t.Fatalf("handle.WriteWithContext() err = %v, want nil", err)
	}
	publicHandle, err := handle.Public()
	if err != nil {
		t.Fatalf("handle.Public() err = %v, want nil", err)
	}
	buf4 := &bytes.Buffer{}
	if err := publicHandle.WriteWithNoSecrets(keyset.NewBinaryWriter(buf4)); err != nil {
		t.Fatalf("publicHandle.WriteWithNoSecrets() err = %v, want nil", err)
	}

	if fakeClient.KeyExportsLogs() != nil {
		t.Errorf("fakeClient.KeyExportsLogs() = %v, want nil", fakeClient.KeyExportsLogs())
	}
}