From 5f212c159be57dc02ffed0caeca965ec53ba0e8e Mon Sep 17 00:00:00 2001
From: Guillaume Desmottes <guillaume.desmottes@collabora.com>
Date: Thu, 16 May 2019 10:50:18 +0200
Subject: [PATCH] omxh265enc: handle CODECCONFIG buffers

Exact same code as omxh264enc.
---
 omx/gstomxh265enc.c | 76 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 omx/gstomxh265enc.h |  2 ++
 2 files changed, 78 insertions(+)

diff --git a/omx/gstomxh265enc.c b/omx/gstomxh265enc.c
index 6cc89d4..e05105f 100644
--- a/omx/gstomxh265enc.c
+++ b/omx/gstomxh265enc.c
@@ -40,6 +40,8 @@ static void gst_omx_h265_enc_set_property (GObject * object, guint prop_id,
     const GValue * value, GParamSpec * pspec);
 static void gst_omx_h265_enc_get_property (GObject * object, guint prop_id,
     GValue * value, GParamSpec * pspec);
+static GstFlowReturn gst_omx_h265_enc_handle_output_frame (GstOMXVideoEnc *
+    self, GstOMXPort * port, GstOMXBuffer * buf, GstVideoCodecFrame * frame);
 
 enum
 {
@@ -106,15 +108,43 @@ gst_omx_h265_enc_loop_filter_mode_get_type (void)
 }
 #endif
 
+static gboolean
+gst_omx_h265_enc_flush (GstVideoEncoder * enc)
+{
+  GstOMXH265Enc *self = GST_OMX_H265_ENC (enc);
+
+  g_list_free_full (self->headers, (GDestroyNotify) gst_buffer_unref);
+  self->headers = NULL;
+
+  return GST_VIDEO_ENCODER_CLASS (parent_class)->flush (enc);
+}
+
+static gboolean
+gst_omx_h265_enc_stop (GstVideoEncoder * enc)
+{
+  GstOMXH265Enc *self = GST_OMX_H265_ENC (enc);
+
+  g_list_free_full (self->headers, (GDestroyNotify) gst_buffer_unref);
+  self->headers = NULL;
+
+  return GST_VIDEO_ENCODER_CLASS (parent_class)->stop (enc);
+}
+
 static void
 gst_omx_h265_enc_class_init (GstOMXH265EncClass * klass)
 {
   GObjectClass *gobject_class = G_OBJECT_CLASS (klass);
   GstElementClass *element_class = GST_ELEMENT_CLASS (klass);
+  GstVideoEncoderClass *basevideoenc_class = GST_VIDEO_ENCODER_CLASS (klass);
   GstOMXVideoEncClass *videoenc_class = GST_OMX_VIDEO_ENC_CLASS (klass);
 
   videoenc_class->set_format = GST_DEBUG_FUNCPTR (gst_omx_h265_enc_set_format);
   videoenc_class->get_caps = GST_DEBUG_FUNCPTR (gst_omx_h265_enc_get_caps);
+  videoenc_class->handle_output_frame =
+      GST_DEBUG_FUNCPTR (gst_omx_h265_enc_handle_output_frame);
+
+  basevideoenc_class->flush = gst_omx_h265_enc_flush;
+  basevideoenc_class->stop = gst_omx_h265_enc_stop;
 
   gobject_class->set_property = gst_omx_h265_enc_set_property;
   gobject_class->get_property = gst_omx_h265_enc_get_property;
@@ -664,3 +694,49 @@ gst_omx_h265_enc_get_caps (GstOMXVideoEnc * enc, GstOMXPort * port,
 
   return caps;
 }
+
+static GstFlowReturn
+gst_omx_h265_enc_handle_output_frame (GstOMXVideoEnc * enc, GstOMXPort * port,
+    GstOMXBuffer * buf, GstVideoCodecFrame * frame)
+{
+  GstOMXH265Enc *self = GST_OMX_H265_ENC (enc);
+
+  if (buf->omx_buf->nFlags & OMX_BUFFERFLAG_CODECCONFIG) {
+    /* The codec data is SPS/PPS but our output is stream-format=byte-stream.
+     * For bytestream stream format the SPS/PPS is only in-stream and not
+     * in the caps!
+     */
+    GstBuffer *hdrs;
+    GstMapInfo map = GST_MAP_INFO_INIT;
+    GstFlowReturn flow_ret;
+
+    GST_DEBUG_OBJECT (self, "got codecconfig in byte-stream format");
+
+    hdrs = gst_buffer_new_and_alloc (buf->omx_buf->nFilledLen);
+    GST_BUFFER_FLAG_SET (hdrs, GST_BUFFER_FLAG_HEADER);
+
+    gst_buffer_map (hdrs, &map, GST_MAP_WRITE);
+    memcpy (map.data,
+        buf->omx_buf->pBuffer + buf->omx_buf->nOffset,
+        buf->omx_buf->nFilledLen);
+    gst_buffer_unmap (hdrs, &map);
+    self->headers = g_list_append (self->headers, gst_buffer_ref (hdrs));
+
+    frame->output_buffer = hdrs;
+    flow_ret =
+        gst_video_encoder_finish_subframe (GST_VIDEO_ENCODER (self), frame);
+
+    if (frame)
+      gst_video_codec_frame_unref (frame);
+
+    return flow_ret;
+  } else if (self->headers) {
+    gst_video_encoder_set_headers (GST_VIDEO_ENCODER (self), self->headers);
+    self->headers = NULL;
+  }
+
+  return
+      GST_OMX_VIDEO_ENC_CLASS
+      (gst_omx_h265_enc_parent_class)->handle_output_frame (enc, port, buf,
+      frame);
+}
diff --git a/omx/gstomxh265enc.h b/omx/gstomxh265enc.h
index c53be18..b67fa1f 100644
--- a/omx/gstomxh265enc.h
+++ b/omx/gstomxh265enc.h
@@ -55,6 +55,8 @@ struct _GstOMXH265Enc
   gboolean constrained_intra_prediction;
   guint32 loop_filter_mode;
 #endif
+
+  GList *headers;
 };
 
 struct _GstOMXH265EncClass
-- 
2.7.4