From cd77270a669fa827912efd86a0723d85e9e388ed Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Thu, 23 Aug 2018 19:17:08 +0100
Subject: [PATCH] Implement trail users

---
 synapse/api/auth.py                     |  6 +++++-
 synapse/config/server.py                |  6 ++++++
 synapse/storage/monthly_active_users.py |  5 +++++
 synapse/storage/registration.py         | 28 ++++++++++++++++++++++++-
 tests/storage/test_registration.py      |  1 +
 5 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/synapse/api/auth.py b/synapse/api/auth.py
index 4ca40a0f71..8d2aa5870a 100644
--- a/synapse/api/auth.py
+++ b/synapse/api/auth.py
@@ -797,11 +797,15 @@ class Auth(object):
                 limit_type=self.hs.config.hs_disabled_limit_type
             )
         if self.hs.config.limit_usage_by_mau is True:
-            # If the user is already part of the MAU cohort
+            # If the user is already part of the MAU cohort or a trial user
             if user_id:
                 timestamp = yield self.store.user_last_seen_monthly_active(user_id)
                 if timestamp:
                     return
+
+                is_trial = yield self.store.is_trial_user(user_id)
+                if is_trial:
+                    return
             # Else if there is no room in the MAU bucket, bail
             current_mau = yield self.store.get_monthly_active_count()
             if current_mau >= self.hs.config.max_mau_value:
diff --git a/synapse/config/server.py b/synapse/config/server.py
index 68a612e594..8eecd28e7d 100644
--- a/synapse/config/server.py
+++ b/synapse/config/server.py
@@ -77,10 +77,15 @@ class ServerConfig(Config):
             self.max_mau_value = config.get(
                 "max_mau_value", 0,
             )
+
         self.mau_limits_reserved_threepids = config.get(
             "mau_limit_reserved_threepids", []
         )
 
+        self.mau_trial_days = config.get(
+            "mau_trial_days", 0,
+        )
+
         # Options to disable HS
         self.hs_disabled = config.get("hs_disabled", False)
         self.hs_disabled_message = config.get("hs_disabled_message", "")
@@ -365,6 +370,7 @@ class ServerConfig(Config):
           # Enables monthly active user checking
           # limit_usage_by_mau: False
           # max_mau_value: 50
+          # mau_trial_days: 2
           #
           # Sometimes the server admin will want to ensure certain accounts are
           # never blocked by mau checking. These accounts are specified here.
diff --git a/synapse/storage/monthly_active_users.py b/synapse/storage/monthly_active_users.py
index fd3b630bd2..d178f5c5ba 100644
--- a/synapse/storage/monthly_active_users.py
+++ b/synapse/storage/monthly_active_users.py
@@ -201,6 +201,11 @@ class MonthlyActiveUsersStore(SQLBaseStore):
             user_id(str): the user_id to query
         """
         if self.hs.config.limit_usage_by_mau:
+            is_trial = yield self.is_trial_user(user_id)
+            if is_trial:
+                # we don't track trial users in the MAU table.
+                return
+
             last_seen_timestamp = yield self.user_last_seen_monthly_active(user_id)
             now = self.hs.get_clock().time_msec()
 
diff --git a/synapse/storage/registration.py b/synapse/storage/registration.py
index 07333f777d..26b429e307 100644
--- a/synapse/storage/registration.py
+++ b/synapse/storage/registration.py
@@ -26,6 +26,11 @@ from synapse.util.caches.descriptors import cached, cachedInlineCallbacks
 
 
 class RegistrationWorkerStore(SQLBaseStore):
+    def __init__(self, db_conn, hs):
+        super(RegistrationWorkerStore, self).__init__(db_conn, hs)
+
+        self.config = hs.config
+
     @cached()
     def get_user_by_id(self, user_id):
         return self._simple_select_one(
@@ -36,12 +41,33 @@ class RegistrationWorkerStore(SQLBaseStore):
             retcols=[
                 "name", "password_hash", "is_guest",
                 "consent_version", "consent_server_notice_sent",
-                "appservice_id",
+                "appservice_id", "creation_ts",
             ],
             allow_none=True,
             desc="get_user_by_id",
         )
 
+    @defer.inlineCallbacks
+    def is_trial_user(self, user_id):
+        """Checks if user is in the "trial" period, i.e. within the first
+        N days of registration defined by `mau_trial_days` config
+
+        Args:
+            user_id (str)
+
+        Returns:
+            Deferred[bool]
+        """
+
+        info = yield self.get_user_by_id(user_id)
+        if not info:
+            defer.returnValue(False)
+
+        now = self.clock.time_msec()
+        trial_duration_ms = self.config.mau_trial_days * 24 * 60 * 60 * 1000
+        is_trial = (now - info["creation_ts"] * 1000) < trial_duration_ms
+        defer.returnValue(is_trial)
+
     @cached()
     def get_user_by_access_token(self, token):
         """Get a user from the given access token.
diff --git a/tests/storage/test_registration.py b/tests/storage/test_registration.py
index 4eda122edc..3dfb7b903a 100644
--- a/tests/storage/test_registration.py
+++ b/tests/storage/test_registration.py
@@ -46,6 +46,7 @@ class RegistrationStoreTestCase(unittest.TestCase):
                 "consent_version": None,
                 "consent_server_notice_sent": None,
                 "appservice_id": None,
+                "creation_ts": 1000,
             },
             (yield self.store.get_user_by_id(self.user_id)),
         )