Skip to main content

photon_ui/
image.rs

1use base64::{
2    Engine as _,
3    engine::general_purpose::STANDARD,
4};
5
6/// Supported terminal image protocols.
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
8pub enum ImageProtocol {
9    /// Kitty graphics protocol.
10    #[default]
11    Kitty,
12    /// iTerm2 inline image protocol.
13    Iterm2,
14}
15
16/// Encode image data for the Kitty graphics protocol.
17///
18/// The payload is treated as a PNG image (`f=100`), base64-encoded and
19/// transmitted in chunks. `q=1` suppresses terminal responses, `m=0` marks
20/// the final chunk, and a trailing `a=p` command places the image by id.
21///
22/// `cols` and `rows` tell the terminal how many cells the image should occupy
23/// on screen. The image is scaled to fit that cell rectangle, which lets the
24/// layout engine reserve the matching amount of space.
25///
26/// Chunking follows the spec: only the first chunk carries the full set of
27/// keys (`a`, `f`, `i`, `q`, `m`); continuation chunks only carry `q` and `m`.
28pub 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
56/// Encode image data for the iTerm2 inline image protocol.
57pub 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
62/// Generate a Kitty graphics protocol delete command for the given image id.
63pub fn delete_kitty_image(id: u32) -> String {
64    format!("\x1b_Ga=d,d=I,i={},q=1\x1b\\", id)
65}
66
67/// Parse width and height from PNG file header bytes.
68///
69/// Returns `None` if the data is too short or the PNG signature is missing.
70pub 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
79/// Parse width and height from JPEG SOF0 / SOF2 marker segments.
80///
81/// Searches for `0xFFC0` (baseline) or `0xFFC2` (progressive) markers.
82pub 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
95/// Parse width and height from GIF logical screen descriptor.
96pub 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
105/// Parse width and height from a VP8X WebP chunk.
106///
107/// Returns `None` if the RIFF/WEBP signatures are missing or the VP8X chunk
108/// is not present.
109pub 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        // parts ends with an empty string after the final terminator.
157        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; // width - 1 = 99 -> width = 100
255        data[27] = 49; // height - 1 = 49 -> height = 50
256        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}