-
-
Notifications
You must be signed in to change notification settings - Fork 2.1k
Save login tokens in database #13844
Changes from 2 commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1 @@ | ||
| Save login tokens in database and prevent login token reuse. |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -115,6 +115,20 @@ class RefreshTokenLookupResult: | |
| If None, the session can be refreshed indefinitely.""" | ||
|
|
||
|
|
||
| @attr.s(auto_attribs=True, frozen=True, slots=True) | ||
| class LoginTokenLookupResult: | ||
| """Result of looking up a login token.""" | ||
|
|
||
| user_id: str | ||
| """The user this token belongs to.""" | ||
|
|
||
| auth_provider_id: Optional[str] | ||
| """The SSO Identity Provider that the user authenticated with, to get this token.""" | ||
|
|
||
| auth_provider_session_id: Optional[str] | ||
| """The session ID advertised by the SSO Identity Provider.""" | ||
|
|
||
|
|
||
| class RegistrationWorkerStore(CacheInvalidationWorkerStore): | ||
| def __init__( | ||
| self, | ||
|
|
@@ -1789,6 +1803,114 @@ def _replace_refresh_token_txn(txn: LoggingTransaction) -> None: | |
| "replace_refresh_token", _replace_refresh_token_txn | ||
| ) | ||
|
|
||
| async def add_login_token_to_user( | ||
| self, | ||
| user_id: str, | ||
| token: str, | ||
| expiry_ts: int, | ||
| auth_provider_id: Optional[str], | ||
| auth_provider_session_id: Optional[str], | ||
| ) -> None: | ||
| """Adds a short-term login token for the given user. | ||
|
|
||
| Args: | ||
| user_id: The user ID. | ||
| token: The new login token to add. | ||
| expiry_ts (milliseconds since the epoch): Time after which the login token | ||
| cannot be used. | ||
| auth_provider_id: The SSO Identity Provider that the user authenticated with | ||
| to get this token, if any | ||
| auth_provider_session_id: The session ID advertised by the SSO Identity | ||
| Provider, if any. | ||
| """ | ||
| await self.db_pool.simple_insert( | ||
| "login_tokens", | ||
| { | ||
| "token": token, | ||
| "user_id": user_id, | ||
| "expiry_ts": expiry_ts, | ||
| "auth_provider_id": auth_provider_id, | ||
| "auth_provider_session_id": auth_provider_session_id, | ||
| }, | ||
| desc="add_login_token_to_user", | ||
| ) | ||
|
|
||
| def _consume_login_token( | ||
| self, | ||
| txn: LoggingTransaction, | ||
| token: str, | ||
| ts: int, | ||
| ) -> Optional[LoginTokenLookupResult]: | ||
| if self.database_engine.supports_returning: | ||
| # If the database engine supports the `RETURNING` keyword, delete and fetch | ||
| # the token with one query | ||
| txn.execute( | ||
| """ | ||
| DELETE FROM login_tokens | ||
|
||
| WHERE token = ? | ||
| RETURNING user_id, expiry_ts, auth_provider_id, auth_provider_session_id | ||
| """, | ||
| (token,), | ||
| ) | ||
| ret = txn.fetchone() | ||
| if ret is None: | ||
| return None | ||
|
|
||
| user_id, expiry_ts, auth_provider_id, auth_provider_session_id = ret | ||
| else: | ||
| values = self.db_pool.simple_select_one_txn( | ||
| txn, | ||
| "login_tokens", | ||
| keyvalues={"token": token}, | ||
| retcols=( | ||
| "user_id", | ||
| "expiry_ts", | ||
| "auth_provider_id", | ||
| "auth_provider_session_id", | ||
| ), | ||
| allow_none=True, | ||
| ) | ||
|
|
||
| if values is None: | ||
| return None | ||
|
|
||
| self.db_pool.simple_delete_one_txn( | ||
| txn, | ||
| "login_tokens", | ||
| keyvalues={"token": token}, | ||
| ) | ||
| user_id = values["user_id"] | ||
| expiry_ts = values["expiry_ts"] | ||
| auth_provider_id = values["auth_provider_id"] | ||
| auth_provider_session_id = values["auth_provider_session_id"] | ||
|
|
||
| # Token expired | ||
| if ts > int(expiry_ts): | ||
| return None | ||
|
|
||
| return LoginTokenLookupResult( | ||
| user_id=user_id, | ||
| auth_provider_id=auth_provider_id, | ||
| auth_provider_session_id=auth_provider_session_id, | ||
| ) | ||
|
|
||
| async def consume_login_token(self, token: str) -> Optional[LoginTokenLookupResult]: | ||
| """Lookup a login token and consume it. | ||
|
|
||
| Args: | ||
| token: The login token. | ||
|
|
||
| Returns: | ||
| The data stored with that token, including the `user_id`. Returns `None` if | ||
| the token does not exist or if it expired. | ||
| """ | ||
| return await self.db_pool.runInteraction( | ||
| "consume_login_token", | ||
| self._consume_login_token, | ||
| token, | ||
| self._clock.time_msec(), | ||
| ) | ||
|
|
||
| @cached() | ||
| async def is_guest(self, user_id: str) -> bool: | ||
| res = await self.db_pool.simple_select_one_onecol( | ||
|
|
@@ -2019,6 +2141,12 @@ def __init__( | |
| and hs.config.experimental.msc3866.require_approval_for_new_accounts | ||
| ) | ||
|
|
||
| # Create a background job for removing expired login tokens | ||
| if hs.config.worker.run_background_tasks: | ||
| self._clock.looping_call( | ||
| self._delete_expired_login_tokens, THIRTY_MINUTES_IN_MS | ||
| ) | ||
|
|
||
| async def add_access_token_to_user( | ||
| self, | ||
| user_id: str, | ||
|
|
@@ -2617,6 +2745,20 @@ async def update_user_approval_status( | |
| approved, | ||
| ) | ||
|
|
||
| @wrap_as_background_process("delete_expired_login_tokens") | ||
| async def _delete_expired_login_tokens(self) -> None: | ||
| """Remove login tokens with expiry dates that have passed.""" | ||
|
|
||
| def _delete_expired_login_tokens_txn(txn: LoggingTransaction, ts: int) -> None: | ||
| sql = "DELETE FROM login_tokens WHERE expiry_ts <= ?" | ||
| txn.execute(sql, (ts,)) | ||
|
|
||
| await self.db_pool.runInteraction( | ||
| "delete_expired_login_tokens", | ||
| _delete_expired_login_tokens_txn, | ||
| self._clock.time_msec(), | ||
| ) | ||
|
|
||
|
|
||
| def find_max_generated_user_id_localpart(cur: Cursor) -> int: | ||
| """ | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,34 @@ | ||
| /* | ||
| * Copyright 2022 The Matrix.org Foundation C.I.C. | ||
| * | ||
| * Licensed under the Apache License, Version 2.0 (the "License"); | ||
| * you may not use this file except in compliance with the License. | ||
| * You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
|
|
||
| -- Login tokens are short-lived tokens that are used for the m.login.token | ||
| -- login method, mainly during SSO logins | ||
| CREATE TABLE login_tokens ( | ||
| token TEXT PRIMARY KEY, | ||
| user_id TEXT NOT NULL, | ||
| expiry_ts BIGINT NOT NULL, | ||
| auth_provider_id TEXT, | ||
| auth_provider_session_id TEXT | ||
| ); | ||
|
|
||
| -- We're sometimes querying them by their session ID we got from their IDP | ||
| CREATE INDEX login_tokens_auth_provider_idx | ||
| ON login_tokens (auth_provider_id, auth_provider_session_id); | ||
|
|
||
| -- We're deleting them by their expiration time | ||
| CREATE INDEX login_tokens_expiry_time_idx | ||
| ON login_tokens (expiry_ts); | ||
|
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is not needed, because the
check_auth_blockingis called by the caller. So before we effectively checked that twice.