Address PR comments

Modify implementation of `PersistICDKey` not to depend on type
of ICD key handle (AES, HMAC)

src/crypto/CHIPCryptoPALPSA.h

@@ -71,7 +71,8 @@ namespace Crypto {
 static_assert(PSA_KEY_ID_USER_MIN <= CHIP_CONFIG_CRYPTO_PSA_KEY_ID_BASE && CHIP_CONFIG_CRYPTO_PSA_KEY_ID_END <= PSA_KEY_ID_USER_MAX,
               "Matter specific PSA key range doesn't fit within PSA allowed range");
 
-static constexpr uint32_t kMaxICDClientKeys = CHIP_CONFIG_ICD_MAX_CLIENTS;
+// Each ICD client requires storing two keys- AES and HMAC
+static constexpr uint32_t kMaxICDClientKeys = 2 * CHIP_CONFIG_CRYPTO_PSA_ICD_MAX_CLIENTS;
 
 static_assert(kMaxICDClientKeys >= CHIP_CONFIG_ICD_CLIENTS_SUPPORTED_PER_FABRIC * CHIP_CONFIG_MAX_FABRICS,
               "Number of allocated ICD key slots is lower than maximum number of supported ICD clients");
@@ -81,12 +82,11 @@ static_assert(kMaxICDClientKeys >= CHIP_CONFIG_ICD_CLIENTS_SUPPORTED_PER_FABRIC
  */
 enum class KeyIdBase : psa_key_id_t
 {
-    Minimum              = CHIP_CONFIG_CRYPTO_PSA_KEY_ID_BASE,
-    Operational          = Minimum, ///< Base of the PSA key ID range for Node Operational Certificate private keys
-    DACPrivKey           = Operational + kMaxValidFabricIndex + 1,
-    ICDHmacKeyRangeStart = DACPrivKey + 1,
-    ICDAesKeyRangeStart  = ICDHmacKeyRangeStart + kMaxICDClientKeys,
-    Maximum              = ICDAesKeyRangeStart + kMaxICDClientKeys,
+    Minimum          = CHIP_CONFIG_CRYPTO_PSA_KEY_ID_BASE,
+    Operational      = Minimum, ///< Base of the PSA key ID range for Node Operational Certificate private keys
+    DACPrivKey       = Operational + kMaxValidFabricIndex + 1,
+    ICDKeyRangeStart = DACPrivKey + 1,
+    Maximum          = ICDKeyRangeStart + kMaxICDClientKeys,
 };
 
 static_assert(to_underlying(KeyIdBase::Minimum) >= CHIP_CONFIG_CRYPTO_PSA_KEY_ID_BASE &&

src/crypto/PSASessionKeystore.cpp

@@ -35,17 +35,6 @@ class KeyAttributesBase
         psa_set_key_bits(&mAttrs, bits);
     }
 
-    CHIP_ERROR SetKeyPersistence(psa_key_id_t keyId)
-    {
-        VerifyOrReturnError(to_underlying(KeyIdBase::Maximum) >= keyId && keyId >= to_underlying(KeyIdBase::Minimum),
-                            CHIP_ERROR_INVALID_ARGUMENT);
-
-        psa_set_key_lifetime(&mAttrs, PSA_KEY_LIFETIME_PERSISTENT);
-        psa_set_key_id(&mAttrs, keyId);
-
-        return CHIP_NO_ERROR;
-    }
-
     ~KeyAttributesBase() { psa_reset_key_attributes(&mAttrs); }
 
     const psa_key_attributes_t & Get() { return mAttrs; }
@@ -79,6 +68,12 @@ class HkdfKeyAttributes : public KeyAttributesBase
     HkdfKeyAttributes() : KeyAttributesBase(PSA_KEY_TYPE_DERIVE, PSA_ALG_HKDF(PSA_ALG_SHA_256), PSA_KEY_USAGE_DERIVE, 0) {}
 };
 
+void SetKeyId(Symmetric128BitsKeyHandle & key, psa_key_id_t newKeyId)
+{
+    auto & KeyId = key.AsMutable<psa_key_id_t>();
+
+    KeyId = newKeyId;
+}
 } // namespace
 
 CHIP_ERROR PSASessionKeystore::CreateKey(const Symmetric128BitsKeyByteArray & keyMaterial, Aes128KeyHandle & key)
@@ -190,66 +185,33 @@ void PSASessionKeystore::DestroyKey(HkdfKeyHandle & key)
 }
 
 #if CHIP_CONFIG_ENABLE_ICD_CIP
