259 lines
8.9 KiB
Rust
259 lines
8.9 KiB
Rust
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();
|
|
}
|