Skip to main content

vector_core/
blossom_capabilities.rs

1//! Per-(server, mime, encrypted) Blossom capability cache and routing.
2//!
3//! Uploads report outcomes here; `rank_servers()` reorders the enabled
4//! list so known-good servers are tried first.
5
6use std::collections::HashMap;
7use serde::Serialize;
8
9pub const OUTCOME_ACCEPTED: u8 = 1;
10pub const OUTCOME_REJECTED_MIME: u8 = 2;
11/// Seeded by a 413-only rejection (no successful upload yet). Kept
12/// distinct from ACCEPTED so the UI doesn't claim an empty accepted state.
13pub const OUTCOME_SIZE_ONLY: u8 = 3;
14
15/// Rows older than this are routed as "unknown" and re-probed. Server
16/// policies drift (limit bumps, MIME allow-list edits) so we refresh.
17pub const STALE_AFTER_SECS: i64 = 4 * 24 * 3600;
18
19#[derive(Clone, Debug, Serialize)]
20pub struct CapabilityEntry {
21    pub mime_type: String,
22    /// Encrypted ciphertext rarely passes a server's content-sniff even
23    /// when the declared MIME is allowed. Rows split on this flag so the
24    /// two contexts learn independently.
25    pub is_encrypted: bool,
26    pub outcome: u8,            // 1=accepted, 2=rejected_mime, 3=size_only
27    pub max_accepted_size: u64,
28    /// Smallest size we've seen this server reject with HTTP 413, if any.
29    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
43// ============================================================================
44// Storage
45// ============================================================================
46
47/// Record a successful upload. Bumps `max_accepted_size`; clears
48/// `min_rejected_size` if reality contradicts it (the old rejection
49/// was a flake, not a policy). `session` pins the write to the
50/// account that started the upload so a mid-flight swap can't bleed.
51pub 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
80/// Mark this `(server, mime, encrypted)` triple as MIME-rejected.
81/// Future uploads route around the server.
82pub 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
104/// Record an HTTP 413. Tracks the smallest rejected size; combined with
105/// `max_accepted_size` this gives pre-flight "too large" feedback.
106/// Outcome is `SIZE_ONLY` (not `REJECTED_MIME`) — smaller blobs of the
107/// same MIME may still succeed.
108pub 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    // ON CONFLICT keeps any existing `outcome` (especially ACCEPTED) — a
120    // 413 above the known accepted size still leaves smaller blobs viable.
121    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
133/// True iff a row exists for `(server, mime, encrypted)` and is younger
134/// than `STALE_AFTER_SECS`. Used by the probe scheduler to skip rows
135/// we already have current data for.
136pub 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
153/// Drop every cached row for `server_url`. Called on hard-remove so a
154/// later re-add starts with a clean slate.
155pub 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
165/// All rows for `server_url`, most-recent first. Renders the info dialog.
166pub 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// ============================================================================
195// Routing
196// ============================================================================
197
198/// Per-server snapshot used for tier classification.
199#[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
206/// Pure tier-classification (DB-free, unit-testable). Reorders into
207/// known-good → unknown → too-large → MIME-rejected, stable within tier.
208pub 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                // At or above the known 413 ceiling, demote behind unknown.
223                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
244/// Reorder `servers` for an upload of `(mime, encrypted, size_bytes)`.
245pub 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
251/// Pre-flight check: is there any enabled server we haven't already
252/// learned will reject this size/MIME/context? Unknown servers count
253/// as "likely accepts" so we stay optimistic.
254pub 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    // Stale rows route as unknown.
279    let cutoff = now_secs().saturating_sub(STALE_AFTER_SECS);
280    // rusqlite doesn't accept slices in IN — build the clause manually.
281    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
314// ============================================================================
315// Outcome classification
316// ============================================================================
317
318/// HTTP 413 = "blob exceeds size cap". Drives `min_rejected_size`.
319pub fn is_size_rejection(http_status: Option<u16>) -> bool {
320    matches!(http_status, Some(413))
321}
322
323/// Classify an upload error as a permanent MIME rejection.
324///
325/// 415 is BUD-02's canonical signal. 401/402 are also treated permanent
326/// (Vector's auth is fixed-shape, so 401 won't change; 402 means paid
327/// server, we don't pay). 408/429/413/409 are explicitly NOT permanent.
328/// For 5xx and other 4xx with no clear status signal we fall back to a
329/// narrow body-keyword check — non-compliant servers (e.g. nostrcheck
330/// returning `500 "could not be processed"`, blossom.band's 400 sniff
331/// mismatch) would otherwise force a re-discover on every upload.
332pub 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            // Only the "could not process" idiom is permanent on 5xx —
341            // everything else stays transient (could be a temporary outage).
342            let lower = error_msg.to_ascii_lowercase();
343            return lower.contains("could not be processed")
344                || lower.contains("cannot be processed");
345        }
346    }
347    // Other 4xx or unknown status: substring-match body hints.
348    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        // Full fn needs DB access; this just sanity-checks the data shape.
459        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        // nostrcheck-shape: 500 + "could not be processed".
496        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        // blossom.band-shape: 400 + body says the bytes didn't match the declared MIME.
505        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        // "file type" / "content-type" in a transient body must not demote permanently.
520        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}