-CHIP_ERROR PSASessionKeystore::PersistICDKey(Aes128KeyHandle & key)
+CHIP_ERROR PSASessionKeystore::PersistICDKey(Symmetric128BitsKeyHandle & key)
 {
     CHIP_ERROR err;
-    AesKeyAttributes attrs;
-    psa_key_id_t previousKeyId = key.As<psa_key_id_t>();
-    psa_key_attributes_t previousKeyAttrs;
-
-    psa_get_key_attributes(previousKeyId, &previousKeyAttrs);
-    // Exit early if key is already persistent
-    if (psa_get_key_lifetime(&previousKeyAttrs) == PSA_KEY_LIFETIME_PERSISTENT)
-    {
-        ExitNow(err = CHIP_NO_ERROR);
-    }
-
-    SuccessOrExit(err = Crypto::FindFreeKeySlotInRange(key.AsMutable<psa_key_id_t>(), to_underlying(KeyIdBase::ICDAesKeyRangeStart),
-                                                       kMaxICDClientKeys));
-
-    SuccessOrExit(err = attrs.SetKeyPersistence(key.As<psa_key_id_t>()));
-    VerifyOrExit(psa_copy_key(previousKeyId, &attrs.Get(), &key.AsMutable<psa_key_id_t>()) == PSA_SUCCESS,
-                 err = CHIP_ERROR_INTERNAL);
+    psa_key_id_t newKeyId = PSA_KEY_ID_NULL;
+    psa_key_attributes_t attrs;
 
-    psa_destroy_key(previousKeyId);
+    psa_get_key_attributes(key.As<psa_key_id_t>(), &attrs);
 
-exit:
-    if (err != CHIP_NO_ERROR)
-    {
-        psa_destroy_key(previousKeyId);
-        psa_destroy_key(key.As<psa_key_id_t>());
-    }
-
-    return err;
-}
-
-CHIP_ERROR PSASessionKeystore::PersistICDKey(Hmac128KeyHandle & key)
-{
-    CHIP_ERROR err;
-    HmacKeyAttributes attrs;
-    psa_key_id_t previousKeyId = key.As<psa_key_id_t>();
-    psa_key_attributes_t previousKeyAttrs;
-
-    psa_get_key_attributes(previousKeyId, &previousKeyAttrs);
     // Exit early if key is already persistent
-    if (psa_get_key_lifetime(&previousKeyAttrs) == PSA_KEY_LIFETIME_PERSISTENT)
+    if (psa_get_key_lifetime(&attrs) == PSA_KEY_LIFETIME_PERSISTENT)
     {
-        ExitNow(err = CHIP_NO_ERROR);
+        psa_reset_key_attributes(&attrs);
+        return CHIP_NO_ERROR;
     }
 
-    SuccessOrExit(err = Crypto::FindFreeKeySlotInRange(key.AsMutable<psa_key_id_t>(),
-                                                       to_underlying(KeyIdBase::ICDHmacKeyRangeStart), kMaxICDClientKeys));
-    SuccessOrExit(err = attrs.SetKeyPersistence(key.As<psa_key_id_t>()));
-    VerifyOrExit(psa_copy_key(previousKeyId, &attrs.Get(), &key.AsMutable<psa_key_id_t>()) == PSA_SUCCESS,
-                 err = CHIP_ERROR_INTERNAL);
-
-    psa_destroy_key(previousKeyId);
+    SuccessOrExit(err = Crypto::FindFreeKeySlotInRange(newKeyId, to_underlying(KeyIdBase::ICDKeyRangeStart), kMaxICDClientKeys));
+    psa_set_key_lifetime(&attrs, PSA_KEY_LIFETIME_PERSISTENT);
+    psa_set_key_id(&attrs, newKeyId);
+    VerifyOrExit(psa_copy_key(key.As<psa_key_id_t>(), &attrs, &newKeyId) == PSA_SUCCESS, err = CHIP_ERROR_INTERNAL);
 
 exit:
-    if (err != CHIP_NO_ERROR)
+    DestroyKey(key);
+    psa_reset_key_attributes(&attrs);
+
+    if(err == CHIP_NO_ERROR)
     {
-        psa_destroy_key(previousKeyId);
-        psa_destroy_key(key.As<psa_key_id_t>());
+        SetKeyId(key, newKeyId);
     }
 
     return err;

src/crypto/PSASessionKeystore.h

@@ -17,7 +17,6 @@
 
 #pragma once
 
-#include <app/icd/server/ICDServerConfig.h>
 #include <crypto/CHIPCryptoPALPSA.h>
 #include <crypto/SessionKeystore.h>
 
@@ -40,8 +39,7 @@ class PSASessionKeystore : public SessionKeystore
     void DestroyKey(Symmetric128BitsKeyHandle & key) override;
     void DestroyKey(HkdfKeyHandle & key) override;
 #if CHIP_CONFIG_ENABLE_ICD_CIP
-    CHIP_ERROR PersistICDKey(Aes128KeyHandle & key) override;
-    CHIP_ERROR PersistICDKey(Hmac128KeyHandle & key) override;
+    CHIP_ERROR PersistICDKey(Symmetric128BitsKeyHandle & key) override;
 #endif
 
 private:

src/crypto/SessionKeystore.h

@@ -17,6 +17,7 @@
 
 #pragma once
 
+#include <app/icd/server/ICDServerConfig.h>
 #include <crypto/CHIPCryptoPAL.h>
 #include <lib/support/Span.h>
 
@@ -141,21 +142,18 @@ class SessionKeystore
                                          Aes128KeyHandle & i2rKey, Aes128KeyHandle & r2iKey,
                                          AttestationChallenge & attestationChallenge) = 0;
 
