teamcraft/services/
quotas.rs

1//! Storage quotas: an optional limit per user (`users.storage_quota_bytes`)
2//! and one for the whole vault (`vault_quota.total_bytes`). Both count the
3//! sizes of the stored file versions – identical content stored twice counts
4//! twice, as it does for the user who uploads it; purged versions free their
5//! share. Without a limit (the default) nothing is checked.
6//!
7//! Every path that stores versions goes through `files::store_uploads`,
8//! which calls `ensure_room` in its transaction; chunked uploads check
9//! their announced size up front, before any bytes are sent.
10
11use serde::Serialize;
12use serde_json::json;
13use sqlx::{PgConnection, PgPool};
14use uuid::Uuid;
15
16use crate::{
17    auth::Actor,
18    domain::Role,
19    error::{AppError, AppResult},
20    services::audit,
21    util,
22};
23
24/// Limit and usage of one user or of the vault, in bytes.
25#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
26pub struct Usage {
27    pub limit: Option<i64>,
28    pub used: i64,
29}
30
31/// Refuses (413) when `pending` more bytes would exceed the actor's quota or
32/// the vault's. Called inside the transaction that stores the versions:
33/// `pending` = 0 there, because the new versions are already inserted and
34/// counted. The quota rows are locked, so concurrent uploads are checked
35/// one after the other (the user row only `FOR NO KEY UPDATE`, which does
36/// not block the foreign-key checks of other inserts).
37pub(crate) async fn ensure_room(
38    conn: &mut PgConnection,
39    user_id: Uuid,
40    pending: i64,
41) -> AppResult<()> {
42    let user_limit: Option<Option<i64>> =
43        sqlx::query_scalar("SELECT storage_quota_bytes FROM users WHERE id = $1 FOR NO KEY UPDATE")
44            .bind(user_id)
45            .fetch_optional(&mut *conn)
46            .await?;
47    if let Some(Some(limit)) = user_limit {
48        let used = user_used(&mut *conn, user_id).await?;
49        check(limit, used, pending, "your storage quota")?;
50    }
51    // Unlocked first: without a vault limit (the common case) uploads are
52    // not serialised at all.
53    let total: Option<i64> = sqlx::query_scalar("SELECT total_bytes FROM vault_quota")
54        .fetch_optional(&mut *conn)
55        .await?
56        .flatten();
57    if total.is_some() {
58        let limit: Option<i64> =
59            sqlx::query_scalar("SELECT total_bytes FROM vault_quota FOR UPDATE")
60                .fetch_optional(&mut *conn)
61                .await?
62                .flatten();
63        if let Some(limit) = limit {
64            let used = vault_used(&mut *conn).await?;
65            check(limit, used, pending, "the storage quota of the vault")?;
66        }
67    }
68    Ok(())
69}
70
71/// The decision of [`ensure_room`]: over the limit after adding `pending`.
72fn check(limit: i64, used: i64, pending: i64, what: &str) -> AppResult<()> {
73    let after = used.saturating_add(pending.max(0));
74    if after > limit {
75        return Err(AppError::PayloadTooLarge(format!(
76            "The upload exceeds {what}: {} of {} would be used. Ask an administrator for more \
77             space or remove old versions.",
78            util::format_bytes(after),
79            util::format_bytes(limit)
80        )));
81    }
82    Ok(())
83}
84
85async fn user_used(conn: &mut PgConnection, user_id: Uuid) -> AppResult<i64> {
86    Ok(sqlx::query_scalar(
87        "SELECT COALESCE(SUM(size_bytes), 0)::BIGINT FROM file_versions WHERE uploaded_by = $1",
88    )
89    .bind(user_id)
90    .fetch_one(&mut *conn)
91    .await?)
92}
93
94async fn vault_used(conn: &mut PgConnection) -> AppResult<i64> {
95    Ok(
96        sqlx::query_scalar("SELECT COALESCE(SUM(size_bytes), 0)::BIGINT FROM file_versions")
97            .fetch_one(&mut *conn)
98            .await?,
99    )
100}
101
102/// Limit and usage of a user.
103pub async fn user_usage(pool: &PgPool, user_id: Uuid) -> AppResult<Usage> {
104    let mut conn = pool.acquire().await?;
105    let limit: Option<i64> =
106        sqlx::query_scalar("SELECT storage_quota_bytes FROM users WHERE id = $1")
107            .bind(user_id)
108            .fetch_optional(&mut *conn)
109            .await?
110            .ok_or_else(|| AppError::NotFound("User not found.".into()))?;
111    Ok(Usage {
112        limit,
113        used: user_used(&mut conn, user_id).await?,
114    })
115}
116
117/// Limit and usage of the whole vault.
118pub async fn vault_usage(pool: &PgPool) -> AppResult<Usage> {
119    let mut conn = pool.acquire().await?;
120    let limit: Option<i64> = sqlx::query_scalar("SELECT total_bytes FROM vault_quota")
121        .fetch_optional(&mut *conn)
122        .await?
123        .flatten();
124    Ok(Usage {
125        limit,
126        used: vault_used(&mut conn).await?,
127    })
128}
129
130fn validate_limit(limit: Option<i64>) -> AppResult<()> {
131    if limit.is_some_and(|l| l < 0) {
132        return Err(AppError::BadRequest(
133            "A storage quota cannot be negative.".into(),
134        ));
135    }
136    Ok(())
137}
138
139/// The actor for the audit entry: none for the CLI (the system actor has no
140/// user row).
141fn audit_actor(actor: &Actor) -> Option<&Actor> {
142    Some(actor).filter(|a| !a.id.is_nil())
143}
144
145fn describe(limit: Option<i64>) -> String {
146    limit.map_or_else(|| "no limit".to_string(), util::format_bytes)
147}
148
149/// Sets (or with `None` removes) the storage quota of a user. Admins only;
150/// stored versions are never removed – a user over the new limit just
151/// cannot store more.
152pub async fn set_user_quota(
153    pool: &PgPool,
154    actor: &Actor,
155    user_id: Uuid,
156    limit: Option<i64>,
157) -> AppResult<()> {
158    actor.require(Role::Admin)?;
159    validate_limit(limit)?;
160    let mut tx = pool.begin().await?;
161    let username: String = sqlx::query_scalar(
162        "UPDATE users SET storage_quota_bytes = $2 WHERE id = $1 RETURNING username",
163    )
164    .bind(user_id)
165    .bind(limit)
166    .fetch_optional(&mut *tx)
167    .await?
168    .ok_or_else(|| AppError::NotFound("User not found.".into()))?;
169    audit::record(
170        &mut tx,
171        audit_actor(actor),
172        "user.quota",
173        "user",
174        Some(user_id),
175        format!("Storage quota of {username} set to {}", describe(limit)),
176        json!({ "storage_quota_bytes": limit }),
177    )
178    .await?;
179    tx.commit().await?;
180    Ok(())
181}
182
183/// Sets (or with `None` removes) the storage quota of the whole vault.
184/// Admins only.
185pub async fn set_vault_quota(pool: &PgPool, actor: &Actor, limit: Option<i64>) -> AppResult<()> {
186    actor.require(Role::Admin)?;
187    validate_limit(limit)?;
188    let mut tx = pool.begin().await?;
189    sqlx::query(
190        "INSERT INTO vault_quota (id, total_bytes) VALUES (TRUE, $1)
191         ON CONFLICT (id) DO UPDATE SET total_bytes = EXCLUDED.total_bytes",
192    )
193    .bind(limit)
194    .execute(&mut *tx)
195    .await?;
196    audit::record(
197        &mut tx,
198        audit_actor(actor),
199        "vault.quota",
200        "vault",
201        None,
202        format!("Storage quota of the vault set to {}", describe(limit)),
203        json!({ "total_bytes": limit }),
204    )
205    .await?;
206    tx.commit().await?;
207    Ok(())
208}
209
210/// Parses the `--mb` argument of the CLI: a number of MiB or `none`.
211pub fn parse_limit_mb(raw: &str) -> Result<Option<i64>, String> {
212    let raw = raw.trim();
213    if raw.eq_ignore_ascii_case("none") || raw.eq_ignore_ascii_case("unlimited") {
214        return Ok(None);
215    }
216    let mb: i64 = raw
217        .parse()
218        .map_err(|_| format!("'{raw}' is neither a number of MiB nor 'none'"))?;
219    if mb < 0 {
220        return Err("a storage quota cannot be negative".into());
221    }
222    mb.checked_mul(1024 * 1024)
223        .map(Some)
224        .ok_or_else(|| format!("{mb} MiB is too large"))
225}
226
227#[cfg(test)]
228mod tests {
229    use super::*;
230
231    #[test]
232    fn limit_is_checked_after_adding_the_pending_bytes() {
233        assert!(check(100, 60, 40, "q").is_ok(), "exactly at the limit");
234        assert!(check(100, 60, 41, "q").is_err());
235        assert!(check(100, 101, 0, "q").is_err(), "already over");
236        assert!(check(100, 100, 0, "q").is_ok());
237        assert!(
238            check(100, 10, -50, "q").is_ok(),
239            "negative pending is ignored"
240        );
241        assert!(check(i64::MAX, i64::MAX, 1, "q").is_ok(), "saturates");
242        match check(1024, 2048, 0, "your storage quota") {
243            Err(AppError::PayloadTooLarge(m)) => assert!(m.contains("your storage quota"), "{m}"),
244            other => panic!("{other:?}"),
245        }
246    }
247
248    #[test]
249    fn cli_limits() {
250        assert_eq!(parse_limit_mb("none"), Ok(None));
251        assert_eq!(parse_limit_mb(" Unlimited "), Ok(None));
252        assert_eq!(parse_limit_mb("0"), Ok(Some(0)));
253        assert_eq!(parse_limit_mb("10"), Ok(Some(10 * 1024 * 1024)));
254        assert!(parse_limit_mb("-1").is_err());
255        assert!(parse_limit_mb("ten").is_err());
256        assert!(parse_limit_mb(&i64::MAX.to_string()).is_err());
257    }
258
259    #[test]
260    fn negative_limits_are_refused() {
261        assert!(validate_limit(Some(-1)).is_err());
262        assert!(validate_limit(Some(0)).is_ok());
263        assert!(validate_limit(None).is_ok());
264        assert_eq!(describe(None), "no limit");
265    }
266}
teamcraft/services/quotas.rs

teamcraft/services/
quotas.rs
