fix: revoke and cleanup

supertokens · Sep 24, 2024 · 7c2d06e · 7c2d06e
1 parent caf5adb
commit 7c2d06e
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 6 deletions.
diff --git a/src/main/java/io/supertokens/storage/postgresql/Start.java b/src/main/java/io/supertokens/storage/postgresql/Start.java
@@ -3120,10 +3120,10 @@ public List<String> listClientsForApp(AppIdentifier appIdentifier) throws Storag
     }
 
     @Override
-    public void revoke(AppIdentifier appIdentifier, String targetType, String targetValue)
+    public void revoke(AppIdentifier appIdentifier, String targetType, String targetValue, long exp)
             throws StorageQueryException {
         try {
-            OAuthQueries.revoke(this, appIdentifier, targetType, targetValue);
+            OAuthQueries.revoke(this, appIdentifier, targetType, targetValue, exp);
         } catch (SQLException e) {
             throw new StorageQueryException(e);
         }
@@ -3149,6 +3149,15 @@ public void addM2MToken(AppIdentifier appIdentifier, String clientId, long iat,
         }
     }
 
+    @Override
+    public void cleanUpExpiredAndRevokedTokens(AppIdentifier appIdentifier) throws StorageQueryException {
+        try {
+            OAuthQueries.cleanUpExpiredAndRevokedTokens(this, appIdentifier);
+        } catch (SQLException e) {
+            throw new StorageQueryException(e);
+        }
+    }
+
     @Override
     public int countTotalNumberOfM2MTokensAlive(AppIdentifier appIdentifier) throws StorageQueryException {
         try {

diff --git a/src/main/java/io/supertokens/storage/postgresql/queries/GeneralQueries.java b/src/main/java/io/supertokens/storage/postgresql/queries/GeneralQueries.java
@@ -563,6 +563,7 @@ public static void createTablesIfNotExists(Start start, Connection con) throws S
 
                     // index
                     update(con, OAuthQueries.getQueryToCreateOAuthRevokeTimestampIndex(start), NO_OP_SETTER);
+                    update(con, OAuthQueries.getQueryToCreateOAuthRevokeExpIndex(start), NO_OP_SETTER);
                 }
 
                 if (!doesTableExists(start, con, Config.getConfig(start).getOAuthM2MTokensTable())) {

diff --git a/src/main/java/io/supertokens/storage/postgresql/queries/OAuthQueries.java b/src/main/java/io/supertokens/storage/postgresql/queries/OAuthQueries.java
@@ -41,6 +41,7 @@ public static String getQueryToCreateOAuthRevokeTable(Start start) {
                 + "target_type VARCHAR(16) NOT NULL,"
                 + "target_value VARCHAR(128) NOT NULL,"
                 + "timestamp BIGINT NOT NULL,"
+                + "exp BIGINT NOT NULL,"
                 + "CONSTRAINT " + Utils.getConstraintName(schema, oAuth2ClientTable, "client_id", "pkey")
                 + " PRIMARY KEY (app_id, target_type, target_value),"
                 + "CONSTRAINT " + Utils.getConstraintName(schema, oAuth2ClientTable, "app_id", "fkey")
@@ -56,6 +57,12 @@ public static String getQueryToCreateOAuthRevokeTimestampIndex(Start start) {
                 + oAuth2ClientTable + "(timestamp DESC, app_id DESC);";
     }
 
+    public static String getQueryToCreateOAuthRevokeExpIndex(Start start) {
+        String oAuth2ClientTable = Config.getConfig(start).getOAuthRevokeTable();
+        return "CREATE INDEX IF NOT EXISTS oauth_revoke_exp_index ON "
+                + oAuth2ClientTable + "(exp DESC, app_id DESC);";
+    }
+
     public static String getQueryToCreateOAuthM2MTokensTable(Start start) {
         String schema = Config.getConfig(start).getTableSchema();
         String oAuth2ClientTable = Config.getConfig(start).getOAuthM2MTokensTable();
@@ -137,19 +144,21 @@ public static boolean deleteClientIdForAppId(Start start, String clientId, AppId
         return numberOfRow > 0;
     }
 
-    public static void revoke(Start start, AppIdentifier appIdentifier, String targetType, String targetValue)
+    public static void revoke(Start start, AppIdentifier appIdentifier, String targetType, String targetValue, long exp)
             throws SQLException, StorageQueryException {
         String INSERT = "INSERT INTO " + Config.getConfig(start).getOAuthRevokeTable()
-                + "(app_id, target_type, target_value, timestamp) VALUES (?, ?, ?, ?) "
-                + "ON CONFLICT (app_id, target_type, target_value) DO UPDATE SET timestamp = ?";
+                + "(app_id, target_type, target_value, timestamp, exp) VALUES (?, ?, ?, ?, ?) "
+                + "ON CONFLICT (app_id, target_type, target_value) DO UPDATE SET timestamp = ?, exp = ?";
 
         long currentTime = System.currentTimeMillis() / 1000;
         update(start, INSERT, pst -> {
             pst.setString(1, appIdentifier.getAppId());
             pst.setString(2, targetType);
             pst.setString(3, targetValue);
             pst.setLong(4, currentTime);
-            pst.setLong(5, currentTime);
+            pst.setLong(5, exp);
+            pst.setLong(6, currentTime);
+            pst.setLong(7, exp);
         });
     }
 
@@ -252,4 +261,30 @@ public static void addM2MToken(Start start, AppIdentifier appIdentifier, String
             pst.setLong(4, exp);
         });
     }
+
+    public static void cleanUpExpiredAndRevokedTokens(Start start, AppIdentifier appIdentifier) throws SQLException, StorageQueryException {
+        {
+            // delete expired M2M tokens
+            String QUERY = "DELETE FROM " + Config.getConfig(start).getOAuthM2MTokensTable() +
+                    " WHERE app_id = ? AND exp < ?";
+
+            long timestamp = System.currentTimeMillis() / 1000 - 3600 * 24 * 31; // expired 31 days ago
+            update(start, QUERY, pst -> {
+                pst.setString(1, appIdentifier.getAppId());
+                pst.setLong(2, timestamp);
+            });
+        }
+
+        {
+            // delete expired revoked tokens
+            String QUERY = "DELETE FROM " + Config.getConfig(start).getOAuthRevokeTable() +
+                    " WHERE app_id = ? AND exp < ?";
+
+            long timestamp = System.currentTimeMillis() / 1000 - 3600 * 24 * 31; // expired 31 days ago
+            update(start, QUERY, pst -> {
+                pst.setString(1, appIdentifier.getAppId());
+                pst.setLong(2, timestamp);
+            });
+        }
+    }
 }