mpegtsmux: generate SPS/PPS header once and fix overflow

Some H264 packets can be as small as 5 bytes for repeated frames.
In such a situation the output buffer size was not big enough (5*2) to fit the
SPS/PPS header and the start codes. This corrupts the ES stream.
We now generate the SPS/PPS only once which is much more optimal and we now
know the size of the header to calculate the output buffer size more safely.

gst/mpegtsmux/mpegtsmux.c

@@ -1,5 +1,5 @@
 /* 
- * Copyright 2006, 2007, 2008 Fluendo S.A. 
+ * Copyright 2006, 2007, 2008, 2009, 2010 Fluendo S.A. 
  *  Authors: Jan Schmidt <jan@fluendo.com>
  *           Kapil Agrawal <kapil@fluendo.com>
  *           Julien Moutte <julien@fluendo.com>
@@ -350,6 +350,7 @@ mpegtsmux_create_stream (MpegTsMux * mux, MpegTsPadData * ts_data, GstPad * pad)
       GST_DEBUG_OBJECT (pad, "we have additional codec data (%d bytes)",
           GST_BUFFER_SIZE (ts_data->codec_data));
       ts_data->prepare_func = mpegtsmux_prepare_h264;
+      ts_data->free_func = mpegtsmux_free_h264;
     } else {
       ts_data->codec_data = NULL;
     }
@@ -713,7 +714,9 @@ mpegtsmux_request_new_pad (GstElement * element,
   pad_data->pid = pid;
   pad_data->last_ts = GST_CLOCK_TIME_NONE;
   pad_data->codec_data = NULL;
+  pad_data->prepare_data = NULL;
   pad_data->prepare_func = NULL;
+  pad_data->free_func = NULL;
   pad_data->prog_id = -1;
   pad_data->prog = NULL;
 
@@ -757,6 +760,10 @@ mpegtsmux_release_pad (GstElement * element, GstPad * pad)
       gst_buffer_unref (pad_data->codec_data);
       pad_data->codec_data = NULL;
     }
+    if (pad_data->prepare_data && pad_data->free_func) {
+      pad_data->free_func (pad_data->prepare_data);
+      pad_data->prepare_data = pad_data->free_func = NULL;
+    }
   }
   GST_OBJECT_UNLOCK (pad);
 

gst/mpegtsmux/mpegtsmux.h

@@ -1,5 +1,5 @@
 /* 
- * Copyright 2006, 2007, 2008 Fluendo S.A. 
+ * Copyright 2006, 2007, 2008, 2009, 2010 Fluendo S.A. 
  *  Authors: Jan Schmidt <jan@fluendo.com>
  *           Kapil Agrawal <kapil@fluendo.com>
  *           Julien Moutte <julien@fluendo.com>
@@ -101,6 +101,8 @@ typedef struct MpegTsPadData MpegTsPadData;
 typedef GstBuffer * (*MpegTsPadDataPrepareFunction) (GstBuffer * buf,
     MpegTsPadData * data, MpegTsMux * mux);
 
+typedef void (*MpegTsPadDataFreePrepareDataFunction) (gpointer prepare_data);
+
 struct MpegTsMux {
   GstElement parent;
 
@@ -140,7 +142,11 @@ struct MpegTsPadData {
 
   GstBuffer * codec_data; /* Optional codec data available in the caps */
 
+  gpointer prepare_data; /* Opaque data pointer to a structure used by the
+                            prepare function */
+
   MpegTsPadDataPrepareFunction prepare_func; /* Handler to prepare input data */
+  MpegTsPadDataFreePrepareDataFunction free_func; /* Handler to free the private data */
 
   gboolean eos;
 

gst/mpegtsmux/mpegtsmux_h264.c

