[ARCH-333] Update import to override key name if it is set. (#1682)

Author: pavel-raykov

keystore/admin.go

@@ -56,9 +56,9 @@ type ImportKeysRequest struct {
 }
 
 type ImportKeyRequest struct {
-	KeyName  string
-	Data     []byte
-	Password string
+	NewKeyName string
+	Data       []byte
+	Password   string
 }
 
 type ImportKeysResponse struct{}
@@ -93,10 +93,30 @@ type SetMetadataUpdate struct {
 type SetMetadataResponse struct{}
 
 type Admin interface {
+	// CreateKeys creates multiple keys in a single atomic operation.
+	// The response preserves the order of the keys in the request.
+	// Returns ErrKeyAlreadyExists if any key name already exists,
+	// ErrInvalidKeyName if any key name is invalid,
+	// ErrUnsupportedKeyType if any key type is not supported.
 	CreateKeys(ctx context.Context, req CreateKeysRequest) (CreateKeysResponse, error)
+
+	// DeleteKeys deletes multiple keys in a single atomic operation.
+	// Returns ErrKeyNotFound if any key does not exist.
 	DeleteKeys(ctx context.Context, req DeleteKeysRequest) (DeleteKeysResponse, error)
+
+	// ImportKeys imports multiple encrypted keys in a single atomic operation.
+	// Keys can be renamed during import using NewKeyName.
+	// Returns ErrKeyAlreadyExists if a key with the target name exists,
+	// ErrInvalidKeyName if the target name is invalid.
 	ImportKeys(ctx context.Context, req ImportKeysRequest) (ImportKeysResponse, error)
+
+	// ExportKeys exports multiple keys in encrypted format.
+	// Each key is encrypted using the parameters specified in the request.
+	// Returns ErrKeyNotFound if any key does not exist.
 	ExportKeys(ctx context.Context, req ExportKeysRequest) (ExportKeysResponse, error)
+
+	// SetMetadata updates metadata for multiple keys in a single atomic operation.
+	// Returns ErrKeyNotFound if any key does not exist.
 	SetMetadata(ctx context.Context, req SetMetadataRequest) (SetMetadataResponse, error)
 }
 
@@ -237,40 +257,45 @@ func (ks *keystore) ImportKeys(ctx context.Context, req ImportKeysRequest) (Impo
 	defer ks.mu.Unlock()
 
 	ksCopy := maps.Clone(ks.keystore)
-	for _, keyReq := range req.Keys {
-		if err := ValidKeyName(keyReq.KeyName); err != nil {
-			return ImportKeysResponse{}, fmt.Errorf("%w: %s", ErrInvalidKeyName, err)
-		}
-		if _, ok := ksCopy[keyReq.KeyName]; ok {
-			return ImportKeysResponse{}, fmt.Errorf("%w: %s", ErrKeyAlreadyExists, keyReq.KeyName)
-		}
+	for i, keyReq := range req.Keys {
 		encData := gethkeystore.CryptoJSON{}
 		err := json.Unmarshal(keyReq.Data, &encData)
 		if err != nil {
-			return ImportKeysResponse{}, fmt.Errorf("key = %s, failed to unmarshal encrypted import data: %w", keyReq.KeyName, err)
+			return ImportKeysResponse{}, fmt.Errorf("key num = %d, failed to unmarshal encrypted import data: %w", i, err)
 		}
 		decData, err := gethkeystore.DecryptDataV3(encData, keyReq.Password)
 		if err != nil {
-			return ImportKeysResponse{}, fmt.Errorf("key = %s, failed to decrypt key: %w", keyReq.KeyName, err)
+			return ImportKeysResponse{}, fmt.Errorf("key num = %d, failed to decrypt key: %w", i, err)
 		}
 		keypb := &serialization.Key{}
 		err = proto.Unmarshal(decData, keypb)
 		if err != nil {
-			return ImportKeysResponse{}, fmt.Errorf("key = %s, failed to unmarshal key: %w", keyReq.KeyName, err)
+			return ImportKeysResponse{}, fmt.Errorf("key num = %d, failed to unmarshal key: %w", i, err)
 		}
 		pkRaw := internal.NewRaw(keypb.PrivateKey)
 		keyType := KeyType(keypb.KeyType)
 		publicKey, err := publicKeyFromPrivateKey(pkRaw, keyType)
 		if err != nil {
-			return ImportKeysResponse{}, fmt.Errorf("key = %s, failed to get public key from private key: %w", keyReq.KeyName, err)
+			return ImportKeysResponse{}, fmt.Errorf("key num = %d, failed to get public key from private key: %w", i, err)
 		}
 		metadata := keypb.Metadata
 		// The proto compiler sets empty slices to nil during the serialization (https://github.com/golang/protobuf/issues/1348).
 		// We set metadata back to empty slice to be consistent with the Create method which initializes it as such.
 		if metadata == nil {
 			metadata = []byte{}
 		}
-		ksCopy[keyReq.KeyName] = newKey(keyType, pkRaw, publicKey, time.Unix(keypb.CreatedAt, 0), metadata)
+
+		keyName := keyReq.NewKeyName
+		if keyName == "" {
+			keyName = keypb.Name
+		}
+		if err := ValidKeyName(keyName); err != nil {
+			return ImportKeysResponse{}, fmt.Errorf("key name = %s, %w: %s", keyName, ErrInvalidKeyName, err)
+		}
+		if _, ok := ksCopy[keyName]; ok {
+			return ImportKeysResponse{}, fmt.Errorf("key name = %s, %w: %s", keyName, ErrKeyAlreadyExists, keyName)
+		}
+		ksCopy[keyName] = newKey(keyType, pkRaw, publicKey, time.Unix(keypb.CreatedAt, 0), metadata)
 	}
 	// Persist it to storage.
 	if err := ks.save(ctx, ksCopy); err != nil {

keystore/admin_test.go

@@ -231,17 +231,35 @@ func TestKeystore_ExportImport(t *testing.T) {
 		require.Len(t, exportResponse.Keys, 1)
 		_, err = ks2.ImportKeys(t.Context(), keystore.ImportKeysRequest{
 			Keys: []keystore.ImportKeyRequest{
-				{KeyName: "key1", Password: exportParams.Password, Data: exportResponse.Keys[0].Data},
+				{Password: exportParams.Password, Data: exportResponse.Keys[0].Data},
 			},
 		})
 		require.NoError(t, err)
+
+		// Importing a key with the same name again fails.
+		_, err = ks2.ImportKeys(t.Context(), keystore.ImportKeysRequest{
+			Keys: []keystore.ImportKeyRequest{
+				{Password: exportParams.Password, Data: exportResponse.Keys[0].Data},
+			},
+		})
+		require.ErrorIs(t, err, keystore.ErrKeyAlreadyExists)
+
+		// Importing a key with a new name is allowed.
+		_, err = ks2.ImportKeys(t.Context(), keystore.ImportKeysRequest{
+			Keys: []keystore.ImportKeyRequest{
+				{NewKeyName: "new-name", Password: exportParams.Password, Data: exportResponse.Keys[0].Data},
+			},
+		})
+		require.NoError(t, err)
+
+		// Verify that imported key matches exported key.
+		// We cannot compare private keys directly, so we test that signing with key1 from ks1 and verifying
+		// with key1 from ks2 works as if the two keys are the same.
 		key1ks1, err := ks1.GetKeys(t.Context(), keystore.GetKeysRequest{KeyNames: []string{"key1"}})
 		require.NoError(t, err)
 		key1ks2, err := ks2.GetKeys(t.Context(), keystore.GetKeysRequest{KeyNames: []string{"key1"}})
 		require.Equal(t, key1ks1, key1ks2)
 
-		// We cannot compare private keys directly, so we test that signing with key1 from ks1 and verifying
-		// with key1 from ks2 works as if the two keys are the same.
 		testData := []byte("hello world")
 		signature, err := ks2.Sign(t.Context(), keystore.SignRequest{
 			KeyName: "key1",
@@ -266,15 +284,6 @@ func TestKeystore_ExportImport(t *testing.T) {
 		})
 		require.ErrorIs(t, err, keystore.ErrKeyNotFound)
 	})
-
-	t.Run("import existing key", func(t *testing.T) {
-		_, err = ks2.ImportKeys(t.Context(), keystore.ImportKeysRequest{
-			Keys: []keystore.ImportKeyRequest{
-				{KeyName: "key1", Password: "", Data: []byte{}},
-			},
-		})
-		require.ErrorIs(t, err, keystore.ErrKeyAlreadyExists)
-	})
 }
 
 func TestKeystore_SetMetadata(t *testing.T) {