1use base64::{
2 Engine as _,
3 engine::general_purpose::STANDARD,
4};
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
8pub enum ImageProtocol {
9 #[default]
11 Kitty,
12 Iterm2,
14}
15
16pub fn encode_kitty(id: u32, data: &[u8], cols: u16, rows: u16) -> String {
29 let b64 = STANDARD.encode(data);
30 if b64.is_empty() {
31 return String::new();
32 }
33 let mut seq = String::new();
34 const CHUNK_SIZE: usize = 4096;
35 let chunks: Vec<&[u8]> = b64.as_bytes().chunks(CHUNK_SIZE).collect();
36 let last = chunks.len() - 1;
37 for (i, chunk) in chunks.iter().enumerate() {
38 let more = if i == last { 0 } else { 1 };
39 if i == 0 {
40 seq.push_str(&format!("\x1b_Ga=t,f=100,i={},q=1,m={};", id, more));
41 } else {
42 seq.push_str(&format!("\x1b_Gq=1,m={};", more));
43 }
44 if let Ok(s) = std::str::from_utf8(chunk) {
45 seq.push_str(s);
46 }
47 seq.push_str("\x1b\\");
48 }
49 seq.push_str(&format!(
50 "\x1b_Ga=p,i={},c={},r={},q=1\x1b\\",
51 id, cols, rows
52 ));
53 seq
54}
55
56pub fn encode_iterm2(data: &[u8], _mime_type: &str) -> String {
58 let b64 = STANDARD.encode(data);
59 format!("\x1b]1337;File=inline=1:{}\x07", b64)
60}
61
62pub fn delete_kitty_image(id: u32) -> String {
64 format!("\x1b_Ga=d,d=I,i={},q=1\x1b\\", id)
65}
66
67pub fn get_png_dimensions(data: &[u8]) -> Option<(u32, u32)> {
71 if data.len() < 24 || &data[0..8] != b"\x89PNG\r\n\x1a\n" {
72 return None;
73 }
74 let width = u32::from_be_bytes([data[16], data[17], data[18], data[19]]);
75 let height = u32::from_be_bytes([data[20], data[21], data[22], data[23]]);
76 Some((width, height))
77}
78
79pub fn get_jpeg_dimensions(data: &[u8]) -> Option<(u32, u32)> {
83 let mut i = 2;
84 while i < data.len().saturating_sub(9) {
85 if data[i] == 0xff && (data[i + 1] == 0xc0 || data[i + 1] == 0xc2) {
86 let h = u16::from_be_bytes([data[i + 5], data[i + 6]]) as u32;
87 let w = u16::from_be_bytes([data[i + 7], data[i + 8]]) as u32;
88 return Some((w, h));
89 }
90 i += 1;
91 }
92 None
93}
94
95pub fn get_gif_dimensions(data: &[u8]) -> Option<(u32, u32)> {
97 if data.len() < 10 {
98 return None;
99 }
100 let w = u16::from_le_bytes([data[6], data[7]]) as u32;
101 let h = u16::from_le_bytes([data[8], data[9]]) as u32;
102 Some((w, h))
103}
104
105pub fn get_webp_dimensions(data: &[u8]) -> Option<(u32, u32)> {
110 if data.len() < 30 || &data[0..4] != b"RIFF" || &data[8..12] != b"WEBP" {
111 return None;
112 }
113 if &data[12..16] == b"VP8X" {
114 let w = u32::from_le_bytes([data[24], data[25], data[26], 0]) + 1;
115 let h = u32::from_le_bytes([data[27], data[28], data[29], 0]) + 1;
116 return Some((w, h));
117 }
118 None
119}
120
121#[cfg(test)]
122mod tests {
123 use super::*;
124
125 #[test]
126 fn image_protocol_default_is_kitty() {
127 assert_eq!(ImageProtocol::default(), ImageProtocol::Kitty);
128 }
129
130 #[test]
131 fn image_protocol_traits() {
132 let proto = ImageProtocol::Kitty;
133 let cloned = proto;
134 assert_eq!(proto, cloned);
135 assert_ne!(ImageProtocol::Kitty, ImageProtocol::Iterm2);
136 }
137
138 #[test]
139 fn kitty_encode_empty_data_returns_empty() {
140 assert!(encode_kitty(1, &[], 1, 1).is_empty());
141 }
142
143 #[test]
144 fn kitty_encode_single_chunk() {
145 let seq = encode_kitty(7, b"png", 10, 5);
146 assert!(seq.starts_with("\x1b_Ga=t,f=100,i=7,q=1,m=0;"));
147 assert!(seq.contains("a=p,i=7,c=10,r=5,q=1\x1b\\"));
148 assert!(seq.contains("cG5n"));
149 }
150
151 #[test]
152 fn kitty_encode_multiple_chunks_only_first_has_full_keys() {
153 let data = vec![0u8; 10000];
154 let seq = encode_kitty(3, &data, 20, 10);
155 let parts: Vec<&str> = seq.split("\x1b\\").collect();
156 assert!(parts[0].starts_with("\x1b_Ga=t,f=100,i=3,q=1,m=1;"));
158 let continuation = parts
159 .iter()
160 .find(|p| p.starts_with("\x1b_Gq=1,m=1;"))
161 .expect("expected an intermediate continuation chunk");
162 assert!(!continuation.contains("a=t"));
163 assert!(!continuation.contains("f=100"));
164 assert!(!continuation.contains("i=3"));
165 assert!(seq.contains("\x1b_Ga=p,i=3,c=20,r=10,q=1\x1b\\"));
166 }
167
168 #[test]
169 fn iterm2_encode_produces_sequence() {
170 let seq = encode_iterm2(b"data", "image/png");
171 assert!(seq.starts_with("\x1b]1337;File=inline=1:"));
172 assert!(seq.ends_with("\x07"));
173 }
174
175 #[test]
176 fn iterm2_encode_empty_data() {
177 let seq = encode_iterm2(&[], "image/png");
178 assert_eq!(seq, "\x1b]1337;File=inline=1:\x07");
179 }
180
181 #[test]
182 fn delete_kitty_image_produces_sequence() {
183 let seq = delete_kitty_image(42);
184 assert_eq!(seq, "\x1b_Ga=d,d=I,i=42,q=1\x1b\\");
185 }
186
187 #[test]
188 fn png_dimensions_valid() {
189 let mut data = vec![0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a];
190 data.extend_from_slice(&[0; 8]);
191 data.extend_from_slice(&100u32.to_be_bytes());
192 data.extend_from_slice(&200u32.to_be_bytes());
193 assert_eq!(get_png_dimensions(&data), Some((100, 200)));
194 }
195
196 #[test]
197 fn png_dimensions_too_short() {
198 assert_eq!(get_png_dimensions(&[0u8; 10]), None);
199 }
200
201 #[test]
202 fn png_dimensions_invalid_signature() {
203 assert_eq!(get_png_dimensions(b"NOTPNG"), None);
204 }
205
206 #[test]
207 fn jpeg_dimensions_valid_baseline() {
208 let mut data = vec![0xff, 0xd8];
209 data.extend_from_slice(&[0xff, 0xc0]);
210 data.extend_from_slice(&[0x00, 0x0b]);
211 data.extend_from_slice(&[0x08]);
212 data.extend_from_slice(&[0x00, 0x10]);
213 data.extend_from_slice(&[0x00, 0x20]);
214 data.extend_from_slice(&[0x01, 0x01, 0x11, 0x00]);
215 assert_eq!(get_jpeg_dimensions(&data), Some((32, 16)));
216 }
217
218 #[test]
219 fn jpeg_dimensions_progressive() {
220 let mut data = vec![0xff, 0xd8];
221 data.extend_from_slice(&[0xff, 0xc2]);
222 data.extend_from_slice(&[0x00, 0x0b]);
223 data.extend_from_slice(&[0x08]);
224 data.extend_from_slice(&[0x00, 0x20]);
225 data.extend_from_slice(&[0x00, 0x10]);
226 data.extend_from_slice(&[0x01, 0x01, 0x11, 0x00]);
227 assert_eq!(get_jpeg_dimensions(&data), Some((16, 32)));
228 }
229
230 #[test]
231 fn jpeg_dimensions_no_sof_marker() {
232 assert_eq!(get_jpeg_dimensions(b"\xff\xd8\xff\xe0"), None);
233 }
234
235 #[test]
236 fn gif_dimensions_valid() {
237 let mut data = b"GIF89a".to_vec();
238 data.extend_from_slice(&100u16.to_le_bytes());
239 data.extend_from_slice(&50u16.to_le_bytes());
240 assert_eq!(get_gif_dimensions(&data), Some((100, 50)));
241 }
242
243 #[test]
244 fn gif_dimensions_too_short() {
245 assert_eq!(get_gif_dimensions(&[0u8; 5]), None);
246 }
247
248 #[test]
249 fn webp_dimensions_vp8x() {
250 let mut data = vec![0u8; 30];
251 data[0..4].copy_from_slice(b"RIFF");
252 data[8..12].copy_from_slice(b"WEBP");
253 data[12..16].copy_from_slice(b"VP8X");
254 data[24] = 99; data[27] = 49; assert_eq!(get_webp_dimensions(&data), Some((100, 50)));
257 }
258
259 #[test]
260 fn webp_dimensions_too_short() {
261 assert_eq!(get_webp_dimensions(&[0u8; 10]), None);
262 }
263
264 #[test]
265 fn webp_dimensions_missing_riff() {
266 let mut data = vec![0u8; 30];
267 data[8..12].copy_from_slice(b"WEBP");
268 assert_eq!(get_webp_dimensions(&data), None);
269 }
270
271 #[test]
272 fn webp_dimensions_missing_webp() {
273 let mut data = vec![0u8; 30];
274 data[0..4].copy_from_slice(b"RIFF");
275 assert_eq!(get_webp_dimensions(&data), None);
276 }
277
278 #[test]
279 fn webp_dimensions_missing_vp8x() {
280 let mut data = vec![0u8; 30];
281 data[0..4].copy_from_slice(b"RIFF");
282 data[8..12].copy_from_slice(b"WEBP");
283 assert_eq!(get_webp_dimensions(&data), None);
284 }
285}