max / makenotwork
2 files changed,
+36 insertions,
-21 deletions
| @@ -71,20 +71,41 @@ pub use models::*; | |||
| 71 | 71 | use crate::error::Result; | |
| 72 | 72 | use sqlx::PgPool; | |
| 73 | 73 | ||
| 74 | - | /// Acquire a session-level PostgreSQL advisory lock (blocks until available). | |
| 75 | - | pub async fn advisory_lock(pool: &PgPool, key: i64) -> Result<()> { | |
| 74 | + | /// Check the sandbox per-IP cap under an advisory lock on a single connection. | |
| 75 | + | /// | |
| 76 | + | /// Acquires a session-level advisory lock, runs the count query, and unlocks — | |
| 77 | + | /// all on the same connection. Returns the active sandbox count. | |
| 78 | + | /// | |
| 79 | + | /// This avoids the bug where `advisory_lock` + `advisory_unlock` through a pool | |
| 80 | + | /// use different connections, leaving locks permanently held. | |
| 81 | + | pub async fn check_sandbox_cap(pool: &PgPool, lock_key: i64, ip: &str) -> Result<i64> { | |
| 82 | + | let mut conn = pool.acquire().await.map_err(|e| { | |
| 83 | + | crate::error::AppError::Internal(anyhow::anyhow!("pool acquire: {}", e)) | |
| 84 | + | })?; | |
| 85 | + | ||
| 86 | + | // Lock, count, unlock — all on the same connection | |
| 76 | 87 | sqlx::query("SELECT pg_advisory_lock($1)") | |
| 77 | - | .bind(key) | |
| 78 | - | .execute(pool) | |
| 88 | + | .bind(lock_key) | |
| 89 | + | .execute(&mut *conn) | |
| 79 | 90 | .await?; | |
| 80 | - | Ok(()) | |
| 81 | - | } | |
| 82 | 91 | ||
| 83 | - | /// Release a session-level PostgreSQL advisory lock. | |
| 84 | - | pub async fn advisory_unlock(pool: &PgPool, key: i64) -> Result<()> { | |
| 92 | + | let count: i64 = sqlx::query_scalar( | |
| 93 | + | r#" | |
| 94 | + | SELECT COUNT(*) FROM users u | |
| 95 | + | JOIN user_sessions us ON us.user_id = u.id | |
| 96 | + | WHERE u.is_sandbox = TRUE | |
| 97 | + | AND u.sandbox_expires_at > NOW() | |
| 98 | + | AND us.ip_address = $1 | |
| 99 | + | "#, | |
| 100 | + | ) | |
| 101 | + | .bind(ip) | |
| 102 | + | .fetch_one(&mut *conn) | |
| 103 | + | .await?; | |
| 104 | + | ||
| 85 | 105 | sqlx::query("SELECT pg_advisory_unlock($1)") | |
| 86 | - | .bind(key) | |
| 87 | - | .execute(pool) | |
| 106 | + | .bind(lock_key) | |
| 107 | + | .execute(&mut *conn) | |
| 88 | 108 | .await?; | |
| 89 | - | Ok(()) | |
| 109 | + | ||
| 110 | + | Ok(count) | |
| 90 | 111 | } |
| @@ -58,14 +58,12 @@ pub(super) async fn create_sandbox( | |||
| 58 | 58 | AppError::BadRequest("Could not determine client address".to_string()) | |
| 59 | 59 | })?; | |
| 60 | 60 | ||
| 61 | - | // Enforce per-IP concurrent sandbox cap under an advisory lock to prevent | |
| 62 | - | // concurrent requests from both passing the count check before either inserts. | |
| 61 | + | // Enforce per-IP concurrent sandbox cap under an advisory lock. | |
| 62 | + | // Uses a single connection for lock + count + unlock to avoid the pool | |
| 63 | + | // connection mismatch bug with session-level advisory locks. | |
| 63 | 64 | let lock_key = crate::helpers::ip_advisory_lock_key(&ip); | |
| 64 | - | db::advisory_lock(&state.db, lock_key).await?; | |
| 65 | - | ||
| 66 | - | let active = db::users::count_active_sandboxes_by_ip(&state.db, &ip).await?; | |
| 65 | + | let active = db::check_sandbox_cap(&state.db, lock_key, &ip).await?; | |
| 67 | 66 | if active >= constants::SANDBOX_MAX_PER_IP { | |
| 68 | - | db::advisory_unlock(&state.db, lock_key).await?; | |
| 69 | 67 | return Err(AppError::BadRequest( | |
| 70 | 68 | "Too many active sandboxes from this address".to_string(), | |
| 71 | 69 | )); | |
| @@ -111,10 +109,6 @@ pub(super) async fn create_sandbox( | |||
| 111 | 109 | auth::login_user(&session, session_user).await?; | |
| 112 | 110 | auth::track_session(&session, &state.db, user.id, &headers).await?; | |
| 113 | 111 | ||
| 114 | - | // Release the advisory lock now that the session row exists and the cap | |
| 115 | - | // query will see this sandbox account. | |
| 116 | - | db::advisory_unlock(&state.db, lock_key).await?; | |
| 117 | - | ||
| 118 | 112 | // Session ends when the browser closes; the scheduler handles DB cleanup | |
| 119 | 113 | session.set_expiry(Some(tower_sessions::Expiry::OnSessionEnd)); | |
| 120 | 114 |