From cc1fd004ea1ae8a919192a4c4b18b623254734fb Mon Sep 17 00:00:00 2001
From: Joose Sainio <joose.sainio@tuni.fi>
Date: Mon, 22 Jul 2024 12:57:07 +0300
Subject: [PATCH] [10bit] don't overwrite heap

---
 src/encoder.c                       | 16 +++++++++----
 src/strategies/strategies-picture.h | 37 +++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/src/encoder.c b/src/encoder.c
index 56d03305..4b6b42d0 100644
--- a/src/encoder.c
+++ b/src/encoder.c
@@ -141,12 +141,16 @@ static int get_max_parallelism(const encoder_control_t *const encoder)
 static int8_t* derive_chroma_QP_mapping_table(const uvg_config* const cfg, int i)
 {
   const int MAX_QP = 63;
+  const int qpBdOffsetC = (cfg->input_bitdepth - 8) * 6;
 
   int8_t qpInVal[16], qpOutVal[16];
-  int8_t* table = calloc(MAX_QP + 1, sizeof(int8_t));
+  int8_t* table = calloc(MAX_QP + 1+qpBdOffsetC, sizeof(int8_t));
+  if (table == NULL) {
+    fprintf(stderr, "Failed to allocate memory for chroma QP mapping table.\n");
+    return NULL;
+  }
+  table += qpBdOffsetC;
 
-
-  const int qpBdOffsetC = (cfg->input_bitdepth - 8) * 6;
   const int numPtsInCQPTableMinus1 = cfg->qp_table_length_minus1[i];
 
   qpInVal[0] = cfg->qp_table_start_minus26[i] + 26;
@@ -652,7 +656,7 @@ encoder_control_t* uvg_encoder_control_init(const uvg_config *const cfg)
   } else {
     encoder->cfg.vps_period = -1;
   }
-
+  
   for (int i = 0; i < cfg->num_used_table; i++) {
     encoder->qp_map[i] = derive_chroma_QP_mapping_table(cfg, i);
   }
@@ -670,6 +674,7 @@ init_failed:
 void uvg_encoder_control_free(encoder_control_t *const encoder)
 {
   if (!encoder) return;
+  const int qpBdOffsetC = (encoder->cfg.input_bitdepth - 8) * 6;
 
   //Slices
   FREE_POINTER(encoder->slice_addresses_in_ts);
@@ -695,7 +700,8 @@ void uvg_encoder_control_free(encoder_control_t *const encoder)
   uvg_threadqueue_free(encoder->threadqueue);
   encoder->threadqueue = NULL;
   for (int i = 0; i < encoder->cfg.num_used_table; i++) {
-    if (encoder->qp_map[i]) FREE_POINTER(encoder->qp_map[i]);
+    int8_t *temp = encoder->qp_map[i] - qpBdOffsetC;
+    if (encoder->qp_map[i] - qpBdOffsetC) FREE_POINTER(temp);
   }
 
   uvg_close_rdcost_outfiles();
diff --git a/src/strategies/strategies-picture.h b/src/strategies/strategies-picture.h
index cd4e2ec5..da8f9be8 100644
--- a/src/strategies/strategies-picture.h
+++ b/src/strategies/strategies-picture.h
@@ -207,6 +207,7 @@ int uvg_strategy_register_picture(void* opaque, uint8_t bitdepth);
 cost_pixel_nxn_multi_func * uvg_pixels_get_satd_dual_func(unsigned width, unsigned height);
 cost_pixel_nxn_multi_func * uvg_pixels_get_sad_dual_func(unsigned width, unsigned height);
 
+#if UVG_BIT_DEPTH == 8 
 #define STRATEGIES_PICTURE_EXPORTS \
   {"crc32c_4x4", (void**) &uvg_crc32c_4x4}, \
   {"crc32c_8x8", (void **)&uvg_crc32c_8x8}, \
@@ -242,6 +243,42 @@ cost_pixel_nxn_multi_func * uvg_pixels_get_sad_dual_func(unsigned width, unsigne
   {"pixel_var", (void**) &uvg_pixel_var}, \
   {"generate_residual", (void**) &uvg_generate_residual}, \
 
+#else
+#define STRATEGIES_PICTURE_EXPORTS \
+  {"reg_sad", (void**) &uvg_reg_sad}, \
+  {"sad_4x4", (void**) &uvg_sad_4x4}, \
+  {"sad_8x8", (void**) &uvg_sad_8x8}, \
+  {"sad_16x16", (void**) &uvg_sad_16x16}, \
+  {"sad_32x32", (void**) &uvg_sad_32x32}, \
+  {"sad_64x64", (void**) &uvg_sad_64x64}, \
+  {"satd_4x4", (void**) &uvg_satd_4x4}, \
+  {"satd_8x8", (void**) &uvg_satd_8x8}, \
+  {"satd_16x16", (void**) &uvg_satd_16x16}, \
+  {"satd_32x32", (void**) &uvg_satd_32x32}, \
+  {"satd_64x64", (void**) &uvg_satd_64x64}, \
+  {"satd_any_size", (void**) &uvg_satd_any_size}, \
+  {"satd_any_size_vtm", (void**) &uvg_satd_any_size_vtm}, \
+  {"sad_4x4_dual", (void**) &uvg_sad_4x4_dual}, \
+  {"sad_8x8_dual", (void**) &uvg_sad_8x8_dual}, \
+  {"sad_16x16_dual", (void**) &uvg_sad_16x16_dual}, \
+  {"sad_32x32_dual", (void**) &uvg_sad_32x32_dual}, \
+  {"sad_64x64_dual", (void**) &uvg_sad_64x64_dual}, \
+  {"satd_4x4_dual", (void**) &uvg_satd_4x4_dual}, \
+  {"satd_8x8_dual", (void**) &uvg_satd_8x8_dual}, \
+  {"satd_16x16_dual", (void**) &uvg_satd_16x16_dual}, \
+  {"satd_32x32_dual", (void**) &uvg_satd_32x32_dual}, \
+  {"satd_64x64_dual", (void**) &uvg_satd_64x64_dual}, \
+  {"satd_any_size_quad", (void**) &uvg_satd_any_size_quad}, \
+  {"pixels_calc_ssd", (void**) &uvg_pixels_calc_ssd}, \
+  {"bipred_average", (void**) &uvg_bipred_average}, \
+  {"get_optimized_sad", (void**) &uvg_get_optimized_sad}, \
+  {"ver_sad", (void**) &uvg_ver_sad}, \
+  {"hor_sad", (void**) &uvg_hor_sad}, \
+  {"pixel_var", (void**) &uvg_pixel_var}, \
+  {"generate_residual", (void**) &uvg_generate_residual}, \
+
+#endif 
+
 
 
 #endif //STRATEGIES_PICTURE_H_