| 1 | use super::*; |
| 2 | use codewhale_models::Role; |
| 3 | |
| 4 | /// A 1x1 PNG, as bytes rather than a fixture file so the encoding tests |
| 5 | /// have no filesystem dependency. |
| 6 | pub(crate) const PNG_1X1: &[u8] = &[ |
| 7 | 0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 0x00, 0x00, 0x00, 0x0d, 0x49, 0x48, 0x44, 0x52, |
| 8 | 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x06, 0x00, 0x00, 0x00, 0x1f, 0x15, 0xc4, |
| 9 | 0x89, 0x00, 0x00, 0x00, 0x0a, 0x49, 0x44, 0x41, 0x54, 0x78, 0x9c, 0x63, 0x00, 0x01, 0x00, 0x00, |
| 10 | 0x05, 0x00, 0x01, 0x0d, 0x0a, 0x2d, 0xb4, 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4e, 0x44, 0xae, |
| 11 | 0x42, 0x60, 0x82, |
| 12 | ]; |
| 13 | |
| 14 | #[test] |
| 15 | fn sniffs_every_accepted_format_from_magic_bytes() { |
| 16 | assert_eq!(sniff_media_type(PNG_1X1), Some("image/png")); |
| 17 | assert_eq!( |
| 18 | sniff_media_type(&[0xff, 0xd8, 0xff, 0xe0, 0x00]), |
| 19 | Some("image/jpeg") |
| 20 | ); |
| 21 | assert_eq!(sniff_media_type(b"GIF89a....."), Some("image/gif")); |
| 22 | assert_eq!(sniff_media_type(b"GIF87a....."), Some("image/gif")); |
| 23 | assert_eq!( |
| 24 | sniff_media_type(b"RIFF\x00\x00\x00\x00WEBPVP8 "), |
| 25 | Some("image/webp") |
| 26 | ); |
| 27 | } |
| 28 | |
| 29 | #[test] |
| 30 | fn sniffing_ignores_the_extension_and_believes_the_bytes() { |
| 31 | // A JPEG named .png must be declared image/jpeg, or the provider |
| 32 | // rejects the media-type mismatch. |
| 33 | let jpeg = [0xff, 0xd8, 0xff, 0xe0, 0x11, 0x22]; |
| 34 | let attached = encode_image_bytes(&jpeg, "screenshot.png").expect("attach"); |
| 35 | assert_eq!(attached.media_type, "image/jpeg"); |
| 36 | assert!(attached.data_url.starts_with("data:image/jpeg;base64,")); |
| 37 | } |
| 38 | |
| 39 | #[test] |
| 40 | fn rich_tool_images_reject_corruption_mime_spoofing_and_decode_bombs() { |
| 41 | use crate::tools::spec::{RichToolResult, ToolResult}; |
| 42 | use codewhale_tools::ToolResultContentBlock; |
| 43 | let mut oversized_header = PNG_1X1.to_vec(); |
| 44 | oversized_header[16..20].copy_from_slice(&(MAX_IMAGE_DIMENSION + 1).to_be_bytes()); |
| 45 | for (mime, bytes) in [ |
| 46 | ("image/png", &PNG_1X1[..8]), |
| 47 | ("image/jpeg", PNG_1X1), |
| 48 | ("image/png", oversized_header.as_slice()), |
| 49 | ] { |
| 50 | assert!(prepare_tool_image_bytes(bytes, mime).block.is_none()); |
| 51 | let rich = bound_rich_tool_result(RichToolResult::with_content_blocks( |
| 52 | ToolResult::success("capture receipt"), |
| 53 | vec![ToolResultContentBlock::Image { |
| 54 | mime_type: mime.into(), |
| 55 | data: STANDARD.encode(bytes), |
| 56 | }], |
| 57 | )); |
| 58 | assert!(rich.result.success); |
| 59 | assert!(rich.content_blocks.is_empty()); |
| 60 | assert!(rich.result.content.starts_with("capture receipt")); |
| 61 | assert!( |
| 62 | rich.result |
| 63 | .content |
| 64 | .contains("1 tool-result image block(s) omitted") |
| 65 | ); |
| 66 | let stored = vec![ |
| 67 | serde_json::json!({"type":"image","mime_type":mime,"data":STANDARD.encode(bytes)}), |
| 68 | ]; |
| 69 | let (image, omitted) = provider_tool_result_image_refs(Some(&stored)); |
| 70 | assert!( |
| 71 | image.is_none(), |
| 72 | "restored history must not bypass validation" |
| 73 | ); |
| 74 | assert_eq!(omitted, 1); |
| 75 | } |
| 76 | } |
| 77 | |
| 78 | #[test] |
| 79 | fn lowercase_read_image_preparation_is_typed_and_bounded() { |
| 80 | let prepared = prepare_tool_image_bytes(PNG_1X1, "image/png"); |
| 81 | assert_eq!(prepared.note, "Read image file [image/png]"); |
| 82 | let codewhale_tools::ToolResultContentBlock::Image { mime_type, data } = |
| 83 | prepared.block.expect("typed image"); |
| 84 | assert_eq!(mime_type, "image/png"); |
| 85 | assert_eq!(STANDARD.decode(data).expect("base64"), PNG_1X1); |
| 86 | |
| 87 | let omitted = prepare_tool_image_bytes(b"BMnot-a-safe-bitmap", "image/bmp"); |
| 88 | assert!(omitted.block.is_none()); |
| 89 | assert!(omitted.note.contains("Image omitted"), "{}", omitted.note); |
| 90 | } |
| 91 | |
| 92 | #[test] |
| 93 | fn blind_route_removes_nested_tool_result_image() { |
| 94 | let mut messages = vec![codewhale_models::Message { |
| 95 | role: Role::User, |
| 96 | content: vec![ContentBlock::ToolResult { |
| 97 | tool_use_id: "call-image".to_string(), |
| 98 | content: "Read image file [image/png]".to_string(), |
| 99 | is_error: None, |
| 100 | content_blocks: Some(vec![serde_json::json!({ |
| 101 | "type": "image", |
| 102 | "mime_type": "image/png", |
| 103 | "data": "QUJD", |
| 104 | })]), |
| 105 | }], |
| 106 | }]; |
| 107 | |
| 108 | assert_eq!( |
| 109 | strip_images_when_unsupported(&mut messages, SupportState::Unsupported, "text-only",), |
| 110 | 1 |
| 111 | ); |
| 112 | let ContentBlock::ToolResult { |
| 113 | content, |
| 114 | content_blocks, |
| 115 | .. |
| 116 | } = &messages[0].content[0] |
| 117 | else { |
| 118 | panic!("tool result") |
| 119 | }; |
| 120 | assert!(content_blocks.is_none()); |
| 121 | assert!(content.contains("text-only"), "{content}"); |
| 122 | } |
| 123 | |
| 124 | #[test] |
| 125 | fn copy_projection_never_contains_inline_image_bytes() { |
| 126 | const SENTINEL: &str = "U0VOU0lUSVZFX0JBU0U2NA=="; |
| 127 | let projected = safe_tool_result_content_blocks(Some(&[serde_json::json!({ |
| 128 | "type": "image", |
| 129 | "mime_type": "image/png", |
| 130 | "data": SENTINEL, |
| 131 | })])) |
| 132 | .expect("projection"); |
| 133 | let encoded = serde_json::to_string(&projected).expect("json"); |
| 134 | assert!(!encoded.contains(SENTINEL), "{encoded}"); |
| 135 | assert!( |
| 136 | encoded.contains("inline_or_local_image_payload"), |
| 137 | "{encoded}" |
| 138 | ); |
| 139 | } |
| 140 | |
| 141 | #[test] |
| 142 | fn riff_that_is_not_webp_is_not_an_image() { |
| 143 | // A WAV file is also RIFF. Matching on "RIFF" alone would attach audio. |
| 144 | assert_eq!(sniff_media_type(b"RIFF\x00\x00\x00\x00WAVEfmt "), None); |
| 145 | } |
| 146 | |
| 147 | #[test] |
| 148 | fn encodes_a_png_to_a_data_url_that_round_trips() { |
| 149 | let attached = encode_image_bytes(PNG_1X1, "shot.png").expect("attach"); |
| 150 | assert_eq!(attached.media_type, "image/png"); |
| 151 | assert_eq!(attached.source_bytes, PNG_1X1.len()); |
| 152 | |
| 153 | let (media_type, payload) = parse_data_url(&attached.data_url).expect("parse"); |
| 154 | assert_eq!(media_type, "image/png"); |
| 155 | assert_eq!(STANDARD.decode(payload).expect("decode"), PNG_1X1); |
| 156 | } |
| 157 | |
| 158 | #[test] |
| 159 | fn rejects_a_file_over_the_size_limit() { |
| 160 | let oversized = vec![0u8; MAX_IMAGE_BYTES + 1]; |
| 161 | let error = encode_image_bytes(&oversized, "huge.png").expect_err("must reject"); |
| 162 | assert!( |
| 163 | matches!(error, ImageAttachError::TooLarge { .. }), |
| 164 | "got {error:?}" |
| 165 | ); |
| 166 | let rendered = error.to_string(); |
| 167 | assert!(rendered.contains("5.0 MB"), "{rendered}"); |
| 168 | assert!(rendered.contains("huge.png"), "{rendered}"); |
| 169 | } |
| 170 | |
| 171 | #[test] |
| 172 | fn accepts_a_file_exactly_at_the_size_limit() { |
| 173 | // The boundary is inclusive; an off-by-one here would reject images |
| 174 | // the providers accept. |
| 175 | let mut at_limit = PNG_1X1.to_vec(); |
| 176 | at_limit.resize(MAX_IMAGE_BYTES, 0); |
| 177 | assert!(encode_image_bytes(&at_limit, "edge.png").is_ok()); |
| 178 | } |
| 179 | |
| 180 | #[test] |
| 181 | fn rejects_a_real_image_in_an_unsupported_format_by_name() { |
| 182 | for (bytes, name) in [ |
| 183 | (b"BM\x00\x00\x00\x00".as_slice(), "BMP"), |
| 184 | (b"II\x2a\x00extra".as_slice(), "TIFF"), |
| 185 | (b"MM\x00\x2aextra".as_slice(), "TIFF"), |
| 186 | (b"<svg xmlns=".as_slice(), "SVG"), |
| 187 | (b"%PDF-1.7".as_slice(), "PDF"), |
| 188 | ] { |
| 189 | let error = encode_image_bytes(bytes, "f").expect_err("must reject"); |
| 190 | match error { |
| 191 | ImageAttachError::UnsupportedFormat { detected, .. } => { |
| 192 | assert_eq!(detected, name); |
| 193 | } |
| 194 | other => panic!("expected UnsupportedFormat for {name}, got {other:?}"), |
| 195 | } |
| 196 | } |
| 197 | } |
| 198 | |
| 199 | #[test] |
| 200 | fn rejects_a_file_that_is_not_an_image_at_all() { |
| 201 | let error = encode_image_bytes(b"#!/bin/sh\necho hi\n", "script.png").expect_err("must reject"); |
| 202 | assert!( |
| 203 | matches!(error, ImageAttachError::NotAnImage { .. }), |
| 204 | "got {error:?}" |
| 205 | ); |
| 206 | } |
| 207 | |
| 208 | #[test] |
| 209 | fn rejects_an_empty_file() { |
| 210 | let error = encode_image_bytes(b"", "empty.png").expect_err("must reject"); |
| 211 | assert!( |
| 212 | matches!(error, ImageAttachError::Empty { .. }), |
| 213 | "got {error:?}" |
| 214 | ); |
| 215 | } |
| 216 | |
| 217 | #[test] |
| 218 | fn parses_and_rejects_data_urls() { |
| 219 | assert_eq!( |
| 220 | parse_data_url("data:image/png;base64,QUJD"), |
| 221 | Some(("image/png", "QUJD")) |
| 222 | ); |
| 223 | // Not base64-tagged: Anthropic has no shape for a raw data URL. |
| 224 | assert_eq!(parse_data_url("data:image/png,QUJD"), None); |
| 225 | // Remote URLs are a different source type, not a malformed data URL. |
| 226 | assert_eq!(parse_data_url("https://example.com/a.png"), None); |
| 227 | // Degenerate forms must not produce an empty base64 payload that the |
| 228 | // provider would reject with an opaque error. |
| 229 | assert_eq!(parse_data_url("data:;base64,QUJD"), None); |
| 230 | assert_eq!(parse_data_url("data:image/png;base64,"), None); |
| 231 | assert_eq!(parse_data_url("data:image/png;base64"), None); |
| 232 | } |
| 233 | |
| 234 | #[test] |
| 235 | fn classifies_remote_urls() { |
| 236 | assert!(is_remote_image_url("https://example.com/a.png")); |
| 237 | assert!(is_remote_image_url("http://example.com/a.png")); |
| 238 | assert!(!is_remote_image_url("data:image/png;base64,QUJD")); |
| 239 | assert!(!is_remote_image_url("file:///tmp/a.png")); |
| 240 | } |
| 241 | |
| 242 | fn message_with_image(url: &str) -> codewhale_models::Message { |
| 243 | codewhale_models::Message { |
| 244 | role: Role::User, |
| 245 | content: vec![ |
| 246 | ContentBlock::ImageUrl { |
| 247 | image_url: ImageUrlContent { |
| 248 | url: url.to_string(), |
| 249 | }, |
| 250 | }, |
| 251 | ContentBlock::Text { |
| 252 | text: "what is this?".to_string(), |
| 253 | cache_control: None, |
| 254 | }, |
| 255 | ], |
| 256 | } |
| 257 | } |
| 258 | |
| 259 | fn write_png(dir: &std::path::Path, name: &str) -> std::path::PathBuf { |
| 260 | let path = dir.join(name); |
| 261 | std::fs::write(&path, PNG_1X1).expect("write fixture"); |
| 262 | path |
| 263 | } |
| 264 | |
| 265 | #[test] |
| 266 | fn expands_a_placeholder_into_an_image_block() { |
| 267 | let dir = tempfile::tempdir().expect("tempdir"); |
| 268 | let path = write_png(dir.path(), "shot.png"); |
| 269 | let text = format!("look at this\n[Attached image: {}]", path.display()); |
| 270 | |
| 271 | let expanded = expand_attachment_blocks(&text); |
| 272 | |
| 273 | assert!(expanded.notices.is_empty(), "{expanded:?}"); |
| 274 | // Bracketed: open tag naming the path, the image, close tag. |
| 275 | assert_eq!(expanded.blocks.len(), 3, "{expanded:?}"); |
| 276 | match &expanded.blocks[0] { |
| 277 | ContentBlock::Text { text, .. } => { |
| 278 | assert!(text.starts_with("<image path=\""), "{text}"); |
| 279 | assert!(text.contains("shot.png"), "{text}"); |
| 280 | } |
| 281 | other => panic!("expected an opening tag, got {other:?}"), |
| 282 | } |
| 283 | match &expanded.blocks[1] { |
| 284 | ContentBlock::ImageUrl { image_url } => { |
| 285 | assert!(image_url.url.starts_with("data:image/png;base64,")); |
| 286 | } |
| 287 | other => panic!("expected an image block, got {other:?}"), |
| 288 | } |
| 289 | assert_eq!( |
| 290 | expanded.blocks[2], |
| 291 | ContentBlock::Text { |
| 292 | text: "</image>".to_string(), |
| 293 | cache_control: None |
| 294 | } |
| 295 | ); |
| 296 | } |
| 297 | |
| 298 | #[test] |
| 299 | fn expands_multiple_placeholders_in_order() { |
| 300 | let dir = tempfile::tempdir().expect("tempdir"); |
| 301 | let first = write_png(dir.path(), "one.png"); |
| 302 | let second = write_png(dir.path(), "two.png"); |
| 303 | std::fs::write(&second, [0xff, 0xd8, 0xff, 0xe0, 0x01]).expect("write jpeg"); |
| 304 | let text = format!( |
| 305 | "[Attached image: {}]\nand\n[Attached image: {}]", |
| 306 | first.display(), |
| 307 | second.display() |
| 308 | ); |
| 309 | |
| 310 | let expanded = expand_attachment_blocks(&text); |
| 311 | |
| 312 | let media: Vec<_> = expanded |
| 313 | .blocks |
| 314 | .iter() |
| 315 | .filter_map(|block| match block { |
| 316 | ContentBlock::ImageUrl { image_url } => Some( |
| 317 | parse_data_url(&image_url.url) |
| 318 | .expect("data url") |
| 319 | .0 |
| 320 | .to_string(), |
| 321 | ), |
| 322 | _ => None, |
| 323 | }) |
| 324 | .collect(); |
| 325 | assert_eq!(media, vec!["image/png", "image/jpeg"]); |
| 326 | |
| 327 | // Each image carries its own path tag, so the model can tell two |
| 328 | // screenshots in one turn apart. |
| 329 | let tags: Vec<_> = expanded |
| 330 | .blocks |
| 331 | .iter() |
| 332 | .filter_map(|block| match block { |
| 333 | ContentBlock::Text { text, .. } if text.starts_with("<image path=") => { |
| 334 | Some(text.clone()) |
| 335 | } |
| 336 | _ => None, |
| 337 | }) |
| 338 | .collect(); |
| 339 | assert_eq!(tags.len(), 2, "{tags:?}"); |
| 340 | assert!(tags[0].contains("one.png"), "{tags:?}"); |
| 341 | assert!(tags[1].contains("two.png"), "{tags:?}"); |
| 342 | } |
| 343 | |
| 344 | #[test] |
| 345 | fn ingest_does_not_consult_model_capability() { |
| 346 | // Capability is a route property and is re-decided per request. If |
| 347 | // ingest started gating on it, attaching under a text-only model would |
| 348 | // destroy the image for the rest of the session. |
| 349 | let dir = tempfile::tempdir().expect("tempdir"); |
| 350 | let path = write_png(dir.path(), "shot.png"); |
| 351 | let text = format!("[Attached image: {}]", path.display()); |
| 352 | |
| 353 | let expanded = expand_attachment_blocks(&text); |
| 354 | |
| 355 | assert_eq!(expanded.blocks.len(), 3); |
| 356 | assert!(expanded.notices.is_empty()); |
| 357 | } |
| 358 | |
| 359 | #[test] |
| 360 | fn a_blind_route_gets_text_in_place_of_every_image() { |
| 361 | let mut messages = vec![ |
| 362 | message_with_image("data:image/png;base64,QUJD"), |
| 363 | codewhale_models::Message { |
| 364 | role: Role::Assistant, |
| 365 | content: vec![ContentBlock::Text { |
| 366 | text: "sure".to_string(), |
| 367 | cache_control: None, |
| 368 | }], |
| 369 | }, |
| 370 | ]; |
| 371 | |
| 372 | let stripped = |
| 373 | strip_images_when_unsupported(&mut messages, SupportState::Unsupported, "deepseek-chat"); |
| 374 | |
| 375 | assert_eq!(stripped, 1); |
| 376 | assert!( |
| 377 | !messages[0] |
| 378 | .content |
| 379 | .iter() |
| 380 | .any(|block| matches!(block, ContentBlock::ImageUrl { .. })), |
| 381 | "no image may survive to a route that cannot read one" |
| 382 | ); |
| 383 | match &messages[0].content[0] { |
| 384 | ContentBlock::Text { text, .. } => { |
| 385 | assert!(text.contains("deepseek-chat"), "{text}"); |
| 386 | assert!(text.contains("/model"), "{text}"); |
| 387 | assert!(text.contains("omitted"), "{text}"); |
| 388 | } |
| 389 | other => panic!("expected replacement text, got {other:?}"), |
| 390 | } |
| 391 | } |
| 392 | |
| 393 | #[test] |
| 394 | fn a_supported_or_unknown_route_keeps_its_images() { |
| 395 | // Unknown is the common case: models.dev has no modality data for most |
| 396 | // routes. Stripping there would make the feature dead on arrival for |
| 397 | // self-hosted and custom providers. |
| 398 | for vision in [SupportState::Supported, SupportState::Unknown] { |
| 399 | let mut messages = vec![message_with_image("data:image/png;base64,QUJD")]; |
| 400 | |
| 401 | let stripped = strip_images_when_unsupported(&mut messages, vision, "some-model"); |
| 402 | |
| 403 | assert_eq!(stripped, 0, "{vision:?} must not strip"); |
| 404 | assert!( |
| 405 | messages[0] |
| 406 | .content |
| 407 | .iter() |
| 408 | .any(|block| matches!(block, ContentBlock::ImageUrl { .. })), |
| 409 | "{vision:?} must keep the image" |
| 410 | ); |
| 411 | } |
| 412 | } |
| 413 | |
| 414 | #[test] |
| 415 | fn stripping_replaces_every_image_across_every_message() { |
| 416 | let mut messages = vec![ |
| 417 | message_with_image("data:image/png;base64,AAAA"), |
| 418 | message_with_image("data:image/jpeg;base64,BBBB"), |
| 419 | ]; |
| 420 | |
| 421 | let stripped = strip_images_when_unsupported(&mut messages, SupportState::Unsupported, "blind"); |
| 422 | |
| 423 | assert_eq!( |
| 424 | stripped, 2, |
| 425 | "a per-message early return would miss the second" |
| 426 | ); |
| 427 | } |
| 428 | |
| 429 | #[test] |
| 430 | fn a_missing_file_becomes_a_notice_not_a_dropped_turn() { |
| 431 | let text = "[Attached image: /nonexistent/definitely-not-here.png]"; |
| 432 | |
| 433 | let expanded = expand_attachment_blocks(text); |
| 434 | |
| 435 | assert!(expanded.blocks.is_empty()); |
| 436 | assert_eq!(expanded.notices.len(), 1); |
| 437 | assert!( |
| 438 | expanded.notices[0].contains("definitely-not-here.png"), |
| 439 | "{:?}", |
| 440 | expanded.notices |
| 441 | ); |
| 442 | } |
| 443 | |
| 444 | #[test] |
| 445 | fn one_bad_attachment_does_not_suppress_a_good_one() { |
| 446 | let dir = tempfile::tempdir().expect("tempdir"); |
| 447 | let good = write_png(dir.path(), "good.png"); |
| 448 | let text = format!( |
| 449 | "[Attached image: /nope/missing.png]\n[Attached image: {}]", |
| 450 | good.display() |
| 451 | ); |
| 452 | |
| 453 | let expanded = expand_attachment_blocks(&text); |
| 454 | |
| 455 | assert_eq!(expanded.blocks.len(), 3); |
| 456 | assert_eq!(expanded.notices.len(), 1); |
| 457 | } |
| 458 | |
| 459 | #[test] |
| 460 | fn video_attachments_are_left_as_text() { |
| 461 | let text = "[Attached video: /tmp/clip.mp4]"; |
| 462 | |
| 463 | let expanded = expand_attachment_blocks(text); |
| 464 | |
| 465 | assert!(expanded.blocks.is_empty(), "{expanded:?}"); |
| 466 | assert!(expanded.notices.is_empty(), "{expanded:?}"); |
| 467 | } |
| 468 | |
| 469 | #[test] |
| 470 | fn text_with_no_attachments_produces_nothing() { |
| 471 | let expanded = expand_attachment_blocks("just a normal question"); |
| 472 | assert!(expanded.blocks.is_empty()); |
| 473 | assert!(expanded.notices.is_empty()); |
| 474 | } |
| 475 | |
| 476 | #[test] |
| 477 | fn notice_block_names_the_failure_and_forbids_guessing() { |
| 478 | assert_eq!(notice_block(&[]), None); |
| 479 | let block = |
| 480 | notice_block(&["Cannot attach a.png: the file is empty".to_string()]).expect("block"); |
| 481 | match block { |
| 482 | ContentBlock::Text { text, .. } => { |
| 483 | assert!(text.contains("<attachment_notice>"), "{text}"); |
| 484 | assert!(text.contains("a.png"), "{text}"); |
| 485 | assert!(text.contains("Do not describe"), "{text}"); |
| 486 | } |
| 487 | other => panic!("expected text, got {other:?}"), |
| 488 | } |
| 489 | } |
| 490 | |
| 491 | #[test] |
| 492 | fn attach_from_path_reports_an_unreadable_file() { |
| 493 | let error = attach_image_from_path(Path::new("/nonexistent/x.png")).expect_err("must fail"); |
| 494 | assert!( |
| 495 | matches!(error, ImageAttachError::Unreadable { .. }), |
| 496 | "got {error:?}" |
| 497 | ); |
| 498 | } |
| 499 | |
| 500 | pub(crate) fn runtime_image_fixture(color: u8) -> codewhale_protocol::runtime::RuntimeImageInput { |
| 501 | let image = image::DynamicImage::ImageRgba8(image::RgbaImage::from_pixel( |
| 502 | 2, |
| 503 | 2, |
| 504 | image::Rgba([color, 31, 99, 255]), |
| 505 | )); |
| 506 | let mut bytes = std::io::Cursor::new(Vec::new()); |
| 507 | image.write_to(&mut bytes, image::ImageFormat::Png).unwrap(); |
| 508 | codewhale_protocol::runtime::RuntimeImageInput { |
| 509 | mime: "image/png".into(), |
| 510 | data_base64: STANDARD.encode(bytes.into_inner()), |
| 511 | } |
| 512 | } |
| 513 | |
| 514 | #[test] |
| 515 | fn runtime_image_validation_preserves_exact_bytes_and_rejects_corruption() { |
| 516 | let input = runtime_image_fixture(7); |
| 517 | let blocks = prepare_runtime_images(std::slice::from_ref(&input)).unwrap(); |
| 518 | assert_eq!( |
| 519 | runtime_images_from_blocks(&blocks).unwrap().as_slice(), |
| 520 | std::slice::from_ref(&input) |
| 521 | ); |
| 522 | for bad in [ |
| 523 | codewhale_protocol::runtime::RuntimeImageInput { |
| 524 | mime: "image/jpeg".into(), |
| 525 | ..input.clone() |
| 526 | }, |
| 527 | codewhale_protocol::runtime::RuntimeImageInput { |
| 528 | data_base64: "not base64".into(), |
| 529 | ..input.clone() |
| 530 | }, |
| 531 | codewhale_protocol::runtime::RuntimeImageInput { |
| 532 | data_base64: String::new(), |
| 533 | ..input.clone() |
| 534 | }, |
| 535 | // Existing signature sniffing alone accepted this truncated PNG. |
| 536 | codewhale_protocol::runtime::RuntimeImageInput { |
| 537 | data_base64: STANDARD.encode(b"\x89PNG\r\n\x1a\n"), |
| 538 | ..input.clone() |
| 539 | }, |
| 540 | ] { |
| 541 | assert!(prepare_runtime_images(&[bad]).is_err()); |
| 542 | } |
| 543 | } |
| 544 | |
| 545 | #[test] |
| 546 | fn runtime_image_validation_bounds_count_encoded_size_and_decode_dimensions() { |
| 547 | let input = runtime_image_fixture(7); |
| 548 | assert!(prepare_runtime_images(&vec![input.clone(); 11]).is_err()); |
| 549 | assert!( |
| 550 | prepare_runtime_images(&[codewhale_protocol::runtime::RuntimeImageInput { |
| 551 | data_base64: "A".repeat(MAX_IMAGE_BYTES.div_ceil(3) * 4 + 1), |
| 552 | ..input |
| 553 | }]) |
| 554 | .is_err() |
| 555 | ); |
| 556 | let wide = image::DynamicImage::ImageRgba8(image::RgbaImage::new(MAX_IMAGE_DIMENSION + 1, 1)); |
| 557 | let mut bytes = std::io::Cursor::new(Vec::new()); |
| 558 | wide.write_to(&mut bytes, image::ImageFormat::Png).unwrap(); |
| 559 | let input = codewhale_protocol::runtime::RuntimeImageInput { |
| 560 | mime: "image/png".into(), |
| 561 | data_base64: STANDARD.encode(bytes.into_inner()), |
| 562 | }; |
| 563 | assert!(prepare_runtime_images(&[input]).is_err()); |
| 564 | } |
| 565 | |
| 566 | pub(crate) fn runtime_image_fixture_bytes( |
| 567 | size: usize, |
| 568 | ) -> codewhale_protocol::runtime::RuntimeImageInput { |
| 569 | let mut image = runtime_image_fixture(42); |
| 570 | let mut bytes = STANDARD.decode(&image.data_base64).unwrap(); |
| 571 | bytes.resize(size, 0); |
| 572 | image.data_base64 = STANDARD.encode(bytes); |
| 573 | image |
| 574 | } |
| 575 | |
| 576 | #[test] |
| 577 | fn runtime_image_network_four_mib_and_historical_five_mib_bounds_are_distinct() { |
| 578 | let at_network_limit = runtime_image_fixture_bytes(MAX_RUNTIME_IMAGE_BYTES); |
| 579 | assert!(prepare_runtime_images(&[at_network_limit]).is_ok()); |
| 580 | let historical = runtime_image_fixture_bytes(MAX_RUNTIME_IMAGE_BYTES + 1); |
| 581 | assert!(prepare_runtime_images(std::slice::from_ref(&historical)).is_err()); |
| 582 | let stored = prepare_stored_images(std::slice::from_ref(&historical)).unwrap(); |
| 583 | assert_eq!(runtime_images_from_blocks(&stored).unwrap(), [historical]); |
| 584 | assert!(prepare_stored_images(&[runtime_image_fixture_bytes(MAX_IMAGE_BYTES)]).is_ok()); |
| 585 | assert!(prepare_stored_images(&[runtime_image_fixture_bytes(MAX_IMAGE_BYTES + 1)]).is_err()); |
| 586 | let three_mib = runtime_image_fixture_bytes(3 * 1024 * 1024); |
| 587 | assert!(prepare_runtime_images(&[three_mib.clone(), three_mib.clone()]).is_err()); |
| 588 | assert!(prepare_stored_images(&[three_mib.clone(), three_mib]).is_ok()); |
| 589 | assert!(prepare_stored_images(&vec![runtime_image_fixture(1); 11]).is_ok()); |
| 590 | } |
| 591 |