1use std::collections::HashMap;
7use serde::Serialize;
8
9pub const OUTCOME_ACCEPTED: u8 = 1;
10pub const OUTCOME_REJECTED_MIME: u8 = 2;
11pub const OUTCOME_SIZE_ONLY: u8 = 3;
14
15pub const STALE_AFTER_SECS: i64 = 4 * 24 * 3600;
18
19#[derive(Clone, Debug, Serialize)]
20pub struct CapabilityEntry {
21 pub mime_type: String,
22 pub is_encrypted: bool,
26 pub outcome: u8, pub max_accepted_size: u64,
28 pub min_rejected_size: Option<u64>,
30 pub updated_at: i64,
31}
32
33fn now_secs() -> i64 {
34 std::time::SystemTime::now()
35 .duration_since(std::time::UNIX_EPOCH)
36 .map(|d| d.as_secs() as i64).unwrap_or(0)
37}
38
39fn norm_url(url: &str) -> String {
40 url.trim().trim_end_matches('/').to_lowercase()
41}
42
43pub fn record_accepted(
52 server_url: &str,
53 mime_type: &str,
54 is_encrypted: bool,
55 size_bytes: u64,
56) -> Result<(), String> {
57 let conn = crate::db::get_write_connection_guard_static()?;
58 let server = norm_url(server_url);
59 let mime = mime_type.to_lowercase();
60 let enc = if is_encrypted { 1i64 } else { 0i64 };
61 let now = now_secs();
62 conn.execute(
63 "INSERT INTO blossom_server_capabilities
64 (server_url, mime_type, is_encrypted, outcome, max_accepted_size, updated_at)
65 VALUES (?1, ?2, ?3, ?4, ?5, ?6)
66 ON CONFLICT(server_url, mime_type, is_encrypted) DO UPDATE SET
67 outcome = ?4,
68 max_accepted_size = MAX(max_accepted_size, ?5),
69 min_rejected_size = CASE
70 WHEN min_rejected_size IS NOT NULL AND min_rejected_size <= ?5
71 THEN NULL
72 ELSE min_rejected_size
73 END,
74 updated_at = ?6",
75 rusqlite::params![server, mime, enc, OUTCOME_ACCEPTED as i64, size_bytes as i64, now],
76 ).map_err(|e| format!("Failed to record accepted capability: {}", e))?;
77 Ok(())
78}
79
80pub fn record_rejected_mime(
83 server_url: &str,
84 mime_type: &str,
85 is_encrypted: bool,
86) -> Result<(), String> {
87 let conn = crate::db::get_write_connection_guard_static()?;
88 let server = norm_url(server_url);
89 let mime = mime_type.to_lowercase();
90 let enc = if is_encrypted { 1i64 } else { 0i64 };
91 let now = now_secs();
92 conn.execute(
93 "INSERT INTO blossom_server_capabilities
94 (server_url, mime_type, is_encrypted, outcome, max_accepted_size, updated_at)
95 VALUES (?1, ?2, ?3, ?4, 0, ?5)
96 ON CONFLICT(server_url, mime_type, is_encrypted) DO UPDATE SET
97 outcome = ?4,
98 updated_at = ?5",
99 rusqlite::params![server, mime, enc, OUTCOME_REJECTED_MIME as i64, now],
100 ).map_err(|e| format!("Failed to record rejected capability: {}", e))?;
101 Ok(())
102}
103
104pub fn record_rejected_size(
109 server_url: &str,
110 mime_type: &str,
111 is_encrypted: bool,
112 size_bytes: u64,
113) -> Result<(), String> {
114 let conn = crate::db::get_write_connection_guard_static()?;
115 let server = norm_url(server_url);
116 let mime = mime_type.to_lowercase();
117 let enc = if is_encrypted { 1i64 } else { 0i64 };
118 let now = now_secs();
119 conn.execute(
122 "INSERT INTO blossom_server_capabilities
123 (server_url, mime_type, is_encrypted, outcome, max_accepted_size, min_rejected_size, updated_at)
124 VALUES (?1, ?2, ?3, ?4, 0, ?5, ?6)
125 ON CONFLICT(server_url, mime_type, is_encrypted) DO UPDATE SET
126 min_rejected_size = MIN(COALESCE(min_rejected_size, ?5), ?5),
127 updated_at = ?6",
128 rusqlite::params![server, mime, enc, OUTCOME_SIZE_ONLY as i64, size_bytes as i64, now],
129 ).map_err(|e| format!("Failed to record size rejection: {}", e))?;
130 Ok(())
131}
132
133pub fn has_fresh_capability_for(server_url: &str, mime_type: &str, is_encrypted: bool) -> bool {
137 let conn = match crate::db::get_db_connection_guard_static() {
138 Ok(c) => c,
139 Err(_) => return false,
140 };
141 let server = norm_url(server_url);
142 let mime = mime_type.to_lowercase();
143 let enc = if is_encrypted { 1i64 } else { 0i64 };
144 let cutoff = now_secs().saturating_sub(STALE_AFTER_SECS);
145 conn.query_row(
146 "SELECT 1 FROM blossom_server_capabilities
147 WHERE server_url = ?1 AND mime_type = ?2 AND is_encrypted = ?3 AND updated_at >= ?4",
148 rusqlite::params![server, mime, enc, cutoff],
149 |_| Ok(()),
150 ).is_ok()
151}
152
153pub fn purge_server(server_url: &str) -> Result<usize, String> {
156 let conn = crate::db::get_write_connection_guard_static()?;
157 let server = norm_url(server_url);
158 let n = conn.execute(
159 "DELETE FROM blossom_server_capabilities WHERE server_url = ?1",
160 rusqlite::params![server],
161 ).map_err(|e| format!("Failed to purge capabilities for {}: {}", server, e))?;
162 Ok(n)
163}
164
165pub fn list_for_server(server_url: &str) -> Result<Vec<CapabilityEntry>, String> {
167 let conn = crate::db::get_db_connection_guard_static()?;
168 let server = norm_url(server_url);
169 let mut stmt = conn.prepare(
170 "SELECT mime_type, is_encrypted, outcome, max_accepted_size, min_rejected_size, updated_at
171 FROM blossom_server_capabilities
172 WHERE server_url = ?1
173 ORDER BY updated_at DESC"
174 ).map_err(|e| format!("Prepare failed: {}", e))?;
175 let rows = stmt.query_map(rusqlite::params![server], |row| {
176 let enc: i64 = row.get(1)?;
177 let min_rej: Option<i64> = row.get(4)?;
178 Ok(CapabilityEntry {
179 mime_type: row.get(0)?,
180 is_encrypted: enc != 0,
181 outcome: row.get::<_, i64>(2)? as u8,
182 max_accepted_size: row.get::<_, i64>(3)? as u64,
183 min_rejected_size: min_rej.map(|v| v as u64),
184 updated_at: row.get(5)?,
185 })
186 }).map_err(|e| format!("Query failed: {}", e))?;
187 let mut out = Vec::new();
188 for r in rows {
189 if let Ok(entry) = r { out.push(entry); }
190 }
191 Ok(out)
192}
193
194#[derive(Clone, Copy, Debug, PartialEq)]
200pub struct CapabilityState {
201 pub outcome: u8,
202 pub max_accepted_size: u64,
203 pub min_rejected_size: Option<u64>,
204}
205
206pub fn classify(
209 cache: &HashMap<String, CapabilityState>,
210 servers: Vec<String>,
211 size_bytes: u64,
212) -> Vec<String> {
213 let mut known_good = Vec::new();
214 let mut unknown = Vec::new();
215 let mut too_large = Vec::new();
216 let mut mime_rejected = Vec::new();
217 for s in servers {
218 let key = norm_url(&s);
219 match cache.get(&key) {
220 Some(st) if st.outcome == OUTCOME_REJECTED_MIME => mime_rejected.push(s),
221 Some(st) => {
222 if let Some(min_rej) = st.min_rejected_size {
224 if size_bytes >= min_rej {
225 too_large.push(s);
226 continue;
227 }
228 }
229 if size_bytes <= st.max_accepted_size {
230 known_good.push(s);
231 } else {
232 unknown.push(s);
233 }
234 }
235 None => unknown.push(s),
236 }
237 }
238 known_good.extend(unknown);
239 known_good.extend(too_large);
240 known_good.extend(mime_rejected);
241 known_good
242}
243
244pub fn rank_servers(servers: Vec<String>, mime: &str, is_encrypted: bool, size_bytes: u64) -> Vec<String> {
246 if servers.is_empty() { return servers; }
247 let cache = load_cache_for(&servers, mime, is_encrypted).unwrap_or_default();
248 classify(&cache, servers, size_bytes)
249}
250
251pub fn any_server_likely_accepts(servers: &[String], mime: &str, is_encrypted: bool, size_bytes: u64) -> bool {
255 if servers.is_empty() { return false; }
256 let cache = match load_cache_for(servers, mime, is_encrypted) { Ok(c) => c, Err(_) => return true };
257 for s in servers {
258 let key = norm_url(s);
259 match cache.get(&key) {
260 Some(st) if st.outcome == OUTCOME_REJECTED_MIME => continue,
261 Some(st) => {
262 if let Some(min_rej) = st.min_rejected_size {
263 if size_bytes >= min_rej { continue; }
264 }
265 return true;
266 }
267 None => return true,
268 }
269 }
270 false
271}
272
273fn load_cache_for(servers: &[String], mime: &str, is_encrypted: bool) -> Result<HashMap<String, CapabilityState>, String> {
274 if servers.is_empty() { return Ok(HashMap::new()); }
275 let conn = crate::db::get_db_connection_guard_static()?;
276 let mime_lower = mime.to_lowercase();
277 let enc: i64 = if is_encrypted { 1 } else { 0 };
278 let cutoff = now_secs().saturating_sub(STALE_AFTER_SECS);
280 let placeholders = servers.iter().enumerate()
282 .map(|(i, _)| format!("?{}", i + 4)).collect::<Vec<_>>().join(",");
283 let sql = format!(
284 "SELECT server_url, outcome, max_accepted_size, min_rejected_size
285 FROM blossom_server_capabilities
286 WHERE mime_type = ?1 AND is_encrypted = ?2 AND updated_at >= ?3 AND server_url IN ({})",
287 placeholders,
288 );
289 let mut stmt = conn.prepare(&sql).map_err(|e| format!("Prepare failed: {}", e))?;
290 let cutoff_param: i64 = cutoff;
291 let mut params: Vec<&dyn rusqlite::ToSql> = vec![&mime_lower, &enc, &cutoff_param];
292 let normalized: Vec<String> = servers.iter().map(|s| norm_url(s)).collect();
293 for n in normalized.iter() { params.push(n); }
294 let rows = stmt.query_map(rusqlite::params_from_iter(params.iter()), |row| {
295 let url: String = row.get(0)?;
296 let outcome: i64 = row.get(1)?;
297 let max: i64 = row.get(2)?;
298 let min_rej: Option<i64> = row.get(3)?;
299 Ok((url, CapabilityState {
300 outcome: outcome as u8,
301 max_accepted_size: max as u64,
302 min_rejected_size: min_rej.map(|v| v as u64),
303 }))
304 }).map_err(|e| format!("Query failed: {}", e))?;
305 let mut out = HashMap::new();
306 for r in rows {
307 if let Ok((url, state)) = r {
308 out.insert(url, state);
309 }
310 }
311 Ok(out)
312}
313
314pub fn is_size_rejection(http_status: Option<u16>) -> bool {
320 matches!(http_status, Some(413))
321}
322
323pub fn is_mime_rejection(http_status: Option<u16>, error_msg: &str) -> bool {
333 if matches!(http_status, Some(415)) { return true; }
334 if matches!(http_status, Some(401)) { return true; }
335 if matches!(http_status, Some(402)) { return true; }
336 if matches!(http_status, Some(413) | Some(409)) { return false; }
337 if matches!(http_status, Some(408) | Some(429)) { return false; }
338 if let Some(s) = http_status {
339 if (500..=504).contains(&s) {
340 let lower = error_msg.to_ascii_lowercase();
343 return lower.contains("could not be processed")
344 || lower.contains("cannot be processed");
345 }
346 }
347 let lower = error_msg.to_ascii_lowercase();
349 let hints = [
350 "could not be processed",
351 "cannot be processed",
352 "unsupported",
353 "file type",
354 "mime",
355 "invalid file",
356 "not allowed",
357 "does not match",
358 "doesn't match",
359 "content-type",
360 ];
361 hints.iter().any(|h| lower.contains(h))
362}
363
364#[cfg(test)]
365mod tests {
366 use super::*;
367
368 fn cache_of(entries: &[(&str, u8, u64)]) -> HashMap<String, CapabilityState> {
369 entries.iter()
370 .map(|(url, o, max)| (norm_url(url), CapabilityState {
371 outcome: *o,
372 max_accepted_size: *max,
373 min_rejected_size: None,
374 }))
375 .collect()
376 }
377
378 fn cache_with_size_cap(entries: &[(&str, u8, u64, u64)]) -> HashMap<String, CapabilityState> {
379 entries.iter()
380 .map(|(url, o, max, min_rej)| (norm_url(url), CapabilityState {
381 outcome: *o,
382 max_accepted_size: *max,
383 min_rejected_size: Some(*min_rej),
384 }))
385 .collect()
386 }
387
388 #[test]
389 fn norm_url_handles_case_and_trailing_slash() {
390 assert_eq!(norm_url("HTTPS://Example.COM/"), "https://example.com");
391 assert_eq!(norm_url(" https://example.com "), "https://example.com");
392 assert_eq!(norm_url("https://example.com"), "https://example.com");
393 }
394
395 #[test]
396 fn classify_empty_input_returns_empty() {
397 let cache = HashMap::new();
398 assert!(classify(&cache, vec![], 1024).is_empty());
399 }
400
401 #[test]
402 fn classify_all_unknown_preserves_order() {
403 let cache = HashMap::new();
404 let servers = vec!["https://a".to_string(), "https://b".to_string(), "https://c".to_string()];
405 assert_eq!(classify(&cache, servers.clone(), 1024), servers);
406 }
407
408 #[test]
409 fn classify_known_good_floats_to_top() {
410 let cache = cache_of(&[("https://b", OUTCOME_ACCEPTED, 10_000)]);
411 let servers = vec!["https://a".to_string(), "https://b".to_string(), "https://c".to_string()];
412 assert_eq!(classify(&cache, servers, 5_000), vec!["https://b", "https://a", "https://c"]);
413 }
414
415 #[test]
416 fn classify_known_good_falls_to_unknown_when_over_size_ceiling() {
417 let cache = cache_of(&[("https://b", OUTCOME_ACCEPTED, 10_000)]);
418 let servers = vec!["https://a".to_string(), "https://b".to_string(), "https://c".to_string()];
419 let out = classify(&cache, servers, 20_000);
420 assert_eq!(out, vec!["https://a", "https://b", "https://c"]);
421 }
422
423 #[test]
424 fn classify_mime_rejected_sinks_to_bottom() {
425 let cache = cache_of(&[("https://b", OUTCOME_REJECTED_MIME, 0)]);
426 let servers = vec!["https://a".to_string(), "https://b".to_string(), "https://c".to_string()];
427 let out = classify(&cache, servers, 1024);
428 assert_eq!(out, vec!["https://a", "https://c", "https://b"]);
429 }
430
431 #[test]
432 fn classify_mixed_tiers_full_ordering() {
433 let cache = cache_of(&[
434 ("https://b", OUTCOME_ACCEPTED, 10_000),
435 ("https://c", OUTCOME_REJECTED_MIME, 0),
436 ("https://d", OUTCOME_ACCEPTED, 1_000),
437 ]);
438 let servers = vec![
439 "https://a".to_string(), "https://b".to_string(),
440 "https://c".to_string(), "https://d".to_string(),
441 ];
442 assert_eq!(
443 classify(&cache, servers, 5_000),
444 vec!["https://b", "https://a", "https://d", "https://c"],
445 );
446 }
447
448 #[test]
449 fn classify_demotes_servers_above_known_size_ceiling() {
450 let cache = cache_with_size_cap(&[("https://b", OUTCOME_ACCEPTED, 10_000, 50_000)]);
451 let servers = vec!["https://a".to_string(), "https://b".to_string(), "https://c".to_string()];
452 let out = classify(&cache, servers, 60_000);
453 assert_eq!(out, vec!["https://a", "https://c", "https://b"]);
454 }
455
456 #[test]
457 fn any_server_likely_accepts_smoke_test() {
458 let cache = cache_with_size_cap(&[
460 ("https://a", OUTCOME_ACCEPTED, 1_000, 10_000),
461 ("https://b", OUTCOME_ACCEPTED, 1_000, 10_000),
462 ]);
463 for (_, st) in &cache {
464 assert!(st.min_rejected_size.unwrap() <= 50_000);
465 }
466 }
467
468 #[test]
469 fn classify_keys_are_normalized_against_cache() {
470 let cache = cache_of(&[("https://b.example.com", OUTCOME_ACCEPTED, 10_000)]);
471 let servers = vec!["https://a".to_string(), "https://B.Example.com/".to_string()];
472 let out = classify(&cache, servers, 5_000);
473 assert_eq!(out, vec!["https://B.Example.com/", "https://a"]);
474 }
475
476 #[test]
477 fn mime_rejection_415_is_always_a_reject() {
478 assert!(is_mime_rejection(Some(415), ""));
479 }
480
481 #[test]
482 fn mime_rejection_413_and_409_never_mime_regardless_of_body() {
483 assert!(!is_mime_rejection(Some(413), "Payload Too Large mime"));
484 assert!(!is_mime_rejection(Some(409), "Conflict file type"));
485 }
486
487 #[test]
488 fn mime_rejection_401_treated_as_permanent() {
489 assert!(is_mime_rejection(Some(401), "Unauthorized"));
490 assert!(is_mime_rejection(Some(401), ""));
491 }
492
493 #[test]
494 fn mime_rejection_5xx_with_body_hint_recorded() {
495 assert!(is_mime_rejection(
497 Some(500),
498 r#"Upload failed with status 500 Internal Server Error: {"status":"error","message":"File could not be processed"}"#,
499 ));
500 }
501
502 #[test]
503 fn mime_rejection_content_type_sniff_mismatch_recorded() {
504 assert!(is_mime_rejection(
506 Some(400),
507 "Upload failed with status 400 Bad Request: Content-Type header does not match the file content, expected application/json",
508 ));
509 }
510
511 #[test]
512 fn mime_rejection_transient_5xx_without_hint_not_recorded() {
513 assert!(!is_mime_rejection(Some(500), "Internal Server Error"));
514 assert!(!is_mime_rejection(Some(503), "service unavailable"));
515 }
516
517 #[test]
518 fn mime_rejection_transient_5xx_with_unrelated_body_keywords_not_recorded() {
519 assert!(!is_mime_rejection(Some(503), "file type detection service down"));
521 assert!(!is_mime_rejection(Some(429), "rate limit reached for content-type uploads"));
522 }
523
524 #[test]
525 fn mime_rejection_payment_required_treated_as_permanent() {
526 assert!(is_mime_rejection(Some(402), ""));
527 assert!(is_mime_rejection(Some(402), "Payment Required"));
528 }
529
530 #[test]
531 fn mime_rejection_no_status_with_body_hint_recorded() {
532 assert!(is_mime_rejection(None, "Unsupported file type"));
533 }
534
535 #[test]
536 fn mime_rejection_no_status_no_body_hint_not_recorded() {
537 assert!(!is_mime_rejection(None, "network error"));
538 assert!(!is_mime_rejection(None, "timeout"));
539 }
540}