fix: inememory impl

supertokens · Sep 24, 2024 · c58d162 · c58d162
1 parent 930cd2c
commit c58d162
Show file tree

Hide file tree

Showing 4 changed files with 141 additions and 15 deletions.
diff --git a/src/main/java/io/supertokens/inmemorydb/Start.java b/src/main/java/io/supertokens/inmemorydb/Start.java
@@ -3023,7 +3023,7 @@ public boolean doesClientIdExistForApp(AppIdentifier appIdentifier, String clien
     public void addOrUpdateClientForApp(AppIdentifier appIdentifier, String clientId, boolean isClientCredentialsOnly)
             throws StorageQueryException {
         try {
-            OAuthQueries.insertClientIdForAppId(this, clientId, appIdentifier);
+            OAuthQueries.insertClientIdForAppId(this, appIdentifier, clientId, isClientCredentialsOnly);
         } catch (SQLException e) {
             throw new StorageQueryException(e);
         }
@@ -3071,28 +3071,48 @@ public boolean isRevoked(AppIdentifier appIdentifier, String[] targetTypes, Stri
     @Override
     public void addM2MToken(AppIdentifier appIdentifier, String clientId, long iat, long exp)
             throws StorageQueryException {
-        // TODO
+        try {
+            OAuthQueries.addM2MToken(this, appIdentifier, clientId, iat, exp);
+        } catch (SQLException e) {
+            throw new StorageQueryException(e);
+        }
     }
 
     @Override
     public int countTotalNumberOfClientCredentialsOnlyClientsForApp(AppIdentifier appIdentifier)
             throws StorageQueryException {
-        return 0; // TODO
+        try {
+            return OAuthQueries.countTotalNumberOfClientsForApp(this, appIdentifier, true);
+        } catch (SQLException e) {
+            throw new StorageQueryException(e);
+        }
     }
 
     @Override
     public int countTotalNumberOfClientsForApp(AppIdentifier appIdentifier) throws StorageQueryException {
-        return 0; // TODO
+        try {
+            return OAuthQueries.countTotalNumberOfClientsForApp(this, appIdentifier, false);
+        } catch (SQLException e) {
+            throw new StorageQueryException(e);
+        }
     }
 
     @Override
     public int countTotalNumberOfM2MTokensAlive(AppIdentifier appIdentifier) throws StorageQueryException {
-        return 0; // TODO
+        try {
+            return OAuthQueries.countTotalNumberOfM2MTokensAlive(this, appIdentifier);
+        } catch (SQLException e) {
+            throw new StorageQueryException(e);
+        }
     }
 
     @Override
     public int countTotalNumberOfM2MTokensCreatedSince(AppIdentifier appIdentifier, long since)
             throws StorageQueryException {
-        return 0; // TODO
+        try {
+            return OAuthQueries.countTotalNumberOfM2MTokensCreatedSince(this, appIdentifier, since);
+        } catch (SQLException e) {
+            throw new StorageQueryException(e);
+        }
     }
 }
diff --git a/src/main/java/io/supertokens/inmemorydb/config/SQLiteConfig.java b/src/main/java/io/supertokens/inmemorydb/config/SQLiteConfig.java
@@ -165,11 +165,15 @@ public String getDashboardSessionsTable() {
         return "dashboard_user_sessions";
     }
 
-    public String getOAuthClientTable() { 
+    public String getOAuthClientsTable() {
         return "oauth_clients";
     }
 
     public String getOAuthRevokeTable() {
         return "oauth_revoke";
     }
+
+    public String getOAuthM2MTokensTable() {
+        return "oauth_m2m_tokens";
+    }
 }
diff --git a/src/main/java/io/supertokens/inmemorydb/queries/GeneralQueries.java b/src/main/java/io/supertokens/inmemorydb/queries/GeneralQueries.java
@@ -423,7 +423,7 @@ public static void createTablesIfNotExists(Start start, Main main) throws SQLExc
             update(start, TOTPQueries.getQueryToCreateUsedCodesExpiryTimeIndex(start), NO_OP_SETTER);
         }
 
-        if (!doesTableExists(start, Config.getConfig(start).getOAuthClientTable())) {
+        if (!doesTableExists(start, Config.getConfig(start).getOAuthClientsTable())) {
             getInstance(main).addState(CREATING_NEW_TABLE, null);
             update(start, OAuthQueries.getQueryToCreateOAuthClientTable(start), NO_OP_SETTER);
         }

diff --git a/src/main/java/io/supertokens/inmemorydb/queries/OAuthQueries.java b/src/main/java/io/supertokens/inmemorydb/queries/OAuthQueries.java
@@ -32,11 +32,12 @@
 public class OAuthQueries {
 
     public static String getQueryToCreateOAuthClientTable(Start start) {
-        String oAuth2ClientTable = Config.getConfig(start).getOAuthClientTable();
+        String oAuth2ClientTable = Config.getConfig(start).getOAuthClientsTable();
         // @formatter:off
         return "CREATE TABLE IF NOT EXISTS " + oAuth2ClientTable + " ("
                 + "app_id VARCHAR(64),"
                 + "client_id VARCHAR(128) NOT NULL,"
+                + "is_client_credentials_only BOOLEAN NOT NULL,"
                 + " PRIMARY KEY (app_id, client_id),"
                 + " FOREIGN KEY(app_id) REFERENCES " + Config.getConfig(start).getAppsTable() + "(app_id) ON DELETE CASCADE);";
         // @formatter:on
@@ -63,9 +64,36 @@ public static String getQueryToCreateOAuthRevokeTimestampIndex(Start start) {
                 + oAuth2ClientTable + "(timestamp DESC, app_id DESC);";
     }
 
+    public static String getQueryToCreateOAuthM2MTokensTable(Start start) {
+        String oAuth2ClientTable = Config.getConfig(start).getOAuthM2MTokensTable();
+        // @formatter:off
+        return "CREATE TABLE IF NOT EXISTS " + oAuth2ClientTable + " ("
+                + "app_id VARCHAR(64) DEFAULT 'public',"
+                + "client_id VARCHAR(128) NOT NULL,"
+                + "iat BIGINT NOT NULL,"
+                + "exp BIGINT NOT NULL,"
+                + "PRIMARY KEY (app_id, client_id, iat),"
+                + "FOREIGN KEY(app_id)"
+                + " REFERENCES " + Config.getConfig(start).getAppsTable() + "(app_id) ON DELETE CASCADE"
+                + ");";
+        // @formatter:on
+    }
+
+    public static String getQueryToCreateOAuthM2MTokenIatIndex(Start start) {
+        String oAuth2ClientTable = Config.getConfig(start).getOAuthM2MTokensTable();
+        return "CREATE INDEX IF NOT EXISTS oauth_m2m_token_iat_index ON "
+                + oAuth2ClientTable + "(iat DESC, app_id DESC);";
+    }
+
+    public static String getQueryToCreateOAuthM2MTokenExpIndex(Start start) {
+        String oAuth2ClientTable = Config.getConfig(start).getOAuthM2MTokensTable();
+        return "CREATE INDEX IF NOT EXISTS oauth_m2m_token_exp_index ON "
+                + oAuth2ClientTable + "(exp DESC, app_id DESC);";
+    }
+
     public static boolean isClientIdForAppId(Start start, String clientId, AppIdentifier appIdentifier)
             throws SQLException, StorageQueryException {
-        String QUERY = "SELECT app_id FROM " + Config.getConfig(start).getOAuthClientTable() +
+        String QUERY = "SELECT app_id FROM " + Config.getConfig(start).getOAuthClientsTable() +
             " WHERE client_id = ? AND app_id = ?";
 
         return execute(start, QUERY, pst -> {
@@ -76,7 +104,7 @@ public static boolean isClientIdForAppId(Start start, String clientId, AppIdenti
 
     public static List<String> listClientsForApp(Start start, AppIdentifier appIdentifier)
             throws SQLException, StorageQueryException {
-        String QUERY = "SELECT client_id FROM " + Config.getConfig(start).getOAuthClientTable() +
+        String QUERY = "SELECT client_id FROM " + Config.getConfig(start).getOAuthClientsTable() +
                 " WHERE app_id = ?";
         return execute(start, QUERY, pst -> {
             pst.setString(1, appIdentifier.getAppId());
@@ -89,19 +117,23 @@ public static List<String> listClientsForApp(Start start, AppIdentifier appIdent
         });
     }
 
-    public static void insertClientIdForAppId(Start start, String clientId, AppIdentifier appIdentifier)
+    public static void insertClientIdForAppId(Start start, AppIdentifier appIdentifier, String clientId,
+            boolean isClientCredentialsOnly)
             throws SQLException, StorageQueryException {
-        String INSERT = "INSERT INTO " + Config.getConfig(start).getOAuthClientTable()
-                + "(app_id, client_id) VALUES(?, ?)";
+        String INSERT = "INSERT INTO " + Config.getConfig(start).getOAuthClientsTable()
+                + "(app_id, client_id, is_client_credentials_only) VALUES(?, ?, ?) "
+                + "ON CONFLICT (app_id, client_id) DO UPDATE SET is_client_credentials_only = ?";
         update(start, INSERT, pst -> {
             pst.setString(1, appIdentifier.getAppId());
             pst.setString(2, clientId);
+            pst.setBoolean(3, isClientCredentialsOnly);
+            pst.setBoolean(4, isClientCredentialsOnly);
         });
     }
 
     public static boolean deleteClientIdForAppId(Start start, String clientId, AppIdentifier appIdentifier)
             throws SQLException, StorageQueryException {
-        String DELETE = "DELETE FROM " + Config.getConfig(start).getOAuthClientTable()
+        String DELETE = "DELETE FROM " + Config.getConfig(start).getOAuthClientsTable()
                 + " WHERE app_id = ? AND client_id = ?";
         int numberOfRow = update(start, DELETE, pst -> {
             pst.setString(1, appIdentifier.getAppId());
@@ -154,4 +186,74 @@ public static boolean isRevoked(Start start, AppIdentifier appIdentifier, String
             }
         }, ResultSet::next);
     }
+
+    public static int countTotalNumberOfClientsForApp(Start start, AppIdentifier appIdentifier,
+            boolean filterByClientCredentialsOnly) throws SQLException, StorageQueryException {
+        if (filterByClientCredentialsOnly) {
+            String QUERY = "SELECT COUNT(*) as c FROM " + Config.getConfig(start).getOAuthClientsTable() +
+                    " WHERE app_id = ? AND is_client_credentials_only = ?";
+            return execute(start, QUERY, pst -> {
+                pst.setString(1, appIdentifier.getAppId());
+                pst.setBoolean(2, true);
+            }, result -> {
+                if (result.next()) {
+                    return result.getInt("c");
+                }
+                return 0;
+            });
+        } else {
+            String QUERY = "SELECT COUNT(*) as c FROM " + Config.getConfig(start).getOAuthClientsTable() +
+                    " WHERE app_id = ?";
+            return execute(start, QUERY, pst -> {
+                pst.setString(1, appIdentifier.getAppId());
+            }, result -> {
+                if (result.next()) {
+                    return result.getInt("c");
+                }
+                return 0;
+            });
+        }
+    }
+
+    public static int countTotalNumberOfM2MTokensAlive(Start start, AppIdentifier appIdentifier)
+            throws SQLException, StorageQueryException {
+        String QUERY = "SELECT COUNT(*) as c FROM " + Config.getConfig(start).getOAuthM2MTokensTable() +
+                " WHERE app_id = ? AND exp > ?";
+        return execute(start, QUERY, pst -> {
+            pst.setString(1, appIdentifier.getAppId());
+            pst.setLong(2, System.currentTimeMillis());
+        }, result -> {
+            if (result.next()) {
+                return result.getInt("c");
+            }
+            return 0;
+        });
+    }
+
+    public static int countTotalNumberOfM2MTokensCreatedSince(Start start, AppIdentifier appIdentifier, long since)
+            throws SQLException, StorageQueryException {
+        String QUERY = "SELECT COUNT(*) as c FROM " + Config.getConfig(start).getOAuthM2MTokensTable() +
+                " WHERE app_id = ? AND iat >= ?";
+        return execute(start, QUERY, pst -> {
+            pst.setString(1, appIdentifier.getAppId());
+            pst.setLong(2, since / 1000);
+        }, result -> {
+            if (result.next()) {
+                return result.getInt("c");
+            }
+            return 0;
+        });
+    }
+
+    public static void addM2MToken(Start start, AppIdentifier appIdentifier, String clientId, long iat, long exp)
+            throws SQLException, StorageQueryException {
+        String QUERY = "INSERT INTO " + Config.getConfig(start).getOAuthM2MTokensTable() +
+                " (app_id, client_id, iat, exp) VALUES (?, ?, ?, ?)";
+        update(start, QUERY, pst -> {
+            pst.setString(1, appIdentifier.getAppId());
+            pst.setString(2, clientId);
+            pst.setLong(3, iat);
+            pst.setLong(4, exp);
+        });
+    }
 }