fix(files): check the quota BEFORE the destructive blob write

Self-review of this branch caught a data-loss bug I introduced with the
per-app / group files ceilings: the quota check ran AFTER the blob write.

`write_atomic_at` renames the new bytes over the FINAL path, so by the time
the ceiling refused an update the previous bytes were already gone. The
per-app path then unlinked the blob (row survives, file destroyed — every
read 404s); the group path left the new bytes in place under the old row's
checksum (every read fails `Corrupted`). Either way a user permanently lost
a file merely by exceeding a quota — a strictly worse outcome than the
unchecked-update bypass the ceiling was added to close.

The check now precedes the write on create AND update, per-app and group. As
a bonus this stops an over-quota caller driving unbounded write+unlink disk
churn: the ceiling now bounds I/O, not just stored bytes. The blob still goes
down inside the transaction, under the advisory lock, so a rollback unlinks
it and nothing can reference it in between.

The original test passed against the bug: it asserted the refused update did
not change the stored byte TOTAL — true, while the blob was already
destroyed. The regression test now reads the file back through the
checksum-verifying `FsFilesRepo::get`, and was confirmed to fail with
`Corrupted` against the old ordering.

Also in the KV writer (same file): drop the redundant pre-read on the hottest
write path. `set`/`set_if` did a SELECT purely to learn whether the write
would add a row, when the upsert already returns the previous value. Check
after the insert instead (`>` not `>=`) and let the transaction roll back on
refusal — identical outcome, one round-trip fewer, and an update pays nothing
at all. `kv_repo::get_on` and `FsFilesRepo::final_path` are now unused and
deleted; `FilesRepo::delete` delegates to the `delete_meta_on` + `unlink_blob`
helpers rather than re-implementing them.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
MechaCat02
2026-07-14 21:11:49 +02:00
parent 735b2daedf
commit fa79b4ee45
4 changed files with 103 additions and 99 deletions

View File

