feat: Make TOTP a paid feature and report stats

supertokens · Mar 22, 2023 · 3d50c06 · 3d50c06
1 parent b02a420
commit 3d50c06
Show file tree

Hide file tree

Showing 5 changed files with 118 additions and 10 deletions.
diff --git a/ee/src/main/java/io/supertokens/ee/EEFeatureFlag.java b/ee/src/main/java/io/supertokens/ee/EEFeatureFlag.java
@@ -6,10 +6,8 @@
 import com.auth0.jwt.exceptions.JWTVerificationException;
 import com.auth0.jwt.interfaces.DecodedJWT;
 import com.auth0.jwt.interfaces.RSAKeyProvider;
-import com.google.gson.JsonArray;
-import com.google.gson.JsonObject;
-import com.google.gson.JsonParser;
-import com.google.gson.JsonPrimitive;
+import com.google.gson.*;
+import io.supertokens.ActiveUsers;
 import io.supertokens.Main;
 import io.supertokens.ProcessState;
 import io.supertokens.cronjobs.Cronjobs;
@@ -146,11 +144,36 @@ public Boolean getIsLicenseKeyPresent() {
     public JsonObject getPaidFeatureStats() throws StorageQueryException {
         JsonObject result = new JsonObject();
         EE_FEATURES[] features = getEnabledEEFeaturesFromDbOrCache();
-        if (Arrays.stream(features).anyMatch(t -> t == EE_FEATURES.DASHBOARD_LOGIN)) {
-            JsonObject stats = new JsonObject();
-            int userCount = StorageLayer.getDashboardStorage(main).getAllDashboardUsers().length;
-            stats.addProperty("user_count", userCount);
-            result.add(EE_FEATURES.DASHBOARD_LOGIN.toString(), stats);
+        for (EE_FEATURES feature : features) {
+            if (feature == EE_FEATURES.DASHBOARD_LOGIN) {
+                JsonObject stats = new JsonObject();
+                int userCount = StorageLayer.getDashboardStorage(main).getAllDashboardUsers().length;
+                stats.addProperty("user_count", userCount);
+                result.add(EE_FEATURES.DASHBOARD_LOGIN.toString(), stats);
+            }
+
+            if (feature == EE_FEATURES.TOTP) {
+                JsonObject stats = new JsonObject();
+                JsonArray mauArr = new JsonArray();
+                JsonArray totpMauArr = new JsonArray();
+                for (int i = 0; i < 30; i++) {
+                    long timestamp = System.currentTimeMillis() - i * 24 * 60 * 60 * 1000;
+
+                    int mau = ActiveUsers.countUsersActiveSince(main, timestamp);
+                    mauArr.add(new JsonPrimitive(mau));
+
+                    int totpMau = StorageLayer.getActiveUsersStorage(main).countUsersEnabledTotpAndActiveSince(timestamp);
+                    totpMauArr.add(new JsonPrimitive(totpMau));
+                }
+
+                stats.add("maus", mauArr);
+                stats.add("totp_maus", totpMauArr);
+
+                int totpTotalEnabled = StorageLayer.getActiveUsersStorage(main).countUsersEnabledTotp();
+                stats.addProperty("total_totp_users", totpTotalEnabled);
+
+                result.add(EE_FEATURES.TOTP.toString(), stats);
+            }
         }
         return result;
     }

diff --git a/src/main/java/io/supertokens/featureflag/EE_FEATURES.java b/src/main/java/io/supertokens/featureflag/EE_FEATURES.java
@@ -17,7 +17,8 @@
 package io.supertokens.featureflag;
 
 public enum EE_FEATURES {
-    ACCOUNT_LINKING("account_linking"), MULTI_TENANCY("multi_tenancy"), TEST("test"), DASHBOARD_LOGIN("dashboard_login");
+    ACCOUNT_LINKING("account_linking"), MULTI_TENANCY("multi_tenancy"), TEST("test"), DASHBOARD_LOGIN("dashboard_login"),
+    TOTP("totp");
 
     private final String name;
 

diff --git a/src/main/java/io/supertokens/inmemorydb/Start.java b/src/main/java/io/supertokens/inmemorydb/Start.java
@@ -460,6 +460,24 @@ public int countUsersActiveSince(long time) throws StorageQueryException {
         }
     }
 
+    @Override
+    public int countUsersEnabledTotp() throws StorageQueryException {
+        try {
+            return ActiveUsersQueries.countUsersEnabledTotp(this);
+        } catch (SQLException e) {
+            throw new StorageQueryException(e);
+        }
+    }
+
+    @Override
+    public int countUsersEnabledTotpAndActiveSince(long time) throws StorageQueryException {
+        try {
+            return ActiveUsersQueries.countUsersEnabledTotpAndActiveSince(this, time);
+        } catch (SQLException e) {
+            throw new StorageQueryException(e);
+        }
+    }
+
     @Override
     public SessionInfo getSessionInfo_Transaction(TransactionConnection con, String sessionHandle)
             throws StorageQueryException {

diff --git a/src/main/java/io/supertokens/inmemorydb/queries/ActiveUsersQueries.java b/src/main/java/io/supertokens/inmemorydb/queries/ActiveUsersQueries.java
@@ -28,6 +28,31 @@ public static int countUsersActiveSince(Start start, long sinceTime) throws SQLE
         });
     }
 
