fix potential encoder crash (#419)

strukturag · Sep 1, 2023 · 296503d · 296503d
1 parent 6fc550f
commit 296503d
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 11 deletions.
diff --git a/libde265/cabac.cc b/libde265/cabac.cc
@@ -492,6 +492,8 @@ void CABAC_encoder_bitstream::write_bits(uint32_t bits,int n)
   vlc_buffer |= bits;
   vlc_buffer_len += n;
 
+  // TODO: errors returned by append_byte() are ignored, resulting in a broken output.
+
   while (vlc_buffer_len>=8) {
     append_byte((vlc_buffer >> (vlc_buffer_len-8)) & 0xFF);
     vlc_buffer_len -= 8;
@@ -524,6 +526,8 @@ void CABAC_encoder::write_svlc(int value)
 
 void CABAC_encoder_bitstream::flush_VLC()
 {
+  // TODO: errors returned by append_byte() are ignored, resulting in a broken output.
+
   while (vlc_buffer_len>=8) {
     append_byte((vlc_buffer >> (vlc_buffer_len-8)) & 0xFF);
     vlc_buffer_len -= 8;
@@ -557,7 +561,7 @@ int  CABAC_encoder_bitstream::number_free_bits_in_byte() const
 }
 
 
-void CABAC_encoder_bitstream::check_size_and_resize(int nBytes)
+bool CABAC_encoder_bitstream::check_size_and_resize(int nBytes)
 {
   if (data_size+nBytes > data_capacity) { // 1 extra byte for stuffing
     if (data_capacity==0) {
@@ -566,14 +570,24 @@ void CABAC_encoder_bitstream::check_size_and_resize(int nBytes)
       data_capacity *= 2;
     }
 
-    data_mem = (uint8_t*)realloc(data_mem,data_capacity);
+    uint8_t* new_data_mem = (uint8_t*)realloc(data_mem,data_capacity);
+    if (new_data_mem) {
+      data_mem = new_data_mem;
+    }
+    else {
+      return false;
+    }
   }
+
+  return true;
 }
 
 
-void CABAC_encoder_bitstream::append_byte(int byte)
+bool CABAC_encoder_bitstream::append_byte(int byte)
 {
-  check_size_and_resize(2);
+  if (!check_size_and_resize(2)) {
+    return false;
+  }
 
   // --- emulation prevention ---
 
@@ -603,17 +617,23 @@ void CABAC_encoder_bitstream::append_byte(int byte)
   // write actual data byte
 
   data_mem[ data_size++ ] = byte;
+
+  return true;
 }
 
 
-void CABAC_encoder_bitstream::write_startcode()
+bool CABAC_encoder_bitstream::write_startcode()
 {
-  check_size_and_resize(3);
+  if (!check_size_and_resize(3)) {
+    return false;
+  }
 
   data_mem[ data_size+0 ] = 0;
   data_mem[ data_size+1 ] = 0;
   data_mem[ data_size+2 ] = 1;
   data_size+=3;
+
+  return true;
 }
 
 void CABAC_encoder_bitstream::init_CABAC()
@@ -628,6 +648,8 @@ void CABAC_encoder_bitstream::init_CABAC()
 
 void CABAC_encoder_bitstream::flush_CABAC()
 {
+  // TODO: errors returned by append_byte() are ignored, resulting in a broken output.
+
   if (low >> (32 - bits_left))
     {
       append_byte(buffered_byte + 1);
@@ -661,6 +683,8 @@ void CABAC_encoder_bitstream::flush_CABAC()
 
 void CABAC_encoder_bitstream::write_out()
 {
+  // TODO: errors returned by append_byte() are ignored, resulting in a broken output.
+
   //logtrace(LogCABAC,"low = %08x (bits_left=%d)\n",low,bits_left);
   int leadByte = low >> (24 - bits_left);
   bits_left += 8;

diff --git a/libde265/cabac.h b/libde265/cabac.h
@@ -66,7 +66,7 @@ class CABAC_encoder
   virtual void write_bit(int bit) { write_bits(bit,1); }
   virtual void write_uvlc(int value);
   virtual void write_svlc(int value);
-  virtual void write_startcode() = 0;
+  virtual bool write_startcode() = 0;
   virtual void skip_bits(int nBits) = 0;
 
   virtual void add_trailing_bits();
@@ -113,7 +113,7 @@ class CABAC_encoder_bitstream : public CABAC_encoder
   // --- VLC ---
 
   virtual void write_bits(uint32_t bits,int n);
-  virtual void write_startcode();
+  virtual bool write_startcode();
   virtual void skip_bits(int nBits);
 
   virtual int  number_free_bits_in_byte() const;
@@ -155,10 +155,10 @@ class CABAC_encoder_bitstream : public CABAC_encoder
   uint16_t num_buffered_bytes;
 
 
-  void check_size_and_resize(int nBytes);
+  bool check_size_and_resize(int nBytes);
   void testAndWriteOut();
   void write_out();
-  void append_byte(int byte);
+  bool append_byte(int byte);
 };
 
 
@@ -178,7 +178,7 @@ class CABAC_encoder_estim : public CABAC_encoder
 
   virtual void write_bits(uint32_t bits,int n) { mFracBits += n<<15; }
   virtual void write_bit(int bit) { mFracBits+=1<<15; }
-  virtual void write_startcode() { mFracBits += (1<<15)*8*3; }
+  virtual bool write_startcode() { mFracBits += (1<<15)*8*3; return true; }
   virtual void skip_bits(int nBits) { mFracBits += nBits<<15; }
   virtual int  number_free_bits_in_byte() const { return 0; } // TODO, good enough for now