zmbv: fix out-of-bounds motion compensation

[nihav.git] / nihav-commonfmt / src / codecs / zmbvenc.rs

diff --git a/nihav-commonfmt/src/codecs/zmbvenc.rs b/nihav-commonfmt/src/codecs/zmbvenc.rs
index 23e8aadd893bfd45eb3255b313229d6ba8ea186e..b3f75bf8b26b928208cfa632835c1f08869f4952 100644 (file)
--- a/nihav-commonfmt/src/codecs/zmbvenc.rs
+++ b/nihav-commonfmt/src/codecs/zmbvenc.rs
@@ -35,6 +35,7 @@ struct ZMBVEncoder {
     width:      usize,
     height:     usize,
     range:      usize,
+    sent_pal:   bool,
 }
 
 fn buf_type_to_bpp(buf: &NABufferType) -> u8 {
@@ -134,6 +135,7 @@ impl ZMBVEncoder {
             width:      0,
             height:     0,
             range:      128,
+            sent_pal:   false,
         }
     }
     fn encode_intra(&mut self, bw: &mut ByteWriter, buf: NABufferType) -> EncoderResult<()> {
@@ -182,7 +184,7 @@ impl ZMBVEncoder {
             }
             bw.write_buf(&self.frm1[..self.width * self.height * bm])?;
         } else {
-            self.tmp_buf.truncate(0);
+            self.tmp_buf.clear();
             if bpp == 8 {
                 self.tmp_buf.extend_from_slice(&self.pal);
             }
@@ -191,7 +193,7 @@ impl ZMBVEncoder {
 
             let mut db = Vec::new();
             std::mem::swap(&mut db, &mut self.zbuf);
-            db.truncate(0);
+            db.clear();
             let mut wr = DeflateWriter::new(db);
             self.compr.write_zlib_header(&mut wr);
             self.compr.compress(&self.tmp_buf, &mut wr);
@@ -209,7 +211,7 @@ impl ZMBVEncoder {
             self.frm1.copy_from_slice(&self.frm2);
 
             bw.write_byte(0)?;
-            self.tmp_buf.truncate(0);
+            self.tmp_buf.clear();
             let tile_w = (self.width  + self.tile_w - 1) / self.tile_w;
             let tile_h = (self.height + self.tile_h - 1) / self.tile_h;
             let mv_size = (tile_w * tile_h * 2 + 3) & !3;
@@ -222,7 +224,7 @@ impl ZMBVEncoder {
                 let mut db = Vec::new();
 
                 std::mem::swap(&mut db, &mut self.zbuf);
-                db.truncate(0);
+                db.clear();
                 let mut wr = DeflateWriter::new(db);
                 self.compr.compress(&self.tmp_buf, &mut wr);
                 self.compr.compress_flush(&mut wr);
@@ -238,7 +240,7 @@ impl ZMBVEncoder {
             return Err(EncoderError::FormatError);
         }
 
-        self.tmp_buf.truncate(0);
+        self.tmp_buf.clear();
         if let (NABufferType::Video(ref vbuf), true) = (&buf, bpp == 8) {
             let mut npal = [0; 768];
             let off = vbuf.get_offset(1);
@@ -258,6 +260,8 @@ impl ZMBVEncoder {
                 self.pal = npal;
 
                 bw.write_byte(2)?;
+
+                self.sent_pal = false;
             } else {
                 bw.write_byte(0)?;
             }
@@ -352,7 +356,7 @@ impl ZMBVEncoder {
             let mut db = Vec::new();
 
             std::mem::swap(&mut db, &mut self.zbuf);
-            db.truncate(0);
+            db.clear();
             let mut wr = DeflateWriter::new(db);
             self.compr.compress(&self.tmp_buf, &mut wr);
             self.compr.compress_flush(&mut wr);
@@ -361,7 +365,7 @@ impl ZMBVEncoder {
 
             bw.write_buf(&self.zbuf)?;
         }
-        
+
         Ok(())
     }
 }
@@ -400,7 +404,7 @@ impl NAEncoder for ZMBVEncoder {
             NACodecTypeInfo::Video(vinfo) => {
                 self.width  = vinfo.width;
                 self.height = vinfo.height;
-                        
+
                 let out_info = NAVideoInfo::new(vinfo.width, vinfo.height, false, vinfo.format);
                 let info = NACodecInfo::new("zmbv", NACodecTypeInfo::Video(out_info), None);
                 let mut stream = NAStream::new(StreamType::Video, stream_id, info, encinfo.tb_num, encinfo.tb_den, 0);
@@ -429,6 +433,21 @@ impl NAEncoder for ZMBVEncoder {
                 false
             };
         self.pkt = Some(NAPacket::new(self.stream.clone().unwrap(), frm.ts, is_intra, dbuf));
+        if self.bpp == 8 && !self.sent_pal {
+            if let NABufferType::Video(ref buf) = frm.get_buffer() {
+                let paloff = buf.get_offset(1);
+                let data = buf.get_data();
+                let mut pal = [0; 1024];
+                let srcpal = &data[paloff..][..768];
+                for (dclr, sclr) in pal.chunks_exact_mut(4).zip(srcpal.chunks_exact(3)) {
+                    dclr[..3].copy_from_slice(sclr);
+                }
+                if let Some(ref mut pkt) = &mut self.pkt {
+                    pkt.side_data.push(NASideData::Palette(true, Arc::new(pal)));
+                }
+            }
+            self.sent_pal = true;
+        }
         self.frmcount += 1;
         if self.frmcount == self.key_int {
             self.frmcount = 0;
@@ -530,6 +549,7 @@ mod test {
     use nihav_codec_support::test::enc_video::*;
     use super::{RGB555_FORMAT, RGB24_0_FORMAT};
 
+    // samples are from https://samples.mplayerhq.hu/V-codecs/ZMBV/
     #[test]
     fn test_zmbv_encoder_8bit() {
         let mut dmx_reg = RegisteredDemuxers::new();
@@ -569,9 +589,12 @@ mod test {
                 tb_den:  0,
                 flags:   0,
             };
-        //test_encoding_to_file(&dec_config, &enc_config, enc_params);
-        test_encoding_md5(&dec_config, &enc_config, enc_params,
-                          &[0x50df10e2, 0x606f3268, 0xdd4bc2ff, 0x844e7d87]);
+        let enc_options = &[
+                NAOption { name: "range", value: NAValue::Int(16) },
+            ];
+        //test_encoding_to_file(&dec_config, &enc_config, enc_params, enc_options);
+        test_encoding_md5(&dec_config, &enc_config, enc_params, enc_options,
+                          &[0x08615111, 0x6f644a35, 0xa4e28f32, 0x35d2e66c]);
     }
 
     #[test]
@@ -613,9 +636,12 @@ mod test {
                 tb_den:  0,
                 flags:   0,
             };
-        //test_encoding_to_file(&dec_config, &enc_config, enc_params);
-        test_encoding_md5(&dec_config, &enc_config, enc_params,
-                          &[0x0b4cb528, 0x66c91f6c, 0x1c2187a5, 0x2723a08d]);
+        let enc_options = &[
+                NAOption { name: "range", value: NAValue::Int(16) },
+            ];
+        //test_encoding_to_file(&dec_config, &enc_config, enc_params, enc_options);
+        test_encoding_md5(&dec_config, &enc_config, enc_params, enc_options,
+                          &[0x1cff4116, 0x5926d91b, 0x60aac53f, 0x8a22bc7b]);
     }
 
     #[test]
@@ -657,9 +683,12 @@ mod test {
                 tb_den:  0,
                 flags:   0,
             };
-        //test_encoding_to_file(&dec_config, &enc_config, enc_params);
-        test_encoding_md5(&dec_config, &enc_config, enc_params,
-                          &[0x1a522743, 0x6c320a6e, 0xd08539e1, 0x03fc17ea]);
+        let enc_options = &[
+                NAOption { name: "range", value: NAValue::Int(16) },
+            ];
+        //test_encoding_to_file(&dec_config, &enc_config, enc_params, enc_options);
+        test_encoding_md5(&dec_config, &enc_config, enc_params, enc_options,
+                          &[0xb48aa633, 0x673539fe, 0xa312d45a, 0x76eee134]);
     }
 
     #[test]
@@ -701,8 +730,11 @@ mod test {
                 tb_den:  0,
                 flags:   0,
             };
-        //test_encoding_to_file(&dec_config, &enc_config, enc_params);
-        test_encoding_md5(&dec_config, &enc_config, enc_params,
-                          &[0x3880e045, 0xe6c88dc7, 0x21066058, 0xc789f1e9]);
+        let enc_options = &[
+                NAOption { name: "range", value: NAValue::Int(16) },
+            ];
+        //test_encoding_to_file(&dec_config, &enc_config, enc_params, enc_options);
+        test_encoding_md5(&dec_config, &enc_config, enc_params, enc_options,
+                          &[0x0836152c, 0xfcd7e1fc, 0xf1e2f619, 0x874d3dbc]);
     }
 }