Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
157 changes: 120 additions & 37 deletions desktop/src-tauri/src/commands/media_download.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,13 @@ use tauri::State;
use crate::app_state::AppState;
use crate::commands::export_util::save_bytes_with_dialog;
use crate::commands::media::{detect_and_validate_mime, sanitize_filename};
use crate::commands::personas::{
decode_snapshot_from_bytes, MAX_SNAPSHOT_JSON_BYTES, MAX_SNAPSHOT_PNG_BYTES, PNG_MAGIC,
use crate::commands::{
personas::{
decode_snapshot_from_bytes, MAX_SNAPSHOT_JSON_BYTES, MAX_SNAPSHOT_PNG_BYTES, PNG_MAGIC,
},
team_snapshot::{
decode_team_snapshot_from_bytes, MAX_TEAM_SNAPSHOT_JSON_BYTES, MAX_TEAM_SNAPSHOT_PNG_BYTES,
},
};
use crate::relay::{classify_request_error, relay_api_base_url_with_override, relay_error_message};

Expand Down Expand Up @@ -275,16 +280,35 @@ async fn fetch_blob_bytes_with_cap(
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub(crate) enum SnapshotFileKind {
/// `.agent.json` — plaintext JSON; accepts memory; 5 MiB cap.
Json,
AgentJson,
/// `.agent.png` — PNG with embedded metadata; no memory; 10 MiB cap.
Png,
AgentPng,
/// `.team.json` — team template; 25 MiB cap.
TeamJson,
/// `.team.png` — team template PNG; 50 MiB cap.
TeamPng,
}

impl SnapshotFileKind {
fn cap(self) -> u64 {
match self {
SnapshotFileKind::Json => MAX_SNAPSHOT_JSON_BYTES as u64,
SnapshotFileKind::Png => MAX_SNAPSHOT_PNG_BYTES as u64,
SnapshotFileKind::AgentJson => MAX_SNAPSHOT_JSON_BYTES as u64,
SnapshotFileKind::AgentPng => MAX_SNAPSHOT_PNG_BYTES as u64,
SnapshotFileKind::TeamJson => MAX_TEAM_SNAPSHOT_JSON_BYTES as u64,
SnapshotFileKind::TeamPng => MAX_TEAM_SNAPSHOT_PNG_BYTES as u64,
}
}

fn is_png(self) -> bool {
matches!(self, SnapshotFileKind::AgentPng | SnapshotFileKind::TeamPng)
}

fn label(self) -> &'static str {
match self {
SnapshotFileKind::AgentJson => ".agent.json",
SnapshotFileKind::AgentPng => ".agent.png",
SnapshotFileKind::TeamJson => ".team.json",
SnapshotFileKind::TeamPng => ".team.png",
}
}
}
Expand All @@ -296,14 +320,18 @@ impl SnapshotFileKind {
/// fails closed before any bytes reach the frontend.
fn ensure_bytes_match_kind(bytes: &[u8], kind: SnapshotFileKind) -> Result<(), String> {
let has_png_magic = bytes.len() >= 4 && bytes[..4] == PNG_MAGIC;
match kind {
SnapshotFileKind::Png if !has_png_magic => {
Err("format mismatch: filename is .agent.png but bytes are not a PNG".to_string())
}
SnapshotFileKind::Json if has_png_magic => {
Err("format mismatch: filename is .agent.json but bytes are a PNG".to_string())
}
_ => Ok(()),
if kind.is_png() && !has_png_magic {
Err(format!(
"format mismatch: filename is {} but bytes are not a PNG",
kind.label()
))
} else if !kind.is_png() && has_png_magic {
Err(format!(
"format mismatch: filename is {} but bytes are a PNG",
kind.label()
))
} else {
Ok(())
}
}

Expand All @@ -312,35 +340,60 @@ fn ensure_bytes_match_kind(bytes: &[u8], kind: SnapshotFileKind) -> Result<(), S
fn snapshot_kind_for_filename(filename: &str) -> Result<SnapshotFileKind, String> {
let lower = filename.to_ascii_lowercase();
if lower.ends_with(".agent.json") {
Ok(SnapshotFileKind::Json)
Ok(SnapshotFileKind::AgentJson)
} else if lower.ends_with(".agent.png") {
Ok(SnapshotFileKind::Png)
Ok(SnapshotFileKind::AgentPng)
} else if lower.ends_with(".team.json") {
Ok(SnapshotFileKind::TeamJson)
} else if lower.ends_with(".team.png") {
Ok(SnapshotFileKind::TeamPng)
} else {
Err(format!(
"\"{}\" is not a snapshot filename — expected .agent.json or .agent.png",
"\"{}\" is not a snapshot filename — expected .agent.json, .agent.png, .team.json, or .team.png",
filename
))
}
}

/// Fetch and validate an agent snapshot attachment in memory.
/// Reject a metadata-declared size before opening the media stream.
///
/// Keeping this as a pure helper makes the per-kind cap a testable production
/// boundary, rather than relying on a duplicated test-side comparison.
fn ensure_declared_size_within_cap(
expected_size: usize,
kind: SnapshotFileKind,
) -> Result<(), String> {
if expected_size as u64 > kind.cap() {
return Err(format!(
"declared size {} exceeds the {} MiB cap for this format",
expected_size,
kind.cap() / (1024 * 1024)
));
}
Ok(())
}

/// Fetch and validate an agent or team snapshot attachment in memory.
///
/// Input validation (before HTTP):
/// - URL must be a valid same-relay `/media/` URL.
/// - Filename must end with `.agent.json` or `.agent.png`.
/// - Filename must end case-insensitively with `.agent.json`, `.agent.png`,
/// `.team.json`, or `.team.png`.
/// - `expected_sha256` and `expected_size` must be non-empty strings.
///
/// During fetch:
/// - Enforces a format-specific cap (5 MiB JSON, 10 MiB PNG) via
/// Content-Length header and streamed byte count.
/// During fetch, `SnapshotFileKind::cap()` enforces the kind-specific cap via
/// Content-Length and streamed byte count: 5 MiB JSON / 10 MiB PNG for agents,
/// or 25 MiB JSON / 50 MiB PNG for teams.
///
/// Post-fetch validation (all must pass; returns an error on first failure):
/// 1. Byte length equals `expected_size`.
/// 2. SHA-256 hex of bytes equals `expected_sha256` (lowercase).
/// 3. `decode_snapshot_from_bytes` succeeds — bytes are a well-formed snapshot.
/// 3. The byte magic matches the filename-selected kind.
/// 4. Agent kinds pass `decode_snapshot_from_bytes`; team kinds pass
/// `decode_team_snapshot_from_bytes`.
///
/// Returns `tauri::ipc::Response` so bytes cross IPC as a raw buffer rather
/// than a JSON number array (which would be ~3× the size at the 5–10 MiB cap).
/// than a JSON number array (which would be ~3× the size at the applicable cap).
#[tauri::command]
pub async fn fetch_snapshot_bytes(
url: String,
Expand Down Expand Up @@ -369,13 +422,7 @@ pub async fn fetch_snapshot_bytes(
if expected_size == 0 {
return Err("missing or zero expected size (imeta size field)".to_string());
}
if expected_size as u64 > cap {
return Err(format!(
"declared size {} exceeds the {} MiB cap for this format",
expected_size,
cap / (1024 * 1024)
));
}
ensure_declared_size_within_cap(expected_size, kind)?;

// ── Bounded fetch ─────────────────────────────────────────────────────
let bytes = fetch_blob_bytes_with_cap(&url, &state, cap).await?;
Expand All @@ -401,10 +448,19 @@ pub async fn fetch_snapshot_bytes(
// JSON) and .agent.json delivering PNG bytes.
ensure_bytes_match_kind(&bytes, kind)?;

// 4. Bytes must parse as a valid agent snapshot. This rejects malformed
// payloads, memory-bearing PNGs, JSON with inconsistent memory fields,
// and any format/extension mismatch not caught by magic-byte check.
decode_snapshot_from_bytes(&bytes).map_err(|e| format!("invalid snapshot: {e}"))?;
// 4. Bytes must parse as the snapshot type selected by the filename.
// Team parsing rejects retired flat JSON and persona-pack ZIP inputs
// before anything reaches the frontend.
match kind {
SnapshotFileKind::AgentJson | SnapshotFileKind::AgentPng => {
decode_snapshot_from_bytes(&bytes)
.map_err(|e| format!("invalid agent snapshot: {e}"))?;
}
SnapshotFileKind::TeamJson | SnapshotFileKind::TeamPng => {
decode_team_snapshot_from_bytes(&bytes)
.map_err(|e| format!("invalid team snapshot: {e}"))?;
}
}

Ok(tauri::ipc::Response::new(bytes))
}
Expand All @@ -416,14 +472,14 @@ mod tests {
#[test]
fn snapshot_kind_json_returns_json_kind_and_correct_cap() {
let kind = snapshot_kind_for_filename("analyst.agent.json").unwrap();
assert_eq!(kind, SnapshotFileKind::Json);
assert_eq!(kind, SnapshotFileKind::AgentJson);
assert_eq!(kind.cap(), MAX_SNAPSHOT_JSON_BYTES as u64);
}

#[test]
fn snapshot_kind_png_returns_png_kind_and_correct_cap() {
let kind = snapshot_kind_for_filename("analyst.agent.png").unwrap();
assert_eq!(kind, SnapshotFileKind::Png);
assert_eq!(kind, SnapshotFileKind::AgentPng);
assert_eq!(kind.cap(), MAX_SNAPSHOT_PNG_BYTES as u64);
}

Expand All @@ -449,6 +505,33 @@ mod tests {
assert!(snapshot_kind_for_filename("agentjson").is_err());
}

#[test]
fn snapshot_kind_team_extensions_are_case_insensitive_and_scale_caps() {
let json = snapshot_kind_for_filename("review.TEAM.JSON").unwrap();
let png = snapshot_kind_for_filename("review.TEAM.PNG").unwrap();
assert_eq!(json, SnapshotFileKind::TeamJson);
assert_eq!(png, SnapshotFileKind::TeamPng);
assert_eq!(json.cap(), 25 * 1024 * 1024);
assert_eq!(png.cap(), 50 * 1024 * 1024);
}

#[test]
fn fetch_boundary_team_png_filename_with_json_bytes_rejected() {
let bytes = br#"{"format":"buzz-team-snapshot","version":1}"#;
let kind = snapshot_kind_for_filename("review.team.png").unwrap();
let error = ensure_bytes_match_kind(bytes, kind).unwrap_err();
assert!(error.contains(".team.png") && error.contains("not a PNG"));
}

#[test]
fn fetch_boundary_team_declared_size_over_cap_rejected() {
let kind = snapshot_kind_for_filename("review.team.json").unwrap();
assert!(ensure_declared_size_within_cap(MAX_TEAM_SNAPSHOT_JSON_BYTES, kind).is_ok());
let error =
ensure_declared_size_within_cap(MAX_TEAM_SNAPSHOT_JSON_BYTES + 1, kind).unwrap_err();
assert!(error.contains("25 MiB"));
}

// ── Focused boundary tests: format mismatch and consistency ──────────────
//
// These tests exercise the guard logic that fetch_snapshot_bytes applies
Expand Down
2 changes: 2 additions & 0 deletions desktop/src-tauri/src/commands/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ mod project_terminal;
mod relay_members;
mod relay_reconnect;
mod social;
mod team_snapshot;
mod teams;
mod updater;
mod window_vibrancy;
Expand Down Expand Up @@ -80,6 +81,7 @@ pub use project_terminal::*;
pub use relay_members::*;
pub use relay_reconnect::*;
pub use social::*;
pub use team_snapshot::*;
pub use teams::*;
pub use updater::*;
pub use window_vibrancy::*;
Expand Down
6 changes: 3 additions & 3 deletions desktop/src-tauri/src/commands/personas/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ fn trim_optional(value: Option<String>) -> Option<String> {
}

mod pending;
use pending::retain_persona_pending;
pub(in crate::commands) use pending::retain_persona_pending;
pub(super) use pending::tombstone_persona_pending;

#[tauri::command]
Expand Down Expand Up @@ -973,11 +973,11 @@ pub async fn set_persona_active(
}

pub(crate) const PNG_MAGIC: [u8; 4] = [0x89, 0x50, 0x4E, 0x47];

mod snapshot;
pub use snapshot::encode_agent_snapshot_for_send;
pub use snapshot::export_agent_snapshot;
pub(crate) use snapshot::import::{
decode_snapshot_from_bytes, MAX_SNAPSHOT_JSON_BYTES, MAX_SNAPSHOT_PNG_BYTES,
decode_snapshot_from_bytes, resolve_snapshot_import_behavior, MAX_SNAPSHOT_JSON_BYTES,
MAX_SNAPSHOT_PNG_BYTES,
};
pub use snapshot::{confirm_agent_snapshot_import, preview_agent_snapshot_import};
2 changes: 1 addition & 1 deletion desktop/src-tauri/src/commands/personas/pending.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ use crate::managed_agents::AgentDefinition;
/// does not retain, so the local-only `is_active` toggle never republishes, and
/// a byte-identical user-save republish is harmlessly NIP-33-replaced). The
/// guard is intentionally omitted.
pub(in crate::commands::personas) fn retain_persona_pending(
pub(in crate::commands) fn retain_persona_pending(
app: &AppHandle,
state: &AppState,
persona: &AgentDefinition,
Expand Down
Loading
Loading