+#if CHIP_CONFIG_ENABLE_ICD_CIP
     /**
-     * @brief Store key in persistent PSA storage and return a key handle for an ICD Aes key.
+     * @brief Persistently store an ICD key.
      *
-     * If the method returns no error, the application is responsible for destroying the handle
-     * using the DestroyKey() method when the key is no longer needed.
-     */
-    virtual CHIP_ERROR PersistICDKey(Aes128KeyHandle & key) { return CHIP_NO_ERROR; }
-
-    /**
-     * @brief Store key in persistent PSA storage and return a key handle for an ICD Hmac key.
+     * If input is already a persistent key handle, the function is a no-op and the original handle is returned.
+     * If input is a volatile key handle, key is persisted and the handle may be updated.
      *
      * If the method returns no error, the application is responsible for destroying the handle
      * using the DestroyKey() method when the key is no longer needed.
      */
-    virtual CHIP_ERROR PersistICDKey(Hmac128KeyHandle & key) { return CHIP_NO_ERROR; }
+    virtual CHIP_ERROR PersistICDKey(Symmetric128BitsKeyHandle & key) { return CHIP_NO_ERROR; }
+#endif
 };
 
 /**

src/lib/core/CHIPConfig.h

@@ -1668,7 +1668,7 @@ extern const char CHIP_NON_PRODUCTION_MARKER[];
 #endif
 
 /**
- * @def CHIP_CONFIG_ICD_MAX_CLIENTS
+ * @def CHIP_CONFIG_CRYPTO_PSA_ICD_MAX_CLIENTS
  *
  * @brief
  *   Maximum number of ICD clients. Based on this number, platforms that utilize the
@@ -1679,8 +1679,8 @@ extern const char CHIP_NON_PRODUCTION_MARKER[];
  *   compute the number of PSA key slots. It should remain unchanged during the device's lifetime,
  *   as alterations may lead to issues with backwards compatibility.
  */
-#ifndef CHIP_CONFIG_ICD_MAX_CLIENTS
-#define CHIP_CONFIG_ICD_MAX_CLIENTS 256
+#ifndef CHIP_CONFIG_CRYPTO_PSA_ICD_MAX_CLIENTS
+#define CHIP_CONFIG_CRYPTO_PSA_ICD_MAX_CLIENTS 256
 #endif
 
 /**

src/platform/nrfconnect/CHIPPlatformConfig.h

@@ -146,8 +146,8 @@
 #endif // CONFIG_CHIP_ICD_CLIENTS_PER_FABRIC
 #endif // CHIP_CONFIG_ICD_CLIENTS_SUPPORTED_PER_FABRIC
 
-#ifndef CHIP_CONFIG_ICD_MAX_CLIENTS
-#define CHIP_CONFIG_ICD_MAX_CLIENTS 256
+#ifndef CHIP_CONFIG_CRYPTO_PSA_ICD_MAX_CLIENTS
+#define CHIP_CONFIG_CRYPTO_PSA_ICD_MAX_CLIENTS 256
 #endif
 
 #ifndef CHIP_CONFIG_ENABLE_BDX_LOG_TRANSFER