Add and register parameters serialization for AES-GCM

PiperOrigin-RevId: 684061848 Change-Id: I5fd9fd6c67ac40edb461e7572c8005c33b9ce2df
tink-crypto · Oct 9, 2024 · 45a25fe · 45a25fe
1 parent 3b27b59
commit 45a25fe
Show file tree

Hide file tree

Showing 5 changed files with 225 additions and 28 deletions.
diff --git a/aead/aesgcm/aesgcm.go b/aead/aesgcm/aesgcm.go
@@ -31,13 +31,16 @@ func init() {
 	if err := internalregistry.AllowKeyDerivation(typeURL); err != nil {
 		panic(fmt.Sprintf("aesgcm.init() failed: %v", err))
 	}
-	if err := protoserialization.RegisterKeySerializer[*Key](&serializer{}); err != nil {
+	if err := protoserialization.RegisterKeySerializer[*Key](&keySerializer{}); err != nil {
 		panic(fmt.Sprintf("aesgcm.init() failed: %v", err))
 	}
-	if err := protoserialization.RegisterKeyParser(typeURL, &parser{}); err != nil {
+	if err := protoserialization.RegisterKeyParser(typeURL, &keyParser{}); err != nil {
 		panic(fmt.Sprintf("aesgcm.init() failed: %v", err))
 	}
 	if err := registryconfig.RegisterPrimitiveConstructor[*Key](primitiveConstructor); err != nil {
 		panic(fmt.Sprintf("aesgcm.init() failed: %v", err))
 	}
+	if err := protoserialization.RegisterParametersSerializer[*Parameters](&parametersSerializer{}); err != nil {
+		panic(fmt.Sprintf("aesgcm.init() failed: %v", err))
+	}
 }
diff --git a/aead/aesgcm/aesgcm_test.go b/aead/aesgcm/aesgcm_test.go
@@ -23,7 +23,7 @@ import (
 	"github.com/tink-crypto/tink-go/v2/keyset"
 )
 
-func TestGetAESGCMKeyFromHandle(t *testing.T) {
+func TestGetKeyFromHandle(t *testing.T) {
 	keysetHandle, err := keyset.NewHandle(aead.AES128GCMKeyTemplate())
 	if err != nil {
 		t.Fatalf("keyset.NewHandle(aead.AES128GCMKeyTemplate()) err = %v, want nil", err)
@@ -59,7 +59,7 @@ func TestGetAESGCMKeyFromHandle(t *testing.T) {
 	}
 }
 
-func TestCreateKeysetHandleFromAESGCMKey(t *testing.T) {
+func TestCreateKeysetHandleFromKey(t *testing.T) {
 	keysetHandle, err := keyset.NewHandle(aead.AES128GCMKeyTemplate())
 	if err != nil {
 		t.Fatalf("keyset.NewHandle(aead.AES128GCMKeyTemplate()) err = %v, want nil", err)
@@ -111,3 +111,42 @@ func TestCreateKeysetHandleFromAESGCMKey(t *testing.T) {
 		t.Errorf("decrypt = %v, want %v", decrypt, plaintext)
 	}
 }
+
+func TestCreateKeysetHandleFromParameters(t *testing.T) {
+	params, err := aesgcm.NewParameters(aesgcm.ParametersOpts{
+		KeySizeInBytes: 32,
+		IVSizeInBytes:  12,
+		TagSizeInBytes: 16,
+		Variant:        aesgcm.VariantTink,
+	})
+	if err != nil {
+		t.Fatalf("aesgcm.NewParameters(%v) err = %v, want nil", params, err)
+	}
+	manager := keyset.NewManager()
+	keyID, err := manager.AddNewKeyFromParameters(params)
+	if err != nil {
+		t.Fatalf("manager.AddNewKeyFromParameters(%v) err = %v, want nil", params, err)
+	}
+	manager.SetPrimary(keyID)
+	handle, err := manager.Handle()
+	if err != nil {
+		t.Fatalf("manager.Handle() err = %v, want nil", err)
+	}
+	aeadPrimitive, err := aead.New(handle)
+	if err != nil {
+		t.Fatalf("aead.New(handle) err = %v, want nil", err)
+	}
+	plaintext := []byte("plaintext")
+	additionalData := []byte("additionalData")
+	ciphertext, err := aeadPrimitive.Encrypt(plaintext, additionalData)
+	if err != nil {
+		t.Fatalf("aeadPrimitive.Encrypt(%v, %v) err = %v, want nil", plaintext, additionalData, err)
+	}
+	decrypted, err := aeadPrimitive.Decrypt(ciphertext, additionalData)
+	if err != nil {
+		t.Fatalf("aeadPrimitive.Decrypt(%v, %v) err = %v, want nil", ciphertext, additionalData, err)
+	}
+	if !bytes.Equal(decrypted, plaintext) {
+		t.Errorf("decrypted = %v, want %v", decrypted, plaintext)
+	}
+}
diff --git a/aead/aesgcm/protoserialization.go b/aead/aesgcm/protoserialization.go
@@ -33,9 +33,9 @@ const (
 	protoVersion = 0
 )
 
-type serializer struct{}
+type keySerializer struct{}
 
-var _ protoserialization.KeySerializer = (*serializer)(nil)
+var _ protoserialization.KeySerializer = (*keySerializer)(nil)
 
 func protoOutputPrefixTypeFromVariant(variant Variant) (tinkpb.OutputPrefixType, error) {
 	switch variant {
@@ -50,7 +50,7 @@ func protoOutputPrefixTypeFromVariant(variant Variant) (tinkpb.OutputPrefixType,
 	}
 }
 
-func (s *serializer) SerializeKey(key key.Key) (*protoserialization.KeySerialization, error) {
+func (s *keySerializer) SerializeKey(key key.Key) (*protoserialization.KeySerialization, error) {
 	actualKey, ok := key.(*Key)
 	if !ok {
 		return nil, fmt.Errorf("key is not a Key")
@@ -82,9 +82,9 @@ func (s *serializer) SerializeKey(key key.Key) (*protoserialization.KeySerializa
 	return protoserialization.NewKeySerialization(keyData, outputPrefixType, idRequirement)
 }
 
-type parser struct{}
+type keyParser struct{}
 
-var _ protoserialization.KeyParser = (*parser)(nil)
+var _ protoserialization.KeyParser = (*keyParser)(nil)
 
 func variantFromProto(prefixType tinkpb.OutputPrefixType) (Variant, error) {
 	switch prefixType {
@@ -99,7 +99,7 @@ func variantFromProto(prefixType tinkpb.OutputPrefixType) (Variant, error) {
 	}
 }
 
-func (s *parser) ParseKey(keySerialization *protoserialization.KeySerialization) (key.Key, error) {
+func (s *keyParser) ParseKey(keySerialization *protoserialization.KeySerialization) (key.Key, error) {
 	if keySerialization == nil {
 		return nil, fmt.Errorf("key serialization is nil")
 	}
@@ -132,7 +132,35 @@ func (s *parser) ParseKey(keySerialization *protoserialization.KeySerialization)
 		return nil, err
 	}
 	keyMaterial := secretdata.NewBytesFromData(protoKey.GetKeyValue(), insecuresecretdataaccess.Token{})
-	// keySerialization.IDRequirement() returns zero if the key doesn't have a key requirement.
+	// keySerialization.IDRequirement() returns zero if the key doesn't have a
+	// key requirement.
 	keyID, _ := keySerialization.IDRequirement()
 	return NewKey(keyMaterial, keyID, params)
 }
+
+type parametersSerializer struct{}
+
+var _ protoserialization.ParametersSerializer = (*parametersSerializer)(nil)
+
+func (s *parametersSerializer) Serialize(parameters key.Parameters) (*tinkpb.KeyTemplate, error) {
+	actualParameters, ok := parameters.(*Parameters)
+	if !ok {
+		return nil, fmt.Errorf("invalid parameters type: got %T, want *aesgcm.Parameters", parameters)
+	}
+	outputPrefixType, err := protoOutputPrefixTypeFromVariant(actualParameters.Variant())
+	if err != nil {
+		return nil, err
+	}
+	format := &gcmpb.AesGcmKeyFormat{
+		KeySize: uint32(actualParameters.KeySizeInBytes()),
+	}
+	serializedFormat, err := proto.Marshal(format)
+	if err != nil {
+		return nil, err
+	}
+	return &tinkpb.KeyTemplate{
+		TypeUrl:          typeURL,
+		OutputPrefixType: outputPrefixType,
+		Value:            serializedFormat,
+	}, nil
+}