+    public static int countUsersEnabledTotp(Start start) throws SQLException, StorageQueryException {
+        String QUERY = "SELECT COUNT(*) as total FROM " + Config.getConfig(start).getTotpUsersTable();
+
+        return execute(start, QUERY, null, result -> {
+            if (result.next()) {
+                return result.getInt("total");
+            }
+            return 0;
+        });
+    }
+
+    public static int countUsersEnabledTotpAndActiveSince(Start start, long sinceTime) throws SQLException, StorageQueryException {
+        String QUERY = "SELECT COUNT(*) as total FROM " + Config.getConfig(start).getTotpUsersTable() + " AS totp_users "
+                + "INNER JOIN " + Config.getConfig(start).getUserLastActiveTable() + " AS user_last_active "
+                + "ON totp_users.user_id = user_last_active.user_id "
+                + "WHERE user_last_active.last_active_time >= ?";
+
+        return execute(start, QUERY, pst -> pst.setLong(1, sinceTime), result -> {
+            if (result.next()) {
+                return result.getInt("total");
+            }
+            return 0;
+        });
+    }
+
     public static int updateUserLastActive(Start start, String userId) throws SQLException, StorageQueryException {
         String QUERY = "INSERT INTO " + Config.getConfig(start).getUserLastActiveTable()
                 + "(user_id, last_active_time) VALUES(?, ?) ON CONFLICT(user_id) DO UPDATE SET last_active_time = ?";

diff --git a/src/test/java/io/supertokens/test/FeatureFlagTest.java b/src/test/java/io/supertokens/test/FeatureFlagTest.java
@@ -16,6 +16,7 @@
 
 package io.supertokens.test;
 
+import com.google.gson.JsonArray;
 import com.google.gson.JsonObject;
 import io.supertokens.ProcessState;
 import io.supertokens.featureflag.FeatureFlag;
@@ -111,4 +112,44 @@ public void testThatCallingGetFeatureFlagAPIReturnsEmptyArray() throws Exception
         process.kill();
         Assert.assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STOPPED));
     }
+
+    private final String OPAQUE_KEY_WITH_TOTP_FEATURE = "pXhNK=nYiEsb6gJEOYP2kIR6M0kn4XLvNqcwT1XbX8xHtm44K-lQfGCbaeN0Ieeza39fxkXr=tiiUU=DXxDH40Y=4FLT4CE-rG1ETjkXxO4yucLpJvw3uSegPayoISGL";
+
+    @Test
+    public void testThatCallingGetFeatureFlagAPIReturnsTotpStats() throws Exception {
+        String[] args = {"../"};
+
+        TestingProcessManager.TestingProcess process = TestingProcessManager.start(args);
+        Assert.assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STARTED));
+
+        FeatureFlag.getInstance(process.main).setLicenseKeyAndSyncFeatures(OPAQUE_KEY_WITH_TOTP_FEATURE);
+
+        JsonObject response = HttpRequestForTesting.sendGETRequest(process.getProcess(), "",
+                "http://localhost:3567/ee/featureflag",
+                null, 1000, 1000, null, WebserverAPI.getLatestCDIVersion(), "");
+        Assert.assertEquals("OK", response.get("status").getAsString());
+
+        JsonArray features = response.get("features").getAsJsonArray();
+        JsonObject totpStats = response.get("usageStats").getAsJsonObject().get("totp").getAsJsonObject();
+
+        assert features.size() == 1;
+        assert features.get(0).getAsString().equals("totp");
+
+        JsonArray mau = totpStats.get("maus").getAsJsonArray();
+        JsonArray totpMau = totpStats.get("totp_maus").getAsJsonArray();
+        int totalTotpUsers = totpStats.get("total_totp_users").getAsInt();
+
+        assert mau.size() == 30;
+        assert mau.get(0).getAsInt() == 0;
+        assert mau.get(29).getAsInt() == 0;
+
+        assert totpMau.size() == 30;
+        assert totpMau.get(0).getAsInt() == 0;
+        assert totpMau.get(29).getAsInt() == 0;
+
+        assert totalTotpUsers == 0;
+
+        process.kill();
+        Assert.assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STOPPED));
+    }
 }