@@ -1,5 +1,5 @@
 /* 
- * Copyright 2006, 2007, 2008 Fluendo S.A. 
+ * Copyright 2006, 2007, 2008, 2009, 2010 Fluendo S.A.
  *  Authors: Jan Schmidt <jan@fluendo.com>
  *           Kapil Agrawal <kapil@fluendo.com>
  *           Julien Moutte <julien@fluendo.com>
@@ -90,82 +90,170 @@
 GST_DEBUG_CATEGORY_EXTERN (mpegtsmux_debug);
 #define GST_CAT_DEFAULT mpegtsmux_debug
 
+#define SPS_PPS_PERIOD GST_SECOND
+
+typedef struct PrivDataH264 PrivDataH264;
+
+struct PrivDataH264
+{
+  GstBuffer *last_codec_data;
+  GstClockTime last_resync_ts;
+  GstBuffer *cached_es;
+  guint8 nal_length_size;
+};
+
+void
+mpegtsmux_free_h264 (gpointer prepare_data)
+{
+  PrivDataH264 *h264_data = (PrivDataH264 *) prepare_data;
+  if (h264_data->cached_es) {
+    gst_buffer_unref (h264_data->cached_es);
+    h264_data->cached_es = NULL;
+  }
+  g_free (prepare_data);
+}
+
+static inline gboolean
+mpegtsmux_process_codec_data_h264 (MpegTsPadData * data, MpegTsMux * mux)
+{
+  PrivDataH264 *h264_data;
+  gboolean ret = FALSE;
+
+  /* Initialize our private data structure for caching */
+  if (G_UNLIKELY (!data->prepare_data)) {
+    data->prepare_data = g_new0 (PrivDataH264, 1);
+    h264_data = (PrivDataH264 *) data->prepare_data;
+    h264_data->last_resync_ts = GST_CLOCK_TIME_NONE;
+  }
+
+  h264_data = (PrivDataH264 *) data->prepare_data;
+
+  /* Detect a codec data change */
+  if (h264_data->last_codec_data != data->codec_data) {
+    gst_buffer_unref (h264_data->cached_es);
+    h264_data->cached_es = NULL;
+    ret = TRUE;
+  }
+
+  /* Generate the SPS/PPS ES header that will be prepended regularly */
+  if (G_UNLIKELY (!h264_data->cached_es)) {
+    gint offset = 4, i = 0, nb_sps = 0, nb_pps = 0;
+    gsize out_offset = 0;
+    guint8 startcode[4] = { 0x00, 0x00, 0x00, 0x01 };
+    h264_data->last_codec_data = data->codec_data;
+    h264_data->cached_es =
+        gst_buffer_new_and_alloc (GST_BUFFER_SIZE (data->codec_data) * 10);
+
+    /* Get NAL length size */
+    h264_data->nal_length_size =
+        (GST_READ_UINT8 (GST_BUFFER_DATA (data->codec_data) + offset) & 0x03) +
+        1;
+    GST_LOG_OBJECT (mux, "NAL length will be coded on %u bytes",
+        h264_data->nal_length_size);
+    offset++;
+
+    /* How many SPS */
+    nb_sps =
+        GST_READ_UINT8 (GST_BUFFER_DATA (data->codec_data) + offset) & 0x1f;
+    GST_DEBUG_OBJECT (mux, "we have %d Sequence Parameter Set", nb_sps);
+    offset++;
+
+    /* For each SPS */
+    for (i = 0; i < nb_sps; i++) {
+      guint16 sps_size =
+          GST_READ_UINT16_BE (GST_BUFFER_DATA (data->codec_data) + offset);
+
+      GST_LOG_OBJECT (mux, "Sequence Parameter Set is %d bytes", sps_size);
+
+      /* Jump over SPS size */
+      offset += 2;
+
+      /* Fake a start code */
+      memcpy (GST_BUFFER_DATA (h264_data->cached_es) + out_offset,
+          startcode, 4);
+      out_offset += 4;
+      /* Now push the SPS */
+      memcpy (GST_BUFFER_DATA (h264_data->cached_es) + out_offset,
+          GST_BUFFER_DATA (data->codec_data) + offset, sps_size);
+
+      out_offset += sps_size;
+      offset += sps_size;
+    }
+
+    /* How many PPS */
+    nb_pps = GST_READ_UINT8 (GST_BUFFER_DATA (data->codec_data) + offset);
+    GST_LOG_OBJECT (mux, "we have %d Picture Parameter Set", nb_sps);
+    offset++;
+
+    /* For each PPS */
+    for (i = 0; i < nb_pps; i++) {
+      gint pps_size =
+          GST_READ_UINT16_BE (GST_BUFFER_DATA (data->codec_data) + offset);
+
+      GST_LOG_OBJECT (mux, "Picture Parameter Set is %d bytes", pps_size);
+
+      /* Jump over PPS size */
+      offset += 2;
+
+      /* Fake a start code */
+      memcpy (GST_BUFFER_DATA (h264_data->cached_es) + out_offset,
+          startcode, 4);
+      out_offset += 4;
+      /* Now push the PPS */
+      memcpy (GST_BUFFER_DATA (h264_data->cached_es) + out_offset,
+          GST_BUFFER_DATA (data->codec_data) + offset, pps_size);
+
+      out_offset += pps_size;
+      offset += pps_size;
+    }
+    GST_BUFFER_SIZE (h264_data->cached_es) = out_offset;
+    GST_DEBUG_OBJECT (mux, "generated a %" G_GSIZE_FORMAT
+        " bytes SPS/PPS header", out_offset);
+  }
+  return ret;
+}
+
 GstBuffer *
 mpegtsmux_prepare_h264 (GstBuffer * buf, MpegTsPadData * data, MpegTsMux * mux)
 {
-  guint8 nal_length_size = 0;
   guint8 startcode[4] = { 0x00, 0x00, 0x00, 0x01 };
-  GstBuffer *out_buf = gst_buffer_new_and_alloc (GST_BUFFER_SIZE (buf) * 2);
-  gint offset = 4, i = 0, nb_sps = 0, nb_pps = 0;
   gsize out_offset = 0, in_offset = 0;
+  GstBuffer *out_buf;
+  gboolean changed;
+  PrivDataH264 *h264_data;
+  GstClockTimeDiff diff = GST_CLOCK_TIME_NONE;
 
   GST_DEBUG_OBJECT (mux, "Preparing H264 buffer for output");
 
+  changed = mpegtsmux_process_codec_data_h264 (data, mux);
+  h264_data = (PrivDataH264 *) data->prepare_data;
+
+  if (GST_CLOCK_TIME_IS_VALID (h264_data->last_resync_ts) &&
+      GST_CLOCK_TIME_IS_VALID (GST_BUFFER_TIMESTAMP (buf))) {
+    diff = GST_CLOCK_DIFF (h264_data->last_resync_ts,
+        GST_BUFFER_TIMESTAMP (buf));
+  }
+
+  if (changed || (GST_CLOCK_TIME_IS_VALID (diff) && diff > SPS_PPS_PERIOD)) {
+    out_buf = gst_buffer_new_and_alloc (GST_BUFFER_SIZE (buf) * 2 +
+        GST_BUFFER_SIZE (h264_data->cached_es));
+    h264_data->last_resync_ts = GST_BUFFER_TIMESTAMP (buf);
+    memcpy (GST_BUFFER_DATA (out_buf), GST_BUFFER_DATA (h264_data->cached_es),
+        GST_BUFFER_SIZE (h264_data->cached_es));
+    out_offset = GST_BUFFER_SIZE (h264_data->cached_es);
+    GST_DEBUG_OBJECT (mux, "prepending SPS/PPS information to that packet");
+  } else {
+    out_buf = gst_buffer_new_and_alloc (GST_BUFFER_SIZE (buf) * 2);
+  }
+
   /* We want the same metadata */
   gst_buffer_copy_metadata (out_buf, buf, GST_BUFFER_COPY_ALL);
 
-  /* Get NAL length size */
-  nal_length_size =
-      (GST_READ_UINT8 (GST_BUFFER_DATA (data->codec_data) + offset) & 0x03) + 1;
-  GST_LOG_OBJECT (mux, "NAL length will be coded on %u bytes", nal_length_size);
-  offset++;
-
-  /* Generate SPS */
-  nb_sps = GST_READ_UINT8 (GST_BUFFER_DATA (data->codec_data) + offset) & 0x1f;
-  GST_DEBUG_OBJECT (mux, "we have %d Sequence Parameter Set", nb_sps);
-  offset++;
-
-  /* For each SPS */
-  for (i = 0; i < nb_sps; i++) {
-    guint16 sps_size =
-        GST_READ_UINT16_BE (GST_BUFFER_DATA (data->codec_data) + offset);
-
-    GST_LOG_OBJECT (mux, "Sequence Parameter Set is %d bytes", sps_size);
-
-    /* Jump over SPS size */
-    offset += 2;
-
-    /* Fake a start code */
-    memcpy (GST_BUFFER_DATA (out_buf) + out_offset, startcode, 4);
-    out_offset += 4;
-    /* Now push the SPS */
-    memcpy (GST_BUFFER_DATA (out_buf) + out_offset,
-        GST_BUFFER_DATA (data->codec_data) + offset, sps_size);
-
-    out_offset += sps_size;
-    offset += sps_size;
-  }
-
-  nb_pps = GST_READ_UINT8 (GST_BUFFER_DATA (data->codec_data) + offset);
-  GST_LOG_OBJECT (mux, "we have %d Picture Parameter Set", nb_sps);
-  offset++;
-
-  /* For each PPS */
-  for (i = 0; i < nb_pps; i++) {
-    gint pps_size =
-        GST_READ_UINT16_BE (GST_BUFFER_DATA (data->codec_data) + offset);
-
-    GST_LOG_OBJECT (mux, "Picture Parameter Set is %d bytes", pps_size);
-
-    /* Jump over PPS size */
-    offset += 2;
-
-    /* Fake a start code */
-    memcpy (GST_BUFFER_DATA (out_buf) + out_offset, startcode, 4);
-    out_offset += 4;
-    /* Now push the PPS */
-    memcpy (GST_BUFFER_DATA (out_buf) + out_offset,
-        GST_BUFFER_DATA (data->codec_data) + offset, pps_size);
-
-    out_offset += pps_size;
-    offset += pps_size;
-  }
-
   while (in_offset < GST_BUFFER_SIZE (buf) &&
       out_offset < GST_BUFFER_SIZE (out_buf) - 4) {
     guint32 nal_size = 0;
 
-    switch (nal_length_size) {
+    switch (h264_data->nal_length_size) {
       case 1:
         nal_size = GST_READ_UINT8 (GST_BUFFER_DATA (buf) + in_offset);
         break;
@@ -177,9 +265,9 @@ mpegtsmux_prepare_h264 (GstBuffer * buf, MpegTsPadData * data, MpegTsMux * mux)
         break;
       default:
         GST_WARNING_OBJECT (mux, "unsupported NAL length size %u",
-            nal_length_size);
+            h264_data->nal_length_size);
     }
-    in_offset += nal_length_size;
+    in_offset += h264_data->nal_length_size;
 
     /* Generate an Elementary stream buffer by inserting a startcode */
     memcpy (GST_BUFFER_DATA (out_buf) + out_offset, startcode, 4);

gst/mpegtsmux/mpegtsmux_h264.h

@@ -1,5 +1,5 @@
 /* 
- * Copyright 2006, 2007, 2008 Fluendo S.A. 
+ * Copyright 2006, 2007, 2008, 2009, 2010 Fluendo S.A. 
  *  Authors: Jan Schmidt <jan@fluendo.com>
  *           Kapil Agrawal <kapil@fluendo.com>
  *           Julien Moutte <julien@fluendo.com>
@@ -88,4 +88,6 @@
 GstBuffer * mpegtsmux_prepare_h264 (GstBuffer * buf, MpegTsPadData * data,
     MpegTsMux * mux);
  
+void mpegtsmux_free_h264 (gpointer prepare_data);
+
 #endif /* __MPEGTSMUX_H264_H__ */