feat: Make TOTP a paid feature and report stats

supertokens · Mar 22, 2023 · bdf72d0 · bdf72d0
1 parent cf1dd29
commit bdf72d0
Show file tree

Hide file tree

Showing 4 changed files with 97 additions and 24 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/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
@@ -18,22 +18,16 @@
 
 import com.google.gson.JsonArray;
 import com.google.gson.JsonObject;
-import com.google.gson.JsonPrimitive;
 import io.supertokens.ProcessState;
-import io.supertokens.featureflag.EE_FEATURES;
 import io.supertokens.featureflag.FeatureFlag;
 import io.supertokens.featureflag.FeatureFlagTestContent;
 import io.supertokens.featureflag.exceptions.NoLicenseKeyFoundException;
-import io.supertokens.pluginInterface.KeyValueInfo;
 import io.supertokens.pluginInterface.exceptions.StorageQueryException;
-import io.supertokens.storageLayer.StorageLayer;
 import io.supertokens.test.httpRequest.HttpRequestForTesting;
 import io.supertokens.webserver.WebserverAPI;
 import org.junit.*;
 import org.junit.rules.TestRule;
 
-import java.util.Arrays;
-
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.fail;
 
@@ -97,7 +91,6 @@ public void missingEEFolderShouldBeSameAsNoLicenseKey() throws InterruptedExcept
 
         Assert.assertEquals(FeatureFlag.getInstance(process.getProcess()).getPaidFeatureStats().entrySet().size(), 0);
 
-        // update tests
         process.kill();
         assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STOPPED));
     }
@@ -120,11 +113,7 @@ public void testThatCallingGetFeatureFlagAPIReturnsEmptyArray() throws Exception
         Assert.assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STOPPED));
     }
 
-    private void setFeatureFlag(TestingProcessManager.TestingProcess process, EE_FEATURES[] features) throws StorageQueryException {
-        JsonArray json = new JsonArray();
-        Arrays.stream(features).forEach(ee_features -> json.add(new JsonPrimitive(ee_features.toString())));
-        StorageLayer.getStorage(process.getProcess()).setKeyValue("FEATURE_FLAG", new KeyValueInfo(json.toString()));
-    }
+    private final String OPAQUE_KEY_WITH_TOTP_FEATURE = "pXhNK=nYiEsb6gJEOYP2kIR6M0kn4XLvNqcwT1XbX8xHtm44K-lQfGCbaeN0Ieeza39fxkXr=tiiUU=DXxDH40Y=4FLT4CE-rG1ETjkXxO4yucLpJvw3uSegPayoISGL";
 
     @Test
     public void testThatCallingGetFeatureFlagAPIReturnsTotpStats() throws Exception {
@@ -133,14 +122,32 @@ public void testThatCallingGetFeatureFlagAPIReturnsTotpStats() throws Exception
         TestingProcessManager.TestingProcess process = TestingProcessManager.start(args);
         Assert.assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STARTED));
 
-        setFeatureFlag(process, new EE_FEATURES[]{EE_FEATURES.TOTP});
+        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());
-        Assert.assertNotNull(response.get("features"));
-        Assert.assertEquals(0, response.get("features").getAsJsonArray().size());
+
+        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));