Avoid potential buffer overflow issue

Warning if the slice type is wrong for encoding

Signed-off-by: Xiang, Haihao <haihao.xiang@intel.com>

7 files changed, 57 insertions, 54 deletions

diff --git a/src/gen6_mfc.c b/src/gen6_mfc.c
index ccfdcf6..5f430a9 100644
--- a/src/gen6_mfc.c
+++ b/src/gen6_mfc.c
@@ -352,9 +352,9 @@ gen6_mfc_avc_slice_state(VADriverContextP ctx,
     int beginy = beginmb / width_in_mbs;
     int nextx =  endmb % width_in_mbs;
     int nexty = endmb / width_in_mbs;
-    int slice_type = slice_param->slice_type;
+    int slice_type = intel_avc_enc_slice_type_fixup(slice_param->slice_type);
     int last_slice = (endmb == (width_in_mbs * height_in_mbs));
-    int bit_rate_control_target, maxQpN, maxQpP;
+    int maxQpN, maxQpP;
     unsigned char correct[6], grow, shrink;
     int i;
     int weighted_pred_idc = 0;
@@ -365,12 +365,6 @@ gen6_mfc_avc_slice_state(VADriverContextP ctx,
     if (batch == NULL)
         batch = encoder_context->base.batch;
 
-    bit_rate_control_target = slice_type;
-    if (slice_type == SLICE_TYPE_SP)
-        bit_rate_control_target = SLICE_TYPE_P;
-    else if (slice_type == SLICE_TYPE_SI)
-        bit_rate_control_target = SLICE_TYPE_I;
-
     if (slice_type == SLICE_TYPE_P) {
         weighted_pred_idc = pic_param->pic_fields.bits.weighted_pred_flag;
     } else if (slice_type == SLICE_TYPE_B) {
@@ -384,16 +378,16 @@ gen6_mfc_avc_slice_state(VADriverContextP ctx,
         }
     }
 
-    maxQpN = mfc_context->bit_rate_control_context[bit_rate_control_target].MaxQpNegModifier;
-    maxQpP = mfc_context->bit_rate_control_context[bit_rate_control_target].MaxQpPosModifier;
+    maxQpN = mfc_context->bit_rate_control_context[slice_type].MaxQpNegModifier;
+    maxQpP = mfc_context->bit_rate_control_context[slice_type].MaxQpPosModifier;
 
     for (i = 0; i < 6; i++)
-        correct[i] = mfc_context->bit_rate_control_context[bit_rate_control_target].Correct[i];
+        correct[i] = mfc_context->bit_rate_control_context[slice_type].Correct[i];
 
-    grow = mfc_context->bit_rate_control_context[bit_rate_control_target].GrowInit + 
-        (mfc_context->bit_rate_control_context[bit_rate_control_target].GrowResistance << 4);
-    shrink = mfc_context->bit_rate_control_context[bit_rate_control_target].ShrinkInit + 
-        (mfc_context->bit_rate_control_context[bit_rate_control_target].ShrinkResistance << 4);
+    grow = mfc_context->bit_rate_control_context[slice_type].GrowInit + 
+        (mfc_context->bit_rate_control_context[slice_type].GrowResistance << 4);
+    shrink = mfc_context->bit_rate_control_context[slice_type].ShrinkInit + 
+        (mfc_context->bit_rate_control_context[slice_type].ShrinkResistance << 4);
 
     BEGIN_BCS_BATCH(batch, 11);;
 
@@ -780,7 +774,6 @@ gen6_mfc_avc_pipeline_slice_programing(VADriverContextP ctx,
     VAEncPictureParameterBufferH264 *pPicParameter = (VAEncPictureParameterBufferH264 *)encode_state->pic_param_ext->buffer;
     VAEncSliceParameterBufferH264 *pSliceParameter = (VAEncSliceParameterBufferH264 *)encode_state->slice_params_ext[slice_index]->buffer; 
     unsigned int *msg = NULL, offset = 0;
-    int is_intra = pSliceParameter->slice_type == SLICE_TYPE_I;
     int width_in_mbs = (mfc_context->surface_state.width + 15) / 16;
     int height_in_mbs = (mfc_context->surface_state.height + 15) / 16;
     int last_slice = (pSliceParameter->macroblock_address + pSliceParameter->num_macroblocks) == (width_in_mbs * height_in_mbs);
@@ -790,8 +783,8 @@ gen6_mfc_avc_pipeline_slice_programing(VADriverContextP ctx,
     unsigned char *slice_header = NULL;
     int slice_header_length_in_bits = 0;
     unsigned int tail_data[] = { 0x0, 0x0 };
-    int slice_type = pSliceParameter->slice_type;
-
+    int slice_type = intel_avc_enc_slice_type_fixup(pSliceParameter->slice_type);
+    int is_intra = slice_type == SLICE_TYPE_I;
 
     if (rate_control_mode == VA_RC_CBR) {
         qp = mfc_context->bit_rate_control_context[slice_type].QpPrimeY;
@@ -846,7 +839,7 @@ gen6_mfc_avc_pipeline_slice_programing(VADriverContextP ctx,
             if (msg[0] & INTRA_MB_FLAG_MASK) {
                 gen6_mfc_avc_pak_object_intra(ctx, x, y, last_mb, qp, msg, encoder_context, 0, 0, slice_batch);
             } else {
-                gen6_mfc_avc_pak_object_inter(ctx, x, y, last_mb, qp, msg, offset, encoder_context, 0, 0, pSliceParameter->slice_type, slice_batch);
+                gen6_mfc_avc_pak_object_inter(ctx, x, y, last_mb, qp, msg, offset, encoder_context, 0, 0, slice_type, slice_batch);
             }
 
             msg += INTER_VME_OUTPUT_IN_DWS;
@@ -1167,7 +1160,7 @@ gen6_mfc_avc_batchbuffer_slice(VADriverContextP ctx,
     long head_offset;
     int old_used = intel_batchbuffer_used_size(slice_batch), used;
     unsigned short head_size, tail_size;
-    int slice_type = pSliceParameter->slice_type;
+    int slice_type = intel_avc_enc_slice_type_fixup(pSliceParameter->slice_type);
 
     if (rate_control_mode == VA_RC_CBR) {
         qp = mfc_context->bit_rate_control_context[slice_type].QpPrimeY;
diff --git a/src/gen6_mfc.h b/src/gen6_mfc.h
index 3e44e91..ba366d3 100644
--- a/src/gen6_mfc.h
+++ b/src/gen6_mfc.h
@@ -260,4 +260,6 @@ extern VAStatus intel_mfc_avc_prepare(VADriverContextP ctx,
                                      struct encode_state *encode_state,
                                      struct intel_encoder_context *encoder_context);
 
+extern int intel_avc_enc_slice_type_fixup(int type);
+
 #endif	/* _GEN6_MFC_BCS_H_ */
diff --git a/src/gen6_mfc_common.c b/src/gen6_mfc_common.c
index 1d57075..ed2711b 100644
--- a/src/gen6_mfc_common.c
+++ b/src/gen6_mfc_common.c
@@ -66,6 +66,24 @@
 #define log2f(x) (logf(x)/(float)M_LN2)
 #endif
 
+int intel_avc_enc_slice_type_fixup(int slice_type)
+{
+    if (slice_type == SLICE_TYPE_SP ||
+        slice_type == SLICE_TYPE_P)
+        slice_type = SLICE_TYPE_P;
+    else if (slice_type == SLICE_TYPE_SI ||
+             slice_type == SLICE_TYPE_I)
+        slice_type = SLICE_TYPE_I;
+    else {
+        if (slice_type != SLICE_TYPE_B)
+            WARN_ONCE("Invalid slice type for H.264 encoding!\n");
+
+        slice_type = SLICE_TYPE_B;
+    }
+
+    return slice_type;
+}
+
 static void
 intel_mfc_bit_rate_control_context_init(struct encode_state *encode_state, 
                                        struct gen6_mfc_context *mfc_context)
@@ -200,7 +218,7 @@ int intel_mfc_brc_postpack(struct encode_state *encode_state,
 {
     gen6_brc_status sts = BRC_NO_HRD_VIOLATION;
     VAEncSliceParameterBufferH264 *pSliceParameter = (VAEncSliceParameterBufferH264 *)encode_state->slice_params_ext[0]->buffer; 
-    int slicetype = pSliceParameter->slice_type;
+    int slicetype = intel_avc_enc_slice_type_fixup(pSliceParameter->slice_type);
     int qpi = mfc_context->bit_rate_control_context[SLICE_TYPE_I].QpPrimeY;
     int qpp = mfc_context->bit_rate_control_context[SLICE_TYPE_P].QpPrimeY;
     int qpb = mfc_context->bit_rate_control_context[SLICE_TYPE_B].QpPrimeY;
@@ -216,11 +234,6 @@ int intel_mfc_brc_postpack(struct encode_state *encode_state,
     double x, y;
     double frame_size_alpha;
 
-    if (slicetype == SLICE_TYPE_SP)
-        slicetype = SLICE_TYPE_P;
-    else if (slicetype == SLICE_TYPE_SI)
-        slicetype = SLICE_TYPE_I;
-
     qp = mfc_context->bit_rate_control_context[slicetype].QpPrimeY;
 
     target_frame_size = mfc_context->brc.target_frame_size[slicetype];
@@ -724,18 +737,20 @@ void intel_vme_update_mbmv_cost(VADriverContextP ctx,
     uint8_t *vme_state_message = (uint8_t *)(vme_context->vme_state_message);
     float   lambda, m_costf;
 
+    int slice_type = intel_avc_enc_slice_type_fixup(slice_param->slice_type);
+
+    
     if (encoder_context->rate_control_mode == VA_RC_CQP)
 	qp = pic_param->pic_init_qp + slice_param->slice_qp_delta;
     else
-	qp = mfc_context->bit_rate_control_context[slice_param->slice_type].QpPrimeY;
+	qp = mfc_context->bit_rate_control_context[slice_type].QpPrimeY;
   
     if (vme_state_message == NULL)
 	return;
  
     assert(qp <= QP_MAX); 
     lambda = intel_lambda_qp(qp);
-    if ((slice_param->slice_type == SLICE_TYPE_I) ||
-		(slice_param->slice_type == SLICE_TYPE_SI)) {
+    if (slice_type == SLICE_TYPE_I) {
 	vme_state_message[MODE_INTRA_16X16] = 0;
 	m_cost = lambda * 4;
 	vme_state_message[MODE_INTRA_8X8] = intel_format_lutvalue(m_cost, 0x8f);
@@ -781,8 +796,7 @@ void intel_vme_update_mbmv_cost(VADriverContextP ctx,
 	m_costf = lambda * 3.5;
 	m_cost = m_costf;
 	vme_state_message[MODE_INTRA_NONPRED] = intel_format_lutvalue(m_cost, 0x6f);
-    	if ((slice_param->slice_type == SLICE_TYPE_P) ||
-		(slice_param->slice_type == SLICE_TYPE_SP)) {
+    	if (slice_type == SLICE_TYPE_P) {
 		m_costf = lambda * 2.5;
 		m_cost = m_costf;
 		vme_state_message[MODE_INTER_16X16] = intel_format_lutvalue(m_cost, 0x8f);
diff --git a/src/gen6_vme.c b/src/gen6_vme.c
index 95f2d08..ec70286 100644
--- a/src/gen6_vme.c
+++ b/src/gen6_vme.c
@@ -444,10 +444,11 @@ static void gen6_vme_state_setup_fixup(VADriverContextP ctx,
     if (slice_param->slice_type != SLICE_TYPE_I &&
         slice_param->slice_type != SLICE_TYPE_SI)
         return;
+
     if (encoder_context->rate_control_mode == VA_RC_CQP)
         vme_state_message[16] = intra_mb_mode_cost_table[pic_param->pic_init_qp + slice_param->slice_qp_delta];
     else
-        vme_state_message[16] = intra_mb_mode_cost_table[mfc_context->bit_rate_control_context[slice_param->slice_type].QpPrimeY];
+        vme_state_message[16] = intra_mb_mode_cost_table[mfc_context->bit_rate_control_context[SLICE_TYPE_I].QpPrimeY];
 }
 
 static VAStatus gen6_vme_vme_state_setup(VADriverContextP ctx,
diff --git a/src/gen75_mfc.c b/src/gen75_mfc.c
index e7a3221..f954d7a 100644
--- a/src/gen75_mfc.c
+++ b/src/gen75_mfc.c
@@ -895,9 +895,9 @@ gen75_mfc_avc_slice_state(VADriverContextP ctx,
     int beginy = beginmb / width_in_mbs;
     int nextx =  endmb % width_in_mbs;
     int nexty = endmb / width_in_mbs;
-    int slice_type = slice_param->slice_type;
+    int slice_type = intel_avc_enc_slice_type_fixup(slice_param->slice_type);
     int last_slice = (endmb == (width_in_mbs * height_in_mbs));
-    int bit_rate_control_target, maxQpN, maxQpP;
+    int maxQpN, maxQpP;
     unsigned char correct[6], grow, shrink;
     int i;
     int bslice = 0;
@@ -908,12 +908,6 @@ gen75_mfc_avc_slice_state(VADriverContextP ctx,
     if (batch == NULL)
         batch = encoder_context->base.batch;
 
-    bit_rate_control_target = slice_type;
-    if (slice_type == SLICE_TYPE_SP)
-        bit_rate_control_target = SLICE_TYPE_P;
-    else if (slice_type == SLICE_TYPE_SI)
-        bit_rate_control_target = SLICE_TYPE_I;
-
     if (slice_type == SLICE_TYPE_P) {
         weighted_pred_idc = pic_param->pic_fields.bits.weighted_pred_flag;
     } else if (slice_type == SLICE_TYPE_B) {
@@ -927,16 +921,16 @@ gen75_mfc_avc_slice_state(VADriverContextP ctx,
         }
     }
 
-    maxQpN = mfc_context->bit_rate_control_context[bit_rate_control_target].MaxQpNegModifier;
-    maxQpP = mfc_context->bit_rate_control_context[bit_rate_control_target].MaxQpPosModifier;
+    maxQpN = mfc_context->bit_rate_control_context[slice_type].MaxQpNegModifier;
+    maxQpP = mfc_context->bit_rate_control_context[slice_type].MaxQpPosModifier;
 
     for (i = 0; i < 6; i++)
-        correct[i] = mfc_context->bit_rate_control_context[bit_rate_control_target].Correct[i];
+        correct[i] = mfc_context->bit_rate_control_context[slice_type].Correct[i];
 
-    grow = mfc_context->bit_rate_control_context[bit_rate_control_target].GrowInit + 
-        (mfc_context->bit_rate_control_context[bit_rate_control_target].GrowResistance << 4);
-    shrink = mfc_context->bit_rate_control_context[bit_rate_control_target].ShrinkInit + 
-        (mfc_context->bit_rate_control_context[bit_rate_control_target].ShrinkResistance << 4);
+    grow = mfc_context->bit_rate_control_context[slice_type].GrowInit + 
+        (mfc_context->bit_rate_control_context[slice_type].GrowResistance << 4);
+    shrink = mfc_context->bit_rate_control_context[slice_type].ShrinkInit + 
+        (mfc_context->bit_rate_control_context[slice_type].ShrinkResistance << 4);
 
     BEGIN_BCS_BATCH(batch, 11);;
 
@@ -1170,7 +1164,6 @@ gen75_mfc_avc_pipeline_slice_programing(VADriverContextP ctx,
     VAEncSliceParameterBufferH264 *pSliceParameter = (VAEncSliceParameterBufferH264 *)encode_state->slice_params_ext[slice_index]->buffer; 
     unsigned int *msg = NULL, offset = 0;
     unsigned char *msg_ptr = NULL;
-    int is_intra = pSliceParameter->slice_type == SLICE_TYPE_I;
     int width_in_mbs = (mfc_context->surface_state.width + 15) / 16;
     int height_in_mbs = (mfc_context->surface_state.height + 15) / 16;
     int last_slice = (pSliceParameter->macroblock_address + pSliceParameter->num_macroblocks) == (width_in_mbs * height_in_mbs);
@@ -1180,8 +1173,8 @@ gen75_mfc_avc_pipeline_slice_programing(VADriverContextP ctx,
     unsigned char *slice_header = NULL;
     int slice_header_length_in_bits = 0;
     unsigned int tail_data[] = { 0x0, 0x0 };
-    int slice_type = pSliceParameter->slice_type;
-
+    int slice_type = intel_avc_enc_slice_type_fixup(pSliceParameter->slice_type);
+    int is_intra = slice_type == SLICE_TYPE_I;
 
     if (rate_control_mode == VA_RC_CBR) {
         qp = mfc_context->bit_rate_control_context[slice_type].QpPrimeY;
@@ -1239,7 +1232,7 @@ gen75_mfc_avc_pipeline_slice_programing(VADriverContextP ctx,
                 gen75_mfc_avc_pak_object_intra(ctx, x, y, last_mb, qp, msg, encoder_context, 0, 0, slice_batch);
             } else {
 		msg += AVC_INTER_MSG_OFFSET;
-                gen75_mfc_avc_pak_object_inter(ctx, x, y, last_mb, qp, msg, offset, encoder_context, 0, 0, pSliceParameter->slice_type, slice_batch);
+                gen75_mfc_avc_pak_object_inter(ctx, x, y, last_mb, qp, msg, offset, encoder_context, 0, 0, slice_type, slice_batch);
             }
         }
     }
@@ -1564,7 +1557,7 @@ gen75_mfc_avc_batchbuffer_slice(VADriverContextP ctx,
     long head_offset;
     int old_used = intel_batchbuffer_used_size(slice_batch), used;
     unsigned short head_size, tail_size;
-    int slice_type = pSliceParameter->slice_type;
+    int slice_type = intel_avc_enc_slice_type_fixup(pSliceParameter->slice_type);
 
     if (rate_control_mode == VA_RC_CBR) {
         qp = mfc_context->bit_rate_control_context[slice_type].QpPrimeY;
diff --git a/src/gen75_vme.c b/src/gen75_vme.c
index e0d982b..af002f2 100644
--- a/src/gen75_vme.c
+++ b/src/gen75_vme.c
@@ -452,7 +452,7 @@ static void gen75_vme_state_setup_fixup(VADriverContextP ctx,
     if (encoder_context->rate_control_mode == VA_RC_CQP)
         vme_state_message[0] = intra_mb_mode_cost_table[pic_param->pic_init_qp + slice_param->slice_qp_delta];
     else
-        vme_state_message[0] = intra_mb_mode_cost_table[mfc_context->bit_rate_control_context[slice_param->slice_type].QpPrimeY];
+        vme_state_message[0] = intra_mb_mode_cost_table[mfc_context->bit_rate_control_context[SLICE_TYPE_I].QpPrimeY];
 }
 
 static VAStatus gen75_vme_vme_state_setup(VADriverContextP ctx,
diff --git a/src/gen7_vme.c b/src/gen7_vme.c
index 9c8f4fe..493992f 100644
--- a/src/gen7_vme.c
+++ b/src/gen7_vme.c
@@ -425,7 +425,7 @@ static void gen7_vme_state_setup_fixup(VADriverContextP ctx,
     if (encoder_context->rate_control_mode == VA_RC_CQP)
         vme_state_message[16] = intra_mb_mode_cost_table[pic_param->pic_init_qp + slice_param->slice_qp_delta];
     else
-        vme_state_message[16] = intra_mb_mode_cost_table[mfc_context->bit_rate_control_context[slice_param->slice_type].QpPrimeY];
+        vme_state_message[16] = intra_mb_mode_cost_table[mfc_context->bit_rate_control_context[SLICE_TYPE_I].QpPrimeY];
 }
 
 static VAStatus gen7_vme_avc_state_setup(VADriverContextP ctx,