@@ -135,15 +135,21 @@ impl PostgresKvWriter {
Self { pool, max_rows }
}
/// Refuse a key-adding write once the app is at its ceiling. Runs on the
/// writing transaction's connection.
async fn check_rows(
/// Refuse a key-adding write once the app is over its ceiling.
///
/// Called AFTER the insert, so the count already includes the new key — hence
/// `>` rather than `>=`. Checking afterwards lets us reuse the previous value
/// the upsert already returns, instead of paying a second SELECT just to learn
/// whether the write adds a row; the transaction rolls the insert back on
/// refusal, so the outcome is identical. `kv::set` is the hottest write path in
/// the system and does not need a redundant round-trip.
async fn check_rows_after_insert(
&self,
conn: &mut PgConnection,
app_id: picloud_shared::AppId,
) -> Result<(), KvError> {
let count = kv_repo::count_rows_on(&mut *conn, app_id).await?;
if count >= self.max_rows {
if count > self.max_rows {
return Err(KvError::QuotaExceeded {
limit: self.max_rows,
actual: count,
@@ -163,14 +169,13 @@ impl KvWriter for PostgresKvWriter {
value: Value,
) -> Result<Option<Value>, KvError> {
let mut tx = self.pool.begin().await.map_err(|e| tx_err(&e))?;
// A NEW key consumes a row against the app's ceiling; an update does not.
if kv_repo::get_on(&mut *tx, cx.app_id, collection, key)
.await?
.is_none()
{
self.check_rows(&mut tx, cx.app_id).await?;
}
let previous = kv_repo::set_on(&mut *tx, cx.app_id, collection, key, value.clone()).await?;
// `previous.is_none()` means the upsert INSERTED, i.e. this write consumed
// a row against the app's ceiling. An update is net-zero and pays nothing —
// not even the COUNT.
if previous.is_none() {
self.check_rows_after_insert(&mut tx, cx.app_id).await?;
}
let op = if previous.is_some() {
"update"
} else {
@@ -200,18 +205,16 @@ impl KvWriter for PostgresKvWriter {
"insert"
};
let mut tx = self.pool.begin().await.map_err(|e| tx_err(&e))?;
// A swap only ADDS a key when the precondition was "absent" AND the key
// really is absent — with a `Some` precondition against an absent key the
// swap cannot happen at all, so it consumes nothing.
if expected.is_none()
&& kv_repo::get_on(&mut *tx, cx.app_id, collection, key)
.await?
.is_none()
{
self.check_rows(&mut tx, cx.app_id).await?;
}
// A swap only ADDS a key when the precondition was "absent" AND it actually
// inserted. With a `Some` precondition it is an UPDATE (net-zero rows), and
// against an absent key it cannot swap at all — either way, nothing to
// charge. So the ceiling is checked only on an insert that really happened.
let inserted = expected.is_none();
let swapped =
kv_repo::set_if_on(&mut *tx, cx.app_id, collection, key, expected, new.clone()).await?;
if swapped && inserted {
self.check_rows_after_insert(&mut tx, cx.app_id).await?;
}
if swapped {
let event = kv_event(op, collection, key, Some(new), None);
emit_on(&mut tx, cx, &event)
@@ -245,6 +248,15 @@ impl KvWriter for PostgresKvWriter {
// Best-effort (repo trait + emitter trait)
// ----------------------------------------------------------------------------
/// A `files_event` that failed to serialize. Never silently dropped — that is the
/// exact class of loss this branch exists to remove.
fn log_event_build_failure(op: &str, e: &str) {
tracing::error!(
error = %e, source = "files", op, event_emit_failure = true,
"could not build the file event — the write committed but its triggers will not fire"
);
}
/// Emit, downgrading a failure to an error-level log. The write has already
/// committed by this point, so there is nothing useful to return to the caller —
/// which is exactly the durability hole `PostgresKvWriter` exists to close.
@@ -1447,15 +1459,21 @@ impl FilesWriter for PostgresFilesWriter {
let owner = files_repo::app_owner_dir(cx.app_id);
let id = Uuid::new_v4();
let size = i64::try_from(new.data.len()).unwrap_or(i64::MAX);
// The bytes go down first — they cannot be part of the transaction.
let checksum = files_repo::write_atomic_at(&self.root, &owner, collection, id, &new.data)?;
let committed = async {
let mut tx = self.pool.begin().await.map_err(|e| files_tx_err(&e))?;
lock_app_quota(&mut tx, cx.app_id, "files")
.await
.map_err(|e| files_tx_err(&e))?;
// Quota FIRST, before a single byte reaches the disk. Writing the blob
// and only then refusing would let an over-quota caller drive unbounded
// write+unlink churn — the ceiling would bound storage but not I/O.
self.check_bytes(&mut tx, cx.app_id, None, size).await?;
// The bytes cannot join the transaction, so they go down inside it,
// under the lock: a rollback below unlinks them, and nothing else can
// reference them in the meantime.
let checksum =
files_repo::write_atomic_at(&self.root, &owner, collection, id, &new.data)?;
let meta = files_repo::insert_meta_on(
&mut *tx,
cx.app_id,
@@ -1502,22 +1520,22 @@ impl FilesWriter for PostgresFilesWriter {
return Ok(false);
};
let size = i64::try_from(upd.data.len()).unwrap_or(i64::MAX);
let checksum = files_repo::write_atomic_at(&self.root, &owner, collection, id, &upd.data)?;
let mut tx = self.pool.begin().await.map_err(|e| files_tx_err(&e))?;
lock_app_quota(&mut tx, cx.app_id, "files")
.await
.map_err(|e| files_tx_err(&e))?;
// Growing an existing blob consumes disk too — an update that skipped the
// ceiling would be a free bypass (exactly the group-files bug). The
// projection subtracts the blob being replaced, so a same-size-or-smaller
// update near the cap still goes through.
if let Err(e) = self.check_bytes(&mut tx, cx.app_id, Some(id), size).await {
// The new bytes are already on disk but no row will reference them.
drop(tx);
files_repo::unlink_blob(&self.root, &owner, collection, id);
return Err(e);
}
// ceiling would be a free bypass. The projection subtracts the blob being
// replaced, so a same-size-or-smaller update near the cap still goes through.
//
// **This MUST precede the blob write.** `write_atomic_at` renames over the
// FINAL path, destroying the previous bytes. Refusing afterwards would
// leave the row pointing at a file that is either gone (if we unlink) or
// has the wrong checksum (if we don't) — i.e. a user would permanently
// lose a file merely by exceeding a quota.
self.check_bytes(&mut tx, cx.app_id, Some(id), size).await?;
let checksum = files_repo::write_atomic_at(&self.root, &owner, collection, id, &upd.data)?;
let updated = files_repo::update_meta_on(
&mut *tx,
cx.app_id,
@@ -1591,8 +1609,9 @@ impl FilesWriter for BestEffortFilesWriter {
new: NewFile,
) -> Result<Uuid, FilesError> {
let meta = self.repo.create(cx.app_id, collection, new).await?;
if let Ok(event) = files_event("create", collection, &meta, None) {
best_effort_emit(&*self.events, cx, event).await;
match files_event("create", collection, &meta, None) {
Ok(event) => best_effort_emit(&*self.events, cx, event).await,
Err(e) => log_event_build_failure("create", &e),
}
Ok(meta.id)
}
@@ -1921,8 +1940,9 @@ impl GroupFilesWriter for BestEffortGroupFilesWriter {
});
}
let meta = self.repo.create(group_id, collection, new).await?;
if let Ok(event) = files_event("create", collection, &meta, None) {
self.emit(cx, group_id, event).await;
match files_event("create", collection, &meta, None) {
Ok(event) => self.emit(cx, group_id, event).await,
Err(e) => log_event_build_failure("create", &e),
}
Ok(meta.id)
}
@@ -1934,8 +1954,22 @@ impl GroupFilesWriter for BestEffortGroupFilesWriter {
collection: &str,
id: Uuid,
upd: FileUpdate,
_quota: GroupFilesQuota,
quota: GroupFilesQuota,
) -> Result<bool, GroupFilesError> {
// Mirror the transactional writer's ceiling, so a unit test can catch its
// removal. (Not serialized — that is what the Postgres writer is for.)
if let Some(prev) = self.repo.head(group_id, collection, id).await? {
let used = self.repo.total_bytes(group_id).await?;
let projected = used
.saturating_sub(prev.size)
.saturating_add(upd.data.len() as u64);
if projected > quota.max_total_bytes {
return Err(GroupFilesError::QuotaExceeded {
limit: quota.max_total_bytes,
actual: projected,
});
}
}
let Some(FileUpdated { new, prev }) =
self.repo.update(group_id, collection, id, upd).await?
else {

View File

@@ -209,11 +209,6 @@ impl FsFilesRepo {
}
Ok(())
}
fn final_path(&self, app_id: AppId, collection: &str, id: Uuid) -> PathBuf {
final_path_at(&self.config.root, &app_owner_dir(app_id), collection, id)
}
fn write_atomic(
&self,
app_id: AppId,
@@ -468,42 +463,14 @@ impl FilesRepo for FsFilesRepo {
id: Uuid,
) -> Result<Option<FileMeta>, FilesRepoError> {
Self::guard_collection(collection)?;
// SELECT + DELETE in one tx; unlink afterwards (outside the tx).
let mut tx = self.pool.begin().await?;
let row: Option<FileRow> = sqlx::query_as(
"SELECT id, collection, name, content_type, size_bytes, \
checksum_sha256, created_at, updated_at \
FROM files WHERE app_id = $1 AND collection = $2 AND id = $3 \
FOR UPDATE",
)
.bind(app_id.into_inner())
.bind(collection)
.bind(id)
.fetch_optional(&mut *tx)
.await?;
let Some(row) = row else {
tx.rollback().await?;
return Ok(None);
};
sqlx::query("DELETE FROM files WHERE app_id = $1 AND collection = $2 AND id = $3")
.bind(app_id.into_inner())
.bind(collection)
.bind(id)
.execute(&mut *tx)
.await?;
tx.commit().await?;
// Row is gone; unlink the bytes. A failure here leaves an orphan
// file (reclaimed by a future sweep) — not fatal.
let path = self.final_path(app_id, collection, id);
if let Err(e) = std::fs::remove_file(&path) {
if e.kind() != std::io::ErrorKind::NotFound {
tracing::warn!(path = %path.display(), error = %e, "files: unlink after delete failed (orphan)");
}
// `DELETE ... RETURNING` is one statement, so the old SELECT-FOR-UPDATE
// + DELETE pair is unnecessary. Unlink only AFTER the row is gone: the
// reverse order would destroy the bytes of a row that a failure keeps.
let meta = delete_meta_on(&self.pool, app_id, collection, id).await?;
if meta.is_some() {
unlink_blob(&self.config.root, &app_owner_dir(app_id), collection, id);
}
Ok(Some(row.into_meta()))
Ok(meta)
}
async fn list(

View File

@@ -99,27 +99,6 @@ const KV_LIST_DEFAULT_LIMIT: u32 = 100;
// trigger fan-out on ONE connection inside ONE transaction.
// ----------------------------------------------------------------------------
pub(crate) async fn get_on<'c, E>(
exec: E,
app_id: AppId,
collection: &str,
key: &str,
) -> Result<Option<serde_json::Value>, KvRepoError>
where
E: sqlx::PgExecutor<'c>,
{
let row: Option<(serde_json::Value,)> = sqlx::query_as(
"SELECT value FROM kv_entries \
WHERE app_id = $1 AND collection = $2 AND key = $3",
)
.bind(app_id.into_inner())
.bind(collection)
.bind(key)
.fetch_optional(exec)
.await?;
Ok(row.map(|(v,)| v))
}
/// Upsert. Returns the previous value, so the caller knows whether this was an
/// `insert` or an `update` for the emitted `ServiceEvent`.
pub(crate) async fn set_on<'c, E>(

View File

@@ -20,6 +20,7 @@ use picloud_manager_core::atomic_write::{
GroupKvTarget, GroupKvWriter, KvWriter, PostgresFilesWriter, PostgresGroupDocsWriter,
PostgresGroupFilesWriter, PostgresGroupKvWriter, PostgresKvWriter,
};
use picloud_manager_core::files_repo::{FilesConfig, FilesRepo, FsFilesRepo};
use picloud_manager_core::quota::GroupWriteQuota;
use picloud_shared::{
AppId, ExecutionId, FileUpdate, FilesError, GroupDocsError, GroupFilesError, GroupId,
@@ -734,6 +735,29 @@ async fn concurrent_uploads_cannot_push_an_app_past_its_disk_ceiling() {
stored,
"the refused update must not have changed the stored total"
);
// …and, crucially, the file must SURVIVE. `write_atomic_at` renames over the
// final path, so a quota check that ran AFTER the blob write would have
// destroyed the original bytes — the user would permanently lose a file
// merely by exceeding a quota. Read it back through the repo, which verifies
// the stored checksum, so both "bytes gone" and "bytes replaced but metadata
// still has the old checksum" fail here.
let repo = FsFilesRepo::new(
pool.clone(),
FilesConfig {
root: root.clone(),
max_file_size_bytes: 1024 * 1024,
},
);
let survivor = repo
.get(AppId::from(f.app), "assets", created[0])
.await
.expect("the refused update must not corrupt or destroy the existing file")
.expect("the file must still exist");
assert_eq!(
survivor,
vec![b'x'; 100],
"the original bytes must be intact — a refused update must not touch the blob"
);
cleanup(&pool, f.app).await;
let _ = std::fs::remove_dir_all(&root);