From 6a79bdf8269037dc7f1d3d738967388bca83f395 Mon Sep 17 00:00:00 2001
From: Maciej Baczmanski <maciej.baczmanski@nordicsemi.no>
Date: Tue, 13 Aug 2024 10:02:05 +0200
Subject: [PATCH] Implement `PersistICDKey` API

---
 src/crypto/PSASessionKeystore.cpp | 58 +++++++++++++++++++++++++++++--
 src/crypto/PSASessionKeystore.h   |  5 +++
 src/crypto/SessionKeystore.h      | 16 +++++++++
 3 files changed, 76 insertions(+), 3 deletions(-)

diff --git a/src/crypto/PSASessionKeystore.cpp b/src/crypto/PSASessionKeystore.cpp
index 8fe6a450e2c78b..3cdd637f16bd3b 100644
--- a/src/crypto/PSASessionKeystore.cpp
+++ b/src/crypto/PSASessionKeystore.cpp
@@ -37,7 +37,8 @@ class KeyAttributesBase
 
     CHIP_ERROR SetKeyPersistence(psa_key_id_t keyId)
     {
-        VerifyOrReturnError(to_underlying(KeyIdBase::Maximum) >= keyId && keyId >= to_underlying(KeyIdBase::Minimum),
+        VerifyOrReturnError((to_underlying(KeyIdBase::Maximum) >= keyId && keyId >= to_underlying(KeyIdBase::Minimum)) ||
+                                (to_underlying(KeyIdOptional::Maximum) >= keyId && keyId >= to_underlying(KeyIdOptional::Minimum)),
                             CHIP_ERROR_INVALID_ARGUMENT);
 
         psa_set_key_lifetime(&mAttrs, PSA_KEY_LIFETIME_PERSISTENT);
@@ -84,7 +85,7 @@ class HkdfKeyAttributes : public KeyAttributesBase
 CHIP_ERROR PSASessionKeystore::CreateKey(const Symmetric128BitsKeyByteArray & keyMaterial, Aes128KeyHandle & key)
 {
     // Destroy the old key if already allocated
-    psa_destroy_key(key.As<psa_key_id_t>());
+    DestroyKey(key);
 
     AesKeyAttributes attrs;
     psa_status_t status =
@@ -98,7 +99,7 @@ CHIP_ERROR PSASessionKeystore::CreateKey(const Symmetric128BitsKeyByteArray & ke
 CHIP_ERROR PSASessionKeystore::CreateKey(const Symmetric128BitsKeyByteArray & keyMaterial, Hmac128KeyHandle & key)
 {
     // Destroy the old key if already allocated
-    psa_destroy_key(key.As<psa_key_id_t>());
+    DestroyKey(key);
 
     HmacKeyAttributes attrs;
     psa_status_t status =
@@ -189,5 +190,56 @@ void PSASessionKeystore::DestroyKey(HkdfKeyHandle & key)
     keyId = PSA_KEY_ID_NULL;
 }
 
+#if CHIP_CONFIG_ENABLE_ICD_CIP
+CHIP_ERROR PSASessionKeystore::PersistICDKey(Aes128KeyHandle & key)
+{
+    CHIP_ERROR err;
+    AesKeyAttributes attrs;
+    psa_key_id_t previousKeyId = key.As<psa_key_id_t>();
+
+    SuccessOrExit(err = Crypto::FindFreeKeySlotInRange(key.AsMutable<psa_key_id_t>(),
+                                                       to_underlying(KeyIdOptional::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_destroy_key(previousKeyId);
+
+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>();
+
+    SuccessOrExit(err = Crypto::FindFreeKeySlotInRange(key.AsMutable<psa_key_id_t>(),
+                                                       to_underlying(KeyIdOptional::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);
+
+exit:
+    if (err != CHIP_NO_ERROR)
+    {
+        psa_destroy_key(previousKeyId);
+        psa_destroy_key(key.As<psa_key_id_t>());
+    }
+
+    return err;
+}
+#endif
+
 } // namespace Crypto
 } // namespace chip
diff --git a/src/crypto/PSASessionKeystore.h b/src/crypto/PSASessionKeystore.h
index 1a55713b82d167..32230972c831c8 100644
--- a/src/crypto/PSASessionKeystore.h
+++ b/src/crypto/PSASessionKeystore.h
@@ -17,6 +17,7 @@
 
 #pragma once
 
+#include <app/icd/server/ICDServerConfig.h>
 #include <crypto/CHIPCryptoPALPSA.h>
 #include <crypto/SessionKeystore.h>
 
@@ -38,6 +39,10 @@ class PSASessionKeystore : public SessionKeystore
                                  AttestationChallenge & attestationChallenge) override;
     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;
+#endif
 
 private:
     CHIP_ERROR DeriveSessionKeys(PsaKdf & kdf, Aes128KeyHandle & i2rKey, Aes128KeyHandle & r2iKey,
diff --git a/src/crypto/SessionKeystore.h b/src/crypto/SessionKeystore.h
index d668524ac93aed..a5149f4baa3136 100644
--- a/src/crypto/SessionKeystore.h
+++ b/src/crypto/SessionKeystore.h
@@ -140,6 +140,22 @@ class SessionKeystore
     virtual CHIP_ERROR DeriveSessionKeys(const HkdfKeyHandle & secretKey, const ByteSpan & salt, const ByteSpan & info,
                                          Aes128KeyHandle & i2rKey, Aes128KeyHandle & r2iKey,
                                          AttestationChallenge & attestationChallenge) = 0;
+
+    /**
+     * @brief Store key in persistent PSA storage and return a key handle for an ICD Aes 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 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; }
 };
 
 /**