Create a new version of paste:

From 37f1a07904d164ada5dd0af82cca3742e916bd82 Mon Sep 17 00:00:00 2001
From: Nicolas Gaullier <nicolas.gaullier@cji.paris>
Date: Fri, 2 Feb 2024 12:34:39 +0100
Subject: [PATCH] avformat/mpegts: fix first NAL start code splited in two
 different packets

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.
---
 libavformat/mpegts.c | 45 ++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 43 insertions(+), 2 deletions(-)

diff --git a/libavformat/mpegts.c b/libavformat/mpegts.c
index 836d5dd5a3..8f3bdfc690 100644
--- a/libavformat/mpegts.c
+++ b/libavformat/mpegts.c
@@ -1148,6 +1148,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)
@@ -1155,6 +1156,39 @@ 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)
+                    {
+                        if (pes->PES_packet_length)
+                            pes->PES_packet_length += 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;
@@ -1206,6 +1240,8 @@ static int mpegts_push_data(MpegTSFilter *filter,
                     }
 
                     pes->PES_packet_length = AV_RB16(pes->header + 4);
+                    if (pes->PES_packet_length)
+                        pes->PES_packet_length -= pes_align_shift;
                     /* NOTE: zero length means the PES size is unbounded */
 
                     if (pes->stream_id != STREAM_ID_PROGRAM_STREAM_MAP &&
@@ -1300,6 +1336,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);
@@ -1408,9 +1445,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