chat: attachment cache, download, and transfer hardening (Phase 3)
CI / check (push) Successful in 2m33s

Phase 3 of docs/chat-hardening-plan.md — attachments can no longer turn
into unbounded memory, bandwidth, decoder, or task pressure (S15 closed;
S14's filename half closed).

Cache and image cost (3A): AttachmentCache now carries encoded- and
decoded-byte budgets (96 MiB / 64 MiB) on top of the count cap, with
per-entry weights, replacement accounting, and oldest-first eviction; an
individually over-budget fetch services any pending Save/Play from the
bytes in hand and is exposed as Evicted instead of retained.
validate_image_bytes prechecks header dimensions (per-side AND a new
14 MP total-pixel limit) before any decode; the renderer only ever
receives a ≤1600 px downscaled RGBA preview whose w*h*4 cost counts
against the decoded budget — originals stay encoded-only for Save.
sanitize_filename strips the bidi/zero-width spoofing set (RTL-override
extension spoof).

Download policy and state (3B): images auto-fetch only when roster-
authored AND declared ≤4 MiB, gated by a new deterministic
AutoFetchBudget (per-author and session request+byte token buckets,
check-then-take, bounded author map) alongside the existing dedup and
four-permit bound. Attachment state is now explicit — absence/Loading/
Ready/Failed/Evicted — driven by a new AttachmentFetchStarted event, so
skipped or evicted images render a "Load image" button instead of an
indefinite "loading…", and repeated clicks can never spawn duplicate
fetch tasks.

Exact transfers and serve store (3C): fetch_blob requires the received
length to equal the declared size (short = local error, overlong =
bounded-read reject, empty keeps meaning "sender no longer has it");
the file picker's unbounded read is replaced by a metadata-prechecked
cap+1 bounded reader; one Arc<Vec<u8>> now backs the UI cache, command
queue, and serve store; served_files is a count- and byte-budgeted FIFO
ServeStore (16 entries / 128 MiB).

37 new tests (568 lib total) including a real two-endpoint loopback
exercising exact/short/overlong/unknown-id transfers. Plan checkboxes
ticked and constant deviations decision-logged. Tests-green-only: the
plan's two-machine field-test section remains open.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
2026-07-18 02:09:30 -04:00
co-authored by Claude Fable 5
parent 8898652349
commit 554b613466
9 changed files with 1305 additions and 147 deletions
+442 -82
View File
@@ -321,13 +321,23 @@ struct ChatEntry {
attachment: Option<crate::files::ChatAttachment>,
}
/// Fetch state of a chat attachment's bytes (session-only).
/// Fetch state of a chat attachment's bytes (session-only). Absence from the
/// cache means NOT FETCHED (auto-fetch skipped or never attempted) — the UI
/// renders a Load/Download button for it, never a loading label (Phase 3B).
#[derive(Debug, Clone)]
enum AttachmentState {
/// Bytes in hand (image decoded-valid, or a file ready to save).
Ready(Vec<u8>),
/// A fetch task is running for this key (core announced it, or the user
/// clicked). Suppresses duplicate fetches on repeated clicks.
Loading,
/// Bytes in hand (image decoded-valid, or a file ready to save). Shared so
/// the cache, save tasks, and the serve path hold one allocation.
Ready(std::sync::Arc<Vec<u8>>),
/// Fetch or decode failed; carries a short reason for the UI.
Failed(String),
/// The bytes were fetched but could not be retained within the cache
/// budgets (or were dropped under budget pressure after any pending
/// save/play was serviced). Rendered like NotFetched: load-on-demand.
Evicted,
}
/// Identifies one fetched attachment by BOTH the authoring peer and the
@@ -341,13 +351,31 @@ type AttachmentKey = (EndpointId, crate::files::AttachmentId);
/// above any realistic on-screen image working set.
const ATTACHMENT_CACHE_CAP: usize = 64;
/// Budget for retained ENCODED attachment bytes (the original file bytes kept
/// for Save). The count cap alone would admit 64 × 25 MiB = 1.6 GiB; this keeps
/// the worst case at a few large attachments plus the normal small working set
/// (Phase 3A).
const ATTACHMENT_CACHE_ENCODED_BUDGET: usize = 96 * 1024 * 1024;
/// Budget for DECODED preview bytes (estimated RGBA cost `w*h*4` of each inline
/// preview, counted even though iced may copy/upload internally). Previews are
/// capped at 1600 px per side (≈10 MiB RGBA worst case), so this comfortably
/// covers the on-screen set while bounding an image-spam session (Phase 3A).
const ATTACHMENT_CACHE_DECODED_BUDGET: usize = 64 * 1024 * 1024;
/// One cached attachment result: its fetch state plus, for ready images, the
/// pre-built iced image handle (built once on arrival, not per redraw — the
/// e917c53 flicker fix). `handle` is `None` for files and failures.
///
/// `encoded` / `decoded` are this entry's weights against the cache byte
/// budgets: the retained encoded length and the preview's estimated RGBA cost.
/// Both are zero for non-`Ready` states.
#[derive(Debug, Clone)]
struct AttachmentEntry {
state: AttachmentState,
handle: Option<iced::widget::image::Handle>,
handle: Option<PreviewHandle>,
encoded: usize,
decoded: usize,
}
/// Bounded store of fetched chat-attachment results, keyed by [`AttachmentKey`].
@@ -367,38 +395,100 @@ struct AttachmentCache {
/// exactly the present keys (one entry each), so it is bounded by `cap`.
order: VecDeque<AttachmentKey>,
cap: usize,
/// Byte budgets and their running totals (Phase 3A). Totals are exactly the
/// sum of the present entries' weights — every remove/replace subtracts the
/// old weights before new ones are counted.
encoded_budget: usize,
decoded_budget: usize,
encoded_total: usize,
decoded_total: usize,
}
impl AttachmentCache {
fn new(cap: usize) -> Self {
fn new(cap: usize, encoded_budget: usize, decoded_budget: usize) -> Self {
Self {
entries: HashMap::new(),
order: VecDeque::new(),
cap: cap.max(1),
encoded_budget,
decoded_budget,
encoded_total: 0,
decoded_total: 0,
}
}
/// Insert or replace an entry. A brand-new key evicts the oldest entries
/// until there is room; replacing an existing key keeps its position (and so
/// its age), only updating the value.
/// Insert or replace an entry, keeping the count cap AND both byte budgets
/// satisfied by evicting oldest entries first. Replacing an existing key
/// keeps its position (and so its age); its old weights are subtracted
/// before the new ones are checked.
///
/// Returns `false` if the entry ALONE exceeds a byte budget: it is not
/// retained (and an existing entry under the key is dropped), leaving no
/// entry for `key`. The caller services any immediate need from the bytes
/// in hand and stores an [`AttachmentState::Evicted`] marker instead.
fn insert(
&mut self,
key: AttachmentKey,
state: AttachmentState,
handle: Option<iced::widget::image::Handle>,
) {
if !self.entries.contains_key(&key) {
while self.entries.len() >= self.cap {
match self.order.pop_front() {
Some(old) => {
self.entries.remove(&old);
}
None => break,
}
handle: Option<PreviewHandle>,
) -> bool {
let encoded = match &state {
AttachmentState::Ready(data) => data.len(),
_ => 0,
};
let decoded = handle_decoded_cost(&state, &handle);
// Replacement: subtract the old weights first so they don't count
// against the new entry's fit.
if let Some(old) = self.entries.get(&key) {
self.encoded_total -= old.encoded;
self.decoded_total -= old.decoded;
}
if encoded > self.encoded_budget || decoded > self.decoded_budget {
// Individually overweight: never retained, budgets never exceeded.
if self.entries.remove(&key).is_some() {
self.order.retain(|k| k != &key);
}
return false;
}
let replacing = self.entries.contains_key(&key);
// Evict oldest (skipping the key being replaced, which keeps its slot)
// until the count cap and both budgets accommodate the new entry.
while self.would_overflow(replacing, encoded, decoded) {
let Some(victim) = self.order.iter().find(|k| **k != key).copied() else {
break;
};
self.remove_entry(&victim);
}
if !replacing {
self.order.push_back(key);
}
self.entries.insert(key, AttachmentEntry { state, handle });
self.encoded_total += encoded;
self.decoded_total += decoded;
self.entries.insert(
key,
AttachmentEntry {
state,
handle,
encoded,
decoded,
},
);
true
}
fn would_overflow(&self, replacing: bool, encoded: usize, decoded: usize) -> bool {
let count_full = !replacing && self.entries.len() >= self.cap;
count_full
|| self.encoded_total + encoded > self.encoded_budget
|| self.decoded_total + decoded > self.decoded_budget
}
fn remove_entry(&mut self, key: &AttachmentKey) {
if let Some(old) = self.entries.remove(key) {
self.encoded_total -= old.encoded;
self.decoded_total -= old.decoded;
self.order.retain(|k| k != key);
}
}
fn get(&self, key: &AttachmentKey) -> Option<&AttachmentState> {
@@ -406,7 +496,9 @@ impl AttachmentCache {
}
fn handle(&self, key: &AttachmentKey) -> Option<&iced::widget::image::Handle> {
self.entries.get(key).and_then(|e| e.handle.as_ref())
self.entries
.get(key)
.and_then(|e| e.handle.as_ref().map(|h| &h.handle))
}
fn is_ready(&self, key: &AttachmentKey) -> bool {
@@ -419,9 +511,16 @@ impl AttachmentCache {
)
}
/// Whether a fetch for this key is already running (dedup for clicks).
fn is_loading(&self, key: &AttachmentKey) -> bool {
matches!(self.get(key), Some(AttachmentState::Loading))
}
fn clear(&mut self) {
self.entries.clear();
self.order.clear();
self.encoded_total = 0;
self.decoded_total = 0;
}
#[cfg(test)]
@@ -430,6 +529,38 @@ impl AttachmentCache {
}
}
/// The decoded-budget weight of an entry: the preview's RGBA estimate for a
/// ready image (`w*h*4` via the handle's known preview dimensions), else zero.
/// The handle is built from the downscaled preview, so its dimensions are the
/// preview's — but iced's `Handle` doesn't expose them, so the caller threads
/// the cost through [`PreviewHandle`].
fn handle_decoded_cost(state: &AttachmentState, handle: &Option<PreviewHandle>) -> usize {
match (state, handle) {
(AttachmentState::Ready(_), Some(h)) => h.decoded_cost,
_ => 0,
}
}
/// An iced image handle built from a downscaled preview, paired with the
/// preview's decoded RGBA cost (iced doesn't expose handle dimensions).
#[derive(Debug, Clone)]
struct PreviewHandle {
handle: iced::widget::image::Handle,
decoded_cost: usize,
}
/// Build the inline preview handle (and its decoded weight) for image bytes, or
/// `None` when the bytes don't validate as a displayable image. The renderer
/// only ever receives the ≤1600 px downscale; originals stay encoded-only.
fn build_preview_handle(bytes: &[u8]) -> Option<PreviewHandle> {
let p = crate::files::decode_preview(bytes)?;
let decoded_cost = crate::files::preview_rgba_cost(p.width, p.height);
Some(PreviewHandle {
handle: iced::widget::image::Handle::from_rgba(p.width, p.height, p.rgba),
decoded_cost,
})
}
/// Cap on retained chat history so a long call can't grow it without bound.
const CHAT_HISTORY_MAX: usize = 300;
@@ -612,11 +743,17 @@ pub enum AppMessage {
/// Open the native picker to attach a file to the chat.
PickAttachmentFile,
/// Result of the attach picker: (filename, bytes), or None if cancelled.
AttachmentFilePicked(Option<(String, Vec<u8>)>),
/// Picker result: `None` = cancelled; `Some((name, None))` = file over the
/// attachment cap or unreadable; `Some((name, Some(bytes)))` = bytes in hand.
AttachmentFilePicked(Option<(String, Option<Vec<u8>>)>),
/// Save (downloading first if needed) a received attachment to disk. Carries
/// the full `(author, id)` key so the correct sender's bytes are fetched and
/// saved even if another peer reused the same attachment id (Tier C F-12).
SaveAttachment(AttachmentKey),
/// Fetch an image that wasn't auto-fetched (over the auto-size cap, budget-
/// skipped, or evicted) for inline display (Phase 3B). Manual, so it may use
/// the full attachment cap; duplicates are suppressed via the Loading state.
LoadImage(AttachmentKey),
/// Open the click-to-enlarge image lightbox for this attachment (author + id).
OpenImageLightbox(AttachmentKey),
/// Close the image lightbox overlay.
@@ -1220,7 +1357,11 @@ impl Default for AppState {
recording: false,
recording_started: None,
chat_messages: Vec::new(),
attachments: AttachmentCache::new(ATTACHMENT_CACHE_CAP),
attachments: AttachmentCache::new(
ATTACHMENT_CACHE_CAP,
ATTACHMENT_CACHE_ENCODED_BUDGET,
ATTACHMENT_CACHE_DECODED_BUDGET,
),
image_lightbox: None,
pending_saves: HashSet::new(),
pending_plays: HashSet::new(),
@@ -2021,25 +2162,40 @@ fn update(state: &mut AppState, message: AppMessage) -> Task<AppMessage> {
}
}
UiEvent::AttachmentReady { from, id, data } => {
// Bytes arrived for this specific (author, id). For images we
// can cache the iced handle now (built once, not per redraw).
// If the user was waiting to save this file, the save dialog
// is opened from update() below by checking pending_saves —
// done lazily so this arm stays simple.
// Bytes arrived for this specific (author, id). For images
// the ≤1600px preview handle is built once here (not per
// redraw); only its downscale ever reaches the renderer.
let key = (from, id);
let handle = crate::files::validate_image_bytes(&data)
.is_some()
.then(|| iced::widget::image::Handle::from_bytes(data.clone()));
let handle = build_preview_handle(&data);
let needs_save = state.pending_saves.remove(&key);
let needs_play = state.pending_plays.remove(&id);
state
.attachments
.insert(key, AttachmentState::Ready(data), handle);
let retained =
state
.attachments
.insert(key, AttachmentState::Ready(data.clone()), handle);
if !retained {
// Individually over-budget: any immediate save/play is
// serviced from the bytes in hand below, and the entry
// is exposed as evicted, never exceeding the budgets.
state
.attachments
.insert(key, AttachmentState::Evicted, None);
}
if needs_play {
play_ready_audio(state, key);
play_audio_data(state, id, &data);
}
if needs_save {
return save_attachment_task(state, key);
return save_attachment_data_task(state, key, data);
}
}
UiEvent::AttachmentFetchStarted { from, id } => {
// A fetch task exists for this key; show a real loading
// state (and dedup further clicks). Never demote Ready.
let key = (from, id);
if !state.attachments.is_ready(&key) {
state
.attachments
.insert(key, AttachmentState::Loading, None);
}
}
UiEvent::AttachmentFailed { from, id, error } => {
@@ -2933,7 +3089,9 @@ fn update(state: &mut AppState, message: AppMessage) -> Task<AppMessage> {
}
}
AppMessage::PickAttachmentFile => {
// Native picker off the UI thread; returns (filename, bytes).
// Native picker off the UI thread; returns (filename, bytes). The
// read is BOUNDED (metadata precheck + cap+1 read, Phase 3C) — an
// oversized pick reports `(name, None)` without buffering the file.
return Task::perform(
async {
let handle = rfd::AsyncFileDialog::new()
@@ -2941,7 +3099,18 @@ fn update(state: &mut AppState, message: AppMessage) -> Task<AppMessage> {
.pick_file()
.await;
match handle {
Some(h) => Some((h.file_name(), h.read().await)),
Some(h) => {
let name = h.file_name();
let path = h.path().to_path_buf();
let bytes = tokio::task::spawn_blocking(move || {
crate::files::read_file_capped(&path)
})
.await
.ok()
.and_then(|r| r.ok())
.flatten();
Some((name, bytes))
}
None => None,
}
},
@@ -2950,6 +3119,13 @@ fn update(state: &mut AppState, message: AppMessage) -> Task<AppMessage> {
}
AppMessage::AttachmentFilePicked(picked) => {
if let Some((name, bytes)) = picked {
let Some(bytes) = bytes else {
state.status_message = format!(
"Can't attach {name} — unreadable or larger than {}.",
crate::files::human_size(crate::files::MAX_ATTACHMENT_BYTES)
);
return Task::none();
};
let size = bytes.len() as u64;
if !crate::files::size_within_cap(size) {
state.status_message = format!(
@@ -2958,6 +3134,7 @@ fn update(state: &mut AppState, message: AppMessage) -> Task<AppMessage> {
);
return Task::none();
}
let bytes = Arc::new(bytes);
let kind = crate::files::classify(&bytes);
// Random 32-byte handle for this attachment.
let id: crate::files::AttachmentId = rand::random();
@@ -2970,12 +3147,13 @@ fn update(state: &mut AppState, message: AppMessage) -> Task<AppMessage> {
// Keep our own bytes locally so we see our own attachment inline
// immediately (others fetch it off the file plane). Cache under
// our own (author, id) key so the render path — which keys the
// line by `from` = self_id — finds them.
// line by `from` = self_id — finds them. The same Arc backs the
// cache, the command queue, and the serve store (Phase 3C).
if let Ok(self_eid) = state.self_id.parse::<EndpointId>() {
let key = (self_eid, id);
let handle = (kind == crate::files::AttachmentKind::Image
&& crate::files::validate_image_bytes(&bytes).is_some())
.then(|| iced::widget::image::Handle::from_bytes(bytes.clone()));
let handle = (kind == crate::files::AttachmentKind::Image)
.then(|| build_preview_handle(&bytes))
.flatten();
state
.attachments
.insert(key, AttachmentState::Ready(bytes.clone()), handle);
@@ -2998,13 +3176,20 @@ fn update(state: &mut AppState, message: AppMessage) -> Task<AppMessage> {
}
}
AppMessage::SaveAttachment(key) => {
// If we already have the bytes, save now; otherwise fetch from the
// sender and save when AttachmentReady arrives (pending_saves). The
// key's author half is the exact sender of the clicked line.
// If we already have the bytes, save now; if a fetch is already
// running, just record the save intent (no duplicate fetch task);
// otherwise fetch from the sender and save when AttachmentReady
// arrives (pending_saves). The key's author half is the exact
// sender of the clicked line.
if state.attachments.is_ready(&key) {
return save_attachment_task(state, key);
} else if state.attachments.is_loading(&key) {
state.pending_saves.insert(key);
} else if let Some(att) = find_attachment_source(state, key) {
state.pending_saves.insert(key);
state
.attachments
.insert(key, AttachmentState::Loading, None);
state.status_message = format!("Downloading {}", att.name);
let _ = state.controller.send(CoreCommand::FetchAttachment {
from: key.0,
@@ -3012,6 +3197,23 @@ fn update(state: &mut AppState, message: AppMessage) -> Task<AppMessage> {
});
}
}
AppMessage::LoadImage(key) => {
// Manual inline-image fetch for a skipped/evicted image. The
// Loading state both drives the label and suppresses duplicates.
if !state.attachments.is_ready(&key)
&& !state.attachments.is_loading(&key)
&& let Some(att) = find_attachment_source(state, key)
{
state
.attachments
.insert(key, AttachmentState::Loading, None);
state.status_message = format!("Loading {}", att.name);
let _ = state.controller.send(CoreCommand::FetchAttachment {
from: key.0,
attachment: att,
});
}
}
AppMessage::AttachmentSaved(msg) => {
if let Some(m) = msg {
state.status_message = m;
@@ -3020,11 +3222,18 @@ fn update(state: &mut AppState, message: AppMessage) -> Task<AppMessage> {
AppMessage::PlayAudio(key) => {
if state.attachments.is_ready(&key) {
play_ready_audio(state, key);
} else if state.attachments.is_loading(&key) {
// A fetch is already running (e.g. a Save click); just record
// the play intent for when the bytes arrive.
state.pending_plays.insert(key.1);
} else if let Some(att) = find_attachment_source(state, key) {
// Repeated clicks while the transfer is pending must not launch
// duplicate fetches. (pending_plays is id-keyed — it's coupled to
// the id-keyed clip player; same-id collisions are cosmetic.)
if state.pending_plays.insert(key.1) {
state
.attachments
.insert(key, AttachmentState::Loading, None);
state.status_message = format!("Loading {}", att.name);
let _ = state.controller.send(CoreCommand::FetchAttachment {
from: key.0,
@@ -3650,9 +3859,17 @@ fn play_ready_audio(state: &mut AppState, key: AttachmentKey) {
let Some(AttachmentState::Ready(data)) = state.attachments.get(&key) else {
return;
};
let id = key.1;
let data = data.clone();
play_audio_data(state, key.1, &data);
}
/// Play audio from bytes in hand — used both for cached entries and for a
/// just-fetched attachment too large to retain (Phase 3A overweight path).
fn play_audio_data(state: &mut AppState, id: crate::files::AttachmentId, data: &Arc<Vec<u8>>) {
if crate::files::is_probably_audio(data) {
let bytes = data.clone();
// The clip player's command channel owns its bytes; one copy here at
// the click, the cache keeps the shared original.
let bytes = data.as_ref().clone();
state.invalid_audio.remove(&id);
state.clip_player.play(id, bytes);
// Apply this clip's effective gain; the command lands after Play so it
@@ -4065,6 +4282,16 @@ fn save_attachment_task(state: &AppState, key: AttachmentKey) -> Task<AppMessage
return Task::none();
};
let data = data.clone();
save_attachment_data_task(state, key, data)
}
/// The dialog+write half of a save, from bytes in hand — also used when a
/// just-fetched attachment was too large to retain in the cache (Phase 3A).
fn save_attachment_data_task(
state: &AppState,
key: AttachmentKey,
data: Arc<Vec<u8>>,
) -> Task<AppMessage> {
let default_name = attachment_default_name(&state.chat_messages, key);
Task::perform(
async move {
@@ -4074,7 +4301,7 @@ fn save_attachment_task(state: &AppState, key: AttachmentKey) -> Task<AppMessage
.save_file()
.await;
match handle {
Some(h) => match std::fs::write(h.path(), &data) {
Some(h) => match std::fs::write(h.path(), data.as_slice()) {
Ok(()) => Some(format!("Saved {}", h.path().display())),
Err(e) => Some(format!("Save failed: {e}")),
},
@@ -7124,10 +7351,33 @@ fn view(state: &AppState) -> Element<'_, AppMessage> {
None => img.into(),
}
}
None => text(format!("🖼 {} — loading…", att.name))
// "loading…" ONLY while a fetch task actually runs
// (Phase 3B). Absent/evicted images — over the
// auto-fetch cap, budget-skipped, or dropped by the
// cache — get an explicit Load button instead of an
// indefinite label.
None if matches!(data, Some(AttachmentState::Loading)) => {
text(format!("🖼 {} — loading…", att.name))
.size(12)
.color(color_subtext)
.into()
}
None => row![
text(format!(
"🖼 {} ({})",
att.name,
crate::files::human_size(att.size)
))
.size(12)
.color(color_subtext)
.into(),
.color(color_text),
button(text("Load image").size(12))
.on_press_maybe(key.map(AppMessage::LoadImage))
.style(b_style(color_blue, color_lavender, color_crust, 6.0,))
.padding(6),
]
.spacing(8)
.align_y(iced::alignment::Vertical::Center)
.into(),
}
} else if crate::files::looks_like_audio_name(&att.name)
&& !state.invalid_audio.contains(&att.id)
@@ -7167,10 +7417,10 @@ fn view(state: &AppState) -> Element<'_, AppMessage> {
.size(12)
.color(color_text),
button(
text(if matches!(data, Some(AttachmentState::Ready(_))) {
"Save"
} else {
"Download"
text(match data {
Some(AttachmentState::Ready(_)) => "Save",
Some(AttachmentState::Loading) => "Downloading…",
_ => "Download",
})
.size(12)
)
@@ -7212,8 +7462,14 @@ fn view(state: &AppState) -> Element<'_, AppMessage> {
.spacing(4)
.into()
} else {
let ready = matches!(data, Some(AttachmentState::Ready(_)));
let btn_label = if ready { "Save" } else { "Download" };
// A repeated click during Loading only re-records the
// save intent (the update arm dedups the fetch), so the
// button stays enabled with an honest label.
let btn_label = match data {
Some(AttachmentState::Ready(_)) => "Save",
Some(AttachmentState::Loading) => "Downloading…",
_ => "Download",
};
row![
text(format!(
"📎 {} ({})",
@@ -9018,17 +9274,35 @@ mod tests {
(SecretKey::generate().public(), id)
}
/// A `Ready` state around bytes (tests only).
fn ready(bytes: Vec<u8>) -> AttachmentState {
AttachmentState::Ready(std::sync::Arc::new(bytes))
}
/// A cache with byte budgets too large to bind, isolating count-cap tests.
fn count_only_cache(cap: usize) -> AttachmentCache {
AttachmentCache::new(cap, usize::MAX, usize::MAX)
}
/// A dummy preview handle carrying an explicit decoded weight (tests only).
fn preview(decoded_cost: usize) -> super::PreviewHandle {
super::PreviewHandle {
handle: iced::widget::image::Handle::from_rgba(1, 1, vec![0u8; 4]),
decoded_cost,
}
}
#[test]
fn attachment_cache_evicts_oldest_when_full() {
let mut cache = AttachmentCache::new(2);
let mut cache = count_only_cache(2);
let k1 = peer_key([1u8; 32]);
let k2 = peer_key([2u8; 32]);
let k3 = peer_key([3u8; 32]);
cache.insert(k1, AttachmentState::Ready(vec![1]), None);
cache.insert(k2, AttachmentState::Ready(vec![2]), None);
cache.insert(k1, ready(vec![1]), None);
cache.insert(k2, ready(vec![2]), None);
assert_eq!(cache.len(), 2);
// Inserting a third NEW key evicts the oldest (k1), not the newest.
cache.insert(k3, AttachmentState::Ready(vec![3]), None);
cache.insert(k3, ready(vec![3]), None);
assert_eq!(cache.len(), 2);
assert!(cache.get(&k1).is_none(), "oldest should be evicted");
assert!(cache.get(&k2).is_some());
@@ -9037,17 +9311,17 @@ mod tests {
#[test]
fn attachment_cache_replace_keeps_position_and_count() {
let mut cache = AttachmentCache::new(2);
let mut cache = count_only_cache(2);
let k1 = peer_key([1u8; 32]);
let k2 = peer_key([2u8; 32]);
cache.insert(k1, AttachmentState::Failed("pending".into()), None);
cache.insert(k2, AttachmentState::Ready(vec![2]), None);
cache.insert(k2, ready(vec![2]), None);
// Replacing k1 (Failed -> Ready) must NOT bump it to newest; it stays the
// eviction candidate, and the count is unchanged (no order leak).
cache.insert(k1, AttachmentState::Ready(vec![1]), None);
cache.insert(k1, ready(vec![1]), None);
assert_eq!(cache.len(), 2);
let k3 = peer_key([3u8; 32]);
cache.insert(k3, AttachmentState::Ready(vec![3]), None);
cache.insert(k3, ready(vec![3]), None);
assert!(cache.get(&k1).is_none(), "replaced entry kept its old age");
assert!(cache.get(&k2).is_some());
assert!(cache.get(&k3).is_some());
@@ -9057,26 +9331,26 @@ mod tests {
fn attachment_cache_same_id_distinct_authors_do_not_alias() {
// The F-12 core property: two peers sending the SAME attachment id keep
// separate bytes — one never overwrites or aliases the other.
let mut cache = AttachmentCache::new(8);
let mut cache = count_only_cache(8);
let shared_id = [7u8; 32];
let victim = peer_key(shared_id);
let attacker = peer_key(shared_id);
cache.insert(victim, AttachmentState::Ready(vec![1, 1, 1]), None);
cache.insert(attacker, AttachmentState::Ready(vec![9, 9, 9]), None);
cache.insert(victim, ready(vec![1, 1, 1]), None);
cache.insert(attacker, ready(vec![9, 9, 9]), None);
assert_eq!(cache.len(), 2);
assert!(matches!(cache.get(&victim), Some(AttachmentState::Ready(b)) if b == &[1, 1, 1]));
assert!(matches!(cache.get(&attacker), Some(AttachmentState::Ready(b)) if b == &[9, 9, 9]));
assert!(
matches!(cache.get(&victim), Some(AttachmentState::Ready(b)) if b.as_slice() == [1, 1, 1])
);
assert!(
matches!(cache.get(&attacker), Some(AttachmentState::Ready(b)) if b.as_slice() == [9, 9, 9])
);
}
#[test]
fn attachment_cache_is_ready_and_handle_and_clear() {
let mut cache = AttachmentCache::new(4);
let mut cache = count_only_cache(4);
let k = peer_key([5u8; 32]);
cache.insert(
k,
AttachmentState::Ready(vec![1]),
Some(iced::widget::image::Handle::from_bytes(vec![1])),
);
cache.insert(k, ready(vec![1]), Some(preview(4)));
assert!(cache.is_ready(&k));
assert!(cache.handle(&k).is_some());
let failed = peer_key([6u8; 32]);
@@ -9086,17 +9360,105 @@ mod tests {
cache.clear();
assert_eq!(cache.len(), 0);
assert!(cache.get(&k).is_none());
assert_eq!(cache.encoded_total, 0);
assert_eq!(cache.decoded_total, 0);
}
#[test]
fn attachment_cache_cap_zero_is_clamped_to_one() {
let mut cache = AttachmentCache::new(0);
let mut cache = count_only_cache(0);
let k = peer_key([1u8; 32]);
cache.insert(k, AttachmentState::Ready(vec![1]), None);
cache.insert(k, ready(vec![1]), None);
assert_eq!(cache.len(), 1);
assert!(cache.get(&k).is_some());
}
#[test]
fn attachment_cache_encoded_budget_evicts_oldest() {
// Budget of 100 encoded bytes; entries of 60 + 30 fit, adding 40 must
// evict the oldest (60) — not fail, not exceed the budget.
let mut cache = AttachmentCache::new(64, 100, usize::MAX);
let k1 = peer_key([1u8; 32]);
let k2 = peer_key([2u8; 32]);
let k3 = peer_key([3u8; 32]);
assert!(cache.insert(k1, ready(vec![0; 60]), None));
assert!(cache.insert(k2, ready(vec![0; 30]), None));
assert_eq!(cache.encoded_total, 90);
assert!(cache.insert(k3, ready(vec![0; 40]), None));
assert!(cache.get(&k1).is_none(), "oldest evicted for byte budget");
assert!(cache.get(&k2).is_some());
assert!(cache.get(&k3).is_some());
assert_eq!(cache.encoded_total, 70);
}
#[test]
fn attachment_cache_decoded_budget_evicts_oldest() {
let mut cache = AttachmentCache::new(64, usize::MAX, 100);
let k1 = peer_key([1u8; 32]);
let k2 = peer_key([2u8; 32]);
assert!(cache.insert(k1, ready(vec![1]), Some(preview(80))));
assert!(cache.insert(k2, ready(vec![2]), Some(preview(40))));
assert!(
cache.get(&k1).is_none(),
"oldest evicted for decoded budget"
);
assert!(cache.get(&k2).is_some());
assert_eq!(cache.decoded_total, 40);
}
#[test]
fn attachment_cache_replacement_subtracts_old_weight_first() {
// k1 holds 80 of a 100-byte budget. Replacing it with 90 must fit
// WITHOUT evicting k2 (10): the old 80 is subtracted before the check.
let mut cache = AttachmentCache::new(64, 100, usize::MAX);
let k1 = peer_key([1u8; 32]);
let k2 = peer_key([2u8; 32]);
assert!(cache.insert(k1, ready(vec![0; 80]), None));
assert!(cache.insert(k2, ready(vec![0; 10]), None));
assert!(cache.insert(k1, ready(vec![0; 90]), None));
assert!(cache.get(&k2).is_some(), "replacement must not over-evict");
assert_eq!(cache.encoded_total, 100);
// A non-Ready replacement releases the weight entirely.
cache.insert(k1, AttachmentState::Evicted, None);
assert_eq!(cache.encoded_total, 10);
}
#[test]
fn attachment_cache_overweight_entry_is_not_retained() {
let mut cache = AttachmentCache::new(64, 100, usize::MAX);
let k1 = peer_key([1u8; 32]);
let k2 = peer_key([2u8; 32]);
assert!(cache.insert(k2, ready(vec![0; 50]), None));
// Individually larger than the whole budget: rejected, nothing evicted.
assert!(!cache.insert(k1, ready(vec![0; 101]), None));
assert!(cache.get(&k1).is_none());
assert!(cache.get(&k2).is_some(), "others untouched by the reject");
assert_eq!(cache.encoded_total, 50);
// Rejecting a REPLACEMENT drops the stale entry too (its old bytes
// must not linger under a key the caller will mark Evicted).
assert!(!cache.insert(k2, ready(vec![0; 200]), None));
assert!(cache.get(&k2).is_none());
assert_eq!(cache.encoded_total, 0);
// The caller's follow-up Evicted marker is retained (zero weight).
assert!(cache.insert(k1, AttachmentState::Evicted, None));
assert!(matches!(cache.get(&k1), Some(AttachmentState::Evicted)));
}
#[test]
fn attachment_cache_loading_states_are_weightless_and_tracked() {
let mut cache = AttachmentCache::new(64, 100, 100);
let k = peer_key([1u8; 32]);
assert!(cache.insert(k, AttachmentState::Loading, None));
assert!(cache.is_loading(&k));
assert!(!cache.is_ready(&k));
assert_eq!(cache.encoded_total, 0);
// Ready replaces Loading, weights counted; not loading any more.
assert!(cache.insert(k, ready(vec![0; 10]), None));
assert!(!cache.is_loading(&k));
assert!(cache.is_ready(&k));
assert_eq!(cache.encoded_total, 10);
}
#[test]
fn attachment_default_name_matches_full_key_not_bare_id() {
// Two chat lines carry the SAME attachment id but come from different
@@ -9166,11 +9528,9 @@ mod tests {
});
state.chat_input = "draft".to_string();
let att_key = (peer, attachment_id);
state.attachments.insert(
att_key,
AttachmentState::Ready(vec![1]),
Some(iced::widget::image::Handle::from_bytes(vec![1])),
);
state
.attachments
.insert(att_key, ready(vec![1]), Some(preview(4)));
state.pending_saves.insert(att_key);
state.pending_plays.insert(attachment_id);
state.invalid_audio.insert(attachment_id);
+252
View File
@@ -0,0 +1,252 @@
//! Byte/request budgets for AUTOMATIC chat-attachment fetches (Phase 3B).
//!
//! The four-permit semaphore bounds how many auto-fetch tasks run at once, but
//! not how much a peer can make us download over time: with permits released
//! after each transfer, an insider could stream distinct ≤4 MiB images
//! sequentially forever. This budget adds per-author and session (room-wide)
//! token buckets over both request COUNT and declared BYTES. Like the Phase 2
//! chat gate, time is passed in — never read from a clock — so every refill
//! boundary is unit-testable.
//!
//! Only the automatic path consults this; a user's explicit click (Save /
//! Download / Load image) is human-rate-limited and always allowed through to
//! the fetch (still subject to the transfer cap and cache/decoder budgets).
use iroh::EndpointId;
use std::collections::HashMap;
/// Per-author request burst: how many auto-fetches one author can trigger
/// back-to-back before refill pacing binds.
pub const AUTHOR_REQ_BURST: f64 = 8.0;
/// Per-author request refill: one recovered every 10 s.
pub const AUTHOR_REQ_REFILL_PER_MS: f64 = 1.0 / 10_000.0;
/// Per-author byte burst (declared sizes): a couple of full-size auto images
/// plus a normal working set.
pub const AUTHOR_BYTES_BURST: f64 = (16 * 1024 * 1024) as f64;
/// Per-author byte refill: 64 KiB/s (~one 4 MiB auto image per minute).
pub const AUTHOR_BYTES_REFILL_PER_MS: f64 = (64 * 1024) as f64 / 1000.0;
/// Session-wide request burst across all authors.
pub const SESSION_REQ_BURST: f64 = 16.0;
/// Session-wide request refill: one recovered every 5 s.
pub const SESSION_REQ_REFILL_PER_MS: f64 = 1.0 / 5_000.0;
/// Session-wide byte burst across all authors.
pub const SESSION_BYTES_BURST: f64 = (48 * 1024 * 1024) as f64;
/// Session-wide byte refill: 128 KiB/s.
pub const SESSION_BYTES_REFILL_PER_MS: f64 = (128 * 1024) as f64 / 1000.0;
/// Bound on the per-author bucket map. Authors are roster members (≤32 live),
/// so this tracks the roster plus recently departed; the least-recently-active
/// entry is pruned past the cap.
pub const AUTHOR_MAP_CAP: usize = 64;
/// A deterministic token bucket that can take a WEIGHTED cost (bytes), unlike
/// the unit-cost bucket in the gossip chat gate.
#[derive(Debug, Clone, Copy)]
struct WeightedBucket {
tokens: f64,
last_ms: u64,
}
impl WeightedBucket {
fn full(burst: f64, now_ms: u64) -> Self {
Self {
tokens: burst,
last_ms: now_ms,
}
}
/// Refill for elapsed time (capped at `burst`) without consuming.
fn refill(&mut self, burst: f64, refill_per_ms: f64, now_ms: u64) {
let elapsed = now_ms.saturating_sub(self.last_ms) as f64;
self.tokens = (self.tokens + elapsed * refill_per_ms).min(burst);
self.last_ms = now_ms;
}
fn has(&self, cost: f64) -> bool {
self.tokens >= cost
}
fn take(&mut self, cost: f64) {
self.tokens -= cost;
}
}
/// One author's pair of buckets plus last activity (for idle pruning).
#[derive(Debug)]
struct AuthorBudget {
reqs: WeightedBucket,
bytes: WeightedBucket,
last_seen_ms: u64,
}
/// Admission budget for automatic attachment fetches. All four buckets are
/// checked BEFORE any is consumed, so a rejection never burns tokens (no
/// refund bookkeeping — the check-then-take is atomic within `admit`).
#[derive(Debug)]
pub struct AutoFetchBudget {
session_reqs: WeightedBucket,
session_bytes: WeightedBucket,
authors: HashMap<EndpointId, AuthorBudget>,
}
impl AutoFetchBudget {
pub fn new(now_ms: u64) -> Self {
Self {
session_reqs: WeightedBucket::full(SESSION_REQ_BURST, now_ms),
session_bytes: WeightedBucket::full(SESSION_BYTES_BURST, now_ms),
authors: HashMap::new(),
}
}
/// Whether an auto-fetch of `size` declared bytes for `author` may start
/// now. Consumes one request token and `size` byte tokens from BOTH the
/// author's and the session's buckets — or nothing at all on rejection.
pub fn admit(&mut self, author: EndpointId, size: u64, now_ms: u64) -> bool {
self.prune(author, now_ms);
let entry = self.authors.entry(author).or_insert_with(|| AuthorBudget {
reqs: WeightedBucket::full(AUTHOR_REQ_BURST, now_ms),
bytes: WeightedBucket::full(AUTHOR_BYTES_BURST, now_ms),
last_seen_ms: now_ms,
});
entry.last_seen_ms = now_ms;
entry
.reqs
.refill(AUTHOR_REQ_BURST, AUTHOR_REQ_REFILL_PER_MS, now_ms);
entry
.bytes
.refill(AUTHOR_BYTES_BURST, AUTHOR_BYTES_REFILL_PER_MS, now_ms);
self.session_reqs
.refill(SESSION_REQ_BURST, SESSION_REQ_REFILL_PER_MS, now_ms);
self.session_bytes
.refill(SESSION_BYTES_BURST, SESSION_BYTES_REFILL_PER_MS, now_ms);
let cost = size as f64;
let ok = entry.reqs.has(1.0)
&& entry.bytes.has(cost)
&& self.session_reqs.has(1.0)
&& self.session_bytes.has(cost);
if ok {
let entry = self.authors.get_mut(&author).expect("just inserted");
entry.reqs.take(1.0);
entry.bytes.take(cost);
self.session_reqs.take(1.0);
self.session_bytes.take(cost);
}
ok
}
/// Keep the author map bounded: past the cap, drop the least-recently
/// active entry that isn't the author being admitted. A pruned author
/// returns with full buckets, but authors are roster-gated upstream, so
/// the map can't be churned by strangers.
fn prune(&mut self, keep: EndpointId, _now_ms: u64) {
while self.authors.len() >= AUTHOR_MAP_CAP {
let Some(victim) = self
.authors
.iter()
.filter(|(id, _)| **id != keep)
.min_by_key(|(_, b)| b.last_seen_ms)
.map(|(id, _)| *id)
else {
break;
};
self.authors.remove(&victim);
}
}
#[cfg(test)]
fn author_count(&self) -> usize {
self.authors.len()
}
}
#[cfg(test)]
mod tests {
use super::*;
use iroh::SecretKey;
const T0: u64 = 1_000_000;
const MIB: u64 = 1024 * 1024;
fn author() -> EndpointId {
SecretKey::generate().public()
}
#[test]
fn author_request_burst_then_refill_recovers() {
let mut b = AutoFetchBudget::new(T0);
let a = author();
// Tiny sizes so only the REQUEST buckets can bind.
for _ in 0..AUTHOR_REQ_BURST as usize {
assert!(b.admit(a, 1, T0));
}
assert!(!b.admit(a, 1, T0), "author request burst exhausted");
// One request refills after 10 s.
assert!(b.admit(a, 1, T0 + 10_000));
assert!(!b.admit(a, 1, T0 + 10_000));
}
#[test]
fn author_byte_budget_binds_and_recovers() {
let mut b = AutoFetchBudget::new(T0);
let a = author();
// 4 × 4 MiB = the full 16 MiB author byte burst (well under the
// 8-request burst, so bytes are the binding constraint).
for _ in 0..4 {
assert!(b.admit(a, 4 * MIB, T0));
}
assert!(!b.admit(a, 4 * MIB, T0), "author byte burst exhausted");
// 64 KiB/s → a 4 MiB image is affordable again after 64 s (which also
// refills 6 request tokens, so bytes stay the binding constraint).
assert!(!b.admit(a, 4 * MIB, T0 + 32_000));
assert!(b.admit(a, 4 * MIB, T0 + 64_000));
}
#[test]
fn session_budget_binds_across_authors_without_burning_author_tokens() {
let mut b = AutoFetchBudget::new(T0);
// Three authors × 16 MiB exhausts the 48 MiB session byte burst even
// though each author is within their own budget.
for _ in 0..3 {
let a = author();
for _ in 0..4 {
assert!(b.admit(a, 4 * MIB, T0));
}
}
let fresh = author();
assert!(!b.admit(fresh, 4 * MIB, T0), "session bytes exhausted");
// The rejection consumed NOTHING: once the session refills enough for
// one image (4 MiB / 128 KiB/s = 32 s), the fresh author's own full
// burst is intact and admits immediately.
assert!(b.admit(fresh, 4 * MIB, T0 + 32_000));
}
#[test]
fn session_request_bucket_binds_across_authors() {
let mut b = AutoFetchBudget::new(T0);
// 16 tiny requests from distinct authors exhaust the session request
// burst while every author bucket stays nearly full.
for _ in 0..SESSION_REQ_BURST as usize {
assert!(b.admit(author(), 1, T0));
}
assert!(!b.admit(author(), 1, T0), "session requests exhausted");
assert!(b.admit(author(), 1, T0 + 5_000), "one recovers after 5 s");
}
#[test]
fn author_map_stays_bounded_pruning_least_recent() {
let mut b = AutoFetchBudget::new(T0);
// Session request refill would bind over a naive loop; space the
// admissions out so only the map bound is under test.
let mut t = T0;
let first = author();
assert!(b.admit(first, 1, t));
for _ in 0..(AUTHOR_MAP_CAP + 10) {
t += 10_000;
assert!(b.admit(author(), 1, t));
assert!(b.author_count() <= AUTHOR_MAP_CAP);
}
assert!(b.author_count() <= AUTHOR_MAP_CAP);
}
}
+13 -2
View File
@@ -74,7 +74,10 @@ pub enum CoreCommand {
SendChatFile {
text: String,
attachment: crate::files::ChatAttachment,
data: Vec<u8>,
/// Shared, not owned: the same allocation is retained by the UI cache
/// and handed to the serve store, so a 25 MiB attachment is held once,
/// not copied across UI / command queue / serve store (Phase 3C).
data: std::sync::Arc<Vec<u8>>,
},
/// Fetch a received attachment's bytes from its sender over the file plane
/// (used for on-demand file/chip downloads; images are auto-fetched on
@@ -434,7 +437,15 @@ pub enum UiEvent {
AttachmentReady {
from: EndpointId,
id: crate::files::AttachmentId,
data: Vec<u8>,
data: std::sync::Arc<Vec<u8>>,
},
/// An attachment fetch task was spawned (auto or on demand). Lets the UI
/// show a real "loading" state instead of inferring it from cache absence —
/// absence now means NOT fetched (e.g. auto-fetch was skipped), which
/// renders a Load button rather than an indefinite "loading…" (Phase 3B).
AttachmentFetchStarted {
from: EndpointId,
id: crate::files::AttachmentId,
},
/// An attachment fetch failed (sender gone, too large, decode error, etc.).
AttachmentFailed {
+61 -15
View File
@@ -1,5 +1,6 @@
pub mod chatroster;
pub mod connstats;
pub mod fetchbudget;
pub mod jitter;
pub mod messages;
mod recovery;
@@ -1017,6 +1018,15 @@ const MAX_INFLIGHT_ATTACHMENT_FETCHES: usize = 4;
/// fetches (Tier C F-02). Bounded by [`MAX_INFLIGHT_ATTACHMENT_FETCHES`].
type InflightAttachments = Arc<std::sync::Mutex<HashSet<(EndpointId, crate::files::AttachmentId)>>>;
/// Milliseconds since the Unix epoch — the time source handed to the
/// deterministic auto-fetch budget (mirrors the gossip plane's timestamps).
fn unix_now_ms() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_millis() as u64)
.unwrap_or(0)
}
/// RAII bookkeeping for one bounded auto-fetch: holds the concurrency permit for
/// the task's lifetime and clears the in-flight `(author, id)` marker when the
/// fetch finishes (success OR failure), so the same image can be retried later.
@@ -1033,11 +1043,22 @@ impl Drop for AutoFetchGuard {
}
/// Whether to AUTO-fetch a chat image attachment. Only authenticated roster
/// authors qualify (closing the non-roster injection vector), and a `(author,
/// id)` already being fetched is skipped (dedup). The concurrency bound itself is
/// enforced separately by the permit. Pure → unit-testable (Tier C F-02).
fn should_auto_fetch(is_image: bool, author_in_roster: bool, already_inflight: bool) -> bool {
is_image && author_in_roster && !already_inflight
/// authors qualify (closing the non-roster injection vector), only declared
/// sizes at or under [`crate::files::MAX_AUTO_IMAGE_BYTES`] (larger images get
/// a Load button instead, Phase 3B), and a `(author, id)` already being fetched
/// is skipped (dedup). The concurrency bound (permits) and the byte/request
/// budgets ([`fetchbudget::AutoFetchBudget`]) are enforced separately. Pure →
/// unit-testable (Tier C F-02).
fn should_auto_fetch(
is_image: bool,
author_in_roster: bool,
already_inflight: bool,
declared_size: u64,
) -> bool {
is_image
&& author_in_roster
&& !already_inflight
&& declared_size <= crate::files::MAX_AUTO_IMAGE_BYTES
}
/// Fetch a chat attachment's bytes from `from` over the file plane in a detached
@@ -1060,6 +1081,12 @@ fn spawn_attachment_fetch(
tokio::spawn(async move {
// Held for the whole fetch; dropped here on completion (Tier C F-02).
let _guard = guard;
// Tell the UI a real fetch task exists for this key, so it can show a
// genuine loading state and dedup further clicks (Phase 3B). Sent from
// this task's channel handle, so it always precedes Ready/Failed.
let _ = ui_tx
.send(UiEvent::AttachmentFetchStarted { from, id: att.id })
.await;
match transport.fetch_attachment(from, &att).await {
Ok(data) => {
if is_image && crate::files::validate_image_bytes(&data).is_none() {
@@ -1076,7 +1103,7 @@ fn spawn_attachment_fetch(
.send(UiEvent::AttachmentReady {
from,
id: att.id,
data,
data: Arc::new(data),
})
.await;
}
@@ -2362,6 +2389,10 @@ async fn run_core_loop(
let chat_roster_events = chat_roster.clone();
let event_task = tokio::spawn(async move {
let roster = chat_roster_events;
// Byte/request budgets for automatic attachment fetches
// (Phase 3B). Only this sequential task consults it, so it
// needs no lock; time is passed in for testability.
let mut auto_fetch_budget = fetchbudget::AutoFetchBudget::new(unix_now_ms());
while let Some(event) = room_events.recv().await {
match event {
RoomEvent::PeerJoined(peer_id, state) => {
@@ -2532,7 +2563,9 @@ async fn run_core_loop(
// any attachment handling.
true,
already_inflight,
) {
att.size,
) && auto_fetch_budget.admit(from, att.size, unix_now_ms())
{
// Reserve the dedup slot, then a permit. If the
// pool is exhausted, drop the auto-fetch (and the
// dedup marker) — the descriptor still shows and
@@ -3242,9 +3275,7 @@ async fn run_core_loop(
if let Some(session) = &active_session {
// Make the bytes fetchable by room members, then broadcast the
// descriptor alongside the (possibly empty) caption text.
session
.transport
.serve_attachment(attachment.id, Arc::new(data));
session.transport.serve_attachment(attachment.id, data);
if let Err(e) = session.room_state.send_chat(text, Some(attachment)).await {
crate::log_msg(&format!("Failed to send chat file: {e}"));
}
@@ -3753,15 +3784,30 @@ mod tests {
#[test]
fn auto_fetch_only_for_roster_images_not_already_inflight() {
// The happy path: a roster author's brand-new image attachment.
assert!(should_auto_fetch(true, true, false));
const OK_SIZE: u64 = 1024;
// The happy path: a roster author's brand-new, small-enough image.
assert!(should_auto_fetch(true, true, false, OK_SIZE));
// A non-image (generic file) never auto-fetches — it waits for "Save".
assert!(!should_auto_fetch(false, true, false));
assert!(!should_auto_fetch(false, true, false, OK_SIZE));
// A non-roster author (e.g. a sock puppet that never announced) is rejected,
// closing the F-02 unbounded-task vector.
assert!(!should_auto_fetch(true, false, false));
assert!(!should_auto_fetch(true, false, false, OK_SIZE));
// An identical (author,id) already being fetched is deduped.
assert!(!should_auto_fetch(true, true, true));
assert!(!should_auto_fetch(true, true, true, OK_SIZE));
// The declared-size gate (Phase 3B): at the cap auto-fetches, the first
// byte over requires a click.
assert!(should_auto_fetch(
true,
true,
false,
crate::files::MAX_AUTO_IMAGE_BYTES
));
assert!(!should_auto_fetch(
true,
true,
false,
crate::files::MAX_AUTO_IMAGE_BYTES + 1
));
}
#[test]
+341 -10
View File
@@ -22,6 +22,22 @@ pub const MAX_ATTACHMENT_BYTES: u64 = 25 * 1024 * 1024;
/// the byte cap. Applied via `image::Limits` when validating/decoding.
pub const MAX_IMAGE_PX: u32 = 4096;
/// Max total decoded pixels, applied on top of the per-side [`MAX_IMAGE_PX`]
/// limit. The per-side cap alone still admits a 4096×4096 ≈ 16.8 MP bitmap
/// (~64 MiB transient RGBA); this bounds the worst-case decode allocation while
/// still clearing common 12 MP phone photos (4032×3024 ≈ 12.2 MP).
pub const MAX_IMAGE_TOTAL_PIXELS: u64 = 14_000_000;
/// Max pixels per side of the downscaled inline preview handed to the renderer.
/// Original bytes are kept only for Save; the chat column never needs more than
/// this (it displays at ~260 px, and the lightbox at window size).
pub const IMAGE_PREVIEW_MAX_SIDE: u32 = 1600;
/// Largest declared size an image attachment may auto-fetch at. Anything larger
/// (or any skipped/evicted image) renders a "Load image" button instead; a
/// manual click may use the full [`MAX_ATTACHMENT_BYTES`] cap.
pub const MAX_AUTO_IMAGE_BYTES: u64 = 4 * 1024 * 1024;
/// Longest filename we keep and display. Keeps the gossip descriptor compact and
/// the UI tidy; the real bytes are unaffected.
pub const MAX_FILENAME_LEN: usize = 96;
@@ -75,8 +91,14 @@ pub fn sanitize_filename(raw: &str) -> String {
.unwrap_or("")
.trim();
// Drop control chars; turn other whitespace into single spaces later.
let cleaned: String = base.chars().filter(|c| !c.is_control()).collect();
// Drop control chars and the same bidi/zero-width spoofing format chars
// stripped from display names (a U+202E override can visually reverse an
// extension, e.g. "photo\u{202E}gnp.exe" renders as "photoexe.png").
// Ordinary non-ASCII filenames pass through untouched.
let cleaned: String = base
.chars()
.filter(|c| !c.is_control() && !crate::sanitize::is_spoofing_format_char(*c))
.collect();
let collapsed = cleaned.split_whitespace().collect::<Vec<_>>().join(" ");
let collapsed = collapsed.trim_matches('.').trim();
@@ -164,20 +186,185 @@ pub fn classify(bytes: &[u8]) -> AttachmentKind {
/// regardless of [`MAX_ATTACHMENT_BYTES`]. Only PNG/JPEG are buildable in our
/// `image` feature set; anything else returns `None` and the caller shows a chip.
pub fn validate_image_bytes(bytes: &[u8]) -> Option<(u32, u32)> {
let mut limits = image::Limits::default();
limits.max_image_width = Some(MAX_IMAGE_PX);
limits.max_image_height = Some(MAX_IMAGE_PX);
let img = decode_image_bounded(bytes)?;
Some((img.width(), img.height()))
}
/// Shared bounded decode: header-check the dimensions (per-side AND total-pixel
/// limits) BEFORE decoding, then decode under `image::Limits` as defense in
/// depth. The precheck reads only the container header, so an over-limit bomb is
/// rejected without paying its decode cost.
fn decode_image_bounded(bytes: &[u8]) -> Option<image::DynamicImage> {
let reader = image::ImageReader::new(std::io::Cursor::new(bytes))
.with_guessed_format()
.ok()?;
let mut reader = reader;
reader.limits(limits);
let img = reader.decode().ok()?;
let (w, h) = (img.width(), img.height());
let (w, h) = reader.into_dimensions().ok()?;
if w == 0 || h == 0 || w > MAX_IMAGE_PX || h > MAX_IMAGE_PX {
return None;
}
Some((w, h))
if u64::from(w) * u64::from(h) > MAX_IMAGE_TOTAL_PIXELS {
return None;
}
let mut limits = image::Limits::default();
limits.max_image_width = Some(MAX_IMAGE_PX);
limits.max_image_height = Some(MAX_IMAGE_PX);
let mut reader = image::ImageReader::new(std::io::Cursor::new(bytes))
.with_guessed_format()
.ok()?;
reader.limits(limits);
let img = reader.decode().ok()?;
// Decoded size must match the header the precheck approved.
if img.width() != w || img.height() != h {
return None;
}
Some(img)
}
/// A decoded, display-ready inline preview: RGBA pixels downscaled so neither
/// side exceeds [`IMAGE_PREVIEW_MAX_SIDE`]. `rgba.len() == width * height * 4`,
/// which is also the preview's decoded-budget weight in the attachment cache.
pub struct ImagePreview {
pub width: u32,
pub height: u32,
pub rgba: Vec<u8>,
}
/// Decode image bytes under the same limits as [`validate_image_bytes`] and
/// build the downscaled inline preview. The full-resolution bitmap exists only
/// transiently here; the renderer is never handed more than
/// [`IMAGE_PREVIEW_MAX_SIDE`]² pixels. Returns `None` for anything that fails
/// validation (caller falls back to a chip / failure row).
pub fn decode_preview(bytes: &[u8]) -> Option<ImagePreview> {
let img = decode_image_bounded(bytes)?;
let img = if img.width() > IMAGE_PREVIEW_MAX_SIDE || img.height() > IMAGE_PREVIEW_MAX_SIDE {
// `thumbnail` preserves aspect ratio within the bounding box.
img.thumbnail(IMAGE_PREVIEW_MAX_SIDE, IMAGE_PREVIEW_MAX_SIDE)
} else {
img
};
let rgba = img.into_rgba8();
let (width, height) = (rgba.width(), rgba.height());
Some(ImagePreview {
width,
height,
rgba: rgba.into_raw(),
})
}
/// Estimated decoded RGBA cost of a preview, the weight counted against the
/// attachment cache's decoded-byte budget (`width * height * 4`).
pub fn preview_rgba_cost(width: u32, height: u32) -> usize {
(width as usize)
.saturating_mul(height as usize)
.saturating_mul(4)
}
/// Read at most [`MAX_ATTACHMENT_BYTES`] bytes from `r`. Returns `Ok(None)` if
/// the source holds even one byte more (detected by reading cap + 1), so a huge
/// or unbounded source is never fully buffered. Pure over `Read` for tests; the
/// picker wraps it via [`read_file_capped`].
pub fn read_capped<R: std::io::Read>(r: R) -> std::io::Result<Option<Vec<u8>>> {
use std::io::Read as _;
let mut buf = Vec::new();
let mut limited = r.take(MAX_ATTACHMENT_BYTES + 1);
limited.read_to_end(&mut buf)?;
if buf.len() as u64 > MAX_ATTACHMENT_BYTES {
return Ok(None);
}
Ok(Some(buf))
}
/// Read a picked file, bounded by [`MAX_ATTACHMENT_BYTES`]. Checks metadata
/// first to reject an obviously-oversized file without opening it, but keeps the
/// bounded read regardless — metadata can race (the file can grow after the
/// check) or be unavailable through a portal. `Ok(None)` = over the cap.
pub fn read_file_capped(path: &std::path::Path) -> std::io::Result<Option<Vec<u8>>> {
if let Ok(meta) = std::fs::metadata(path)
&& meta.len() > MAX_ATTACHMENT_BYTES
{
return Ok(None);
}
read_capped(std::fs::File::open(path)?)
}
/// Cap on how many blobs the session serve store retains at once (sent chat
/// attachments plus the current/next broadcast music tracks).
pub const SERVED_FILES_MAX_ENTRIES: usize = 16;
/// Byte budget for the serve store. Without it, a sender's own session could
/// grow unbounded at up to [`MAX_ATTACHMENT_BYTES`] per send (Phase 3C).
pub const SERVED_FILES_MAX_BYTES: usize = 128 * 1024 * 1024;
/// Count- and byte-budgeted FIFO store of blobs we serve to room members over
/// the file plane. Evicting an id makes a later request for it read as an empty
/// body — the existing "sender no longer has the file" response — never stale
/// or aliased bytes. Pure (no locks/IO) so budgets are unit-testable; the
/// transport wraps it in its own mutex.
#[derive(Debug, Default)]
pub struct ServeStore {
entries: std::collections::HashMap<AttachmentId, std::sync::Arc<Vec<u8>>>,
/// Present ids in insertion order; the front is the eviction candidate.
order: std::collections::VecDeque<AttachmentId>,
total_bytes: usize,
}
impl ServeStore {
/// Insert or replace a blob, evicting oldest entries until the count and
/// byte budgets fit. Replacement keeps the id's age and subtracts the old
/// bytes before the new ones are counted. Returns `false` for a blob that
/// alone exceeds the byte budget (not stored; an existing entry under the
/// id is dropped rather than left stale).
pub fn insert(&mut self, id: AttachmentId, bytes: std::sync::Arc<Vec<u8>>) -> bool {
if let Some(old) = self.entries.get(&id) {
self.total_bytes -= old.len();
}
if bytes.len() > SERVED_FILES_MAX_BYTES {
if self.entries.remove(&id).is_some() {
self.order.retain(|k| k != &id);
}
return false;
}
let replacing = self.entries.contains_key(&id);
loop {
let count_full = !replacing && self.entries.len() >= SERVED_FILES_MAX_ENTRIES;
let bytes_full = self.total_bytes + bytes.len() > SERVED_FILES_MAX_BYTES;
if !count_full && !bytes_full {
break;
}
let Some(victim) = self.order.iter().find(|k| **k != id).copied() else {
break;
};
self.remove(&victim);
}
if !replacing {
self.order.push_back(id);
}
self.total_bytes += bytes.len();
self.entries.insert(id, bytes);
true
}
pub fn get(&self, id: &AttachmentId) -> Option<std::sync::Arc<Vec<u8>>> {
self.entries.get(id).cloned()
}
pub fn remove(&mut self, id: &AttachmentId) {
if let Some(old) = self.entries.remove(id) {
self.total_bytes -= old.len();
self.order.retain(|k| k != id);
}
}
pub fn clear(&mut self) {
self.entries.clear();
self.order.clear();
self.total_bytes = 0;
}
#[cfg(test)]
fn len(&self) -> usize {
self.entries.len()
}
}
/// Parse a file-plane request: it must be exactly one [`AttachmentId`] (32
@@ -352,6 +539,87 @@ mod tests {
assert_eq!(validate_image_bytes(&buf.into_inner()), Some((4, 3)));
}
#[test]
fn sanitize_strips_bidi_and_zero_width_spoofing_chars() {
// U+202E would visually reverse the tail, disguising the extension.
assert_eq!(sanitize_filename("photo\u{202E}gnp.exe"), "photognp.exe");
assert_eq!(sanitize_filename("a\u{200B}b\u{FEFF}.txt"), "ab.txt");
// Ordinary Unicode filenames pass through.
assert_eq!(sanitize_filename("família_fotos.png"), "família_fotos.png");
assert_eq!(sanitize_filename("日本語.pdf"), "日本語.pdf");
}
/// Encode a solid PNG of the given dimensions for limit tests.
fn png_bytes(w: u32, h: u32) -> Vec<u8> {
let img = image::RgbImage::from_pixel(w, h, image::Rgb([10, 20, 30]));
let mut buf = std::io::Cursor::new(Vec::new());
image::DynamicImage::ImageRgb8(img)
.write_to(&mut buf, image::ImageFormat::Png)
.unwrap();
buf.into_inner()
}
#[test]
fn validate_image_rejects_excessive_total_pixels() {
// Both sides within MAX_IMAGE_PX, but 4096 * 4096 > MAX_IMAGE_TOTAL_PIXELS.
assert!(u64::from(MAX_IMAGE_PX) * u64::from(MAX_IMAGE_PX) > MAX_IMAGE_TOTAL_PIXELS);
assert_eq!(validate_image_bytes(&png_bytes(4096, 4096)), None);
// A 12 MP phone-photo shape passes both limits.
assert_eq!(
validate_image_bytes(&png_bytes(4032, 3024)),
Some((4032, 3024))
);
}
#[test]
fn preview_downscales_to_max_side_preserving_aspect() {
// Wide: 3200x400 → 1600x200.
let p = decode_preview(&png_bytes(3200, 400)).unwrap();
assert_eq!((p.width, p.height), (1600, 200));
assert_eq!(p.rgba.len(), preview_rgba_cost(1600, 200));
// Tall: 400x3200 → 200x1600.
let p = decode_preview(&png_bytes(400, 3200)).unwrap();
assert_eq!((p.width, p.height), (200, 1600));
// Square over the side cap: 2000x2000 → 1600x1600.
let p = decode_preview(&png_bytes(2000, 2000)).unwrap();
assert_eq!((p.width, p.height), (1600, 1600));
// At/under the cap is untouched.
let p = decode_preview(&png_bytes(1600, 900)).unwrap();
assert_eq!((p.width, p.height), (1600, 900));
let p = decode_preview(&png_bytes(4, 3)).unwrap();
assert_eq!((p.width, p.height), (4, 3));
assert_eq!(p.rgba.len(), preview_rgba_cost(4, 3));
}
#[test]
fn preview_rejects_what_validation_rejects() {
assert!(decode_preview(b"not an image").is_none());
assert!(decode_preview(&png_bytes(4096, 4096)).is_none());
}
#[test]
fn read_capped_stops_at_cap_plus_one() {
// Under the cap: full read.
let small = vec![7u8; 1024];
assert_eq!(
read_capped(std::io::Cursor::new(&small))
.unwrap()
.as_deref(),
Some(&small[..])
);
// Exactly at the cap: accepted. `repeat` is endless, `take` proves the
// reader is bounded rather than draining the source.
let at_cap = std::io::Read::take(std::io::repeat(1), MAX_ATTACHMENT_BYTES);
let got = read_capped(at_cap).unwrap().unwrap();
assert_eq!(got.len() as u64, MAX_ATTACHMENT_BYTES);
// One byte over: rejected, and only cap + 1 bytes were ever buffered
// (an unbounded source returns instead of allocating forever).
let over = std::io::Read::take(std::io::repeat(1), MAX_ATTACHMENT_BYTES + 1);
assert_eq!(read_capped(over).unwrap(), None);
let endless = std::io::repeat(1);
assert_eq!(read_capped(endless).unwrap(), None);
}
#[test]
fn human_size_units() {
assert_eq!(human_size(40), "40 B");
@@ -359,6 +627,69 @@ mod tests {
assert_eq!(human_size(3 * 1024 * 1024 + 300 * 1024), "3.3 MB");
}
#[test]
fn serve_store_count_and_byte_eviction_fifo() {
use std::sync::Arc;
let mut s = ServeStore::default();
let blob = |n: u8, len: usize| ([n; 32], Arc::new(vec![n; len]));
// Count cap: entry 0 is evicted when the 17th arrives.
for n in 0..=SERVED_FILES_MAX_ENTRIES as u8 {
let (id, b) = blob(n, 8);
assert!(s.insert(id, b));
}
assert_eq!(s.len(), SERVED_FILES_MAX_ENTRIES);
assert!(s.get(&[0u8; 32]).is_none(), "oldest evicted by count");
assert!(s.get(&[1u8; 32]).is_some());
// Byte budget: two ~half-budget blobs evict everything older.
let half = SERVED_FILES_MAX_BYTES / 2;
let (a, ab) = blob(100, half);
let (b, bb) = blob(101, half);
assert!(s.insert(a, ab));
assert!(s.insert(b, bb));
assert!(s.get(&a).is_some());
assert!(s.get(&b).is_some());
assert!(s.get(&[1u8; 32]).is_none(), "evicted for byte budget");
// A third half-budget blob evicts `a` (oldest), keeps `b`.
let (c, cb) = blob(102, half);
assert!(s.insert(c, cb));
assert!(s.get(&a).is_none());
assert!(s.get(&b).is_some());
assert!(s.get(&c).is_some());
}
#[test]
fn serve_store_replacement_accounting_and_remove_clear() {
use std::sync::Arc;
let mut s = ServeStore::default();
let id = [9u8; 32];
assert!(s.insert(id, Arc::new(vec![1; SERVED_FILES_MAX_BYTES - 10])));
// Replacing the near-budget blob must subtract its old bytes first —
// otherwise this same-id replacement would evict itself.
assert!(s.insert(id, Arc::new(vec![2; SERVED_FILES_MAX_BYTES - 5])));
assert_eq!(s.get(&id).unwrap()[0], 2);
assert_eq!(s.len(), 1);
s.remove(&id);
assert!(s.get(&id).is_none());
// Removed bytes were released: the budget admits a full-size blob again.
assert!(s.insert(id, Arc::new(vec![3; SERVED_FILES_MAX_BYTES])));
s.clear();
assert_eq!(s.len(), 0);
assert!(s.insert(id, Arc::new(vec![4; SERVED_FILES_MAX_BYTES])));
}
#[test]
fn serve_store_rejects_individually_overweight_blob() {
use std::sync::Arc;
let mut s = ServeStore::default();
let id = [7u8; 32];
assert!(s.insert(id, Arc::new(vec![1; 8])));
assert!(!s.insert(id, Arc::new(vec![2; SERVED_FILES_MAX_BYTES + 1])));
// The stale small blob is gone too — a fetch reads "no longer has it",
// never old bytes under a replaced id.
assert!(s.get(&id).is_none());
assert_eq!(s.len(), 0);
}
#[test]
fn attachment_descriptor_round_trips_json() {
let a = ChatAttachment {
+18 -4
View File
@@ -69,7 +69,9 @@ struct Shared {
/// the random attachment id. Populated when we send a chat file; read by the
/// file protocol handler to answer a member's fetch. Cleared on leave. Each
/// blob is already byte-capped at send time.
served_files: StdMutex<HashMap<crate::files::AttachmentId, Arc<Vec<u8>>>>,
/// Blobs we serve to room members, bounded by count and byte budgets
/// (Phase 3C) — an evicted id reads as "sender no longer has the file".
served_files: StdMutex<crate::files::ServeStore>,
incoming_tx: mpsc::Sender<(EndpointId, Bytes)>,
/// Best-effort link-state notifications for the UI (connecting / connected).
conn_events_tx: mpsc::Sender<ConnEvent>,
@@ -518,7 +520,7 @@ impl iroh::protocol::ProtocolHandler for FileRouter {
let Some(id) = crate::files::parse_request(&req) else {
return Ok(());
};
let blob = shared.served_files.lock().unwrap().get(&id).cloned();
let blob = shared.served_files.lock().unwrap().get(&id);
if let Some(blob) = blob {
let _ = send.write_all(&blob).await;
}
@@ -563,7 +565,7 @@ impl IrohTransport {
peers: tokio::sync::Mutex::new(HashMap::new()),
live_conns: StdMutex::new(HashMap::new()),
admitted_audio: StdMutex::new(HashSet::new()),
served_files: StdMutex::new(HashMap::new()),
served_files: StdMutex::new(crate::files::ServeStore::default()),
incoming_tx,
conn_events_tx,
});
@@ -634,7 +636,9 @@ impl IrohTransport {
/// session (served by the [`FileRouter`] handler). Called by core when we
/// send a chat file. The blob is cleared on leave.
pub fn serve_attachment(&self, id: AttachmentId, bytes: Arc<Vec<u8>>) {
self.shared.served_files.lock().unwrap().insert(id, bytes);
if !self.shared.served_files.lock().unwrap().insert(id, bytes) {
crate::log_msg("Transport: refused to serve an over-budget blob");
}
}
/// Drop a previously-served blob (e.g. a music track no longer current-or-next).
@@ -676,6 +680,8 @@ impl IrohTransport {
send.finish()
.map_err(|e| NetError::Other(format!("file fetch: request finish failed: {e}")))?;
// `read_to_end(size)` errors if the stream exceeds `size`, rejecting an
// overlong transfer; the exact-length check below rejects a short one.
let read = recv.read_to_end(size as usize);
let bytes = tokio::time::timeout(FILE_FETCH_TIMEOUT, read)
.await
@@ -686,6 +692,14 @@ impl IrohTransport {
"file fetch: sender no longer has the file".to_string(),
));
}
// Exact transfer required (Phase 3C): a truncated body must not be
// cached/saved/decoded as if it were the declared attachment.
if bytes.len() as u64 != size {
return Err(NetError::Other(format!(
"file fetch: incomplete transfer ({} of {size} bytes)",
bytes.len()
)));
}
Ok(bytes)
}
+1 -1
View File
@@ -18,7 +18,7 @@ pub const NAME_MAX_CHARS: usize = 48;
/// zero-width / BOM characters (invisible, can hide or fake content). Listed
/// explicitly so the sanitizer stays dependency-free (std exposes no category
/// query). Stripped outright rather than replaced.
fn is_spoofing_format_char(c: char) -> bool {
pub(crate) fn is_spoofing_format_char(c: char) -> bool {
matches!(c,
'\u{200B}'..='\u{200F}' // zero-width space, ZWNJ, ZWJ, LRM, RLM
| '\u{202A}'..='\u{202E}' // LRE, RLE, PDF, LRO, RLO (bidi overrides)