diff --git a/config.toml.example b/config.toml.example index 8ac355ac..00413a44 100644 --- a/config.toml.example +++ b/config.toml.example @@ -126,3 +126,16 @@ error_hold_ms = 2500 # schedule = "0 0 * * 0" # channel = "123456789" # message = "generate weekly status report" + +# --- Outbound file attachments (agent → chat) --- +# When enabled, agents that include `![alt](/path/to/file)` markdown in +# their replies will have the file uploaded as a native chat attachment. +# +# Security: only files under ~/.oab/outgoing/ are permitted. The agent +# must explicitly copy files there before referencing them. +# +# [outbound] +# enabled = false # opt-in, disabled by default +# max_file_size_mb = 25 # 25 MB matches Discord's limit +# max_per_message = 10 # per agent response +# max_per_minute_per_channel = 30 # sliding window flood guard \ No newline at end of file diff --git a/docs/security.md b/docs/security.md new file mode 100644 index 00000000..6f26e74c --- /dev/null +++ b/docs/security.md @@ -0,0 +1,109 @@ +# Security Model — Shared Responsibility + +OpenAB bridges messaging platforms (Discord, Slack, Telegram, LINE) to coding agent CLIs over ACP. This document defines the security boundaries and who is responsible for what. + +## Responsibility Layers + +``` +┌──────────────────────────────────────────────────────────┐ +│ Infrastructure Layer (K8s / VPC / Network) │ +│ Egress/ingress control, network isolation, pod security │ +│ → Responsibility: Platform / Infra team │ +├──────────────────────────────────────────────────────────┤ +│ OpenAB Layer (this project) │ +│ Message routing, outbound content validation, rate limit │ +│ → Responsibility: OpenAB maintainers │ +├──────────────────────────────────────────────────────────┤ +│ Agent CLI Layer (Kiro, Claude Code, Codex, Gemini, etc.) │ +│ Filesystem access, tool permissions, sandbox policy │ +│ → Responsibility: User configuration + agent vendor │ +├──────────────────────────────────────────────────────────┤ +│ User Behavior Layer │ +│ OAuth tokens, API keys, prompts, intentional actions │ +│ → Responsibility: End user │ +└──────────────────────────────────────────────────────────┘ +``` + +## What OpenAB Controls + +**Inbound (user → agent):** +- Channel and user allowlists (`allowed_channels`, `allowed_users`) +- Bot message gating (`allow_bot_messages`, `trusted_bot_ids`) +- Bot turn limits (soft + hard caps to prevent runaway loops) +- Attachment size limits and text file caps + +**Outbound (agent → chat):** +- Outbound file attachments are **opt-in** (`outbound.enabled = false` by default) +- When enabled, only files under `~/.oab/outgoing/` are permitted +- **Only image files** are accepted (validated by magic bytes: PNG, JPEG, GIF, WebP, BMP) +- Per-message and per-channel rate limiting prevents flood +- Path traversal and symlink escape blocked via `canonicalize` + `Path::starts_with` + +**What OpenAB does NOT control:** +- What the agent CLI does with filesystem access (that's the agent's sandbox policy) +- What the agent sends via its own network calls (e.g. `curl`, API calls) +- How the agent's tools are configured (e.g. `--trust-all-tools`) + +## Why Images Only + +When an agent produces a file and sends it back through OpenAB to a chat channel, that file crosses a trust boundary — from the agent's local environment to a shared messaging platform. OpenAB is the gatekeeper at this boundary. + +**Threat:** a prompt-injected agent could dump environment variables, secrets, or sensitive files into the outgoing directory and request OpenAB to deliver them to the chat channel. + +**Mitigation:** OpenAB validates file content via magic bytes. Only files whose headers match known image formats are accepted. This blocks the most common exfiltration vector (text/binary dumps) while preserving the primary use case — screenshots, diagrams, charts, and generated images. + +**What this does not prevent:** +- Data hidden in image metadata (EXIF, PNG tEXt chunks) or steganography +- Agent pasting secrets directly in reply text (already possible without outbound attachments) + +**Why this is acceptable:** the image-only check raises the attack bar from trivial (`env > leak.txt`) to sophisticated (encoding secrets into valid image bytes). The fundamental truth: if you don't trust your agent's text output, you shouldn't trust its file output either — outbound attachments don't make the existing text-exfil risk worse. + +## Non-Image Files: Recommended Pattern + +For documents (PDF, Word, Excel, CSV) and other non-image artifacts, the recommended pattern is: + +1. Agent uploads the file to an **external storage provider** (Google Drive, S3, SharePoint, etc.) +2. Agent returns a **signed URL with TTL** to the chat +3. User clicks the link to access the file + +**Benefits:** +- File never touches the chat platform's servers — no compliance issues +- Access control is handled by the storage provider (OAuth, IAM policies) +- URLs can expire (e.g. S3 presigned URL with 60-second TTL) +- Works across all chat platforms (Discord, Slack, Telegram, LINE) + +**Example with S3:** +``` +Agent: "Here's your report: https://bucket.s3.amazonaws.com/report.pdf?X-Amz-Expires=60&..." +``` + +**Example with Google Drive:** +``` +Agent: "Report uploaded: https://drive.google.com/file/d/abc123/view (shared with your org)" +``` + +## Enterprise Deployment + +### Q: Can an agent exfiltrate data to external services? + +The agent CLI runs inside a container. Network-level controls are the infrastructure team's responsibility: + +- **Kubernetes:** use `NetworkPolicy` to restrict pod egress to specific CIDRs or services +- **AWS VPC:** deploy in a private subnet with no internet gateway; use VPC endpoints for allowed AWS services only +- **Service mesh:** use Istio/Linkerd egress policies to allowlist specific external domains + +OpenAB provides a container-ready architecture (`Dockerfile` + Helm chart + k8s manifests) that integrates with these controls. + +### Q: Can an agent send sensitive files to the chat channel? + +- Outbound attachments are **disabled by default** — operators must explicitly opt in +- When enabled, only **image files** are accepted (magic bytes validation) +- Only files in `~/.oab/outgoing/` are permitted — the agent must explicitly copy files there +- Non-image documents should use the signed-URL pattern described above + +### Q: How do I audit what the agent sends? + +- All outbound attachment operations are logged via `tracing` at INFO/WARN level +- Accepted files: `outbound: attachment accepted path=... size=...` +- Blocked files: `outbound: not an image file`, `outbound: path not in outgoing dir`, `outbound: over size limit` +- Rate limit hits: `outbound: rate-limit hit, dropping excess` diff --git a/src/adapter.rs b/src/adapter.rs index a66b5aa9..7376ca68 100644 --- a/src/adapter.rs +++ b/src/adapter.rs @@ -5,10 +5,11 @@ use std::sync::Arc; use tracing::error; use crate::acp::{classify_notification, AcpEvent, ContentBlock, SessionPool}; -use crate::config::ReactionsConfig; +use crate::config::{OutboundConfig, ReactionsConfig}; use crate::error_display::{format_coded_error, format_user_error}; use crate::format; use crate::markdown::{self, TableMode}; +use crate::outbound_rate::OutboundRateLimiter; use crate::reactions::StatusReactionController; // --- Platform-agnostic types --- @@ -130,6 +131,16 @@ pub trait ChatAdapter: Send + Sync + 'static { /// not be detected until the next message. This is acceptable: the first /// response may stream, but subsequent ones will correctly use send-once. fn use_streaming(&self, other_bot_present: bool) -> bool; + + /// Upload file attachments as follow-up messages. Default no-op so + /// adapters that don't support native file upload silently drop. + async fn send_file_attachments( + &self, + _channel: &ChannelRef, + _paths: &[std::path::PathBuf], + ) -> Result<()> { + Ok(()) + } } // --- AdapterRouter --- @@ -140,14 +151,18 @@ pub struct AdapterRouter { pool: Arc, reactions_config: ReactionsConfig, table_mode: TableMode, + outbound_config: OutboundConfig, + outbound_rate: Arc, } impl AdapterRouter { - pub fn new(pool: Arc, reactions_config: ReactionsConfig, table_mode: TableMode) -> Self { + pub fn new(pool: Arc, reactions_config: ReactionsConfig, table_mode: TableMode, outbound_config: OutboundConfig) -> Self { Self { pool, reactions_config, table_mode, + outbound_config, + outbound_rate: Arc::new(OutboundRateLimiter::new()), } } @@ -274,6 +289,10 @@ impl AdapterRouter { let message_limit = adapter.message_limit(); let streaming = adapter.use_streaming(other_bot_present); let table_mode = self.table_mode; + let outbound_cfg = self.outbound_config.clone(); + let outbound_rate = Arc::clone(&self.outbound_rate); + let outbound_channel_key = + format!("{}:{}", adapter.platform(), &thread_channel.channel_id); self.pool .with_connection(thread_key, |conn| { @@ -426,6 +445,29 @@ impl AdapterRouter { }; let final_content = markdown::convert_tables(&final_content, table_mode); + + // Extract outbound `![alt](/path)` attachment markers. + // No-op when `outbound.enabled` is false (the default). + let (final_content, mut outbound_paths) = + crate::media::extract_outbound_attachments(&final_content, &outbound_cfg); + + if !outbound_paths.is_empty() && outbound_cfg.enabled { + let grant = outbound_rate.admit( + &outbound_channel_key, + outbound_paths.len(), + outbound_cfg.max_per_minute_per_channel, + ); + if grant < outbound_paths.len() { + tracing::warn!( + channel = outbound_channel_key, + requested = outbound_paths.len(), + granted = grant, + "outbound: rate-limit hit, dropping excess" + ); + outbound_paths.truncate(grant); + } + } + let chunks = format::split_message(&final_content, message_limit); if let Some(msg) = placeholder_msg { // Streaming: edit first chunk into placeholder, send rest as new messages @@ -442,6 +484,15 @@ impl AdapterRouter { } } + if !outbound_paths.is_empty() { + if let Err(e) = adapter + .send_file_attachments(&thread_channel, &outbound_paths) + .await + { + tracing::warn!(error = %e, "outbound: send_file_attachments failed"); + } + } + Ok(()) }) }) diff --git a/src/config.rs b/src/config.rs index 68a6fdaf..9df5ed0d 100644 --- a/src/config.rs +++ b/src/config.rs @@ -48,8 +48,61 @@ pub struct Config { pub markdown: MarkdownConfig, #[serde(default)] pub cron: CronConfig, + #[serde(default)] + pub outbound: OutboundConfig, +} + +/// Controls outbound file attachments — the `![alt](/path)` markdown marker +/// in agent responses that instructs the bot to upload a local file as a +/// native chat attachment. Disabled by default; operators must explicitly +/// opt in because this opens a path from the host filesystem to the chat +/// channel. +/// +/// Security: only files under `~/.oab/outgoing/` are permitted. The agent +/// must explicitly copy files there before referencing them. This eliminates +/// path traversal and symlink escape risks by design. +#[derive(Debug, Clone, Deserialize)] +pub struct OutboundConfig { + #[serde(default)] + pub enabled: bool, + #[serde(default = "default_outbound_max_size_mb")] + pub max_file_size_mb: u64, + #[serde(default = "default_outbound_max_per_message")] + pub max_per_message: usize, + #[serde(default = "default_outbound_max_per_minute")] + pub max_per_minute_per_channel: usize, +} + +impl Default for OutboundConfig { + fn default() -> Self { + Self { + enabled: false, + max_file_size_mb: default_outbound_max_size_mb(), + max_per_message: default_outbound_max_per_message(), + max_per_minute_per_channel: default_outbound_max_per_minute(), + } + } +} + +impl OutboundConfig { + pub fn max_size_bytes(&self) -> u64 { + self.max_file_size_mb.saturating_mul(1024 * 1024) + } + + /// The single hardcoded outgoing directory. Agents must copy files here. + pub fn outgoing_dir() -> std::path::PathBuf { + std::env::var("HOME") + .map(std::path::PathBuf::from) + .unwrap_or_else(|_| std::path::PathBuf::from("/tmp")) + .join(".oab") + .join("outgoing") + } } +fn default_outbound_max_size_mb() -> u64 { 25 } +fn default_outbound_max_per_message() -> usize { 10 } +fn default_outbound_max_per_minute() -> usize { 30 } + #[derive(Debug, Clone, Default, Deserialize)] pub struct CronConfig { /// Enable usercron hot-reload (default: false). Must be explicitly set to true. diff --git a/src/discord.rs b/src/discord.rs index 86e12990..2b2a1ff6 100644 --- a/src/discord.rs +++ b/src/discord.rs @@ -7,7 +7,7 @@ use crate::format; use crate::media; use async_trait::async_trait; use std::sync::LazyLock; -use serenity::builder::{CreateActionRow, CreateCommand, CreateInteractionResponse, CreateInteractionResponseMessage, CreateSelectMenu, CreateSelectMenuKind, CreateSelectMenuOption, CreateThread, EditMessage}; +use serenity::builder::{CreateActionRow, CreateAttachment, CreateCommand, CreateInteractionResponse, CreateInteractionResponseMessage, CreateMessage, CreateSelectMenu, CreateSelectMenuKind, CreateSelectMenuOption, CreateThread, EditMessage}; use serenity::http::Http; use serenity::model::application::{ComponentInteractionDataKind, Interaction}; use serenity::model::channel::{AutoArchiveDuration, Message, MessageType, ReactionType}; @@ -128,6 +128,30 @@ impl ChatAdapter for DiscordAdapter { .await?; Ok(()) } + + async fn send_file_attachments( + &self, + channel: &ChannelRef, + paths: &[std::path::PathBuf], + ) -> anyhow::Result<()> { + let ch_id: u64 = Self::resolve_channel(channel).parse()?; + for path in paths { + match CreateAttachment::path(path).await { + Ok(file) => { + let msg = CreateMessage::new().add_file(file); + if let Err(e) = ChannelId::new(ch_id).send_message(&self.http, msg).await { + tracing::warn!(path = %path.display(), error = %e, "outbound: discord upload failed"); + } else { + info!(path = %path.display(), "outbound: attachment sent"); + } + } + Err(e) => { + tracing::warn!(path = %path.display(), error = %e, "outbound: failed to read file"); + } + } + } + Ok(()) + } } // --- Handler: serenity EventHandler that delegates to AdapterRouter --- diff --git a/src/main.rs b/src/main.rs index 8dce88a4..0af83a43 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,6 +8,7 @@ mod error_display; mod format; mod markdown; mod media; +mod outbound_rate; mod reactions; mod setup; mod slack; @@ -136,7 +137,7 @@ async fn main() -> anyhow::Result<()> { info!(model = %cfg.stt.model, base_url = %cfg.stt.base_url, "STT enabled"); } - let router = Arc::new(AdapterRouter::new(pool.clone(), cfg.reactions, cfg.markdown.tables)); + let router = Arc::new(AdapterRouter::new(pool.clone(), cfg.reactions, cfg.markdown.tables, cfg.outbound)); // Shutdown signal for Slack adapter let (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(false); diff --git a/src/media.rs b/src/media.rs index 5e0c057f..c3aaf7b5 100644 --- a/src/media.rs +++ b/src/media.rs @@ -1,11 +1,12 @@ use crate::acp::ContentBlock; -use crate::config::SttConfig; +use crate::config::{OutboundConfig, SttConfig}; use base64::engine::general_purpose::STANDARD as BASE64; use base64::Engine; use image::ImageReader; use std::io::Cursor; +use std::path::PathBuf; use std::sync::LazyLock; -use tracing::{debug, error}; +use tracing::{debug, error, info, warn}; /// Reusable HTTP client for downloading attachments (shared across adapters). pub static HTTP_CLIENT: LazyLock = LazyLock::new(|| { @@ -290,6 +291,122 @@ pub async fn download_and_read_text_file( )) } +// --- Outbound attachments --- +// +// Agent → chat file upload. Agents write `![alt](/path)` markdown in their +// response; this module extracts and validates paths. Only files under +// `~/.oab/outgoing/` are permitted — the agent must explicitly copy files +// there before referencing them. + +/// Regex for outbound attachment markers: `![alt](/path/to/file)`. +static OUTBOUND_RE: LazyLock = LazyLock::new(|| { + regex::Regex::new(r"!\[[^\]]*\]\((/[^\)]+)\)").unwrap() +}); + +/// Check file magic bytes to verify it is an image. Only images are +/// allowed for outbound attachments to prevent data exfiltration via +/// text files. +fn is_image_file(path: &std::path::Path) -> bool { + let Ok(buf) = std::fs::read(path) else { + return false; + }; + // Check magic bytes for common image formats + let header = buf.get(..12).unwrap_or(&buf); + matches!( + header, + [0x89, 0x50, 0x4E, 0x47, ..] // PNG + | [0xFF, 0xD8, 0xFF, ..] // JPEG + | [0x47, 0x49, 0x46, 0x38, ..] // GIF + | [0x52, 0x49, 0x46, 0x46, _, _, _, _, 0x57, 0x45, 0x42, 0x50] // WebP + | [0x42, 0x4D, ..] // BMP + ) +} + +/// Scan agent response `text` for `![alt](/path)` markers, validate each +/// path against `config`, and return `(cleaned_text, list_of_paths)`. +/// +/// Only files under `OutboundConfig::outgoing_dir()` are accepted. +/// Markers for accepted files are stripped; rejected markers stay visible. +pub fn extract_outbound_attachments( + text: &str, + config: &OutboundConfig, +) -> (String, Vec) { + if !config.enabled { + return (text.to_string(), Vec::new()); + } + + let outgoing_dir = OutboundConfig::outgoing_dir(); + if let Err(e) = std::fs::create_dir_all(&outgoing_dir) { + warn!(dir = %outgoing_dir.display(), error = %e, "outbound: cannot create outgoing dir"); + return (text.to_string(), Vec::new()); + } + + let canonical_outgoing = match std::fs::canonicalize(&outgoing_dir) { + Ok(p) => p, + Err(e) => { + warn!(dir = %outgoing_dir.display(), error = %e, "outbound: cannot canonicalize outgoing dir"); + return (text.to_string(), Vec::new()); + } + }; + + let mut attachments = Vec::new(); + let mut markers_to_strip = Vec::new(); + + for cap in OUTBOUND_RE.captures_iter(text) { + if attachments.len() >= config.max_per_message { + warn!(cap = config.max_per_message, "outbound: per-message cap hit"); + break; + } + + let full_match = cap.get(0).unwrap().as_str(); + let path_str = &cap[1]; + let path = PathBuf::from(path_str); + + let canonical = match std::fs::canonicalize(&path) { + Ok(p) => p, + Err(e) => { + debug!(path = %path_str, error = %e, "outbound: cannot canonicalize"); + continue; + } + }; + + if !canonical.starts_with(&canonical_outgoing) { + warn!(path = %path_str, canonical = %canonical.display(), "outbound: path not in outgoing dir"); + continue; + } + + match std::fs::metadata(&canonical) { + Ok(meta) if meta.is_file() && meta.len() <= config.max_size_bytes() => { + if !is_image_file(&canonical) { + warn!(path = %canonical.display(), "outbound: not an image file, only images are allowed"); + continue; + } + info!(path = %canonical.display(), size = meta.len(), "outbound: attachment accepted"); + attachments.push(canonical); + markers_to_strip.push(full_match.to_string()); + } + Ok(meta) if meta.len() > config.max_size_bytes() => { + warn!(path = %canonical.display(), size = meta.len(), limit_mb = config.max_file_size_mb, "outbound: over size limit"); + } + Ok(_) => { + warn!(path = %canonical.display(), "outbound: not a regular file"); + } + Err(e) => { + debug!(path = %canonical.display(), error = %e, "outbound: metadata error"); + } + } + } + + let mut cleaned = text.to_string(); + for marker in &markers_to_strip { + cleaned = cleaned.replace(marker, ""); + } + while cleaned.contains("\n\n\n") { + cleaned = cleaned.replace("\n\n\n", "\n\n"); + } + (cleaned.trim().to_string(), attachments) +} + #[cfg(test)] mod tests { use super::*; @@ -372,3 +489,153 @@ mod tests { assert!(resize_and_compress(&garbage).is_err()); } } + +#[cfg(test)] +mod outbound_tests { + use super::*; + use crate::config::OutboundConfig; + + fn cfg_enabled() -> OutboundConfig { + OutboundConfig { + enabled: true, + ..OutboundConfig::default() + } + } + + fn outgoing_dir() -> PathBuf { + OutboundConfig::outgoing_dir() + } + + #[test] + fn disabled_by_default_is_noop() { + let cfg = OutboundConfig::default(); + assert!(!cfg.enabled); + let text = "![foo](/tmp/does-not-matter.png)"; + let (cleaned, atts) = extract_outbound_attachments(text, &cfg); + assert_eq!(cleaned, text); + assert!(atts.is_empty()); + } + + #[test] + fn enabled_extracts_outgoing_file() { + let dir = outgoing_dir(); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("test_happy.png"); + // Valid PNG magic bytes + std::fs::write(&path, &[0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]).unwrap(); + let text = format!("Here: ![screenshot]({}) done.", path.display()); + let (cleaned, atts) = extract_outbound_attachments(&text, &cfg_enabled()); + assert_eq!(atts.len(), 1); + assert!(!cleaned.contains("test_happy")); + std::fs::remove_file(&path).ok(); + } + + #[test] + fn blocks_path_outside_outgoing() { + let text = "![secret](/etc/passwd)"; + let (cleaned, atts) = extract_outbound_attachments(text, &cfg_enabled()); + assert!(atts.is_empty()); + assert!(cleaned.contains("/etc/passwd")); + } + + #[test] + fn blocks_tmp_path() { + let path = "/tmp/openab_outbound_test.png"; + std::fs::write(path, b"x").unwrap(); + let text = format!("![img]({path})"); + let (_, atts) = extract_outbound_attachments(&text, &cfg_enabled()); + assert!(atts.is_empty(), "/tmp must be blocked, only ~/.oab/outgoing/ allowed"); + std::fs::remove_file(path).ok(); + } + + #[test] + fn blocks_symlink_escape() { + let dir = outgoing_dir(); + std::fs::create_dir_all(&dir).unwrap(); + let link = dir.join("escape.png"); + let _ = std::fs::remove_file(&link); + std::os::unix::fs::symlink("/etc/hosts", &link).unwrap(); + let text = format!("![esc]({})", link.display()); + let (_, atts) = extract_outbound_attachments(&text, &cfg_enabled()); + assert!(atts.is_empty(), "symlink escaping outgoing dir must be blocked"); + std::fs::remove_file(&link).ok(); + } + + #[test] + fn blocks_path_traversal() { + let dir = outgoing_dir(); + let text = format!("![x]({}/../../../etc/hosts)", dir.display()); + let (_, atts) = extract_outbound_attachments(&text, &cfg_enabled()); + assert!(atts.is_empty(), "path traversal must be blocked"); + } + + #[test] + fn enforces_max_per_message() { + let dir = outgoing_dir(); + std::fs::create_dir_all(&dir).unwrap(); + let png = [0x89u8, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]; + let mut text = String::new(); + let mut paths = Vec::new(); + for i in 0..5 { + let p = dir.join(format!("cap_{i}.png")); + std::fs::write(&p, &png).unwrap(); + text.push_str(&format!("![a{i}]({})\n", p.display())); + paths.push(p); + } + let cfg = OutboundConfig { + enabled: true, + max_per_message: 2, + ..OutboundConfig::default() + }; + let (_, atts) = extract_outbound_attachments(&text, &cfg); + assert_eq!(atts.len(), 2, "must cap at max_per_message"); + for p in &paths { std::fs::remove_file(p).ok(); } + } + + #[test] + fn blocks_text_file_exfiltration() { + let dir = outgoing_dir(); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("secrets.txt"); + std::fs::write(&path, b"SECRET_KEY=hunter2").unwrap(); + let text = format!("![leak]({})", path.display()); + let (cleaned, atts) = extract_outbound_attachments(&text, &cfg_enabled()); + assert!(atts.is_empty(), "text files must be blocked"); + assert!(cleaned.contains("secrets.txt"), "blocked marker stays visible"); + std::fs::remove_file(&path).ok(); + } + + #[test] + fn accepts_real_png() { + let dir = outgoing_dir(); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("real.png"); + // Minimal valid PNG header + let png_header: Vec = vec![0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]; + std::fs::write(&path, &png_header).unwrap(); + let text = format!("![img]({})", path.display()); + let (_, atts) = extract_outbound_attachments(&text, &cfg_enabled()); + assert_eq!(atts.len(), 1, "real PNG must be accepted"); + std::fs::remove_file(&path).ok(); + } + + #[test] + fn enforces_max_file_size() { + let dir = outgoing_dir(); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("large.bin"); + // PNG header + padding to exceed size limit + let mut data = vec![0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]; + data.resize(2 * 1024 * 1024, 0); + std::fs::write(&path, &data).unwrap(); + let cfg = OutboundConfig { + enabled: true, + max_file_size_mb: 1, + ..OutboundConfig::default() + }; + let text = format!("![big]({})", path.display()); + let (_, atts) = extract_outbound_attachments(&text, &cfg); + assert!(atts.is_empty(), "file exceeding max_file_size_mb must be blocked"); + std::fs::remove_file(&path).ok(); + } +} \ No newline at end of file diff --git a/src/outbound_rate.rs b/src/outbound_rate.rs new file mode 100644 index 00000000..9584df3c --- /dev/null +++ b/src/outbound_rate.rs @@ -0,0 +1,92 @@ +//! Per-channel sliding-window rate limiter for outbound attachments. + +use std::collections::{HashMap, VecDeque}; +use std::sync::Mutex; +use std::time::{Duration, Instant}; + +const WINDOW: Duration = Duration::from_secs(60); + +pub struct OutboundRateLimiter { + inner: Mutex>>, +} + +impl OutboundRateLimiter { + pub fn new() -> Self { + Self { + inner: Mutex::new(HashMap::new()), + } + } + + /// Request permission to send `requested` attachments. Returns the + /// number actually granted. + pub fn admit(&self, channel_key: &str, requested: usize, limit: usize) -> usize { + if requested == 0 || limit == 0 { + return 0; + } + let now = Instant::now(); + let cutoff = now - WINDOW; + let mut map = self.inner.lock().expect("rate limiter mutex poisoned"); + let entry = map.entry(channel_key.to_string()).or_default(); + while entry.front().is_some_and(|t| *t < cutoff) { + entry.pop_front(); + } + let remaining = limit.saturating_sub(entry.len()); + let grant = requested.min(remaining); + for _ in 0..grant { + entry.push_back(now); + } + grant + } +} + +impl Default for OutboundRateLimiter { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn admits_up_to_limit_then_blocks() { + let rl = OutboundRateLimiter::new(); + assert_eq!(rl.admit("ch1", 5, 5), 5); + assert_eq!(rl.admit("ch1", 1, 5), 0); + } + + #[test] + fn partial_admit() { + let rl = OutboundRateLimiter::new(); + assert_eq!(rl.admit("ch1", 3, 5), 3); + assert_eq!(rl.admit("ch1", 5, 5), 2); + } + + #[test] + fn channels_are_independent() { + let rl = OutboundRateLimiter::new(); + assert_eq!(rl.admit("ch1", 5, 5), 5); + assert_eq!(rl.admit("ch2", 5, 5), 5); + } + + #[test] + fn zero_grants_zero() { + let rl = OutboundRateLimiter::new(); + assert_eq!(rl.admit("ch1", 0, 10), 0); + assert_eq!(rl.admit("ch1", 10, 0), 0); + } + + #[test] + fn prunes_old_entries() { + let rl = OutboundRateLimiter::new(); + { + let mut map = rl.inner.lock().unwrap(); + let entry = map.entry("ch1".to_string()).or_default(); + for _ in 0..5 { + entry.push_back(Instant::now() - Duration::from_secs(120)); + } + } + assert_eq!(rl.admit("ch1", 5, 5), 5); + } +}