Handle race condition for unclaimed mappers

v-model/src/db.rs

@@ -3,7 +3,7 @@
 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
 
 use chrono::{DateTime, Utc};
-use diesel::{Insertable, Queryable};
+use diesel::{Insertable, Queryable, QueryableByName};
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use uuid::Uuid;
@@ -190,7 +190,7 @@ pub struct AccessGroupModel<T> {
     pub deleted_at: Option<DateTime<Utc>>,
 }
 
-#[derive(Debug, Deserialize, Serialize, Queryable, Insertable)]
+#[derive(Debug, Deserialize, Serialize, Queryable, QueryableByName, Insertable)]
 #[diesel(table_name = mapper)]
 pub struct MapperModel {
     pub id: Uuid,

v-model/src/storage/postgres.rs

@@ -1521,23 +1521,54 @@ impl MapperStore for PostgresStore {
             .map(|max| new_mapper.activations.unwrap_or(0) == max)
             .unwrap_or(false);
 
-        let mapper_m: MapperModel = insert_into(mapper::dsl::mapper)
-            .values((
-                mapper::id.eq(new_mapper.id.into_untyped_uuid()),
-                mapper::name.eq(new_mapper.name.clone()),
-                mapper::rule.eq(new_mapper.rule.clone()),
-                mapper::activations.eq(new_mapper.activations),
-                mapper::max_activations.eq(new_mapper.max_activations),
-                mapper::depleted_at.eq(if depleted { Some(Utc::now()) } else { None }),
-            ))
-            .on_conflict(mapper::id)
-            .do_update()
-            .set((
-                mapper::activations.eq(excluded(mapper::activations)),
-                mapper::depleted_at.eq(excluded(mapper::depleted_at)),
-            ))
-            .get_result_async(&*self.pool.get().await?)
-            .await?;
+        // On insert (a brand new mapper) the provided activation count is stored
+        // as-is. On conflict (an existing mapper redeeming an activation) the
+        // count is incremented atomically relative to the *current* row value,
+        // guarded so it can never exceed `max_activations`. Computing the new
+        // value from the committed row (rather than a caller-supplied absolute
+        // value derived from a stale snapshot) closes the read-modify-write race
+        // where two concurrent logins could both redeem a single-use mapper.
+        //
+        // `INSERT ... ON CONFLICT DO UPDATE` takes a row lock, so concurrent
+        // redemptions serialize: the second observes the incremented value and
+        // the `WHERE` guard prevents it from applying. When the guard rejects the
+        // update the statement returns no rows, surfacing as a not-found error
+        // that the caller treats as a depleted mapper.
+        let mapper_m: MapperModel = diesel::sql_query(
+            r#"
+            WITH upserted AS (
+                INSERT INTO mapper (id, name, rule, activations, max_activations, depleted_at)
+                VALUES ($1, $2, $3, $4, $5, $6)
+                ON CONFLICT (id) DO UPDATE
+                SET
+                    activations = COALESCE(mapper.activations, 0) + 1,
+                    depleted_at = CASE
+                        WHEN mapper.max_activations IS NOT NULL
+                             AND COALESCE(mapper.activations, 0) + 1 >= mapper.max_activations
+                        THEN now()
+                        ELSE NULL
+                    END
+                WHERE mapper.max_activations IS NULL
+                   OR COALESCE(mapper.activations, 0) < mapper.max_activations
+                RETURNING *
+            )
+            SELECT * FROM upserted
+            "#,
+        )
+        .bind::<diesel::sql_types::Uuid, _>(new_mapper.id.into_untyped_uuid())
+        .bind::<diesel::sql_types::Text, _>(new_mapper.name.clone())
+        .bind::<diesel::sql_types::Jsonb, _>(new_mapper.rule.clone())
+        .bind::<diesel::sql_types::Nullable<diesel::sql_types::Integer>, _>(new_mapper.activations)
+        .bind::<diesel::sql_types::Nullable<diesel::sql_types::Integer>, _>(
+            new_mapper.max_activations,
+        )
+        .bind::<diesel::sql_types::Nullable<diesel::sql_types::Timestamptz>, _>(if depleted {
+            Some(Utc::now())
+        } else {
+            None
+        })
+        .get_result_async(&*self.pool.get().await?)
+        .await?;
 
         Ok(mapper_m.into())
     }

v-model/tests/postgres.rs

@@ -16,12 +16,12 @@ use serde::{Deserialize, Serialize};
 use std::ops::{Add, Sub};
 use uuid::Uuid;
 use v_model::{
-    NewApiKey, NewApiUser, NewMagicLink, NewMagicLinkAttempt, UserId,
+    NewApiKey, NewApiUser, NewMagicLink, NewMagicLinkAttempt, NewMapper, UserId,
     migrations::run_migrations,
     schema_ext::MagicLinkAttemptState,
     storage::{
         ApiKeyFilter, ApiKeyStore, ApiUserFilter, ApiUserStore, ListPagination,
-        MagicLinkAttemptStore, MagicLinkStore, postgres::PostgresStore,
+        MagicLinkAttemptStore, MagicLinkStore, MapperStore, postgres::PostgresStore,
     },
 };
 
@@ -514,3 +514,87 @@ async fn test_magic_link_attempt() {
 
     assert_eq!(None, expired_attempt_lookup);
 }
+
+// Verifies that a mapper's activation limit is enforced atomically by the
+// store's upsert. A mapper with `max_activations = 1` must be redeemable at
+// most once, even when two redemptions race. This guards against the
+// read-modify-write race where two concurrent logins could both redeem a
+// single-use mapper.
+//
+// Steps:
+//   1. Create a single-use mapper (activations = 0, max_activations = 1)
+//   2. Redeem it once (activations -> 1, depleted_at set)
+//   3. A second sequential redemption is rejected
+//   4. Two concurrent redemptions of a fresh single-use mapper yield exactly
+//      one success, and the persisted counter records a single activation
+#[tokio::test]
+async fn test_mapper_activation_limit() {
+    let db = TestDb::new("test_mapper_activation_limit");
+    let store = PostgresStore::new(&db.url()).await.unwrap();
+
+    // 1. Create a single-use mapper. The insert path stores the provided
+    //    activation count as-is.
+    let mapper_id = TypedUuid::new_v4();
+    let seed = NewMapper {
+        id: mapper_id,
+        name: format!("seq-{}", Uuid::new_v4()),
+        rule: serde_json::json!({ "rule": "default" }),
+        activations: Some(0),
+        max_activations: Some(1),
+    };
+
+    let created = MapperStore::upsert(&store, &seed).await.unwrap();
+    assert_eq!(created.activations, Some(0));
+    assert!(created.depleted_at.is_none());
+
+    // 2. The first redemption increments the counter and depletes the mapper.
+    let consumed = MapperStore::upsert(&store, &seed).await.unwrap();
+    assert_eq!(consumed.activations, Some(1));
+    assert!(consumed.depleted_at.is_some());
+
+    // 3. A second redemption is rejected: the guarded update affects no rows,
+    //    which surfaces as an error.
+    let depleted = MapperStore::upsert(&store, &seed).await;
+    assert!(
+        depleted.is_err(),
+        "a depleted single-use mapper must not be redeemable again, got {:?}",
+        depleted.map(|m| m.activations),
+    );
+
+    // 4. Race two redemptions of a fresh single-use mapper. Regardless of
+    //    interleaving, `INSERT ... ON CONFLICT DO UPDATE` locks the row and the
+    //    `WHERE activations < max_activations` guard ensures exactly one wins.
+    let race_id = TypedUuid::new_v4();
+    let race_seed = NewMapper {
+        id: race_id,
+        name: format!("race-{}", Uuid::new_v4()),
+        rule: serde_json::json!({ "rule": "default" }),
+        activations: Some(0),
+        max_activations: Some(1),
+    };
+    MapperStore::upsert(&store, &race_seed).await.unwrap();
+
+    let (first, second) = tokio::join!(
+        MapperStore::upsert(&store, &race_seed),
+        MapperStore::upsert(&store, &race_seed),
+    );
+
+    let successes = [&first, &second].iter().filter(|r| r.is_ok()).count();
+    assert_eq!(
+        successes,
+        1,
+        "exactly one of two concurrent redemptions of a max_activations=1 mapper should succeed, \
+         got first={:?} second={:?}",
+        first.as_ref().map(|m| m.activations),
+        second.as_ref().map(|m| m.activations),
+    );
+
+    // The persisted counter reflects exactly one activation, and the mapper is
+    // marked depleted.
+    let final_mapper = MapperStore::get(&store, &race_id, true, false)
+        .await
+        .unwrap()
+        .unwrap();
+    assert_eq!(final_mapper.activations, Some(1));
+    assert!(final_mapper.depleted_at.is_some());
+}