From 56cf3f6b0e950a2d4f7d5aee9b183498beb514f2 Mon Sep 17 00:00:00 2001
From: vishesh92 <vishesh92@gmail.com>
Date: Tue, 10 Feb 2026 21:30:06 +0530
Subject: [PATCH] Change wrapping algo for pkcs

---
 .../admin/kms/MigrateVolumesToKMSCmd.java     |   3 +-
 .../command/admin/kms/RotateKMSKeyCmd.java    |   3 +-
 .../provider/pkcs11/PKCS11HSMProvider.java    | 165 ++++++------------
 .../apache/cloudstack/kms/KMSManagerImpl.java |  52 +++---
 .../kms/KMSManagerImplKeyRotationTest.java    | 120 ++++++++-----
 .../utils/exception/CSExceptionErrorCode.java |   1 +
 6 files changed, 157 insertions(+), 187 deletions(-)

diff --git a/api/src/main/java/org/apache/cloudstack/api/command/admin/kms/MigrateVolumesToKMSCmd.java b/api/src/main/java/org/apache/cloudstack/api/command/admin/kms/MigrateVolumesToKMSCmd.java
index 308e1dd38a0..4d76abf4155 100644
--- a/api/src/main/java/org/apache/cloudstack/api/command/admin/kms/MigrateVolumesToKMSCmd.java
+++ b/api/src/main/java/org/apache/cloudstack/api/command/admin/kms/MigrateVolumesToKMSCmd.java
@@ -16,6 +16,7 @@
 // under the License.
 package org.apache.cloudstack.api.command.admin.kms;
 
+import com.cloud.dc.DataCenter;
 import com.cloud.user.Account;
 import org.apache.cloudstack.acl.RoleType;
 import org.apache.cloudstack.api.APICommand;
