3 files changed, 48 insertions, 28 deletions

diff --git a/gst/mpegtsmux/tsmux/tsmux.c b/gst/mpegtsmux/tsmux/tsmux.c
index a51218615..1a50e7224 100644
--- a/gst/mpegtsmux/tsmux/tsmux.c
+++ b/gst/mpegtsmux/tsmux/tsmux.c
@@ -767,8 +767,10 @@ tsmux_write_stream_packet (TsMux * mux, TsMuxStream * stream)
     }
   }
 
-  pi->stream_avail = tsmux_stream_bytes_avail (stream);
   pi->packet_start_unit_indicator = tsmux_stream_at_pes_start (stream);
+  if (pi->packet_start_unit_indicator)
+    tsmux_stream_initialize_pes_packet (stream);
+  pi->stream_avail = tsmux_stream_bytes_avail (stream);
 
   if (!tsmux_write_ts_header (mux->packet_buf, pi, &payload_len, &payload_offs))
     return FALSE;
diff --git a/gst/mpegtsmux/tsmux/tsmuxstream.c b/gst/mpegtsmux/tsmux/tsmuxstream.c
index 3d848b490..62c244fb2 100644
--- a/gst/mpegtsmux/tsmux/tsmuxstream.c
+++ b/gst/mpegtsmux/tsmux/tsmuxstream.c
@@ -332,6 +332,50 @@ tsmux_stream_bytes_in_buffer (TsMuxStream * stream)
 }
 
 /**
+ * tsmux_stream_initialize_pes_packet:
+ * @stream: a #TsMuxStream
+ *
+ * Initializes the PES packet.
+ *
+ * Returns: TRUE if we the packet was initialized.
+ */
+gboolean
+tsmux_stream_initialize_pes_packet (TsMuxStream * stream)
+{
+  if (stream->state != TSMUX_STREAM_STATE_HEADER)
+    return TRUE;
+
+  if (stream->pes_payload_size != 0) {
+    /* Use prescribed fixed PES payload size */
+    stream->cur_pes_payload_size = stream->pes_payload_size;
+    tsmux_stream_find_pts_dts_within (stream, stream->cur_pes_payload_size,
+        &stream->pts, &stream->dts);
+  } else if (stream->is_video_stream) {
+    /* Unbounded for video streams */
+    stream->cur_pes_payload_size = 0;
+    tsmux_stream_find_pts_dts_within (stream,
+        tsmux_stream_bytes_in_buffer (stream), &stream->pts, &stream->dts);
+  } else {
+    /* Output a PES packet of all currently available bytes otherwise */
+    stream->cur_pes_payload_size = tsmux_stream_bytes_in_buffer (stream);
+    tsmux_stream_find_pts_dts_within (stream, stream->cur_pes_payload_size,
+        &stream->pts, &stream->dts);
+  }
+
+  stream->pi.flags &= ~(TSMUX_PACKET_FLAG_PES_WRITE_PTS_DTS |
+      TSMUX_PACKET_FLAG_PES_WRITE_PTS);
+
+  if (stream->pts != -1 && stream->dts != -1)
+    stream->pi.flags |= TSMUX_PACKET_FLAG_PES_WRITE_PTS_DTS;
+  else {
+    if (stream->pts != -1)
+      stream->pi.flags |= TSMUX_PACKET_FLAG_PES_WRITE_PTS;
+  }
+
+  return TRUE;
+}
+
+/**
  * tsmux_stream_get_data:
  * @stream: a #TsMuxStream
  * @buf: a buffer to hold the result
@@ -350,33 +394,6 @@ tsmux_stream_get_data (TsMuxStream * stream, guint8 * buf, guint len)
   if (stream->state == TSMUX_STREAM_STATE_HEADER) {
     guint8 pes_hdr_length;
 
-    if (stream->pes_payload_size != 0) {
-      /* Use prescribed fixed PES payload size */
-      stream->cur_pes_payload_size = stream->pes_payload_size;
-      tsmux_stream_find_pts_dts_within (stream, stream->cur_pes_payload_size,
-          &stream->pts, &stream->dts);
-    } else if (stream->is_video_stream) {
-      /* Unbounded for video streams */
-      stream->cur_pes_payload_size = 0;
-      tsmux_stream_find_pts_dts_within (stream,
-          tsmux_stream_bytes_in_buffer (stream), &stream->pts, &stream->dts);
-    } else {
-      /* Output a PES packet of all currently available bytes otherwise */
-      stream->cur_pes_payload_size = tsmux_stream_bytes_in_buffer (stream);
-      tsmux_stream_find_pts_dts_within (stream, stream->cur_pes_payload_size,
-          &stream->pts, &stream->dts);
-    }
-
-    stream->pi.flags &= ~(TSMUX_PACKET_FLAG_PES_WRITE_PTS_DTS |
-        TSMUX_PACKET_FLAG_PES_WRITE_PTS);
-
-    if (stream->pts != -1 && stream->dts != -1)
-      stream->pi.flags |= TSMUX_PACKET_FLAG_PES_WRITE_PTS_DTS;
-    else {
-      if (stream->pts != -1)
-        stream->pi.flags |= TSMUX_PACKET_FLAG_PES_WRITE_PTS;
-    }
-
     pes_hdr_length = tsmux_stream_pes_header_length (stream);
 
     /* Submitted buffer must be at least as large as the PES header */
diff --git a/gst/mpegtsmux/tsmux/tsmuxstream.h b/gst/mpegtsmux/tsmux/tsmuxstream.h
index 4ee6ae181..b384d39a2 100644
--- a/gst/mpegtsmux/tsmux/tsmuxstream.h
+++ b/gst/mpegtsmux/tsmux/tsmuxstream.h
@@ -212,6 +212,7 @@ void 		tsmux_stream_get_es_descrs 	(TsMuxStream *stream, guint8 *buf, guint16 *l
 
 gint 		tsmux_stream_bytes_in_buffer 	(TsMuxStream *stream);
 gint 		tsmux_stream_bytes_avail 	(TsMuxStream *stream);
+gboolean 	tsmux_stream_initialize_pes_packet (TsMuxStream *stream);
 gboolean 	tsmux_stream_get_data 		(TsMuxStream *stream, guint8 *buf, guint len);
 
 guint64 	tsmux_stream_get_pts 		(TsMuxStream *stream);