1use byteorder_lite::{LittleEndian, WriteBytesExt};
2use std::io::{self, Write};
3
4use crate::error::{
5 EncodingError, ImageError, ImageFormatHint, ImageResult, ParameterError, ParameterErrorKind,
6 UnsupportedError, UnsupportedErrorKind,
7};
8use crate::{DynamicImage, ExtendedColorType, ImageEncoder, ImageFormat};
9
10const BITMAPFILEHEADER_SIZE: u32 = 14;
11const BITMAPINFOHEADER_SIZE: u32 = 40;
12const BITMAPV4HEADER_SIZE: u32 = 108;
13
14pub struct BmpEncoder<'a, W: 'a> {
16 writer: &'a mut W,
17}
18
19impl<'a, W: Write + 'a> BmpEncoder<'a, W> {
20 pub fn new(w: &'a mut W) -> Self {
22 BmpEncoder { writer: w }
23 }
24
25 #[track_caller]
31 pub fn encode(
32 &mut self,
33 image: &[u8],
34 width: u32,
35 height: u32,
36 c: ExtendedColorType,
37 ) -> ImageResult<()> {
38 self.encode_with_palette(image, width, height, c, None)
39 }
40
41 #[track_caller]
48 pub fn encode_with_palette(
49 &mut self,
50 image: &[u8],
51 width: u32,
52 height: u32,
53 color_type: ExtendedColorType,
54 palette: Option<&[[u8; 3]]>,
55 ) -> ImageResult<()> {
56 if palette.is_some()
57 && color_type != ExtendedColorType::L8
58 && color_type != ExtendedColorType::La8
59 {
60 return Err(ImageError::Parameter(ParameterError::from_kind(
61 ParameterErrorKind::Generic(
62 "Palette given which must only be used with L8 or La8 color types".to_string(),
63 ),
64 )));
65 }
66
67 let expected_buffer_len = color_type.buffer_size(width, height);
68 assert_eq!(
69 expected_buffer_len,
70 image.len() as u64,
71 "Invalid buffer length: expected {expected_buffer_len} got {} for {width}x{height} image",
72 image.len(),
73 );
74
75 let bmp_header_size = BITMAPFILEHEADER_SIZE;
76
77 let (dib_header_size, written_pixel_size, palette_color_count) =
78 written_pixel_info(color_type, palette)?;
79
80 let (padded_row, image_size) = width
81 .checked_mul(written_pixel_size)
82 .and_then(|v| v.checked_next_multiple_of(4))
84 .and_then(|v| {
85 let image_bytes = v.checked_mul(height)?;
86 Some((v, image_bytes))
87 })
88 .ok_or_else(|| {
89 ImageError::Parameter(ParameterError::from_kind(
90 ParameterErrorKind::DimensionMismatch,
91 ))
92 })?;
93
94 let row_padding = padded_row - width * written_pixel_size;
95
96 let palette_size = palette_color_count.checked_mul(4).ok_or_else(|| {
98 ImageError::Encoding(EncodingError::new(
99 ImageFormatHint::Exact(ImageFormat::Bmp),
100 "calculated palette size larger than 2^32",
101 ))
102 })?;
103
104 let file_size = bmp_header_size
105 .checked_add(dib_header_size)
106 .and_then(|v| v.checked_add(palette_size))
107 .and_then(|v| v.checked_add(image_size))
108 .ok_or_else(|| {
109 ImageError::Encoding(EncodingError::new(
110 ImageFormatHint::Exact(ImageFormat::Bmp),
111 "calculated BMP header size larger than 2^32",
112 ))
113 })?;
114
115 let image_data_offset = bmp_header_size
116 .checked_add(dib_header_size)
117 .and_then(|v| v.checked_add(palette_size))
118 .ok_or_else(|| {
119 ImageError::Encoding(EncodingError::new(
120 ImageFormatHint::Exact(ImageFormat::Bmp),
121 "calculated BMP size larger than 2^32",
122 ))
123 })?;
124
125 self.writer.write_u8(b'B')?;
127 self.writer.write_u8(b'M')?;
128 self.writer.write_u32::<LittleEndian>(file_size)?; self.writer.write_u16::<LittleEndian>(0)?; self.writer.write_u16::<LittleEndian>(0)?; self.writer.write_u32::<LittleEndian>(image_data_offset)?; self.writer.write_u32::<LittleEndian>(dib_header_size)?;
135 self.writer.write_i32::<LittleEndian>(width as i32)?;
136 self.writer.write_i32::<LittleEndian>(height as i32)?;
137 self.writer.write_u16::<LittleEndian>(1)?; self.writer
139 .write_u16::<LittleEndian>((written_pixel_size * 8) as u16)?; if dib_header_size >= BITMAPV4HEADER_SIZE {
141 self.writer.write_u32::<LittleEndian>(3)?; } else {
144 self.writer.write_u32::<LittleEndian>(0)?; }
146 self.writer.write_u32::<LittleEndian>(image_size)?;
147 self.writer.write_i32::<LittleEndian>(0)?; self.writer.write_i32::<LittleEndian>(0)?; self.writer.write_u32::<LittleEndian>(palette_color_count)?;
150 self.writer.write_u32::<LittleEndian>(0)?; if dib_header_size >= BITMAPV4HEADER_SIZE {
152 self.writer.write_u32::<LittleEndian>(0xff << 16)?; self.writer.write_u32::<LittleEndian>(0xff << 8)?; self.writer.write_u32::<LittleEndian>(0xff)?; self.writer.write_u32::<LittleEndian>(0xff << 24)?; self.writer.write_u32::<LittleEndian>(0x7352_4742)?; for _ in 0..12 {
161 self.writer.write_u32::<LittleEndian>(0)?;
162 }
163 }
164
165 match color_type {
167 ExtendedColorType::Rgb8 => self.encode_rgb(image, width, height, row_padding, 3)?,
168 ExtendedColorType::Rgba8 => self.encode_rgba(image, width, height, row_padding, 4)?,
169 ExtendedColorType::L8 => {
170 self.encode_gray(image, width, height, row_padding, 1, palette)?;
171 }
172 ExtendedColorType::La8 => {
173 self.encode_gray(image, width, height, row_padding, 2, palette)?;
174 }
175 _ => {
176 return Err(ImageError::Unsupported(
177 UnsupportedError::from_format_and_kind(
178 ImageFormat::Bmp.into(),
179 UnsupportedErrorKind::Color(color_type),
180 ),
181 ));
182 }
183 }
184
185 Ok(())
186 }
187
188 fn encode_rgb(
189 &mut self,
190 image: &[u8],
191 width: u32,
192 height: u32,
193 row_padding: u32,
194 bytes_per_pixel: u32,
195 ) -> io::Result<()> {
196 let width = width as usize;
197 let height = height as usize;
198 let x_stride = bytes_per_pixel as usize;
199 let y_stride = width * x_stride;
200 for row in (0..height).rev() {
201 let row_start = row * y_stride;
203 for px in image[row_start..][..y_stride].chunks_exact(x_stride) {
204 let r = px[0];
205 let g = px[1];
206 let b = px[2];
207 self.writer.write_all(&[b, g, r])?;
209 }
210 self.write_row_pad(row_padding)?;
211 }
212
213 Ok(())
214 }
215
216 fn encode_rgba(
217 &mut self,
218 image: &[u8],
219 width: u32,
220 height: u32,
221 row_padding: u32,
222 bytes_per_pixel: u32,
223 ) -> io::Result<()> {
224 let width = width as usize;
225 let height = height as usize;
226 let x_stride = bytes_per_pixel as usize;
227 let y_stride = width * x_stride;
228 for row in (0..height).rev() {
229 let row_start = row * y_stride;
231 for px in image[row_start..][..y_stride].chunks_exact(x_stride) {
232 let r = px[0];
233 let g = px[1];
234 let b = px[2];
235 let a = px[3];
236 self.writer.write_all(&[b, g, r, a])?;
238 }
239 self.write_row_pad(row_padding)?;
240 }
241
242 Ok(())
243 }
244
245 fn encode_gray(
246 &mut self,
247 image: &[u8],
248 width: u32,
249 height: u32,
250 row_padding: u32,
251 bytes_per_pixel: u32,
252 palette: Option<&[[u8; 3]]>,
253 ) -> io::Result<()> {
254 if let Some(palette) = palette {
256 for item in palette {
257 self.writer.write_all(&[item[2], item[1], item[0], 0])?;
259 }
260 } else {
261 for val in 0u8..=255 {
262 self.writer.write_all(&[val, val, val, 0])?;
264 }
265 }
266
267 let x_stride = bytes_per_pixel;
269 let y_stride = width * x_stride;
270 for row in (0..height).rev() {
271 let row_start = row * y_stride;
273
274 if x_stride == 1 {
276 self.writer
278 .write_all(&image[row_start as usize..][..y_stride as usize])?;
279 } else {
280 for col in 0..width {
281 let pixel_start = (row_start + (col * x_stride)) as usize;
282 self.writer.write_u8(image[pixel_start])?;
283 }
285 }
286
287 self.write_row_pad(row_padding)?;
288 }
289
290 Ok(())
291 }
292
293 fn write_row_pad(&mut self, row_pad_size: u32) -> io::Result<()> {
294 for _ in 0..row_pad_size {
295 self.writer.write_u8(0)?;
296 }
297
298 Ok(())
299 }
300}
301
302impl<W: Write> ImageEncoder for BmpEncoder<'_, W> {
303 #[track_caller]
304 fn write_image(
305 mut self,
306 buf: &[u8],
307 width: u32,
308 height: u32,
309 color_type: ExtendedColorType,
310 ) -> ImageResult<()> {
311 self.encode(buf, width, height, color_type)
312 }
313
314 fn make_compatible_img(
315 &self,
316 _: crate::io::encoder::MethodSealedToImage,
317 img: &DynamicImage,
318 ) -> Option<DynamicImage> {
319 crate::io::encoder::dynimage_conversion_8bit(img)
320 }
321}
322
323fn written_pixel_info(
325 c: ExtendedColorType,
326 palette: Option<&[[u8; 3]]>,
327) -> Result<(u32, u32, u32), ImageError> {
328 let (header, color_bytes, palette_count) = match c {
329 ExtendedColorType::Rgb8 => (BITMAPINFOHEADER_SIZE, 3, Some(0)),
330 ExtendedColorType::Rgba8 => (BITMAPV4HEADER_SIZE, 4, Some(0)),
331 ExtendedColorType::L8 => (
332 BITMAPINFOHEADER_SIZE,
333 1,
334 u32::try_from(palette.map(|p| p.len()).unwrap_or(256)).ok(),
335 ),
336 ExtendedColorType::La8 => (
337 BITMAPINFOHEADER_SIZE,
338 1,
339 u32::try_from(palette.map(|p| p.len()).unwrap_or(256)).ok(),
340 ),
341 _ => {
342 return Err(ImageError::Unsupported(
343 UnsupportedError::from_format_and_kind(
344 ImageFormat::Bmp.into(),
345 UnsupportedErrorKind::Color(c),
346 ),
347 ));
348 }
349 };
350
351 let palette_count = palette_count.ok_or_else(|| {
352 ImageError::Encoding(EncodingError::new(
353 ImageFormatHint::Exact(ImageFormat::Bmp),
354 "calculated palette size larger than 2^32",
355 ))
356 })?;
357
358 Ok((header, color_bytes, palette_count))
359}
360
361#[cfg(test)]
362mod tests {
363 use super::super::BmpDecoder;
364 use super::BmpEncoder;
365
366 use crate::ExtendedColorType;
367 use crate::ImageDecoder as _;
368 use std::io::Cursor;
369
370 fn round_trip_image(image: &[u8], width: u32, height: u32, c: ExtendedColorType) -> Vec<u8> {
371 let mut encoded_data = Vec::new();
372 {
373 let mut encoder = BmpEncoder::new(&mut encoded_data);
374 encoder
375 .encode(image, width, height, c)
376 .expect("could not encode image");
377 }
378
379 let decoder = BmpDecoder::new(Cursor::new(&encoded_data)).expect("failed to decode");
380
381 let mut buf = vec![0; decoder.total_bytes() as usize];
382 decoder.read_image(&mut buf).expect("failed to decode");
383 buf
384 }
385
386 #[test]
387 fn round_trip_single_pixel_rgb() {
388 let image = [255u8, 0, 0]; let decoded = round_trip_image(&image, 1, 1, ExtendedColorType::Rgb8);
390 assert_eq!(3, decoded.len());
391 assert_eq!(255, decoded[0]);
392 assert_eq!(0, decoded[1]);
393 assert_eq!(0, decoded[2]);
394 }
395
396 #[test]
397 #[cfg(target_pointer_width = "64")]
398 fn huge_files_return_error() {
399 let mut encoded_data = Vec::new();
400 let image = vec![0u8; 3 * 40_000 * 40_000]; let mut encoder = BmpEncoder::new(&mut encoded_data);
402 let result = encoder.encode(&image, 40_000, 40_000, ExtendedColorType::Rgb8);
403 assert!(result.is_err());
404 }
405
406 #[test]
407 fn round_trip_single_pixel_rgba() {
408 let image = [1, 2, 3, 4];
409 let decoded = round_trip_image(&image, 1, 1, ExtendedColorType::Rgba8);
410 assert_eq!(&decoded[..], &image[..]);
411 }
412
413 #[test]
414 fn round_trip_3px_rgb() {
415 let image = [0u8; 3 * 3 * 3]; let _decoded = round_trip_image(&image, 3, 3, ExtendedColorType::Rgb8);
417 }
418
419 #[test]
420 fn round_trip_gray() {
421 let image = [0u8, 1, 2]; let decoded = round_trip_image(&image, 3, 1, ExtendedColorType::L8);
423 assert_eq!(9, decoded.len());
425 assert_eq!(0, decoded[0]);
426 assert_eq!(0, decoded[1]);
427 assert_eq!(0, decoded[2]);
428 assert_eq!(1, decoded[3]);
429 assert_eq!(1, decoded[4]);
430 assert_eq!(1, decoded[5]);
431 assert_eq!(2, decoded[6]);
432 assert_eq!(2, decoded[7]);
433 assert_eq!(2, decoded[8]);
434 }
435
436 #[test]
437 fn round_trip_graya() {
438 let image = [0u8, 0, 1, 0, 2, 0]; let decoded = round_trip_image(&image, 1, 3, ExtendedColorType::La8);
440 assert_eq!(9, decoded.len());
442 assert_eq!(0, decoded[0]);
443 assert_eq!(0, decoded[1]);
444 assert_eq!(0, decoded[2]);
445 assert_eq!(1, decoded[3]);
446 assert_eq!(1, decoded[4]);
447 assert_eq!(1, decoded[5]);
448 assert_eq!(2, decoded[6]);
449 assert_eq!(2, decoded[7]);
450 assert_eq!(2, decoded[8]);
451 }
452
453 #[test]
454 fn regression_issue_2604() {
455 let mut image = vec![];
456 let mut encoder = BmpEncoder::new(&mut image);
457 encoder
458 .encode(&[], 1 << 31, 0, ExtendedColorType::Rgb8)
459 .unwrap_err();
460 }
461}