forge_primitives/
image.rs1use serde::{Deserialize, Serialize};
5
6use crate::content::ContentBlock;
7use crate::messages::Extras;
8
9pub const SUPPORTED_IMAGE_MIME_TYPES: &[&str] =
11 &["image/png", "image/jpeg", "image/gif", "image/webp"];
12
13#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
15pub struct ImageAttachment {
16 pub data: String,
17 pub mime_type: String,
18}
19
20impl ImageAttachment {
21 fn source(&self, with_payload: bool) -> serde_json::Value {
24 let mut source = serde_json::Map::new();
25 source.insert("type".to_owned(), serde_json::Value::String("base64".to_owned()));
26 source.insert("media_type".to_owned(), serde_json::Value::String(self.mime_type.clone()));
27 if with_payload {
28 source.insert("data".to_owned(), serde_json::Value::String(self.data.clone()));
29 }
30 serde_json::Value::Object(source)
31 }
32
33 pub fn wire_block(&self) -> serde_json::Value {
35 serde_json::json!({
36 "type": "image",
37 "source": self.source(true),
38 })
39 }
40
41 pub fn display_block(&self) -> ContentBlock {
45 let mut extras = Extras::new();
46 extras.insert("forge_bytes".to_owned(), serde_json::json!(payload_bytes(&self.data)));
47 ContentBlock::Image { source: self.source(false), extras }
48 }
49}
50
51pub fn display_blocks(images: &[ImageAttachment]) -> Vec<ContentBlock> {
55 images
56 .iter()
57 .filter(|image| validate_image(&image.data, &image.mime_type).is_ok())
58 .map(ImageAttachment::display_block)
59 .collect()
60}
61
62pub fn payload_bytes(data: &str) -> u64 {
65 let padding = if data.ends_with("==") { 2 } else { u64::from(data.ends_with('=')) };
66 (data.len() as u64 * 3 / 4).saturating_sub(padding)
67}
68
69pub fn is_supported_image_type(mime_type: &str) -> bool {
71 SUPPORTED_IMAGE_MIME_TYPES.contains(&mime_type)
72}
73
74pub fn is_valid_base64(data: &str) -> bool {
77 use base64::Engine as _;
78 if data.is_empty() {
79 return false;
80 }
81 let clean = data.trim();
82 if !clean.len().is_multiple_of(4) {
83 return false;
84 }
85 base64::engine::general_purpose::STANDARD.decode(clean).is_ok()
86}
87
88pub fn validate_image(data: &str, mime_type: &str) -> Result<(), String> {
95 if !is_supported_image_type(mime_type) {
96 return Err(format!(
97 "unsupported image type \"{mime_type}\"; expected one of: {}",
98 SUPPORTED_IMAGE_MIME_TYPES.join(", ")
99 ));
100 }
101 if !is_valid_base64(data) {
102 return Err("image data is not valid base64".to_owned());
103 }
104 Ok(())
105}
106
107#[cfg(test)]
108mod tests {
109 use super::*;
110
111 #[test]
115 fn the_wire_block_carries_the_base64_payload() {
116 let image =
117 ImageAttachment { data: "iVBORw0KGgo=".to_owned(), mime_type: "image/png".to_owned() };
118 assert_eq!(
119 image.wire_block(),
120 serde_json::json!({
121 "type": "image",
122 "source": {
123 "type": "base64",
124 "media_type": "image/png",
125 "data": "iVBORw0KGgo=",
126 },
127 }),
128 "the block the CLI receives keeps the payload it was sent",
129 );
130 }
131}