From 7e39ef8800dfb03726b65b73168e79e67e04892a Mon Sep 17 00:00:00 2001
From: Mithun James <1007084+drtechie@users.noreply.github.com>
Date: Fri, 30 May 2025 19:02:23 +0530
Subject: [PATCH] fix: ensure jwt is not in deny list before further
 authentication

---
 src/main/java/com/iemr/tm/utils/JwtUtil.java  | 62 +++++++++++++------
 .../java/com/iemr/tm/utils/TokenDenylist.java | 40 ++++++++++++
 2 files changed, 83 insertions(+), 19 deletions(-)
 create mode 100644 src/main/java/com/iemr/tm/utils/TokenDenylist.java

diff --git a/src/main/java/com/iemr/tm/utils/JwtUtil.java b/src/main/java/com/iemr/tm/utils/JwtUtil.java
index 162295f..4905794 100644
--- a/src/main/java/com/iemr/tm/utils/JwtUtil.java
+++ b/src/main/java/com/iemr/tm/utils/JwtUtil.java
@@ -1,15 +1,15 @@
 package com.iemr.tm.utils;
 
-import java.security.Key;
-import java.util.Date;
 import java.util.function.Function;
 
+import javax.crypto.SecretKey;
+
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.stereotype.Component;
 
 import io.jsonwebtoken.Claims;
 import io.jsonwebtoken.Jwts;
-import io.jsonwebtoken.SignatureAlgorithm;
 import io.jsonwebtoken.security.Keys;
 
 @Component
@@ -18,46 +18,70 @@ public class JwtUtil {
 	@Value("${jwt.secret}")
 	private String SECRET_KEY;
 
-	private static final long EXPIRATION_TIME = 24L * 60 * 60 * 1000; // 1 day in milliseconds
+	@Autowired
+	private TokenDenylist tokenDenylist;
 
 	// Generate a key using the secret
-	private Key getSigningKey() {
+	private SecretKey getSigningKey() {
 		if (SECRET_KEY == null || SECRET_KEY.isEmpty()) {
 			throw new IllegalStateException("JWT secret key is not set in application.properties");
 		}
 		return Keys.hmacShaKeyFor(SECRET_KEY.getBytes());
 	}
 
-	// Generate JWT Token
-	public String generateToken(String username, String userId) {
-		Date now = new Date();
-		Date expiryDate = new Date(now.getTime() + EXPIRATION_TIME);
-
-		// Include the userId in the JWT claims
-		return Jwts.builder().setSubject(username).claim("userId", userId) // Add userId as a claim
-				.setIssuedAt(now).setExpiration(expiryDate).signWith(getSigningKey(), SignatureAlgorithm.HS256)
-				.compact();
-	}
-
 	// Validate and parse JWT Token
 	public Claims validateToken(String token) {
 		try {
-			return Jwts.parser().setSigningKey(getSigningKey()).build().parseClaimsJws(token).getBody();
+			Claims claims = Jwts.parser()
+				.verifyWith(getSigningKey())
+				.build()
+				.parseSignedClaims(token)
+				.getPayload();
+				
+			String jti = claims.getId();
+			
+			// Check if token is denylisted (only if jti exists)
+			if (jti != null && tokenDenylist.isTokenDenylisted(jti)) {
+				return null;
+			}
+			
+			return claims;
 		} catch (Exception e) {
 			return null; // Handle token parsing/validation errors
 		}
 	}
 
+	// Invalidate a token by adding it to denylist
+	public void invalidateToken(String token) {
+		try {
+			Claims claims = extractAllClaims(token);
+			if (claims != null && claims.getId() != null) {
+				// Get remaining time until expiration
+				long expirationTime = claims.getExpiration().getTime() - System.currentTimeMillis();
+				if (expirationTime > 0) {
+					tokenDenylist.addTokenToDenylist(claims.getId(), expirationTime);
+				}
+			}
+		} catch (Exception e) {
+			// Log error but don't throw - token might be invalid already
+			throw new RuntimeException("Failed to invalidate token", e);
+		}
+	}
+
 	public String extractUsername(String token) {
 		return extractClaim(token, Claims::getSubject);
 	}
 
 	public <T> T extractClaim(String token, Function<Claims, T> claimsResolver) {
 		final Claims claims = extractAllClaims(token);
-		return claimsResolver.apply(claims);
+		return claims != null ? claimsResolver.apply(claims) : null;
 	}
 
 	private Claims extractAllClaims(String token) {
-		return Jwts.parser().setSigningKey(getSigningKey()).build().parseClaimsJws(token).getBody();
+		return Jwts.parser()
+			.verifyWith(getSigningKey())
+			.build()
+			.parseSignedClaims(token)
+			.getPayload();
 	}
 }
diff --git a/src/main/java/com/iemr/tm/utils/TokenDenylist.java b/src/main/java/com/iemr/tm/utils/TokenDenylist.java
new file mode 100644
index 0000000..54a7635
--- /dev/null
+++ b/src/main/java/com/iemr/tm/utils/TokenDenylist.java
@@ -0,0 +1,40 @@
+package com.iemr.tm.utils;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.data.redis.core.RedisTemplate;
+import org.springframework.stereotype.Component;
+import java.util.concurrent.TimeUnit;
+
+@Component
+public class TokenDenylist {
+    private static final String PREFIX = "denied_";
+    
+    @Autowired
+    private RedisTemplate<String, Object> redisTemplate;
+    
+    public void addTokenToDenylist(String jti, Long expirationTime) {
+        if (jti != null && expirationTime != null && expirationTime > 0) {
+            try {
+                String key = PREFIX + jti;
+                redisTemplate.opsForValue().set(key, true);
+                redisTemplate.expire(key, expirationTime, TimeUnit.MILLISECONDS);
+            } catch (Exception e) {
+                throw new RuntimeException("Failed to add token to denylist", e);
+            }
+        }
+    }
+    
+    public boolean isTokenDenylisted(String jti) {
+        if (jti == null) {
+            return false;
+        }
+        try {
+            Object value = redisTemplate.opsForValue().get(PREFIX + jti);
+            return value != null;
+        } catch (Exception e) {
+            // If Redis is down, assume token is valid to prevent service disruption
+            return false;
+        }
+    }
+
+}