Rework how email verification works

crates/data-model/src/lib.rs

@@ -47,6 +47,7 @@ pub use self::{
     user_agent::{DeviceType, UserAgent},
     users::{
         Authentication, AuthenticationMethod, BrowserSession, Password, User, UserEmail,
-        UserEmailVerification, UserEmailVerificationState, UserRecoverySession, UserRecoveryTicket,
+        UserEmailAuthentication, UserEmailAuthenticationCode, UserEmailVerification,
+        UserEmailVerificationState, UserRecoverySession, UserRecoveryTicket,
     },
 };

crates/data-model/src/users.rs

@@ -107,6 +107,26 @@ impl UserRecoveryTicket {
     }
 }
 
+/// A user email authentication session
+#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
+pub struct UserEmailAuthentication {
+    pub id: Ulid,
+    pub user_session_id: Option<Ulid>,
+    pub email: String,
+    pub created_at: DateTime<Utc>,
+    pub completed_at: Option<DateTime<Utc>>,
+}
+
+/// A user email authentication code
+#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
+pub struct UserEmailAuthenticationCode {
+    pub id: Ulid,
+    pub user_email_authentication_id: Ulid,
+    pub code: String,
+    pub created_at: DateTime<Utc>,
+    pub expires_at: DateTime<Utc>,
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Serialize)]
 pub struct BrowserSession {
     pub id: Ulid,

crates/storage-pg/migrations/20250109105709_user_email_authentication_codes.sql

@@ -0,0 +1,29 @@
+-- Copyright 2025 New Vector Ltd.
+--
+-- SPDX-License-Identifier: AGPL-3.0-only
+-- Please see LICENSE in the repository root for full details.
+
+-- Add a table for storing email authentication sessions
+CREATE TABLE "user_email_authentications" (
+  "user_email_authentication_id" UUID PRIMARY KEY,
+  "user_session_id" UUID
+    REFERENCES "user_sessions" ("user_session_id")
+    ON DELETE SET NULL,
+  "email" TEXT NOT NULL,
+  "created_at" TIMESTAMP WITH TIME ZONE NOT NULL,
+  "completed_at" TIMESTAMP WITH TIME ZONE
+);
+
+-- A single authentication session has multiple codes, in case the user ask for re-sending
+CREATE TABLE "user_email_authentication_codes" (
+  "user_email_authentication_code_id" UUID PRIMARY KEY,
+  "user_email_authentication_id" UUID
+    NOT NULL
+    REFERENCES "user_email_authentications" ("user_email_authentication_id")
+    ON DELETE CASCADE,
+  "code" TEXT NOT NULL,
+  "created_at" TIMESTAMP WITH TIME ZONE NOT NULL,
+  "expires_at" TIMESTAMP WITH TIME ZONE NOT NULL,
+  CONSTRAINT "user_email_authentication_codes_auth_id_code_unique"
+    UNIQUE ("user_email_authentication_id", "code")
+);

crates/storage-pg/src/user/email.rs

@@ -1,12 +1,15 @@
-// Copyright 2024 New Vector Ltd.
+// Copyright 2024, 2025 New Vector Ltd.
 // Copyright 2022-2024 The Matrix.org Foundation C.I.C.
 //
 // SPDX-License-Identifier: AGPL-3.0-only
 // Please see LICENSE in the repository root for full details.
 
 use async_trait::async_trait;
 use chrono::{DateTime, Utc};
-use mas_data_model::{User, UserEmail, UserEmailVerification, UserEmailVerificationState};
+use mas_data_model::{
+    BrowserSession, User, UserEmail, UserEmailAuthentication, UserEmailAuthenticationCode,
+    UserEmailVerification, UserEmailVerificationState,
+};
 use mas_storage::{
     user::{UserEmailFilter, UserEmailRepository},
     Clock, Page, Pagination,
@@ -95,6 +98,46 @@ impl UserEmailConfirmationCodeLookup {
     }
 }
 
+struct UserEmailAuthenticationLookup {
+    user_email_authentication_id: Uuid,
+    user_session_id: Option<Uuid>,
+    email: String,
+    created_at: DateTime<Utc>,
+    completed_at: Option<DateTime<Utc>>,
+}
+
+impl From<UserEmailAuthenticationLookup> for UserEmailAuthentication {
+    fn from(value: UserEmailAuthenticationLookup) -> Self {
+        UserEmailAuthentication {
+            id: value.user_email_authentication_id.into(),
+            user_session_id: value.user_session_id.map(Ulid::from),
+            email: value.email,
+            created_at: value.created_at,
+            completed_at: value.completed_at,
+        }
+    }
+}
+
+struct UserEmailAuthenticationCodeLookup {
+    user_email_authentication_code_id: Uuid,
+    user_email_authentication_id: Uuid,
+    code: String,
+    created_at: DateTime<Utc>,
+    expires_at: DateTime<Utc>,
+}
+
+impl From<UserEmailAuthenticationCodeLookup> for UserEmailAuthenticationCode {
+    fn from(value: UserEmailAuthenticationCodeLookup) -> Self {
+        UserEmailAuthenticationCode {
+            id: value.user_email_authentication_code_id.into(),
+            user_email_authentication_id: value.user_email_authentication_id.into(),
+            code: value.code,
+            created_at: value.created_at,
+            expires_at: value.expires_at,
+        }
+    }
+}
+
 impl Filter for UserEmailFilter<'_> {
     fn generate_condition(&self, _has_joins: bool) -> impl sea_query::IntoCondition {
         sea_query::Condition::all()
@@ -544,4 +587,229 @@ impl UserEmailRepository for PgUserEmailRepository<'_> {
 
         Ok(user_email_verification)
     }
+
+    #[tracing::instrument(
+        name = "db.user_email.add_authentication_for_session",
+        skip_all,
+        fields(
+            db.query.text,
+            %session.id,
+            user_email_authentication.id,
+            user_email_authentication.email = email,
+        ),
+        err,
+    )]
+    async fn add_authentication_for_session(
+        &mut self,
+        rng: &mut (dyn RngCore + Send),
+        clock: &dyn Clock,
+        email: String,
+        session: &BrowserSession,
+    ) -> Result<UserEmailAuthentication, Self::Error> {
+        let created_at = clock.now();
+        let id = Ulid::from_datetime_with_source(created_at.into(), rng);
+        tracing::Span::current()
+            .record("user_email_authentication.id", tracing::field::display(id));
+
+        sqlx::query!(
+            r#"
+                INSERT INTO user_email_authentications
+                  ( user_email_authentication_id
+                  , user_session_id
+                  , email
+                  , created_at
+                  )
+                VALUES ($1, $2, $3, $4)
+            "#,
+            Uuid::from(id),
+            Uuid::from(session.id),
+            &email,
+            created_at,
+        )
+        .traced()
+        .execute(&mut *self.conn)
+        .await?;
+
+        Ok(UserEmailAuthentication {
+            id,
+            user_session_id: Some(session.id),
+            email,
+            created_at,
+            completed_at: None,
+        })
+    }
+
+    #[tracing::instrument(
+        name = "db.user_email.add_authentication_code",
+        skip_all,
+        fields(
+            db.query.text,
+            %user_email_authentication.id,
+            %user_email_authentication.email,
+            user_email_authentication_code.id,
+            user_email_authentication_code.code = code,
+        ),
+        err,
+    )]
+    async fn add_authentication_code(
+        &mut self,
+        rng: &mut (dyn RngCore + Send),
+        clock: &dyn Clock,
+        duration: chrono::Duration,
+        user_email_authentication: &UserEmailAuthentication,
+        code: String,
+    ) -> Result<UserEmailAuthenticationCode, Self::Error> {
+        let created_at = clock.now();
+        let expires_at = created_at + duration;
+        let id = Ulid::from_datetime_with_source(created_at.into(), rng);
+        tracing::Span::current().record(
+            "user_email_authentication_code.id",
+            tracing::field::display(id),
+        );
+
+        sqlx::query!(
+            r#"
+                INSERT INTO user_email_authentication_codes
+                  ( user_email_authentication_code_id
+                  , user_email_authentication_id
+                  , code
+                  , created_at
+                  , expires_at
+                  )
+                VALUES ($1, $2, $3, $4, $5)
+            "#,
+            Uuid::from(id),
+            Uuid::from(user_email_authentication.id),
+            &code,
+            created_at,
+            expires_at,
+        )
+        .traced()
+        .execute(&mut *self.conn)
+        .await?;
+
+        Ok(UserEmailAuthenticationCode {
+            id,
+            user_email_authentication_id: user_email_authentication.id,
+            code,
+            created_at,
+            expires_at,
+        })
+    }
+
+    #[tracing::instrument(
+        name = "db.user_email.lookup_authentication",
+        skip_all,
+        fields(
+            db.query.text,
+            user_email_authentication.id = %id,
+        ),
+        err,
+    )]
+    async fn lookup_authentication(
+        &mut self,
+        id: Ulid,
+    ) -> Result<Option<UserEmailAuthentication>, Self::Error> {
+        let res = sqlx::query_as!(
+            UserEmailAuthenticationLookup,
+            r#"
+                SELECT user_email_authentication_id
+                     , user_session_id
+                     , email
+                     , created_at
+                     , completed_at
+                FROM user_email_authentications
+                WHERE user_email_authentication_id = $1
+            "#,
+            Uuid::from(id),
+        )
+        .traced()
+        .fetch_optional(&mut *self.conn)
+        .await?;
+
+        Ok(res.map(UserEmailAuthentication::from))
+    }
+
+    #[tracing::instrument(
+        name = "db.user_email.find_authentication_by_code",
+        skip_all,
+        fields(
+            db.query.text,
+            %authentication.id,
+            user_email_authentication_code.code = code,
+        ),
+        err,
+    )]
+    async fn find_authentication_code(
+        &mut self,
+        authentication: &UserEmailAuthentication,
+        code: &str,
+    ) -> Result<Option<UserEmailAuthenticationCode>, Self::Error> {
+        let res = sqlx::query_as!(
+            UserEmailAuthenticationCodeLookup,
+            r#"
+                SELECT user_email_authentication_code_id
+                     , user_email_authentication_id
+                     , code
+                     , created_at
+                     , expires_at
+                FROM user_email_authentication_codes
+                WHERE user_email_authentication_id = $1
+                  AND code = $2
+            "#,
+            Uuid::from(authentication.id),
+            code,
+        )
+        .traced()
+        .fetch_optional(&mut *self.conn)
+        .await?;
+
+        Ok(res.map(UserEmailAuthenticationCode::from))
+    }
+
+    #[tracing::instrument(
+        name = "db.user_email.complete_email_authentication",
+        skip_all,
+        fields(
+            db.query.text,
+            %user_email_authentication.id,
+            %user_email_authentication.email,
+            %user_email_authentication_code.id,
+            %user_email_authentication_code.code,
+        ),
+        err,
+    )]
+    async fn complete_authentication(
+        &mut self,
+        clock: &dyn Clock,
+        mut user_email_authentication: UserEmailAuthentication,
+        user_email_authentication_code: &UserEmailAuthenticationCode,
+    ) -> Result<UserEmailAuthentication, Self::Error> {
+        // We technically don't use the authentication code here (other than
+        // recording it in the span), but this is to make sure the caller has
+        // fetched one before calling this
+        let completed_at = clock.now();
+
+        // We'll assume the caller has checked that completed_at is None, so in case
+        // they haven't, the update will not affect any rows, which will raise
+        // an error
+        let res = sqlx::query!(
+            r#"
+                UPDATE user_email_authentications
+                SET completed_at = $2
+                WHERE user_email_authentication_id = $1
+                  AND completed_at IS NULL
+            "#,
+            Uuid::from(user_email_authentication.id),
+            completed_at,
+        )
+        .traced()
+        .execute(&mut *self.conn)
+        .await?;
+
+        DatabaseError::ensure_affected_rows(&res, 1)?;
+
+        user_email_authentication.completed_at = Some(completed_at);
+        Ok(user_email_authentication)
+    }
 }

