[FFmpeg-devel] [PATCH] avformat/mpegts: fix first NAL start code splited in two different packets

[Nicolas Gaullier]

When PES are not aligned and a tiny nal-encoded frame is contained in a single TS packet,
the first NAL startcode may be split by the PES boundary, making the packet unworkable.
This patch shift the PES boundaries to avoid this.

Signed-off-by: Nicolas Gaullier <nicolas.gaullier at cji.paris>
---
 libavformat/mpegts.c | 41 +++++++++++++++++++++++++++++++++++++++--
 1 file changed, 39 insertions(+), 2 deletions(-)

diff --git a/libavformat/mpegts.c b/libavformat/mpegts.c
index bef00c88e7..4b4cf82fb9 100644
--- a/libavformat/mpegts.c
+++ b/libavformat/mpegts.c
@@ -1149,6 +1149,7 @@ static int mpegts_push_data(MpegTSFilter *filter,
     PESContext *pes   = filter->u.pes_filter.opaque;
     MpegTSContext *ts = pes->ts;
     const uint8_t *p;
+    int pes_align_shift = 0;
     int ret, len;
 
     if (!ts->pkt)
@@ -1156,6 +1157,37 @@ static int mpegts_push_data(MpegTSFilter *filter,
 
     if (is_start) {
         if (pes->state == MPEGTS_PAYLOAD && pes->data_index > 0) {
+            if ((pes->st->codecpar->codec_id == AV_CODEC_ID_H264
+             || pes->st->codecpar->codec_id == AV_CODEC_ID_HEVC)
+                && pes->data_index < TS_PACKET_SIZE - 4 - PES_HEADER_SIZE
+                && pes->data_index >= 4
+                && buf_size >= 4 ) {
+                    /* check/avoid spliting the start code + first byte of the first NAL unit in two different packets.
+                     * this could happen with a tiny unaligned PES that fits in a single ts packet. */
+                    uint8_t *last_p_end = pes->buffer->data + pes->data_index - 4;
+                    p = buf + PES_HEADER_SIZE + buf[PES_HEADER_SIZE - 1];
+                    if (last_p_end[3] == 0x00 && AV_RB24(p) == 0x000001)
+                        pes_align_shift = 4;
+                    else if (AV_RB16(last_p_end + 2)== 0x0000 && AV_RB16(p) == 0x0001)
+                        pes_align_shift = 3;
+                    else if (AV_RB24(last_p_end + 1)== 0x000000 && *p == 0x01)
+                        pes_align_shift = 2;
+                    else if (AV_RB32(last_p_end) == 0x00000001)
+                        pes_align_shift = 1;
+                    if (pes_align_shift)
+                    {
+                        last_p_end += 4;
+                        if (pes_align_shift > 3)
+                            *last_p_end++ = 0x00;
+                        if (pes_align_shift > 2)
+                            *last_p_end++ = 0x00;
+                        if (pes_align_shift > 1)
+                            *last_p_end++ = 0x01;
+                        *last_p_end = *(p + pes_align_shift - 1);
+                        pes->data_index += pes_align_shift;
+                        buf_size -= pes_align_shift;
+                    }
+                }
             ret = new_pes_packet(pes, ts->pkt);
             if (ret < 0)
                 return ret;
@@ -1301,6 +1333,7 @@ skip:
                 /* we got the full header. We parse it and get the payload */
                 pes->state = MPEGTS_PAYLOAD;
                 pes->data_index = 0;
+                p += pes_align_shift;
                 if (pes->stream_type == 0x12 && buf_size > 0) {
                     int sl_header_bytes = read_sl_header(pes, &pes->sl, p,
                                                          buf_size);
@@ -1409,9 +1442,13 @@ skip:
                 pes->data_index += buf_size;
                 /* emit complete packets with known packet size
                  * decreases demuxer delay for infrequent packets like subtitles from
-                 * a couple of seconds to milliseconds for properly muxed files. */
+                 * a couple of seconds to milliseconds for properly muxed files.
+                 * disabled for video/NALs because at this point it could split/break the first NAL start code.
+                 */
                 if (!ts->stop_parse && pes->PES_packet_length &&
-                    pes->pes_header_size + pes->data_index == pes->PES_packet_length + PES_START_SIZE) {
+                    pes->pes_header_size + pes->data_index == pes->PES_packet_length + PES_START_SIZE &&
+                    pes->st->codecpar->codec_id != AV_CODEC_ID_H264 &&
+                    pes->st->codecpar->codec_id != AV_CODEC_ID_HEVC) {
                     ts->stop_parse = 1;
                     ret = new_pes_packet(pes, ts->pkt);
                     pes->state = MPEGTS_SKIP;
-- 
2.30.2