From 8bcdd0b10ac32c886de64a26c8e30e1b21627f75 Mon Sep 17 00:00:00 2001
From: Thulinma <jaron@vietors.com>
Date: Fri, 17 Mar 2023 18:12:33 +0100
Subject: [PATCH] Fix non-4-byte H264 NAL unit lengths in TS and raw H264
 outputs

---
 src/output/output_h264.cpp    | 32 ++++++++++++++++++++++++++------
 src/output/output_ts_base.cpp | 18 +++++++++++++-----
 2 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/src/output/output_h264.cpp b/src/output/output_h264.cpp
index 5520f7a2..174fd56b 100644
--- a/src/output/output_h264.cpp
+++ b/src/output/output_h264.cpp
@@ -414,13 +414,23 @@ namespace Mist{
         Bit::htobs(webBuf+10, 0);
       }
 
+      size_t lenSize = 4;
+      if (M.getCodec(thisIdx) == "H264"){lenSize = (M.getInit(thisIdx)[4] & 3) + 1;}
       unsigned int i = 0;
+      uint32_t ThisNaluSize;
       while (i + 4 < len){
-        uint32_t ThisNaluSize = Bit::btohl(dataPointer + i);
+        if (lenSize == 4){
+          ThisNaluSize = Bit::btohl(dataPointer + i);
+        }else if (lenSize == 2){
+          ThisNaluSize = Bit::btohs(dataPointer + i);
+        }else{
+          ThisNaluSize = dataPointer[i];
+        }
         webBuf.append("\000\000\000\001", 4);
-        webBuf.append(dataPointer + i + 4, ThisNaluSize);
-        i += ThisNaluSize + 4;
+        webBuf.append(dataPointer + i + lenSize, ThisNaluSize);
+        i += ThisNaluSize + lenSize;
       }
+
       webSock->sendFrame(webBuf, webBuf.size(), 2);
 
       if (stayLive && thisPacket.getFlag("keyframe")){liveSeek();}
@@ -428,13 +438,23 @@ namespace Mist{
       return;
     }
 
+    size_t lenSize = 4;
+    if (M.getCodec(thisIdx) == "H264"){lenSize = (M.getInit(thisIdx)[4] & 3) + 1;}
     unsigned int i = 0;
+    uint32_t ThisNaluSize;
     while (i + 4 < len){
-      uint32_t ThisNaluSize = Bit::btohl(dataPointer + i);
+      if (lenSize == 4){
+        ThisNaluSize = Bit::btohl(dataPointer + i);
+      }else if (lenSize == 2){
+        ThisNaluSize = Bit::btohs(dataPointer + i);
+      }else{
+        ThisNaluSize = dataPointer[i];
+      }
       H.Chunkify("\000\000\000\001", 4, myConn);
-      H.Chunkify(dataPointer + i + 4, ThisNaluSize, myConn);
-      i += ThisNaluSize + 4;
+      H.Chunkify(dataPointer + i + lenSize, ThisNaluSize, myConn);
+      i += ThisNaluSize + lenSize;
     }
+
   }
 
   void OutH264::sendHeader(){
diff --git a/src/output/output_ts_base.cpp b/src/output/output_ts_base.cpp
index ac6dd218..695eab24 100644
--- a/src/output/output_ts_base.cpp
+++ b/src/output/output_ts_base.cpp
@@ -164,16 +164,24 @@ namespace Mist{
           }
           /*LTS-END*/
         }
-        while (i + 4 < (unsigned int)dataLen){
-          ThisNaluSize = Bit::btohl(dataPointer + i);
-          if (ThisNaluSize + i + 4 > dataLen){
+        size_t lenSize = 4;
+        if (codec == "H264"){lenSize = (M.getInit(thisIdx)[4] & 3) + 1;}
+        while (i + lenSize < (unsigned int)dataLen){
+          if (lenSize == 4){
+            ThisNaluSize = Bit::btohl(dataPointer + i);
+          }else if (lenSize == 2){
+            ThisNaluSize = Bit::btohs(dataPointer + i);
+          }else{
+            ThisNaluSize = dataPointer[i];
+          }
+          if (ThisNaluSize + i + lenSize > dataLen){
             WARN_MSG("Too big NALU detected (%" PRIu32 " > %zu) - skipping!",
                      ThisNaluSize + i + 4, dataLen);
             break;
           }
           fillPacket("\000\000\000\001", 4, firstPack, video, keyframe, pkgPid, contPkg);
-          fillPacket(dataPointer + i + 4, ThisNaluSize, firstPack, video, keyframe, pkgPid, contPkg);
-          i += ThisNaluSize + 4;
+          fillPacket(dataPointer + i + lenSize, ThisNaluSize, firstPack, video, keyframe, pkgPid, contPkg);
+          i += ThisNaluSize + lenSize;
         }
       }else{
         uint64_t offset = thisPacket.getInt("offset") * 90;