Clean up orphan images
This commit is contained in:
jude committed
1 parent
8763a8e74a
commit
c3bb015d2f
11 files changed
+496
-2
No files matched your search
@@ -0,0 +1,258 @@
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use chrono::NaiveDateTime;
|
||||
use s3::{creds::Credentials, Bucket, Region};
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
net::TcpListener,
|
||||
sync::Semaphore,
|
||||
task::JoinHandle,
|
||||
};
|
||||
|
||||
use super::*;
|
||||
|
||||
const BASE_URL: &str = "https://images.example.com/bucket";
|
||||
|
||||
fn upload_url(number: u64) -> String {
|
||||
format!("{BASE_URL}/uploads/{number:064x}.png")
|
||||
}
|
||||
|
||||
async fn upload(pool: &Pool<MySql>, number: u64) -> u64 {
|
||||
sqlx::query("INSERT INTO image_uploads (uploader, url) VALUES (123, ?)")
|
||||
.bind(upload_url(number))
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap()
|
||||
.last_insert_id()
|
||||
}
|
||||
|
||||
async fn orphaned_at(pool: &Pool<MySql>, id: u64) -> Option<NaiveDateTime> {
|
||||
sqlx::query_scalar("SELECT orphaned_at FROM image_uploads WHERE id = ?")
|
||||
.bind(id)
|
||||
.fetch_one(pool)
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
async fn expire(pool: &Pool<MySql>, id: u64) {
|
||||
sqlx::query(
|
||||
"UPDATE image_uploads SET orphaned_at = UTC_TIMESTAMP(6) - INTERVAL 2 DAY WHERE id = ?",
|
||||
)
|
||||
.bind(id)
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn reference_tables(pool: &Pool<MySql>) -> (u64, u64) {
|
||||
let guild = sqlx::query("INSERT INTO guilds (guild) VALUES (123)")
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap()
|
||||
.last_insert_id();
|
||||
let channel = sqlx::query("INSERT INTO channels (channel, guild_id) VALUES (123, ?)")
|
||||
.bind(guild)
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap()
|
||||
.last_insert_id();
|
||||
(guild, channel)
|
||||
}
|
||||
|
||||
async fn storage(status: u16) -> (ImageUploadService, JoinHandle<String>) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let endpoint = format!("http://{}", listener.local_addr().unwrap());
|
||||
let task = tokio::spawn(async move {
|
||||
let (mut socket, _) =
|
||||
tokio::time::timeout(Duration::from_secs(5), listener.accept()).await.unwrap().unwrap();
|
||||
let mut received = Vec::new();
|
||||
let mut buffer = [0; 4096];
|
||||
while !received.windows(4).any(|window| window == b"\r\n\r\n") {
|
||||
let count = socket.read(&mut buffer).await.unwrap();
|
||||
assert!(count > 0);
|
||||
received.extend_from_slice(&buffer[..count]);
|
||||
}
|
||||
socket
|
||||
.write_all(
|
||||
format!("HTTP/1.1 {status} Test\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")
|
||||
.as_bytes(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
String::from_utf8(received).unwrap()
|
||||
});
|
||||
let credentials = Credentials::new(Some("test"), Some("test"), None, None, None).unwrap();
|
||||
let mut bucket =
|
||||
Bucket::new("images", Region::Custom { region: "local".into(), endpoint }, credentials)
|
||||
.unwrap()
|
||||
.with_path_style();
|
||||
bucket.set_request_timeout(Some(Duration::from_secs(2)));
|
||||
(
|
||||
ImageUploadService {
|
||||
bucket,
|
||||
public_url: BASE_URL.into(),
|
||||
model: None,
|
||||
slots: Arc::new(Semaphore::new(2)),
|
||||
},
|
||||
task,
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn retention_defaults_and_rejects_invalid_settings() {
|
||||
for value in [None, Some(""), Some(" ")] {
|
||||
assert_eq!(retention_seconds(value).unwrap(), 604800);
|
||||
}
|
||||
assert_eq!(retention_seconds(Some(" 3600 ")).unwrap(), 3600);
|
||||
for value in ["0", "-1", "1.5", "forever", "4294967296"] {
|
||||
assert!(retention_seconds(Some(value)).is_err());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn restricts_deletion_to_generated_keys_in_the_configured_url() {
|
||||
assert_eq!(
|
||||
object_key(BASE_URL, &upload_url(1)),
|
||||
Some(format!("uploads/{:064x}.png", 1).as_str())
|
||||
);
|
||||
for url in [
|
||||
upload_url(1).replace("images.example.com", "other.example.com"),
|
||||
upload_url(1).replace("/bucket/", "/bucket-other/"),
|
||||
upload_url(1).replace("/uploads/", "/private/"),
|
||||
format!("{BASE_URL}/uploads/../private.png"),
|
||||
format!("{BASE_URL}/uploads/{}.png", "g".repeat(64)),
|
||||
format!("{}?download=true", upload_url(1)),
|
||||
format!("{BASE_URL}/uploads/short.png"),
|
||||
] {
|
||||
assert!(object_key(BASE_URL, &url).is_none(), "{url}");
|
||||
}
|
||||
}
|
||||
|
||||
#[sqlx::test]
|
||||
async fn orphan_timestamp_survives_repeated_runs_and_resets_on_reuse(pool: Pool<MySql>) {
|
||||
let id = upload(&pool, 1).await;
|
||||
assert!(orphaned_at(&pool, id).await.is_none());
|
||||
mark_orphans(&pool).await.unwrap();
|
||||
let first = orphaned_at(&pool, id).await.unwrap();
|
||||
mark_orphans(&pool).await.unwrap();
|
||||
assert_eq!(orphaned_at(&pool, id).await, Some(first));
|
||||
|
||||
let (guild, _) = reference_tables(&pool).await;
|
||||
let template =
|
||||
sqlx::query("INSERT INTO reminder_template (guild_id, embed_image_url) VALUES (?, ?)")
|
||||
.bind(guild)
|
||||
.bind(upload_url(1))
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap()
|
||||
.last_insert_id();
|
||||
mark_orphans(&pool).await.unwrap();
|
||||
assert!(orphaned_at(&pool, id).await.is_none());
|
||||
sqlx::query("DELETE FROM reminder_template WHERE id = ?")
|
||||
.bind(template)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
mark_orphans(&pool).await.unwrap();
|
||||
assert!(orphaned_at(&pool, id).await.unwrap() > first);
|
||||
}
|
||||
|
||||
#[sqlx::test]
|
||||
async fn protects_all_image_fields_in_reminders_and_templates(pool: Pool<MySql>) {
|
||||
let (guild, channel) = reference_tables(&pool).await;
|
||||
let mut ids = Vec::new();
|
||||
for (index, column) in IMAGE_COLUMNS.split(", ").enumerate() {
|
||||
for table in ["reminders", "reminder_template"] {
|
||||
let number = ids.len() as u64 + 1;
|
||||
ids.push(upload(&pool, number).await);
|
||||
let query = if table == "reminders" {
|
||||
// Archived/disabled reminders still retain their image URLs.
|
||||
format!(
|
||||
"INSERT INTO reminders (uid, channel_id, `utc_time`, status, enabled, {column})
|
||||
VALUES ('field-{index}', {channel}, UTC_TIMESTAMP(), 'sent', 0, ?)"
|
||||
)
|
||||
} else {
|
||||
format!("INSERT INTO reminder_template (guild_id, {column}) VALUES ({guild}, ?)")
|
||||
};
|
||||
sqlx::query(&query).bind(upload_url(number)).execute(&pool).await.unwrap();
|
||||
}
|
||||
}
|
||||
let unused = upload(&pool, 100).await;
|
||||
// Previously orphaned files are unmarked as soon as cleanup observes a reference.
|
||||
for id in &ids {
|
||||
expire(&pool, *id).await;
|
||||
}
|
||||
mark_orphans(&pool).await.unwrap();
|
||||
for id in ids {
|
||||
assert!(orphaned_at(&pool, id).await.is_none());
|
||||
}
|
||||
assert!(orphaned_at(&pool, unused).await.is_some());
|
||||
}
|
||||
|
||||
#[sqlx::test]
|
||||
async fn only_expired_orphans_are_deleted_and_rows_removed_after_success(pool: Pool<MySql>) {
|
||||
let old = upload(&pool, 1).await;
|
||||
let fresh = upload(&pool, 2).await;
|
||||
expire(&pool, old).await;
|
||||
let (storage, request) = storage(204).await;
|
||||
cleanup(&pool, &storage, 86400).await.unwrap();
|
||||
assert!(request.await.unwrap().starts_with(&format!("DELETE /images/uploads/{:064x}.png ", 1)));
|
||||
let ids: Vec<u64> = sqlx::query_scalar("SELECT id FROM image_uploads ORDER BY id")
|
||||
.fetch_all(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(ids, vec![fresh]);
|
||||
assert!(orphaned_at(&pool, fresh).await.is_some());
|
||||
}
|
||||
|
||||
#[sqlx::test]
|
||||
async fn failed_deletions_keep_the_original_timestamp_and_retry_missing_objects(pool: Pool<MySql>) {
|
||||
let id = upload(&pool, 1).await;
|
||||
expire(&pool, id).await;
|
||||
let original = orphaned_at(&pool, id).await;
|
||||
let (failed_storage, failed_request) = storage(403).await;
|
||||
assert!(cleanup(&pool, &failed_storage, 86400).await.is_err());
|
||||
failed_request.await.unwrap();
|
||||
assert_eq!(orphaned_at(&pool, id).await, original);
|
||||
|
||||
let (retry_storage, retry_request) = storage(404).await;
|
||||
cleanup(&pool, &retry_storage, 86400).await.unwrap();
|
||||
retry_request.await.unwrap();
|
||||
let count: i64 =
|
||||
sqlx::query_scalar("SELECT COUNT(*) FROM image_uploads").fetch_one(&pool).await.unwrap();
|
||||
assert_eq!(count, 0);
|
||||
}
|
||||
|
||||
#[sqlx::test]
|
||||
async fn rechecks_references_before_deleting_a_candidate(pool: Pool<MySql>) {
|
||||
let id = upload(&pool, 1).await;
|
||||
expire(&pool, id).await;
|
||||
let (guild, _) = reference_tables(&pool).await;
|
||||
sqlx::query("INSERT INTO reminder_template (guild_id, avatar) VALUES (?, ?)")
|
||||
.bind(guild)
|
||||
.bind(upload_url(1))
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let (storage, request) = storage(204).await;
|
||||
assert!(!delete_orphan(&pool, &storage, 86400, id).await.unwrap());
|
||||
assert!(orphaned_at(&pool, id).await.is_none());
|
||||
request.abort();
|
||||
}
|
||||
|
||||
#[sqlx::test]
|
||||
async fn foreign_urls_are_retained_without_contacting_s3(pool: Pool<MySql>) {
|
||||
let id = upload(&pool, 1).await;
|
||||
sqlx::query(
|
||||
"UPDATE image_uploads SET url = 'https://elsewhere.example.com/file.png' WHERE id = ?",
|
||||
)
|
||||
.bind(id)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
expire(&pool, id).await;
|
||||
let (storage, request) = storage(204).await;
|
||||
assert!(cleanup(&pool, &storage, 86400).await.is_err());
|
||||
assert!(orphaned_at(&pool, id).await.is_some());
|
||||
request.abort();
|
||||
}
|
||||
Reference in new issue
Block a user