From 8835f42f32ff8bc40b5fc5883cd4afa1cbd39fc7 Mon Sep 17 00:00:00 2001
From: mohamedlajmileanix <amin.lajmi@leanix.net>
Date: Tue, 19 Nov 2024 15:54:10 +0100
Subject: [PATCH] CID-3153: Improve JWT generation failure logging

---
 .../advice/GlobalExceptionHandler.kt          |  2 +-
 .../services/GitHubAuthenticationService.kt   | 45 +++++++++++++------
 2 files changed, 33 insertions(+), 14 deletions(-)

diff --git a/src/main/kotlin/net/leanix/githubagent/controllers/advice/GlobalExceptionHandler.kt b/src/main/kotlin/net/leanix/githubagent/controllers/advice/GlobalExceptionHandler.kt
index 57c8372..c207b10 100644
--- a/src/main/kotlin/net/leanix/githubagent/controllers/advice/GlobalExceptionHandler.kt
+++ b/src/main/kotlin/net/leanix/githubagent/controllers/advice/GlobalExceptionHandler.kt
@@ -39,7 +39,7 @@ class GlobalExceptionHandler(
 
     @ExceptionHandler(Exception::class)
     fun handleUncaughtException(exception: Exception): ProblemDetail {
-        val detail = "An unexpected error occurred ${exception.message}"
+        val detail = "An unexpected error occurred. ${exception.message}"
         val problemDetail = ProblemDetail.forStatusAndDetail(HttpStatus.INTERNAL_SERVER_ERROR, detail)
         problemDetail.title = exception.message
         exceptionLogger.error("Uncaught exception: ${exception.message}", exception)
diff --git a/src/main/kotlin/net/leanix/githubagent/services/GitHubAuthenticationService.kt b/src/main/kotlin/net/leanix/githubagent/services/GitHubAuthenticationService.kt
index f8d5d3f..9525a91 100644
--- a/src/main/kotlin/net/leanix/githubagent/services/GitHubAuthenticationService.kt
+++ b/src/main/kotlin/net/leanix/githubagent/services/GitHubAuthenticationService.kt
@@ -17,7 +17,6 @@ import java.nio.file.Files
 import java.security.KeyFactory
 import java.security.PrivateKey
 import java.security.Security
-import java.security.spec.InvalidKeySpecException
 import java.security.spec.PKCS8EncodedKeySpec
 import java.util.*
 
@@ -49,23 +48,31 @@ class GitHubAuthenticationService(
     fun generateAndCacheJwtToken() {
         runCatching {
             logger.info("Generating JWT token")
+            val privateKey = loadPrivateKey()
+            val jwt = createJwtToken(privateKey)
+            verifyAndCacheJwtToken(jwt)
+        }.onFailure {
+            logger.error("Failed to generate a valid jwt token", it)
+            throw it
+        }
+    }
+
+    private fun loadPrivateKey(): PrivateKey {
+        return runCatching {
             Security.addProvider(BouncyCastleProvider())
             val rsaPrivateKey: String = readPrivateKey()
             val keySpec = PKCS8EncodedKeySpec(Base64.getDecoder().decode(rsaPrivateKey))
-            val privateKey = KeyFactory.getInstance("RSA").generatePrivate(keySpec)
-            val jwt = createJwtToken(privateKey)
-            gitHubEnterpriseService.verifyJwt(jwt.getOrThrow())
-            cachingService.set("jwtToken", jwt.getOrThrow(), JWT_EXPIRATION_DURATION)
-        }.onFailure {
-            if (it is InvalidKeySpecException) {
-                throw IllegalArgumentException("The provided private key is not in a valid PKCS8 format.", it)
-            } else {
-                throw it
-            }
+            KeyFactory.getInstance("RSA").generatePrivate(keySpec)
+        }.getOrElse {
+            logger.error("Failed to load private key", it)
+            throw IllegalArgumentException(
+                "Failed to load private key, " +
+                    "the provided private key is not a valid PKCS8 key."
+            )
         }
     }
 
-    private fun createJwtToken(privateKey: PrivateKey): Result<String> {
+    private fun createJwtToken(privateKey: PrivateKey): String {
         return runCatching {
             Jwts.builder()
                 .setIssuedAt(Date())
@@ -73,8 +80,20 @@ class GitHubAuthenticationService(
                 .setIssuer(cachingService.get("githubAppId").toString())
                 .signWith(privateKey, SignatureAlgorithm.RS256)
                 .compact()
+        }.getOrElse {
+            logger.error("Failed to generate a JWT token", it)
+            throw FailedToCreateJWTException("Failed to generate a JWT token")
+        }
+    }
+
+    private fun verifyAndCacheJwtToken(jwt: String) {
+        runCatching {
+            gitHubEnterpriseService.verifyJwt(jwt)
+            cachingService.set("jwtToken", jwt, JWT_EXPIRATION_DURATION)
+            logger.info("JWT token generated and cached successfully")
         }.onFailure {
-            throw FailedToCreateJWTException("Failed to generate a valid JWT token")
+            logger.error("Failed to verify and cache JWT token", it)
+            throw it
         }
     }