crates/storage-pg/src/user/tests.rs

@@ -1,4 +1,4 @@
-// Copyright 2024 New Vector Ltd.
+// Copyright 2024, 2025 New Vector Ltd.
 // Copyright 2023, 2024 The Matrix.org Foundation C.I.C.
 //
 // SPDX-License-Identifier: AGPL-3.0-only
@@ -11,7 +11,7 @@ use mas_storage::{
         BrowserSessionFilter, BrowserSessionRepository, UserEmailFilter, UserEmailRepository,
         UserFilter, UserPasswordRepository, UserRepository,
     },
-    Pagination, RepositoryAccess,
+    Clock, Pagination, RepositoryAccess,
 };
 use rand::SeedableRng;
 use rand_chacha::ChaChaRng;
@@ -399,6 +399,121 @@ async fn test_user_email_repo(pool: PgPool) {
     repo.save().await.unwrap();
 }
 
+/// Test the authentication codes methods in the user email repository
+#[sqlx::test(migrator = "crate::MIGRATOR")]
+async fn test_user_email_repo_authentications(pool: PgPool) {
+    let mut repo = PgRepository::from_pool(&pool).await.unwrap().boxed();
+    let mut rng = ChaChaRng::seed_from_u64(42);
+    let clock = MockClock::default();
+
+    // Create a user and a user session so that we can create an authentication
+    let user = repo
+        .user()
+        .add(&mut rng, &clock, "alice".to_owned())
+        .await
+        .unwrap();
+
+    let browser_session = repo
+        .browser_session()
+        .add(&mut rng, &clock, &user, None)
+        .await
+        .unwrap();
+
+    // Create an authentication session
+    let authentication = repo
+        .user_email()
+        .add_authentication_for_session(
+            &mut rng,
+            &clock,
+            "alice@example.com".to_owned(),
+            &browser_session,
+        )
+        .await
+        .unwrap();
+
+    assert_eq!(authentication.email, "alice@example.com");
+    assert_eq!(authentication.user_session_id, Some(browser_session.id));
+    assert_eq!(authentication.created_at, clock.now());
+    assert_eq!(authentication.completed_at, None);
+
+    // Check that we can find the authentication by its ID
+    let lookup = repo
+        .user_email()
+        .lookup_authentication(authentication.id)
+        .await
+        .unwrap()
+        .unwrap();
+    assert_eq!(lookup.id, authentication.id);
+    assert_eq!(lookup.email, "alice@example.com");
+    assert_eq!(lookup.user_session_id, Some(browser_session.id));
+    assert_eq!(lookup.created_at, clock.now());
+    assert_eq!(lookup.completed_at, None);
+
+    // Add a code to the session
+    let code = repo
+        .user_email()
+        .add_authentication_code(
+            &mut rng,
+            &clock,
+            Duration::minutes(5),
+            &authentication,
+            "123456".to_owned(),
+        )
+        .await
+        .unwrap();
+
+    assert_eq!(code.code, "123456");
+    assert_eq!(code.created_at, clock.now());
+    assert_eq!(code.expires_at, clock.now() + Duration::minutes(5));
+
+    // Check that we can find the code by its ID
+    let id = code.id;
+    let lookup = repo
+        .user_email()
+        .find_authentication_code(&authentication, "123456")
+        .await
+        .unwrap()
+        .unwrap();
+
+    assert_eq!(lookup.id, id);
+    assert_eq!(lookup.code, "123456");
+    assert_eq!(lookup.created_at, clock.now());
+    assert_eq!(lookup.expires_at, clock.now() + Duration::minutes(5));
+
+    // Complete the authentication
+    let authentication = repo
+        .user_email()
+        .complete_authentication(&clock, authentication, &code)
+        .await
+        .unwrap();
+
+    assert_eq!(authentication.id, authentication.id);
+    assert_eq!(authentication.email, "alice@example.com");
+    assert_eq!(authentication.user_session_id, Some(browser_session.id));
+    assert_eq!(authentication.created_at, clock.now());
+    assert_eq!(authentication.completed_at, Some(clock.now()));
+
+    // Check that we can find the completed authentication by its ID
+    let lookup = repo
+        .user_email()
+        .lookup_authentication(authentication.id)
+        .await
+        .unwrap()
+        .unwrap();
+    assert_eq!(lookup.id, authentication.id);
+    assert_eq!(lookup.email, "alice@example.com");
+    assert_eq!(lookup.user_session_id, Some(browser_session.id));
+    assert_eq!(lookup.created_at, clock.now());
+    assert_eq!(lookup.completed_at, Some(clock.now()));
+
+    // Completing a second time should fail
+    let res = repo
+        .user_email()
+        .complete_authentication(&clock, authentication, &code)
+        .await;
+    assert!(res.is_err());
+}
+
 /// Test the user password repository implementation.
 #[sqlx::test(migrator = "crate::MIGRATOR")]
 async fn test_user_password_repo(pool: PgPool) {

crates/storage/src/user/email.rs

@@ -1,11 +1,14 @@
-// Copyright 2024 New Vector Ltd.
+// Copyright 2024, 2025 New Vector Ltd.
 // Copyright 2022-2024 The Matrix.org Foundation C.I.C.
 //
 // SPDX-License-Identifier: AGPL-3.0-only
 // Please see LICENSE in the repository root for full details.
 
 use async_trait::async_trait;
-use mas_data_model::{User, UserEmail, UserEmailVerification};
+use mas_data_model::{
+    BrowserSession, User, UserEmail, UserEmailAuthentication, UserEmailAuthenticationCode,
+    UserEmailVerification,
+};
 use rand_core::RngCore;
 use ulid::Ulid;
 
@@ -281,6 +284,102 @@ pub trait UserEmailRepository: Send + Sync {
         clock: &dyn Clock,
         verification: UserEmailVerification,
     ) -> Result<UserEmailVerification, Self::Error>;
+
+    /// Add a new [`UserEmailAuthentication`] for a [`BrowserSession`]
+    ///
+    /// # Parameters
+    ///
+    /// * `rng`: The random number generator to use
+    /// * `clock`: The clock to use
+    /// * `email`: The email address to add
+    /// * `session`: The [`BrowserSession`] for which to add the
+    ///   [`UserEmailAuthentication`]
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the underlying repository fails
+    async fn add_authentication_for_session(
+        &mut self,
+        rng: &mut (dyn RngCore + Send),
+        clock: &dyn Clock,
+        email: String,
+        session: &BrowserSession,
+    ) -> Result<UserEmailAuthentication, Self::Error>;
+
+    /// Add a new [`UserEmailAuthenticationCode`] for a
+    /// [`UserEmailAuthentication`]
+    ///
+    /// # Parameters
+    ///
+    /// * `rng`: The random number generator to use
+    /// * `clock`: The clock to use
+    /// * `duration`: The duration for which the code is valid
+    /// * `authentication`: The [`UserEmailAuthentication`] for which to add the
+    ///   [`UserEmailAuthenticationCode`]
+    /// * `code`: The code to add
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the underlying repository fails or if the code
+    /// already exists for this session
+    async fn add_authentication_code(
+        &mut self,
+        rng: &mut (dyn RngCore + Send),
+        clock: &dyn Clock,
+        duration: chrono::Duration,
+        authentication: &UserEmailAuthentication,
+        code: String,
+    ) -> Result<UserEmailAuthenticationCode, Self::Error>;
+
+    /// Lookup a [`UserEmailAuthentication`]
+    ///
+    /// # Parameters
+    ///
+    /// * `id`: The ID of the [`UserEmailAuthentication`] to lookup
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the underlying repository fails
+    async fn lookup_authentication(
+        &mut self,
+        id: Ulid,
+    ) -> Result<Option<UserEmailAuthentication>, Self::Error>;
+
+    /// Find a [`UserEmailAuthenticationCode`] by its code and session
+    ///
+    /// # Parameters
+    ///
+    /// * `authentication`: The [`UserEmailAuthentication`] to find the code for
+    /// * `code`: The code of the [`UserEmailAuthentication`] to lookup
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the underlying repository fails
+    async fn find_authentication_code(
+        &mut self,
+        authentication: &UserEmailAuthentication,
+        code: &str,
+    ) -> Result<Option<UserEmailAuthenticationCode>, Self::Error>;
+
+    /// Complete a [`UserEmailAuthentication`] by using the given code
+    ///
+    /// Returns the completed [`UserEmailAuthentication`]
+    ///
+    /// # Parameters
+    ///
+    /// * `clock`: The clock to use to generate timestamps
+    /// * `authentication`: The [`UserEmailAuthentication`] to complete
+    /// * `code`: The [`UserEmailAuthenticationCode`] to use
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the underlying repository fails
+    async fn complete_authentication(
+        &mut self,
+        clock: &dyn Clock,
+        authentication: UserEmailAuthentication,
+        code: &UserEmailAuthenticationCode,
+    ) -> Result<UserEmailAuthentication, Self::Error>;
 }
 
 repository_impl!(UserEmailRepository:
@@ -331,4 +430,39 @@ repository_impl!(UserEmailRepository:
         clock: &dyn Clock,
         verification: UserEmailVerification,
     ) -> Result<UserEmailVerification, Self::Error>;
+
+    async fn add_authentication_for_session(
+        &mut self,
+        rng: &mut (dyn RngCore + Send),
+        clock: &dyn Clock,
+        email: String,
+        session: &BrowserSession,
+    ) -> Result<UserEmailAuthentication, Self::Error>;
+
+    async fn add_authentication_code(
+        &mut self,
+        rng: &mut (dyn RngCore + Send),
+        clock: &dyn Clock,
+        duration: chrono::Duration,
+        authentication: &UserEmailAuthentication,
+        code: String,
+    ) -> Result<UserEmailAuthenticationCode, Self::Error>;
+
+    async fn lookup_authentication(
+        &mut self,
+        id: Ulid,
+    ) -> Result<Option<UserEmailAuthentication>, Self::Error>;
+
+    async fn find_authentication_code(
+        &mut self,
+        authentication: &UserEmailAuthentication,
+        code: &str,
+    ) -> Result<Option<UserEmailAuthenticationCode>, Self::Error>;
+
+    async fn complete_authentication(
+        &mut self,
+        clock: &dyn Clock,
+        authentication: UserEmailAuthentication,
+        code: &UserEmailAuthenticationCode,
+    ) -> Result<UserEmailAuthentication, Self::Error>;
 );