@@ -127,7 +128,7 @@ public class MigrateVolumesToKMSCmd extends BaseAsyncCmd {
 
     @Override
     public String getEventDescription() {
-        return "Migrating volumes to KMS for zone: " + _uuidMgr.getUuid(ZoneResponse.class, zoneId);
+        return "Migrating volumes to KMS for zone: " + _uuidMgr.getUuid(DataCenter.class, zoneId);
     }
 
     @Override
diff --git a/api/src/main/java/org/apache/cloudstack/api/command/admin/kms/RotateKMSKeyCmd.java b/api/src/main/java/org/apache/cloudstack/api/command/admin/kms/RotateKMSKeyCmd.java
index ffe0bc32ab3..7ab47a65351 100644
--- a/api/src/main/java/org/apache/cloudstack/api/command/admin/kms/RotateKMSKeyCmd.java
+++ b/api/src/main/java/org/apache/cloudstack/api/command/admin/kms/RotateKMSKeyCmd.java
@@ -29,6 +29,7 @@ import org.apache.cloudstack.api.response.AsyncJobResponse;
 import org.apache.cloudstack.api.response.HSMProfileResponse;
 import org.apache.cloudstack.api.response.KMSKeyResponse;
 import org.apache.cloudstack.framework.kms.KMSException;
+import org.apache.cloudstack.kms.KMSKey;
 import org.apache.cloudstack.kms.KMSManager;
 
 import javax.inject.Inject;
@@ -103,7 +104,7 @@ public class RotateKMSKeyCmd extends BaseAsyncCmd {
 
     @Override
     public String getEventDescription() {
-        return "Rotating KMS key: " + _uuidMgr.getUuid(KMSKeyResponse.class, id);
+        return "Rotating KMS key: " + _uuidMgr.getUuid(KMSKey.class, id);
     }
 
     @Override
diff --git a/plugins/kms/pkcs11/src/main/java/org/apache/cloudstack/kms/provider/pkcs11/PKCS11HSMProvider.java b/plugins/kms/pkcs11/src/main/java/org/apache/cloudstack/kms/provider/pkcs11/PKCS11HSMProvider.java
index e1d059258bd..487a753889c 100644
--- a/plugins/kms/pkcs11/src/main/java/org/apache/cloudstack/kms/provider/pkcs11/PKCS11HSMProvider.java
+++ b/plugins/kms/pkcs11/src/main/java/org/apache/cloudstack/kms/provider/pkcs11/PKCS11HSMProvider.java
@@ -39,7 +39,7 @@ import javax.crypto.Cipher;
 import javax.crypto.IllegalBlockSizeException;
 import javax.crypto.NoSuchPaddingException;
 import javax.crypto.SecretKey;
-import javax.crypto.spec.GCMParameterSpec;
+import javax.crypto.spec.IvParameterSpec;
 import javax.crypto.spec.SecretKeySpec;
 import javax.inject.Inject;
 import java.io.Closeable;
@@ -564,9 +564,10 @@ public class PKCS11HSMProvider extends AdapterBase implements KMSProvider {
      * {@link KMSException.ErrorType} values for proper retry logic and error reporting.
      */
     private static class PKCS11Session {
-        private static final String ALGORITHM = "AES/GCM/NoPadding";
-        private static final int GCM_IV_LENGTH = 12; // 96 bits
-        private static final int GCM_TAG_LENGTH = 16; // 128 bits
+        // Use AES-CBC with PKCS5Padding for key wrapping
+        // This is FIPS-compliant (NIST SP 800-38A) and has universal PKCS#11 support
+        private static final String CIPHER_ALGORITHM = "AES/CBC/PKCS5Padding";
+        private static final int IV_LENGTH = 16; // 128 bits for CBC
         private static final String PROVIDER_PREFIX = "CloudStackPKCS11-";
 
         private final Map<String, String> config;
@@ -658,7 +659,7 @@ public class PKCS11HSMProvider extends AdapterBase implements KMSProvider {
                 // Zeroize PIN from memory
                 Arrays.fill(pinChars, '\0');
 
-                logger.debug("Successfully connected to PKCS#11 HSM at {}", config.get("library"));
+                logger.debug("aSuccessfully connected to PKCS#11 HSM at {}", config.get("library"));
             } catch (KeyStoreException | CertificateException | NoSuchAlgorithmException e) {
                 handlePKCS11Exception(e, "Failed to initialize PKCS#11 connection");
             } catch (IOException e) {
@@ -937,19 +938,18 @@ public class PKCS11HSMProvider extends AdapterBase implements KMSProvider {
         /**
          * Wraps (encrypts) a plaintext DEK using a KEK stored in the HSM.
          *
-         * <p>Uses AES-GCM for authenticated encryption:
+         * <p>Uses AES-CBC with PKCS5Padding (FIPS 197 + NIST SP 800-38A):
          * <ul>
-         *   <li>Generates a random 96-bit IV</li>
-         *   <li>Encrypts the DEK using the KEK from HSM</li>
-         *   <li>Appends a 128-bit authentication tag</li>
-         *   <li>Returns format: [IV (12 bytes)][ciphertext+tag]</li>
+         *   <li>Generates a random 128-bit IV</li>
+         *   <li>Encrypts the DEK using AES-CBC with the KEK from HSM</li>
+         *   <li>Returns format: [IV (16 bytes)][ciphertext]</li>
          * </ul>
          *
          * <p>Security: The plaintext DEK should be zeroized by the caller after wrapping.
          *
          * @param plainDek Plaintext DEK to wrap (will be encrypted)
          * @param kekLabel Label of the KEK stored in the HSM
-         * @return Wrapped blob: [IV][ciphertext+tag]
+         * @return Wrapped key blob: [IV][ciphertext]
          * @throws KMSException with appropriate ErrorType:
          *                      <ul>
          *                        <li>{@code INVALID_PARAMETER} if plainDek is null or empty</li>
@@ -966,23 +966,30 @@ public class PKCS11HSMProvider extends AdapterBase implements KMSProvider {
             try {
                 kek = getKekFromKeyStore(kekLabel);
 
-                // Generate random IV for GCM
-                byte[] iv = new byte[GCM_IV_LENGTH];
+                // Generate random IV for CBC
+                byte[] iv = new byte[IV_LENGTH];
                 new SecureRandom().nextBytes(iv);
 
-                // Create and initialize AES-GCM cipher in ENCRYPT_MODE
-                Cipher cipher = createGCMCipher(kek, iv, Cipher.ENCRYPT_MODE);
+                // Create cipher with AES-CBC
+                Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM, provider);
+                cipher.init(Cipher.ENCRYPT_MODE, kek, new IvParameterSpec(iv));
 
-                // Encrypt the plaintext DEK using doFinal (GCM includes authentication tag)
-                byte[] wrappedBlob = cipher.doFinal(plainDek);
+                // Encrypt the plaintext DEK
+                byte[] ciphertext = cipher.doFinal(plainDek);
 
-                // Prepend IV to wrapped blob: [IV][ciphertext+tag]
-                byte[] result = prependIV(iv, wrappedBlob);
+                // Prepend IV to ciphertext: [IV][ciphertext]
+                byte[] result = new byte[IV_LENGTH + ciphertext.length];
+                System.arraycopy(iv, 0, result, 0, IV_LENGTH);
+                System.arraycopy(ciphertext, 0, result, IV_LENGTH, ciphertext.length);
 
-                logger.debug("Wrapped key with KEK '{}'", kekLabel);
+                logger.debug("Wrapped key with KEK '{}' using AES-CBC", kekLabel);
                 return result;
-            } catch (IllegalBlockSizeException e) {
-                handlePKCS11Exception(e, "Invalid block size for wrapping");
+            } catch (IllegalBlockSizeException | BadPaddingException | InvalidKeyException e) {
+                handlePKCS11Exception(e, "Invalid key or data for wrapping");
+            } catch (NoSuchAlgorithmException | NoSuchPaddingException e) {
+                handlePKCS11Exception(e, "AES-CBC not supported by HSM");
+            } catch (InvalidAlgorithmParameterException e) {
+                handlePKCS11Exception(e, "Invalid IV for CBC mode");
             } catch (Exception e) {
                 handlePKCS11Exception(e, "Failed to wrap key with HSM");
             } finally {
@@ -1019,71 +1026,30 @@ public class PKCS11HSMProvider extends AdapterBase implements KMSProvider {
             return null; // Unreachable
         }
 
-        /**
-         * Prepends IV to data, creating a new byte array.
-         *
-         * @param iv   Initialization vector
-         * @param data Data to prepend IV to
-         * @return Combined array: [IV][data]
-         */
-        private byte[] prependIV(byte[] iv, byte[] data) {
-            byte[] result = new byte[GCM_IV_LENGTH + data.length];
-            System.arraycopy(iv, 0, result, 0, GCM_IV_LENGTH);
-            System.arraycopy(data, 0, result, GCM_IV_LENGTH, data.length);
-            return result;
-        }
-
-        /**
-         * Creates and initializes an AES-GCM cipher.
-         *
-         * @param kek  Key Encryption Key
-         * @param iv   Initialization vector
-         * @param mode Cipher mode (ENCRYPT_MODE or DECRYPT_MODE)
-         * @return Initialized Cipher instance
-         * @throws KMSException if cipher creation or initialization fails
-         */
-        private Cipher createGCMCipher(SecretKey kek, byte[] iv, int mode) throws KMSException {
-            try {
-                Cipher cipher = Cipher.getInstance(ALGORITHM, provider);
-                GCMParameterSpec gcmSpec = new GCMParameterSpec(GCM_TAG_LENGTH * 8, iv);
-                cipher.init(mode, kek, gcmSpec);
-                return cipher;
-            } catch (NoSuchPaddingException e) {
-                handlePKCS11Exception(e, "GCM padding not supported");
-            } catch (InvalidKeyException e) {
-                handlePKCS11Exception(e, "Invalid KEK");
-            } catch (InvalidAlgorithmParameterException e) {
-                handlePKCS11Exception(e, "Invalid GCM parameters");
-            } catch (NoSuchAlgorithmException e) {
-                handlePKCS11Exception(e, String.format("Algorithm %s not supported.", ALGORITHM));
-            }
-            return null; // Unreachable
-        }
 
         /**
          * Unwraps (decrypts) a wrapped DEK using a KEK stored in the HSM.
          *
-         * <p>Process:
+         * <p>Uses AES-CBC with PKCS5Padding (FIPS 197 + NIST SP 800-38A):
          * <ol>
          *   <li>Extracts IV from the wrapped blob</li>
          *   <li>Retrieves KEK from HSM using the label</li>
-         *   <li>Decrypts using AES-GCM (verifies authentication tag)</li>
+         *   <li>Decrypts using AES-CBC</li>
          *   <li>Returns plaintext DEK</li>
          * </ol>
          *
          * <p>Security: The returned plaintext DEK must be zeroized by the caller after use.
          *
-         * <p>Expected format: [IV (12 bytes)][ciphertext+tag]
+         * <p>Expected format: [IV (16 bytes)][ciphertext]
          *
-         * @param wrappedBlob Wrapped DEK blob (IV + ciphertext + tag)
+         * @param wrappedBlob Wrapped DEK blob (IV + ciphertext)
          * @param kekLabel    Label of the KEK stored in the HSM
          * @return Plaintext DEK
          * @throws KMSException with appropriate ErrorType:
          *                      <ul>
          *                        <li>{@code INVALID_PARAMETER} if wrappedBlob is null, empty, or too short</li>
          *                        <li>{@code KEK_NOT_FOUND} if KEK with label doesn't exist or is not accessible</li>
-         *                        <li>{@code WRAP_UNWRAP_FAILED} if unwrapping fails (e.g., authentication tag
-         *                        verification fails)</li>
+         *                        <li>{@code WRAP_UNWRAP_FAILED} if unwrapping fails</li>
          *                      </ul>
          */
         byte[] unwrapKey(byte[] wrappedBlob, String kekLabel) throws KMSException {
@@ -1091,9 +1057,10 @@ public class PKCS11HSMProvider extends AdapterBase implements KMSProvider {
                 throw KMSException.invalidParameter("Wrapped blob cannot be null or empty");
             }
 
-            if (wrappedBlob.length < GCM_IV_LENGTH + GCM_TAG_LENGTH) {
+            // Minimum size: IV (16) + at least one block of ciphertext (16)
+            if (wrappedBlob.length < IV_LENGTH + 16) {
                 throw KMSException.invalidParameter("Wrapped blob too short: expected at least " +
-                                                    (GCM_IV_LENGTH + GCM_TAG_LENGTH) + " bytes");
+                        (IV_LENGTH + 16) + " bytes");
             }
 
             SecretKey kek = null;
@@ -1101,22 +1068,29 @@ public class PKCS11HSMProvider extends AdapterBase implements KMSProvider {
                 kek = getKekFromKeyStore(kekLabel);
 
                 // Extract IV and ciphertext from wrapped blob
-                IVAndCiphertext extracted = extractIVAndCiphertext(wrappedBlob);
+                byte[] iv = new byte[IV_LENGTH];
+                System.arraycopy(wrappedBlob, 0, iv, 0, IV_LENGTH);
+                byte[] ciphertext = new byte[wrappedBlob.length - IV_LENGTH];
+                System.arraycopy(wrappedBlob, IV_LENGTH, ciphertext, 0, ciphertext.length);
 
-                // Create and initialize AES-GCM cipher in DECRYPT_MODE
-                Cipher cipher = createGCMCipher(kek, extracted.iv, Cipher.DECRYPT_MODE);
+                // Create cipher with AES-CBC
+                Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM, provider);
+                cipher.init(Cipher.DECRYPT_MODE, kek, new IvParameterSpec(iv));
 
-                // Decrypt the ciphertext to get plaintext DEK (GCM verifies authentication tag)
-                byte[] plainDek = cipher.doFinal(extracted.ciphertextWithTag);
+                // Decrypt the ciphertext to get plaintext DEK
+                byte[] plainDek = cipher.doFinal(ciphertext);
 
-                logger.debug("Unwrapped key with KEK '{}'", kekLabel);
+                logger.debug("Unwrapped key with KEK '{}' using AES-CBC", kekLabel);
                 return plainDek;
             } catch (BadPaddingException e) {
-                // GCM authentication tag verification failed
                 throw KMSException.wrapUnwrapFailed(
-                        "Authentication failed: wrapped key may be corrupted or KEK is incorrect", e);
-            } catch (IllegalBlockSizeException e) {
-                handlePKCS11Exception(e, "Invalid block size for unwrapping");
+                        "Decryption failed: wrapped key may be corrupted or KEK is incorrect", e);
+            } catch (IllegalBlockSizeException | InvalidKeyException e) {
+                handlePKCS11Exception(e, "Invalid key or data for unwrapping");
+            } catch (NoSuchAlgorithmException | NoSuchPaddingException e) {
+                handlePKCS11Exception(e, "AES-CBC not supported by HSM");
+            } catch (InvalidAlgorithmParameterException e) {
+                handlePKCS11Exception(e, "Invalid IV for CBC mode");
             } catch (Exception e) {
                 handlePKCS11Exception(e, "Failed to unwrap key with HSM");
             } finally {
@@ -1126,24 +1100,6 @@ public class PKCS11HSMProvider extends AdapterBase implements KMSProvider {
             return null; // Unreachable
         }
 
-        /**
-         * Extracts IV and ciphertext from a wrapped blob.
-         *
-         * @param wrappedBlob Wrapped blob containing IV and ciphertext
-         * @return IVAndCiphertext containing extracted IV and ciphertext
-         * @throws KMSException if wrapped blob is too short
-         */
-        private IVAndCiphertext extractIVAndCiphertext(byte[] wrappedBlob) throws KMSException {
-            if (wrappedBlob.length < GCM_IV_LENGTH + GCM_TAG_LENGTH) {
-                throw KMSException.invalidParameter("Wrapped blob too short: expected at least " +
-                                                    (GCM_IV_LENGTH + GCM_TAG_LENGTH) + " bytes");
-            }
-            byte[] iv = new byte[GCM_IV_LENGTH];
-            System.arraycopy(wrappedBlob, 0, iv, 0, GCM_IV_LENGTH);
-            byte[] ciphertextWithTag = new byte[wrappedBlob.length - GCM_IV_LENGTH];
-            System.arraycopy(wrappedBlob, GCM_IV_LENGTH, ciphertextWithTag, 0, ciphertextWithTag.length);
-            return new IVAndCiphertext(iv, ciphertextWithTag);
-        }
 
         /**
          * Deletes a key from the HSM.
@@ -1212,18 +1168,5 @@ public class PKCS11HSMProvider extends AdapterBase implements KMSProvider {
                 return false;
             }
         }
-
-        /**
-         * Helper class to hold IV and ciphertext extracted from wrapped blob.
-         */
-        private static class IVAndCiphertext {
-            final byte[] iv;
-            final byte[] ciphertextWithTag;
-
-            IVAndCiphertext(byte[] iv, byte[] ciphertextWithTag) {
-                this.iv = iv;
-                this.ciphertextWithTag = ciphertextWithTag;
-            }
-        }
     }
 }
diff --git a/server/src/main/java/org/apache/cloudstack/kms/KMSManagerImpl.java b/server/src/main/java/org/apache/cloudstack/kms/KMSManagerImpl.java
index 958e4f94d34..09fd6838a0e 100644
--- a/server/src/main/java/org/apache/cloudstack/kms/KMSManagerImpl.java
+++ b/server/src/main/java/org/apache/cloudstack/kms/KMSManagerImpl.java
@@ -147,49 +147,45 @@ public class KMSManagerImpl extends ManagerBase implements KMSManager, Pluggable
     /**
      * Internal method to rotate a KEK (create new version and update KMS key state)
      */
-    String rotateKek(Long zoneId, KeyPurpose purpose, String oldKekLabel,
-            String newKekLabel, int keyBits, Long newProfileId) throws KMSException {
-        validateKmsEnabled(zoneId);
+    String rotateKek(KMSKeyVO kmsKey, String oldKekLabel,
+            String newKekLabel, int keyBits, HSMProfileVO newHSMProfile) throws KMSException {
+
+        validateKmsEnabled(kmsKey.getZoneId());
 
         if (StringUtils.isEmpty(oldKekLabel)) {
             throw KMSException.invalidParameter("oldKekLabel must be specified");
         }
 
-        KMSProvider provider = getKMSProviderForZone(zoneId);
+        if (newHSMProfile == null) {
+            newHSMProfile = hsmProfileDao.findById(kmsKey.getHsmProfileId());
+        }
+
+
+        KMSProvider provider = getKMSProvider(newHSMProfile.getProtocol());
 
         try {
-            logger.info("Starting KEK rotation from {} to {} for zone {} and purpose {}",
-                    oldKekLabel, newKekLabel, zoneId, purpose);
+            logger.info("Starting KEK rotation from {} to {} for kms key {}", oldKekLabel, newKekLabel, kmsKey);
 
             // Find KMS key by old KEK label
-            KMSKeyVO kmsKey = kmsKeyDao.findByKekLabel(oldKekLabel, provider.getProviderName());
-            if (kmsKey == null) {
-                throw KMSException.kekNotFound("KMS key not found for KEK label: " + oldKekLabel);
-            }
 
             // Generate new KEK label if not provided
             if (StringUtils.isEmpty(newKekLabel)) {
-                newKekLabel = purpose.getName() + "-kek-" + UUID.randomUUID().toString().substring(0, 8);
-            }
-
-            // Resolve profile ID if not provided (use current profile from key)
-            if (newProfileId == null) {
-                newProfileId = kmsKey.getHsmProfileId();
+                newKekLabel = kmsKey.getPurpose().getName() + "-kek-" + UUID.randomUUID().toString().substring(0, 8);
             }
 
             // Create new KEK in provider
             String finalNewKekLabel = newKekLabel;
-            Long finalProfileId = newProfileId;
-            String newKekId = retryOperation(() -> provider.createKek(purpose, finalNewKekLabel, keyBits, finalProfileId));
+            Long newProfileId = newHSMProfile.getId();
+            String newKekId = retryOperation(() -> provider.createKek(kmsKey.getPurpose(), finalNewKekLabel, keyBits, newProfileId));
 
             // Create new KEK version (marks old as Previous, new as Active)
-            KMSKekVersionVO newVersion = createKekVersion(kmsKey.getId(), newKekId, finalProfileId);
+            KMSKekVersionVO newVersion = createKekVersion(kmsKey.getId(), newKekId, newProfileId);
 
             // Update KMS Key with new profile if different
-            if (finalProfileId != null && !finalProfileId.equals(kmsKey.getHsmProfileId())) {
-                kmsKey.setHsmProfileId(finalProfileId);
+            if (!newProfileId.equals(kmsKey.getHsmProfileId())) {
+                kmsKey.setHsmProfileId(newProfileId);
                 kmsKeyDao.update(kmsKey.getId(), kmsKey);
-                logger.info("Updated KMS key {} to use HSM profile ID {}", kmsKey.getUuid(), finalProfileId);
+                logger.info("Updated KMS key {} to use HSM profile {}", kmsKey, newHSMProfile);
             }
 
             logger.info("KEK rotation: KMS key {} now has {} versions (active: v{}, previous: v{})",
@@ -201,7 +197,7 @@ public class KMSManagerImpl extends ManagerBase implements KMSManager, Pluggable
             return newKekId;
 
         } catch (Exception e) {
-            logger.error("KEK rotation failed for zone {}: {}", zoneId, e.getMessage());
+            logger.error("KEK rotation failed for kmsKey {}: {}", kmsKey, e.getMessage());
             throw handleKmsException(e);
         }
     }
@@ -693,9 +689,9 @@ public class KMSManagerImpl extends ManagerBase implements KMSManager, Pluggable
         }
 
         // Validate and resolve target HSM profile if provided
-        Long targetProfileId = null;
+        HSMProfileVO profile = null;
         if (hsmProfileName != null) {
-            HSMProfileVO profile = hsmProfileDao.findByName(hsmProfileName);
+            profile = hsmProfileDao.findByName(hsmProfileName);
             if (profile == null) {
                 throw KMSException.invalidParameter("Target HSM Profile not found: " + hsmProfileName);
             }
@@ -704,7 +700,6 @@ public class KMSManagerImpl extends ManagerBase implements KMSManager, Pluggable
                  // Warn or fail - admin can migrate to any profile really, but key owner should have access ideally.
                  // For now allow admin to do anything.
             }
-            targetProfileId = profile.getId();
         }
 
         // Get current active version to determine key bits if not provided
@@ -712,12 +707,11 @@ public class KMSManagerImpl extends ManagerBase implements KMSManager, Pluggable
         KMSKekVersionVO currentActive = getActiveKekVersion(kmsKey.getId());
 
         rotateKek(
-                kmsKey.getZoneId(),
-                kmsKey.getPurpose(),
+                kmsKey,
                 currentActive.getKekLabel(),
                 null, // auto-generate new label
                 newKeyBits,
-                targetProfileId
+                profile
         );
 
         KMSKekVersionVO newVersion = getActiveKekVersion(kmsKey.getId());
diff --git a/server/src/test/java/org/apache/cloudstack/kms/KMSManagerImplKeyRotationTest.java b/server/src/test/java/org/apache/cloudstack/kms/KMSManagerImplKeyRotationTest.java
index 34a76fa6ef3..a256f6dd46a 100644
--- a/server/src/test/java/org/apache/cloudstack/kms/KMSManagerImplKeyRotationTest.java
+++ b/server/src/test/java/org/apache/cloudstack/kms/KMSManagerImplKeyRotationTest.java
@@ -17,20 +17,6 @@
 
 package org.apache.cloudstack.kms;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotNull;
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.ArgumentMatchers.anyString;
-import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.doReturn;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
-
-import java.util.Arrays;
-
 import org.apache.cloudstack.framework.kms.KMSException;
 import org.apache.cloudstack.framework.kms.KMSProvider;
 import org.apache.cloudstack.framework.kms.KeyPurpose;
@@ -48,6 +34,20 @@ import org.mockito.Mock;
 import org.mockito.Spy;
 import org.mockito.junit.MockitoJUnitRunner;
 
+import java.util.Arrays;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
 /**
  * Unit tests for KMS key rotation logic in KMSManagerImpl
  * Tests key rotation within same HSM and cross-HSM migration
@@ -79,7 +79,6 @@ public class KMSManagerImplKeyRotationTest {
 
     @Before
     public void setUp() {
-        when(kmsProvider.getProviderName()).thenReturn(testProviderName);
     }
 
     /**
@@ -95,8 +94,14 @@ public class KMSManagerImplKeyRotationTest {
 
         KMSKeyVO kmsKey = mock(KMSKeyVO.class);
         when(kmsKey.getId()).thenReturn(kmsKeyId);
+        when(kmsKey.getZoneId()).thenReturn(testZoneId);
         when(kmsKey.getHsmProfileId()).thenReturn(oldProfileId);
-        when(kmsKeyDao.findByKekLabel(oldKekLabel, testProviderName)).thenReturn(kmsKey);
+        when(kmsKey.getPurpose()).thenReturn(KeyPurpose.VOLUME_ENCRYPTION);
+
+        HSMProfileVO hsmProfile = mock(HSMProfileVO.class);
+        when(hsmProfile.getId()).thenReturn(oldProfileId);
+        when(hsmProfile.getProtocol()).thenReturn(testProviderName);
+        when(hsmProfileDao.findById(oldProfileId)).thenReturn(hsmProfile);
 
         // Old version should be marked as Previous
         KMSKekVersionVO oldVersion = mock(KMSKekVersionVO.class);
@@ -107,17 +112,16 @@ public class KMSManagerImplKeyRotationTest {
 
         // Provider creates new KEK
         when(kmsProvider.createKek(any(KeyPurpose.class), eq(newKekLabel), anyInt(), eq(oldProfileId)))
-            .thenReturn("new-kek-id");
+                .thenReturn("new-kek-id");
 
         KMSKekVersionVO newVersion = mock(KMSKekVersionVO.class);
         when(newVersion.getVersionNumber()).thenReturn(2);
         when(kmsKekVersionDao.persist(any(KMSKekVersionVO.class))).thenReturn(newVersion);
 
-        doReturn(kmsProvider).when(kmsManager).getKMSProviderForZone(testZoneId);
+        doReturn(kmsProvider).when(kmsManager).getKMSProvider(testProviderName);
         doReturn(true).when(kmsManager).isKmsEnabled(testZoneId);
 
-        String result = kmsManager.rotateKek(testZoneId, KeyPurpose.VOLUME_ENCRYPTION,
-            oldKekLabel, newKekLabel, 256, null);
+        String result = kmsManager.rotateKek(kmsKey, oldKekLabel, newKekLabel, 256, null);
 
         // Verify new KEK was created in same HSM
         assertNotNull(result);
@@ -149,8 +153,13 @@ public class KMSManagerImplKeyRotationTest {
 
         KMSKeyVO kmsKey = mock(KMSKeyVO.class);
         when(kmsKey.getId()).thenReturn(kmsKeyId);
+        when(kmsKey.getZoneId()).thenReturn(testZoneId);
         when(kmsKey.getHsmProfileId()).thenReturn(oldProfileId);
-        when(kmsKeyDao.findByKekLabel(oldKekLabel, testProviderName)).thenReturn(kmsKey);
+        when(kmsKey.getPurpose()).thenReturn(KeyPurpose.VOLUME_ENCRYPTION);
+
+        HSMProfileVO hsmProfile = mock(HSMProfileVO.class);
+        when(hsmProfile.getId()).thenReturn(newProfileId);
+        when(hsmProfile.getProtocol()).thenReturn(testProviderName);
 
         KMSKekVersionVO oldVersion = mock(KMSKekVersionVO.class);
         when(oldVersion.getVersionNumber()).thenReturn(1);
@@ -160,17 +169,16 @@ public class KMSManagerImplKeyRotationTest {
 
         // Provider creates new KEK in different HSM
         when(kmsProvider.createKek(any(KeyPurpose.class), eq(newKekLabel), anyInt(), eq(newProfileId)))
-            .thenReturn("new-kek-id");
+                .thenReturn("new-kek-id");
 
         KMSKekVersionVO newVersion = mock(KMSKekVersionVO.class);
         when(newVersion.getVersionNumber()).thenReturn(2);
         when(kmsKekVersionDao.persist(any(KMSKekVersionVO.class))).thenReturn(newVersion);
 
-        doReturn(kmsProvider).when(kmsManager).getKMSProviderForZone(testZoneId);
+        doReturn(kmsProvider).when(kmsManager).getKMSProvider(testProviderName);
         doReturn(true).when(kmsManager).isKmsEnabled(testZoneId);
 
-        String result = kmsManager.rotateKek(testZoneId, KeyPurpose.VOLUME_ENCRYPTION,
-            oldKekLabel, newKekLabel, 256, newProfileId);
+        String result = kmsManager.rotateKek(kmsKey, oldKekLabel, newKekLabel, 256, hsmProfile);
 
         // Verify new KEK was created in new HSM
         assertNotNull(result);
@@ -219,7 +227,7 @@ public class KMSManagerImplKeyRotationTest {
         WrappedKey newWrappedKey = mock(WrappedKey.class);
         when(newWrappedKey.getWrappedKeyMaterial()).thenReturn("new-wrapped-blob".getBytes());
         when(kmsProvider.wrapKey(plainDek, KeyPurpose.VOLUME_ENCRYPTION, "new-kek-label", newProfileId))
-            .thenReturn(newWrappedKey);
+                .thenReturn(newWrappedKey);
 
         kmsManager.rewrapSingleKey(wrappedKeyVO, kmsKey, newVersion, kmsProvider);
 
@@ -248,7 +256,13 @@ public class KMSManagerImplKeyRotationTest {
         KMSKeyVO kmsKey = mock(KMSKeyVO.class);
         when(kmsKey.getId()).thenReturn(kmsKeyId);
         when(kmsKey.getHsmProfileId()).thenReturn(oldProfileId);
-        when(kmsKeyDao.findByKekLabel(oldKekLabel, testProviderName)).thenReturn(kmsKey);
+        when(kmsKey.getPurpose()).thenReturn(KeyPurpose.VOLUME_ENCRYPTION);
+        when(kmsKey.getZoneId()).thenReturn(testZoneId);
+
+        HSMProfileVO hsmProfile = mock(HSMProfileVO.class);
+        when(hsmProfile.getId()).thenReturn(oldProfileId);
+        when(hsmProfile.getProtocol()).thenReturn(testProviderName);
+        when(hsmProfileDao.findById(oldProfileId)).thenReturn(hsmProfile);
 
         KMSKekVersionVO oldVersion = mock(KMSKekVersionVO.class);
         when(oldVersion.getVersionNumber()).thenReturn(1);
@@ -256,24 +270,24 @@ public class KMSManagerImplKeyRotationTest {
         when(kmsKekVersionDao.getActiveVersion(kmsKeyId)).thenReturn(oldVersion);
         when(kmsKekVersionDao.listByKmsKeyId(kmsKeyId)).thenReturn(Arrays.asList(oldVersion));
 
-        // Provider creates new KEK - capture the generated label
-        ArgumentCaptor<String> labelCaptor = ArgumentCaptor.forClass(String.class);
-        when(kmsProvider.createKek(any(KeyPurpose.class), labelCaptor.capture(), anyInt(), eq(oldProfileId)))
-            .thenReturn("new-kek-id");
+        // Provider creates new KEK - will accept any label
+        when(kmsProvider.createKek(any(KeyPurpose.class), anyString(), anyInt(), eq(oldProfileId)))
+                .thenReturn("new-kek-id");
 
         KMSKekVersionVO newVersion = mock(KMSKekVersionVO.class);
+        when(newVersion.getVersionNumber()).thenReturn(2);
         when(kmsKekVersionDao.persist(any(KMSKekVersionVO.class))).thenReturn(newVersion);
 
-        doReturn(kmsProvider).when(kmsManager).getKMSProviderForZone(testZoneId);
+        doReturn(kmsProvider).when(kmsManager).getKMSProvider(testProviderName);
         doReturn(true).when(kmsManager).isKmsEnabled(testZoneId);
 
-        kmsManager.rotateKek(testZoneId, KeyPurpose.VOLUME_ENCRYPTION,
-            oldKekLabel, null, 256, null);
+        kmsManager.rotateKek(kmsKey, oldKekLabel, null, 256, null);
 
-        // Verify a label was generated
+        // Verify a label was generated and passed to createKek
+        ArgumentCaptor<String> labelCaptor = ArgumentCaptor.forClass(String.class);
+        verify(kmsProvider).createKek(any(KeyPurpose.class), labelCaptor.capture(), eq(256), eq(oldProfileId));
         String generatedLabel = labelCaptor.getValue();
         assertNotNull("Label should be generated", generatedLabel);
-        verify(kmsProvider).createKek(any(KeyPurpose.class), eq(generatedLabel), eq(256), eq(oldProfileId));
     }
 
     /**
@@ -282,13 +296,23 @@ public class KMSManagerImplKeyRotationTest {
     @Test(expected = KMSException.class)
     public void testRotateKek_ThrowsExceptionWhenOldKekNotFound() throws KMSException {
         // Setup: Old KEK doesn't exist
-        when(kmsKeyDao.findByKekLabel("non-existent-label", testProviderName)).thenReturn(null);
+        Long oldProfileId = 10L;
 
-        doReturn(kmsProvider).when(kmsManager).getKMSProviderForZone(testZoneId);
+        KMSKeyVO kmsKey = mock(KMSKeyVO.class);
+        when(kmsKey.getZoneId()).thenReturn(testZoneId);
+        when(kmsKey.getHsmProfileId()).thenReturn(oldProfileId);
+        when(kmsKey.getPurpose()).thenReturn(KeyPurpose.VOLUME_ENCRYPTION);
+
+        HSMProfileVO hsmProfile = mock(HSMProfileVO.class);
+        when(hsmProfile.getId()).thenReturn(oldProfileId);
+        when(hsmProfile.getProtocol()).thenReturn(testProviderName);
+        when(hsmProfileDao.findById(oldProfileId)).thenReturn(hsmProfile);
+
+        doReturn(kmsProvider).when(kmsManager).getKMSProvider(testProviderName);
         doReturn(true).when(kmsManager).isKmsEnabled(testZoneId);
 
-        kmsManager.rotateKek(testZoneId, KeyPurpose.VOLUME_ENCRYPTION,
-            "non-existent-label", "new-label", 256, null);
+        kmsManager.rotateKek(kmsKey,
+                "non-existent-label", "new-label", 256, null);
     }
 
     /**
@@ -303,8 +327,14 @@ public class KMSManagerImplKeyRotationTest {
 
         KMSKeyVO kmsKey = mock(KMSKeyVO.class);
         when(kmsKey.getId()).thenReturn(kmsKeyId);
+        when(kmsKey.getZoneId()).thenReturn(testZoneId);
         when(kmsKey.getHsmProfileId()).thenReturn(currentProfileId);
-        when(kmsKeyDao.findByKekLabel(oldKekLabel, testProviderName)).thenReturn(kmsKey);
+        when(kmsKey.getPurpose()).thenReturn(KeyPurpose.VOLUME_ENCRYPTION);
+
+        HSMProfileVO hsmProfile = mock(HSMProfileVO.class);
+        when(hsmProfile.getId()).thenReturn(currentProfileId);
+        when(hsmProfile.getProtocol()).thenReturn(testProviderName);
+        when(hsmProfileDao.findById(currentProfileId)).thenReturn(hsmProfile);
 
         KMSKekVersionVO oldVersion = mock(KMSKekVersionVO.class);
         when(oldVersion.getVersionNumber()).thenReturn(1);
@@ -313,16 +343,16 @@ public class KMSManagerImplKeyRotationTest {
         when(kmsKekVersionDao.listByKmsKeyId(kmsKeyId)).thenReturn(Arrays.asList(oldVersion));
 
         when(kmsProvider.createKek(any(KeyPurpose.class), anyString(), anyInt(), eq(currentProfileId)))
-            .thenReturn("new-kek-id");
+                .thenReturn("new-kek-id");
 
         KMSKekVersionVO newVersion = mock(KMSKekVersionVO.class);
+        when(newVersion.getVersionNumber()).thenReturn(2);
         when(kmsKekVersionDao.persist(any(KMSKekVersionVO.class))).thenReturn(newVersion);
 
-        doReturn(kmsProvider).when(kmsManager).getKMSProviderForZone(testZoneId);
+        doReturn(kmsProvider).when(kmsManager).getKMSProvider(testProviderName);
         doReturn(true).when(kmsManager).isKmsEnabled(testZoneId);
 
-        kmsManager.rotateKek(testZoneId, KeyPurpose.VOLUME_ENCRYPTION,
-            oldKekLabel, "new-label", 256, null);
+        kmsManager.rotateKek(kmsKey, oldKekLabel, "new-label", 256, null);
 
         // Verify current profile was used (not a different one)
         verify(kmsProvider).createKek(any(KeyPurpose.class), anyString(), eq(256), eq(currentProfileId));
diff --git a/utils/src/main/java/com/cloud/utils/exception/CSExceptionErrorCode.java b/utils/src/main/java/com/cloud/utils/exception/CSExceptionErrorCode.java
index fd31db8a5a3..8b25d77e7f3 100644
--- a/utils/src/main/java/com/cloud/utils/exception/CSExceptionErrorCode.java
+++ b/utils/src/main/java/com/cloud/utils/exception/CSExceptionErrorCode.java
@@ -74,6 +74,7 @@ public class CSExceptionErrorCode {
             ExceptionErrorCodeMap.put("com.cloud.exception.StorageConflictException", 4550);
             ExceptionErrorCodeMap.put("com.cloud.exception.UnavailableCommandException", 4555);
             ExceptionErrorCodeMap.put("com.cloud.exception.OperationTimedoutException", 4560);
+            ExceptionErrorCodeMap.put("org.apache.cloudstack.framework.kms.KMSException", 4561);
 
             // Have a special error code for ServerApiException when it is
             // thrown in a standalone manner when failing to detect any of the above