Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 43 additions & 19 deletions src/main/java/com/iemr/tm/utils/JwtUtil.java
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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);
}
Comment on lines +65 to +68
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Fix inconsistency between comment and implementation

The comment states "Log error but don't throw" but the code throws a RuntimeException.

Either update the comment or change the implementation:

         } catch (Exception e) {
-            // Log error but don't throw - token might be invalid already
+            // Propagate error - token invalidation should be reliable
             throw new RuntimeException("Failed to invalidate token", e);
         }

Or if you want to follow the comment's intent:

         } catch (Exception e) {
-            // Log error but don't throw - token might be invalid already
-            throw new RuntimeException("Failed to invalidate token", e);
+            // Log error but don't throw - token might be invalid already
+            // logger.error("Failed to invalidate token", e);
         }
πŸ“ Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
} catch (Exception e) {
// Log error but don't throw - token might be invalid already
throw new RuntimeException("Failed to invalidate token", e);
}
} catch (Exception e) {
// Propagate error - token invalidation should be reliable
throw new RuntimeException("Failed to invalidate token", e);
}
πŸ€– Prompt for AI Agents
In src/main/java/com/iemr/tm/utils/JwtUtil.java around lines 65 to 68, the
comment says to log the error without throwing, but the code throws a
RuntimeException. To fix this inconsistency, either update the comment to
reflect that an exception is thrown or modify the catch block to log the error
without throwing any exception, following the comment's intent.

}

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();
}
Comment on lines 80 to 86
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

πŸ› οΈ Refactor suggestion

Consider adding exception handling for invalid tokens

The extractAllClaims method will throw exceptions for invalid tokens. While validateToken handles this properly, invalidateToken expects this method to return null for invalid tokens based on the null check at line 58.

Consider wrapping in try-catch to align with the usage pattern:

 private Claims extractAllClaims(String token) {
+    try {
         return Jwts.parser()
             .verifyWith(getSigningKey())
             .build()
             .parseSignedClaims(token)
             .getPayload();
+    } catch (Exception e) {
+        return null;
+    }
 }
πŸ“ Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
private Claims extractAllClaims(String token) {
return Jwts.parser().setSigningKey(getSigningKey()).build().parseClaimsJws(token).getBody();
return Jwts.parser()
.verifyWith(getSigningKey())
.build()
.parseSignedClaims(token)
.getPayload();
}
private Claims extractAllClaims(String token) {
try {
return Jwts.parser()
.verifyWith(getSigningKey())
.build()
.parseSignedClaims(token)
.getPayload();
} catch (Exception e) {
return null;
}
}
πŸ€– Prompt for AI Agents
In src/main/java/com/iemr/tm/utils/JwtUtil.java around lines 80 to 86, the
extractAllClaims method currently throws exceptions for invalid tokens, but
invalidateToken expects it to return null on invalid tokens. To fix this, wrap
the method's logic in a try-catch block that catches relevant exceptions (e.g.,
JwtException) and returns null when an exception occurs, ensuring consistency
with the null check usage.

}
40 changes: 40 additions & 0 deletions src/main/java/com/iemr/tm/utils/TokenDenylist.java
Original file line number Diff line number Diff line change
@@ -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);
}
}
}
Comment on lines +15 to +25
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Use atomic operation for setting key with expiration

The current implementation sets the value and expiration in two separate operations. If the expire call fails, the key will remain in Redis indefinitely without expiration.

Apply this diff to use an atomic operation:

 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);
+            redisTemplate.opsForValue().set(key, true, expirationTime, TimeUnit.MILLISECONDS);
         } catch (Exception e) {
             throw new RuntimeException("Failed to add token to denylist", e);
         }
     }
 }
πŸ“ Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
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 void addTokenToDenylist(String jti, Long expirationTime) {
if (jti != null && expirationTime != null && expirationTime > 0) {
try {
String key = PREFIX + jti;
redisTemplate.opsForValue().set(key, true, expirationTime, TimeUnit.MILLISECONDS);
} catch (Exception e) {
throw new RuntimeException("Failed to add token to denylist", e);
}
}
}
πŸ€– Prompt for AI Agents
In src/main/java/com/iemr/tm/utils/TokenDenylist.java between lines 15 and 25,
the method addTokenToDenylist sets the Redis key and expiration in two separate
calls, risking keys without expiration if expire fails. Fix this by using a
single atomic Redis operation that sets the key with its expiration time in one
call, such as redisTemplate.opsForValue().set(key, true, expirationTime,
TimeUnit.MILLISECONDS), to ensure the key always has the correct TTL.


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;
}
}

}