Merge pull request #5 from vstroebel/bad_color_conversion

Fix integer overflow in rgb to ycbcr conversion.

Author: vstroebel

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "jpeg-encoder"
-version = "0.5.0"
+version = "0.5.1"
 authors = ["Volker Ströbel <[email protected]>"]
 edition = "2018"
 license = "MIT OR Apache-2.0"

src/avx2/ycbcr.rs

@@ -70,7 +70,7 @@ macro_rules! ycbcr_image_avx2 {
                     let yb = _mm256_mullo_epi32(ymulb, b);
 
                     let y = _mm256_add_epi32(_mm256_add_epi32(yr, yg), yb);
-                    let y = _mm256_add_epi32(y, _mm256_set1_epi32(1 << 15));
+                    let y = _mm256_add_epi32(y, _mm256_set1_epi32(0x7FFF));
                     let y = _mm256_srli_epi32(y, 16);
                     let y: [i32; 8] = core::mem::transmute(y);
 
@@ -80,7 +80,7 @@ macro_rules! ycbcr_image_avx2 {
 
                     let cb = _mm256_add_epi32(_mm256_sub_epi32(cbr, cbg), cbb);
                     let cb = _mm256_add_epi32(cb, _mm256_set1_epi32(128 << 16));
-                    let cb = _mm256_add_epi32(cb, _mm256_set1_epi32(1 << 15));
+                    let cb = _mm256_add_epi32(cb, _mm256_set1_epi32(0x7FFF));
                     let cb = _mm256_srli_epi32(cb, 16);
                     let cb: [i32; 8] = core::mem::transmute(cb);
 
@@ -90,7 +90,7 @@ macro_rules! ycbcr_image_avx2 {
 
                     let cr = _mm256_sub_epi32(_mm256_sub_epi32(crr, crg), crb);
                     let cr = _mm256_add_epi32(cr, _mm256_set1_epi32(128 << 16));
-                    let cr = _mm256_add_epi32(cr, _mm256_set1_epi32(1 << 15));
+                    let cr = _mm256_add_epi32(cr, _mm256_set1_epi32(0x7FFF));
                     let cr = _mm256_srli_epi32(cr, 16);
                     let cr: [i32; 8] = core::mem::transmute(cr);
 

src/image_buffer.rs

@@ -23,9 +23,9 @@ pub fn rgb_to_ycbcr(r: u8, g: u8, b: u8) -> (u8, u8, u8) {
     let cb = -11059 * r - 21709 * g + 32768 * b + (128 << 16);
     let cr = 32768 * r - 27439 * g - 5329 * b + (128 << 16);
 
-    let y = (y + (1 << 15)) >> 16;
-    let cb = (cb + (1 << 15)) >> 16;
-    let cr = (cr + (1 << 15)) >> 16;
+    let y = (y + 0x7FFF) >> 16;
+    let cb = (cb + 0x7FFF) >> 16;
+    let cr = (cr + 0x7FFF) >> 16;
 
     (y as u8, cb as u8, cr as u8)
 }
@@ -287,6 +287,13 @@ mod tests {
 
     #[test]
     fn test_rgb_to_ycbcr() {
+
+        assert_rgb_to_ycbcr([0, 0, 0], [0, 128, 128]);
+        assert_rgb_to_ycbcr([255, 255, 255], [255, 128, 128]);
+        assert_rgb_to_ycbcr([255, 0, 0], [76, 85, 255]);
+        assert_rgb_to_ycbcr([0, 255, 0], [150, 44, 21]);
+        assert_rgb_to_ycbcr([0, 0, 255], [29, 255, 107]);
+
         // Values taken from libjpeg for a common image
 
         assert_rgb_to_ycbcr([59, 109, 6], [82, 85, 111]);

src/lib.rs

@@ -77,6 +77,26 @@ mod tests {
         (data, width as u16, height as u16)
     }
 
+    fn create_test_img_rgba() -> (Vec<u8>, u16, u16) {
+        // Ensure size which which ensures an odd MCU count per row to test chroma subsampling
+        let width = 258;
+        let height = 128;
+
+        let mut data = Vec::with_capacity(width * height * 3);
+
+        for y in 0..height {
+            for x in 0..width {
+                let x = x.min(255);
+                data.push(x as u8);
+                data.push((y * 2) as u8);
+                data.push(((x + y * 2) / 2) as u8);
+                data.push(x as u8);
+            }
+        }
+
+        (data, width as u16, height as u16)
+    }
+
     fn create_test_img_gray() -> (Vec<u8>, u16, u16) {
         let width = 258;
         let height = 128;
@@ -184,6 +204,21 @@ mod tests {
         check_result(data, width, height, &mut result, PixelFormat::RGB24);
     }
 
+    #[test]
+    fn test_rgba_80() {
+        let (data, width, height) = create_test_img_rgba();
+
+        let mut result = Vec::new();
+        let encoder = Encoder::new(&mut result, 80);
+        encoder
+            .encode(&data, width, height, ColorType::Rgba)
+            .unwrap();
+
+        let (data, width, height) = create_test_img_rgb();
+
+        check_result(data, width, height, &mut result, PixelFormat::RGB24);
+    }
+
     #[test]
     fn test_rgb_custom_q_table() {
         let (data, width, height) = create_test_img_rgb();