2fa

idp/migrations/20250712182250_sessions.sql

@@ -4,6 +4,7 @@ CREATE TABLE IF NOT EXISTS sessions (
     user_id INTEGER NOT NULL,
     created INTEGER NOT NULL,
     user_agent TEXT NOT NULL,
+    locked_2fa INTEGER NOT NULL, -- whether the session is currently locked and needs a 2FA code to unlock it
 
     FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
 ) STRICT;

idp/migrations/20250712211103_totp.sql

@@ -1,8 +1,8 @@
 CREATE TABLE IF NOT EXISTS used_totp (
     user_id INTEGER,
-    code TEXT,
+    time_step INTEGER,
 
-    PRIMARY KEY (user_id, code),
+    PRIMARY KEY (user_id, time_step),
     FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
 ) STRICT;
 

idp/src/main.rs

@@ -59,6 +59,7 @@ async fn main() -> Result<()> {
         .route("/", get(root))
         .route("/signup", get(signup).post(signup_post))
         .route("/login", get(login).post(login_post))
+        .route("/login-2fa", get(login_2fa).post(login_2fa_post))
         .route("/2fa", get(list_2fa))
         .route("/sessions", get(list_sessions))
         .route("/sessions/delete", post(delete_session))
@@ -134,6 +135,25 @@ async fn login() -> impl IntoResponse {
     Html(LoginTemplate { error: false }.render().unwrap())
 }
 
+#[derive(askama::Template)]
+#[template(path = "login-2fa.html")]
+struct Login2faTemplate {
+    error: bool,
+    reuse: bool,
+}
+
+#[axum::debug_handler]
+async fn login_2fa() -> impl IntoResponse {
+    Html(
+        Login2faTemplate {
+            error: false,
+            reuse: false,
+        }
+        .render()
+        .unwrap(),
+    )
+}
+
 async fn list_2fa(user: UserSession, State(db): State<Db>) -> Result<impl IntoResponse, Response> {
     let Some(user) = user.0 else {
         return Err(Redirect::to("/").into_response());
@@ -217,6 +237,8 @@ async fn delete_2fa(
         return Err(Redirect::to("/").into_response());
     };
 
+    // TODO: This should require 2FA authentication
+
     totp::delete_totp_device(&db, user.user_id, form.device_id)
         .await
         .map_err(|err| {
@@ -257,7 +279,10 @@ async fn add_totp_post(
         return Err(Redirect::to("/").into_response());
     };
 
-    let computed = totp::Totp::compute(&form.secret, jiff::Timestamp::now().as_second() as u64);
+    let computed = totp::Totp::compute(
+        &form.secret,
+        totp::Totp::time_step(jiff::Timestamp::now().as_second()),
+    );
 
     if computed.digits != form.code.trim() {
         return Err(Html(
@@ -307,8 +332,9 @@ async fn make_session_cookie_for_user(
     db: &Db,
     user_id: i64,
     user_agent: &str,
+    locked_2fa: bool,
 ) -> Result<Cookie<'static>, Response> {
-    let session = session::create_session(&db, user_id, user_agent)
+    let session = session::create_session(&db, user_id, user_agent, locked_2fa)
         .await
         .map_err(|err| {
             error!(?err, "Failed to create session for user");
@@ -356,7 +382,7 @@ async fn signup_post(
         .and_then(|v| str::from_utf8(v.as_bytes()).ok())
         .unwrap_or("unknown");
 
-    let session_id = make_session_cookie_for_user(&db, user.id, user_agent).await?;
+    let session_id = make_session_cookie_for_user(&db, user.id, user_agent, false).await?;
 
     Ok((jar.add(session_id), Redirect::to("/")))
 }
@@ -378,12 +404,107 @@ async fn login_post(
         return Err(Html(LoginTemplate { error: true }.render().unwrap()).into_response());
     };
 
+    let totp_devices = totp::list_totp_devices(&db, user.id).await.map_err(|err| {
+        error!(?err, "Failed to list totp devices");
+        StatusCode::INTERNAL_SERVER_ERROR.into_response()
+    })?;
+
+    let locked_2fa = !totp_devices.is_empty();
+
     let user_agent = headers
         .get(axum::http::header::USER_AGENT)
         .and_then(|v| str::from_utf8(v.as_bytes()).ok())
         .unwrap_or("unknown");
 
-    let session_id = make_session_cookie_for_user(&db, user.id, user_agent).await?;
+    let session_id = make_session_cookie_for_user(&db, user.id, user_agent, locked_2fa).await?;
 
-    Ok((jar.add(session_id), Redirect::to("/")))
+    let redirect_target = if locked_2fa { "/login-2fa" } else { "/" };
+
+    Ok((jar.add(session_id), Redirect::to(redirect_target)))
+}
+
+#[derive(Deserialize)]
+struct Login2faForm {
+    totp_code: String,
+}
+
+async fn login_2fa_post(
+    State(db): State<Db>,
+    jar: CookieJar,
+    Form(form): Form<Login2faForm>,
+) -> Result<impl IntoResponse, Response> {
+    let now = jiff::Timestamp::now();
+
+    let Some(session_id) = jar.get(crate::SESSION_ID_COOKIE_NAME) else {
+        return Err(Redirect::to("/").into_response());
+    };
+
+    let session = session::find_locked_session(&db, session_id.value())
+        .await
+        .map_err(|err| {
+            error!(?err, "Failed to find locked session");
+            StatusCode::INTERNAL_SERVER_ERROR.into_response()
+        })?;
+
+    let Some(session) = session else {
+        return Err(Redirect::to("/").into_response());
+    };
+
+    let totp_devices = totp::list_totp_devices(&db, session.user_id)
+        .await
+        .map_err(|err| {
+            error!(?err, "Failed to list totp devices");
+            StatusCode::INTERNAL_SERVER_ERROR.into_response()
+        })?;
+
+    let time_step = now.as_second() / 30;
+    let used = totp::find_used_totp(&db, session.user_id, time_step)
+        .await
+        .map_err(|err| {
+            error!(?err, "Failed to find used totp");
+            StatusCode::INTERNAL_SERVER_ERROR.into_response()
+        })?;
+
+    if used.has_used > 0 {
+        return Err(Html(
+            Login2faTemplate {
+                error: false,
+                reuse: true,
+            }
+            .render()
+            .unwrap(),
+        )
+        .into_response());
+    }
+
+    let code_matches = totp_devices.iter().any(|device| {
+        totp::Totp::compute(&device.secret, time_step).digits == form.totp_code.trim()
+    });
+
+    if code_matches {
+        totp::insert_used_totp_code(&db, session.user_id, time_step)
+            .await
+            .map_err(|err| {
+                error!(?err, "Failed to insert used totp");
+                StatusCode::INTERNAL_SERVER_ERROR.into_response()
+            })?;
+        session::unlock_session(&db, &session.session_id)
+            .await
+            .map_err(|err| {
+                error!(?err, "Failed to unlock session");
+                StatusCode::INTERNAL_SERVER_ERROR.into_response()
+            })?;
+    } else {
+        return Err(Html(
+            Login2faTemplate {
+                error: true,
+                reuse: false,
+            }
+            .render()
+            .unwrap(),
+        )
+        .into_response());
+    }
+
+    Ok(Redirect::to("/"))
 }

idp/src/session.rs

@@ -22,7 +22,7 @@ pub struct SessionId(pub String);
 
 pub async fn find_session(db: &Db, session_id: &str) -> Result<Option<SessionWithUser>> {
     let result = sqlx::query_as::<_, SessionWithUser>(
-        "select user_id, created, username from sessions left join users on sessions.user_id = users.id where session_id = ?",
+        "select user_id, created, username from sessions left join users on sessions.user_id = users.id where session_id = ? and locked_2fa = false",
     )
     .bind(session_id)
     .fetch_one(&db.pool)
@@ -34,6 +34,35 @@ pub async fn find_session(db: &Db, session_id: &str) -> Result<Option<SessionWit
     }
 }
 
+#[derive(Debug, sqlx::FromRow)]
+pub struct LockedSession {
+    pub user_id: i64,
+    pub session_id: String,
+}
+
+pub async fn find_locked_session(db: &Db, session_id: &str) -> Result<Option<LockedSession>> {
+    let result = sqlx::query_as::<_, LockedSession>(
+        "select user_id, session_id from sessions where session_id = ? and locked_2fa = true",
+    )
+    .bind(session_id)
+    .fetch_one(&db.pool)
+    .await;
+    match result {
+        Ok(session) => Ok(Some(session)),
+        Err(sqlx::Error::RowNotFound) => Ok(None),
+        Err(e) => Err(e).wrap_err("failed to fetch locked session"),
+    }
+}
+
+pub async fn unlock_session(db: &Db, session_id: &str) -> Result<()> {
+    sqlx::query("update sessions set locked_2fa = false where session_id = ?")
+        .bind(session_id)
+        .execute(&db.pool)
+        .await
+        .wrap_err("unlocking session")?;
+    Ok(())
+}
+
 #[derive(Debug, sqlx::FromRow)]
 pub struct SessionForList {
     pub created: i64,
@@ -61,18 +90,24 @@ pub async fn delete_session(db: &Db, user_id: i64, session_public_id: i64) -> Re
     Ok(())
 }
 
-pub async fn create_session(db: &Db, user_id: i64, user_agent: &str) -> Result<SessionId> {
+pub async fn create_session(
+    db: &Db,
+    user_id: i64,
+    user_agent: &str,
+    locked_2fa: bool,
+) -> Result<SessionId> {
     let mut session_id = [0_u8; 32];
     rand_core::OsRng.fill_bytes(&mut session_id);
     let session_id = format!("idpsess_{}", hex::encode(session_id));
 
     sqlx::query(
-        "insert into sessions (session_id, user_id, created, user_agent) values (?, ?, ?, ?)",
+        "insert into sessions (session_id, user_id, created, user_agent, locked_2fa) values (?, ?, ?, ?, ?)",
     )
     .bind(&session_id)
     .bind(user_id)
     .bind(jiff::Timestamp::now().as_millisecond())
     .bind(user_agent)
+    .bind(locked_2fa)
     .execute(&db.pool)
     .await
     .wrap_err("inserting new session")?;

idp/src/totp.rs

@@ -13,7 +13,6 @@ mod compute {
     }
 
     struct TotpConfig {
-        time_step_s: u64,
         digits: u32,
     }
 
@@ -42,23 +41,19 @@ mod compute {
     }
 
     impl Totp {
-        pub fn compute(secret: &str, unix_seconds: u64) -> Self {
+        pub fn time_step(unix_seconds: i64) -> i64 {
+            unix_seconds / 30
+        }
+
+        pub fn compute(secret: &str, time_step: i64) -> Self {
             let secret = base32::decode(base32::Alphabet::Rfc4648 { padding: false }, &secret)
                 .unwrap_or_default(); // nonsense secret will result in the wrong code as intended
 
-            Self::compute_inner(
+            Self::compute_inner(&secret, time_step, TotpConfig { digits: 6 })
-                &secret,
-                unix_seconds,
-                TotpConfig {
-                    time_step_s: 30,
-                    digits: 6,
-                },
-            )
         }
 
-        fn compute_inner(secret: &[u8], unix_seconds: u64, config: TotpConfig) -> Self {
+        fn compute_inner(secret: &[u8], time_step: i64, config: TotpConfig) -> Self {
-            let time_step = unix_seconds / config.time_step_s;
+            let code = hotp(secret, time_step as u64, config.digits);
-            let code = hotp(secret, time_step, config.digits);
             Totp { digits: code }
         }
     }
@@ -80,11 +75,8 @@ mod compute {
             for test in tests {
                 let totp = super::Totp::compute_inner(
                     secret,
-                    test.0,
+                    super::Totp::time_step(test.0),
-                    TotpConfig {
+                    TotpConfig { digits: 8 },
-                        time_step_s: 30,
-                        digits: 8,
-                    },
                 );
                 assert_eq!(totp.digits, test.1);
             }
@@ -118,11 +110,12 @@ pub struct TotpDevice {
     pub id: i64,
     pub created_time: i64,
     pub name: String,
+    pub secret: String,
 }
 
 pub async fn list_totp_devices(db: &Db, user_id: i64) -> Result<Vec<TotpDevice>> {
     sqlx::query_as::<_, TotpDevice>(
-        "select id, created_time, name from totp_devices where user_id = ?",
+        "select id, created_time, name, secret from totp_devices where user_id = ?",
     )
     .bind(user_id)
     .fetch_all(&db.pool)
@@ -139,3 +132,29 @@ pub async fn delete_totp_device(db: &Db, user_id: i64, totp_device_id: i64) -> R
         .wrap_err("failed to delete totp device")?;
     Ok(())
 }
+
+#[derive(sqlx::FromRow)]
+pub struct UsedTotp {
+    pub has_used: i64,
+}
+
+pub async fn find_used_totp(db: &Db, user_id: i64, time_step: i64) -> Result<UsedTotp> {
+    sqlx::query_as::<_, UsedTotp>(
+        "select COUNT(time_step) as has_used from used_totp where user_id = ? and time_step = ?",
+    )
+    .bind(user_id)
+    .bind(time_step)
+    .fetch_one(&db.pool)
+    .await
+    .wrap_err("fetching totp devices")
+}
+
+pub async fn insert_used_totp_code(db: &Db, user_id: i64, time_step: i64) -> Result<()> {
+    sqlx::query("insert into used_totp (user_id, time_step) values (?, ?)")
+        .bind(user_id)
+        .bind(time_step)
+        .execute(&db.pool)
+        .await
+        .wrap_err("failed to insert used totp code")?;
+    Ok(())
+}

idp/templates/add-totp.html

@@ -9,7 +9,7 @@
   <body>
     <h1>Add 2FA to your account</h1>
     <a href="/">home</a>
-    <form method="post" action="/add-totp">
+    <form method="post">
       <p>
         Copy this secret into your authenticator app and enter the code from it
         below.

idp/templates/login-2fa.html

@@ -0,0 +1,29 @@
+<!doctype html>
+<html lang="en">
+  <head>
+    <meta charset="UTF-8" />
+    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+    <title>Login - IDP</title>
+    <link rel="stylesheet" href="/style.css" />
+  </head>
+  <body>
+    <h1>Log into your beautiful account</h1>
+    <a href="/">home</a>
+    <form method="post">
+      <div>
+        <label for="totp_code">2FA code</label>
+        <input id="totp_code" name="totp_code" />
+      </div>
+
+      <button type="submit">Login</button>
+
+      {% if error %}
+      <p>Incorrect 2FA code.</p>
+      {% endif %}
+      <!---->
+      {% if reuse %}
+      <p>2FA code reused, wait 30s for the next one</p>
+      {% endif %}
+    </form>
+  </body>
+</html>

idp/templates/login.html

@@ -9,7 +9,7 @@
   <body>
     <h1>Log into your beautiful account</h1>
     <a href="/">home</a>
-    <form method="post" action="/login">
+    <form method="post">
       <div>
         <label for="username">Username</label>
         <input id="username" name="username" />

idp/templates/signup.html

@@ -9,7 +9,7 @@
   <body>
     <h1>Create a new account</h1>
     <a href="/">home</a>
-    <form method="post" action="/signup">
+    <form method="post">
       <div>
         <label for="username">Username</label>
         <input id="username" name="username" />

idp/templates/style.css

@@ -44,7 +44,7 @@ a[href="/"] {
   margin-bottom: 20px;
 }
 
-form + :not(.fake-form) {
+form:not(.fake-form) {
   display: flex;
   flex-direction: column;
   gap: 10px;