From 52b0c0dc506faec9952de64cccb2633488caee60 Mon Sep 17 00:00:00 2001
From: Robert Mathew <robert.c.mathew@gmail.com>
Date: Wed, 28 Aug 2024 12:42:31 +0530
Subject: [PATCH] fix: added missing function in LocalKmsDatabaseConnection

---
 .../src/main/kotlin/KeyRepository.kt          |  3 +--
 .../main/kotlin/LocalKmsDatabaseConnection.kt | 21 ++++++++++++++++---
 .../src/main/resources/application.properties |  5 +++++
 .../src/main/resources/{1.sql => schema.sql}  |  0
 4 files changed, 24 insertions(+), 5 deletions(-)
 rename modules/local-kms/src/main/resources/{1.sql => schema.sql} (100%)

diff --git a/modules/local-kms/src/main/kotlin/KeyRepository.kt b/modules/local-kms/src/main/kotlin/KeyRepository.kt
index b95f1269..b1ce65e7 100644
--- a/modules/local-kms/src/main/kotlin/KeyRepository.kt
+++ b/modules/local-kms/src/main/kotlin/KeyRepository.kt
@@ -5,8 +5,7 @@ import org.springframework.data.repository.CrudRepository
 import org.springframework.stereotype.Repository
 
 @Repository
-interface KeyRepository : CrudRepository<Key,
-        String> {
+interface KeyRepository : CrudRepository<Key, String> {
 
     @Query("SELECT * FROM keys WHERE id = :keyId")
     fun findByKeyId(keyId: String): Key?
diff --git a/modules/local-kms/src/main/kotlin/LocalKmsDatabaseConnection.kt b/modules/local-kms/src/main/kotlin/LocalKmsDatabaseConnection.kt
index b96e3d15..b79d307d 100644
--- a/modules/local-kms/src/main/kotlin/LocalKmsDatabaseConnection.kt
+++ b/modules/local-kms/src/main/kotlin/LocalKmsDatabaseConnection.kt
@@ -1,6 +1,7 @@
 package com.sphereon.oid.fed.kms.local
 
 import org.springframework.beans.factory.annotation.Autowired
+import org.springframework.dao.EmptyResultDataAccessException
 import org.springframework.stereotype.Component
 
 @Component
@@ -12,8 +13,22 @@ class LocalKmsDatabaseConnection @Autowired constructor(private val keyRepositor
     }
 
     fun getKey(keyId: String): Key {
-        return keyRepository.findByKeyId(keyId) ?: throw Exception("Key not found")
+        return keyRepository.findByKeyId(keyId) ?: throw KeyNotFoundException("Key with ID $keyId not found")
     }
 
-    // ... (Implement other methods like updateKey, deleteKey as needed)
-}
\ No newline at end of file
+    fun updateKey(keyId: String, privateKey: ByteArray, publicKey: ByteArray, algorithm: String) {
+        val existingKey = keyRepository.findByKeyId(keyId) ?: throw KeyNotFoundException("Key with ID $keyId not found")
+        val updatedKey = existingKey.copy(privateKey = privateKey, publicKey = publicKey, algorithm = algorithm)
+        keyRepository.save(updatedKey)
+    }
+
+    fun deleteKey(keyId: String) {
+        try {
+            keyRepository.deleteById(keyId)
+        } catch (e: EmptyResultDataAccessException) {
+            throw KeyNotFoundException("Key with ID $keyId not found")
+        }
+    }
+}
+
+class KeyNotFoundException(message: String) : Exception(message)
\ No newline at end of file
diff --git a/modules/local-kms/src/main/resources/application.properties b/modules/local-kms/src/main/resources/application.properties
index 0ac4201e..98062da0 100644
--- a/modules/local-kms/src/main/resources/application.properties
+++ b/modules/local-kms/src/main/resources/application.properties
@@ -3,7 +3,12 @@ spring.application.name=OpenID Federation Server
 spring.datasource.url=${DATASOURCE_URL}
 spring.datasource.username=${DATASOURCE_USER}
 spring.datasource.password=${DATASOURCE_PASSWORD}
+
 # Mapping /actuator/health to /status
 management.endpoints.web.base-path=/
 management.endpoints.web.path-mapping.health=status
 server.port=8080
+
+# Spring Boot to execute the schema.sql script on startup
+spring.sql.init.mode=always
+spring.datasource.schema=classpath:schema.sql
diff --git a/modules/local-kms/src/main/resources/1.sql b/modules/local-kms/src/main/resources/schema.sql
similarity index 100%
rename from modules/local-kms/src/main/resources/1.sql
rename to modules/local-kms/src/main/resources/schema.sql