diff options
| -rw-r--r-- | synapse/handlers/register.py | 9 | ||||
| -rw-r--r-- | synapse/storage/__init__.py | 26 | ||||
| -rw-r--r-- | tests/handlers/test_auth.py | 39 | ||||
| -rw-r--r-- | tests/handlers/test_register.py | 10 |
4 files changed, 46 insertions, 38 deletions
diff --git a/synapse/handlers/register.py b/synapse/handlers/register.py index f46b8355c0..cc935a5e84 100644 --- a/synapse/handlers/register.py +++ b/synapse/handlers/register.py @@ -144,7 +144,7 @@ class RegistrationHandler(BaseHandler): Raises: RegistrationError if there was a problem registering. """ - self._check_mau_limits() + yield self._check_mau_limits() password_hash = None if password: password_hash = yield self.auth_handler().hash(password) @@ -289,7 +289,7 @@ class RegistrationHandler(BaseHandler): 400, "User ID can only contain characters a-z, 0-9, or '=_-./'", ) - self._check_mau_limits() + yield self._check_mau_limits() user = UserID(localpart, self.hs.hostname) user_id = user.to_string() @@ -439,7 +439,7 @@ class RegistrationHandler(BaseHandler): """ if localpart is None: raise SynapseError(400, "Request must include user id") - self._check_mau_limits() + yield self._check_mau_limits() need_register = True try: @@ -534,13 +534,14 @@ class RegistrationHandler(BaseHandler): action="join", ) + @defer.inlineCallbacks def _check_mau_limits(self): """ Do not accept registrations if monthly active user limits exceeded and limiting is enabled """ if self.hs.config.limit_usage_by_mau is True: - current_mau = self.store.count_monthly_users() + current_mau = yield self.store.count_monthly_users() if current_mau >= self.hs.config.max_mau_value: raise RegistrationError( 403, "MAU Limit Exceeded", Codes.MAU_LIMIT_EXCEEDED diff --git a/synapse/storage/__init__.py b/synapse/storage/__init__.py index 4747118ed7..f9682832ca 100644 --- a/synapse/storage/__init__.py +++ b/synapse/storage/__init__.py @@ -273,24 +273,24 @@ class DataStore(RoomMemberStore, RoomStore, This method should be refactored with count_daily_users - the only reason not to is waiting on definition of mau returns: - int: count of current monthly active users + defered: resolves to int """ + def _count_monthly_users(txn): + thirty_days_ago = int(self._clock.time_msec()) - (1000 * 60 * 60 * 24 * 30) + sql = """ + SELECT COALESCE(count(*), 0) FROM ( + SELECT user_id FROM user_ips + WHERE last_seen > ? + GROUP BY user_id + ) u + """ - thirty_days_ago = int(self._clock.time_msec()) - (1000 * 60 * 60 * 24 * 30) - sql = """ - SELECT COALESCE(count(*), 0) FROM ( - SELECT user_id FROM user_ips - WHERE last_seen > ? - GROUP BY user_id - ) u - """ - try: - txn = self.db_conn.cursor() txn.execute(sql, (thirty_days_ago,)) count, = txn.fetchone() + print "Count is %d" % (count,) return count - finally: - txn.close() + + return self.runInteraction("count_monthly_users", _count_monthly_users) def count_r30_users(self): """ diff --git a/tests/handlers/test_auth.py b/tests/handlers/test_auth.py index e01f14a10a..440a453082 100644 --- a/tests/handlers/test_auth.py +++ b/tests/handlers/test_auth.py @@ -77,38 +77,37 @@ class AuthTestCase(unittest.TestCase): v.satisfy_general(verify_nonce) v.verify(macaroon, self.hs.config.macaroon_secret_key) + @defer.inlineCallbacks def test_short_term_login_token_gives_user_id(self): self.hs.clock.now = 1000 token = self.macaroon_generator.generate_short_term_login_token( "a_user", 5000 ) - - self.assertEqual( - "a_user", - self.auth_handler.validate_short_term_login_token_and_get_user_id( - token - ) + user_id = yield self.auth_handler.validate_short_term_login_token_and_get_user_id( + token ) + self.assertEqual("a_user", user_id) # when we advance the clock, the token should be rejected self.hs.clock.now = 6000 with self.assertRaises(synapse.api.errors.AuthError): - self.auth_handler.validate_short_term_login_token_and_get_user_id( + yield self.auth_handler.validate_short_term_login_token_and_get_user_id( token ) + @defer.inlineCallbacks def test_short_term_login_token_cannot_replace_user_id(self): token = self.macaroon_generator.generate_short_term_login_token( "a_user", 5000 ) macaroon = pymacaroons.Macaroon.deserialize(token) + user_id = yield self.auth_handler.validate_short_term_login_token_and_get_user_id( + macaroon.serialize() + ) self.assertEqual( - "a_user", - self.auth_handler.validate_short_term_login_token_and_get_user_id( - macaroon.serialize() - ) + "a_user", user_id ) # add another "user_id" caveat, which might allow us to override the @@ -116,7 +115,7 @@ class AuthTestCase(unittest.TestCase): macaroon.add_first_party_caveat("user_id = b_user") with self.assertRaises(synapse.api.errors.AuthError): - self.auth_handler.validate_short_term_login_token_and_get_user_id( + yield self.auth_handler.validate_short_term_login_token_and_get_user_id( macaroon.serialize() ) @@ -126,7 +125,7 @@ class AuthTestCase(unittest.TestCase): # Ensure does not throw exception yield self.auth_handler.get_access_token_for_user_id('user_a') - self.auth_handler.validate_short_term_login_token_and_get_user_id( + yield self.auth_handler.validate_short_term_login_token_and_get_user_id( self._get_macaroon().serialize() ) @@ -134,24 +133,30 @@ class AuthTestCase(unittest.TestCase): def test_mau_limits_exceeded(self): self.hs.config.limit_usage_by_mau = True self.hs.get_datastore().count_monthly_users = Mock( - return_value=self.large_number_of_users + return_value=defer.succeed(self.large_number_of_users) ) + with self.assertRaises(AuthError): yield self.auth_handler.get_access_token_for_user_id('user_a') + + self.hs.get_datastore().count_monthly_users = Mock( + return_value=defer.succeed(self.large_number_of_users) + ) with self.assertRaises(AuthError): - self.auth_handler.validate_short_term_login_token_and_get_user_id( + yield self.auth_handler.validate_short_term_login_token_and_get_user_id( self._get_macaroon().serialize() ) @defer.inlineCallbacks def test_mau_limits_not_exceeded(self): self.hs.config.limit_usage_by_mau = True + self.hs.get_datastore().count_monthly_users = Mock( - return_value=self.small_number_of_users + return_value=defer.succeed(self.small_number_of_users) ) # Ensure does not raise exception yield self.auth_handler.get_access_token_for_user_id('user_a') - self.auth_handler.validate_short_term_login_token_and_get_user_id( + yield self.auth_handler.validate_short_term_login_token_and_get_user_id( self._get_macaroon().serialize() ) diff --git a/tests/handlers/test_register.py b/tests/handlers/test_register.py index a5a8e7c954..0937d71cf6 100644 --- a/tests/handlers/test_register.py +++ b/tests/handlers/test_register.py @@ -90,7 +90,7 @@ class RegistrationTestCase(unittest.TestCase): lots_of_users = 100 small_number_users = 1 - store.count_monthly_users = Mock(return_value=lots_of_users) + store.count_monthly_users = Mock(return_value=defer.succeed(lots_of_users)) # Ensure does not throw exception yield self.handler.get_or_create_user(requester, 'a', display_name) @@ -100,7 +100,7 @@ class RegistrationTestCase(unittest.TestCase): with self.assertRaises(RegistrationError): yield self.handler.get_or_create_user(requester, 'b', display_name) - store.count_monthly_users = Mock(return_value=small_number_users) + store.count_monthly_users = Mock(return_value=defer.succeed(small_number_users)) self._macaroon_mock_generator("another_secret") @@ -108,12 +108,14 @@ class RegistrationTestCase(unittest.TestCase): yield self.handler.get_or_create_user("@neil:matrix.org", 'c', "Neil") self._macaroon_mock_generator("another another secret") - store.count_monthly_users = Mock(return_value=lots_of_users) + store.count_monthly_users = Mock(return_value=defer.succeed(lots_of_users)) + with self.assertRaises(RegistrationError): yield self.handler.register(localpart=local_part) self._macaroon_mock_generator("another another secret") - store.count_monthly_users = Mock(return_value=lots_of_users) + store.count_monthly_users = Mock(return_value=defer.succeed(lots_of_users)) + with self.assertRaises(RegistrationError): yield self.handler.register_saml2(local_part) |