#24 Implement token invalidation

Author: JavaCoDED78

src/main/java/io/github/javacoded78/jwthumble/storage/TokenStorageImpl.java

@@ -1,9 +1,15 @@
 package io.github.javacoded78.jwthumble.storage;
 
 import io.github.javacoded78.jwthumble.config.TokenParameters;
+import lombok.Getter;
 
+import java.util.Date;
 import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * Basic implementation of TokenStorage.
@@ -13,13 +19,32 @@ public class TokenStorageImpl implements TokenStorage {
     /**
      * Inner map of key-value pairs.
      */
-    private final ConcurrentHashMap<String, String> tokens;
+    private final ConcurrentHashMap<String, TokenEntry> tokens;
+
+    /**
+     * Scheduled executor for cleanup tokens.
+     */
+    private final ScheduledExecutorService scheduler;
 
     /**
      * Creates an object.
      */
     public TokenStorageImpl() {
         this.tokens = new ConcurrentHashMap<>();
+        this.scheduler = Executors.newSingleThreadScheduledExecutor();
+        scheduler.scheduleAtFixedRate(
+                () -> {
+                    Date date = new Date();
+                    for (Map.Entry<String, TokenEntry> entry : tokens.entrySet()) {
+                        if (entry.getValue().isAfter(date)) {
+                            tokens.remove(entry.getKey());
+                        }
+                    }
+                },
+                0,
+                1,
+                TimeUnit.SECONDS
+        );
     }
 
     private String subjectTokenKey(final String subject,
@@ -34,7 +59,8 @@ public void save(final String token,
                 params.getSubject(),
                 params.getType()
         );
-        tokens.put(tokenKey, token);
+        TokenEntry entry = new TokenEntry(token, params.getExpiredAt());
+        tokens.put(tokenKey, entry);
     }
 
     @Override
@@ -44,7 +70,11 @@ public boolean exists(final String token,
                 params.getSubject(),
                 params.getType()
         );
-        return token.equals(tokens.get(tokenKey));
+        TokenEntry entry = tokens.get(tokenKey);
+        if (entry == null) {
+            return false;
+        }
+        return token.equals(entry.token);
     }
 
     @Override
@@ -53,14 +83,18 @@ public String get(final TokenParameters params) {
                 params.getSubject(),
                 params.getType()
         );
-        return tokens.get(tokenKey);
+        TokenEntry entry = tokens.get(tokenKey);
+        if (entry == null) {
+            return null;
+        }
+        return entry.token;
     }
 
     @Override
     public boolean remove(final String token) {
         AtomicBoolean deleted = new AtomicBoolean(false);
-        for (Map.Entry<String, String> entry : tokens.entrySet()) {
-            if (entry.getValue().equals(token)) {
+        for (Map.Entry<String, TokenEntry> entry : tokens.entrySet()) {
+            if (entry.getValue().token.equals(token)) {
                 tokens.remove(entry.getKey());
                 deleted.set(true);
                 return true;
@@ -78,4 +112,29 @@ public boolean remove(final TokenParameters params) {
         return tokens.remove(tokenKey) != null;
     }
 
+    @Getter
+    private static class TokenEntry {
+
+        /**
+         * Token.
+         */
+        private final String token;
+
+        /**
+         * Expiration date.
+         */
+        private final Date expiredAt;
+
+        TokenEntry(final String token,
+                   final Date expiredAt) {
+            this.token = token;
+            this.expiredAt = expiredAt;
+        }
+
+        public boolean isAfter(final Date date) {
+            return expiredAt.after(date);
+        }
+
+    }
+
 }