From 0cf67c3be7b70f6da5a06a9a34f3621b9c3d9204 Mon Sep 17 00:00:00 2001
From: Seungha Yang <seungha.yang@navercorp.com>
Date: Thu, 14 Nov 2019 19:00:51 +0900
Subject: [PATCH] nvenc: Fix crash when nvenc was reused then freed without
 encoding

GstNvBaseEnc::n_bufs was set from the previous encoding session
but it wasn't cleared after stop. That might result to invalid memory
access at the next start (no encoded data) and then stop sequence.
Instead of defining a variable for array length, use GArray::len
directly to avoid such confusion.
---
 sys/nvcodec/gstnvbaseenc.c | 18 +++++++++++-------
 sys/nvcodec/gstnvbaseenc.h |  1 -
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/sys/nvcodec/gstnvbaseenc.c b/sys/nvcodec/gstnvbaseenc.c
index 8cfee75..317ebaa 100644
--- a/sys/nvcodec/gstnvbaseenc.c
+++ b/sys/nvcodec/gstnvbaseenc.c
@@ -1272,8 +1272,11 @@ gst_nv_base_enc_free_buffers (GstNvBaseEnc * nvenc)
 
   gst_nv_base_enc_reset_queues (nvenc);
 
+  if (!nvenc->items || !nvenc->items->len)
+    return;
+
   gst_cuda_context_push (nvenc->cuda_ctx);
-  for (i = 0; i < nvenc->n_bufs; ++i) {
+  for (i = 0; i < nvenc->items->len; ++i) {
     NV_ENC_OUTPUT_PTR out_buf =
         g_array_index (nvenc->items, GstNvEncFrameState, i).out_buf;
     GstNvEncInputResource *in_buf =
@@ -1448,11 +1451,11 @@ gst_nv_base_enc_setup_rate_control (GstNvBaseEnc * nvenc,
   }
 }
 
-static gint
+static guint
 gst_nv_base_enc_calculate_num_prealloc_buffers (GstNvBaseEnc * enc,
     NV_ENC_CONFIG * config)
 {
-  gint num_buffers;
+  guint num_buffers;
 
   /* At least 4 surfaces are required as documented by Nvidia Encoder guide */
   num_buffers = 4;
@@ -1712,16 +1715,17 @@ gst_nv_base_enc_set_format (GstVideoEncoder * enc, GstVideoCodecState * state)
 #endif
     guint i;
     guint input_width, input_height;
+    guint n_bufs;
 
     input_width = GST_VIDEO_INFO_WIDTH (info);
     input_height = GST_VIDEO_INFO_HEIGHT (info);
 
-    nvenc->n_bufs =
+    n_bufs =
         gst_nv_base_enc_calculate_num_prealloc_buffers (nvenc,
         params->encodeConfig);
 
     /* input buffers */
-    g_array_set_size (nvenc->items, nvenc->n_bufs);
+    g_array_set_size (nvenc->items, n_bufs);
 
 #if HAVE_NVCODEC_GST_GL
     features = gst_caps_get_features (state->caps, 0);
@@ -1732,7 +1736,7 @@ gst_nv_base_enc_set_format (GstVideoEncoder * enc, GstVideoCodecState * state)
 #endif
 
     gst_cuda_context_push (nvenc->cuda_ctx);
-    for (i = 0; i < nvenc->n_bufs; ++i) {
+    for (i = 0; i < nvenc->items->len; ++i) {
       GstNvEncInputResource *resource = g_new0 (GstNvEncInputResource, 1);
       CUresult cu_ret;
 
@@ -1773,7 +1777,7 @@ gst_nv_base_enc_set_format (GstVideoEncoder * enc, GstVideoCodecState * state)
     gst_cuda_context_pop (NULL);
 
     /* output buffers */
-    for (i = 0; i < nvenc->n_bufs; ++i) {
+    for (i = 0; i < nvenc->items->len; ++i) {
       NV_ENC_CREATE_BITSTREAM_BUFFER cout_buf = { 0, };
 
       cout_buf.version = gst_nvenc_get_create_bitstream_buffer_version ();
diff --git a/sys/nvcodec/gstnvbaseenc.h b/sys/nvcodec/gstnvbaseenc.h
index 7541f11..e9b2fe1 100644
--- a/sys/nvcodec/gstnvbaseenc.h
+++ b/sys/nvcodec/gstnvbaseenc.h
@@ -117,7 +117,6 @@ typedef struct {
   /* array of allocated input/output buffers (GstNvEncFrameState),
    * and hold the ownership of the GstNvEncFrameState. */
   GArray            *items;
-  guint              n_bufs;
 
   /* (GstNvEncFrameState) available empty items which could be submitted
    * to encoder */
-- 
2.7.4