From e57cecdcf87303eb4777e45f80816c3f9addacbc Mon Sep 17 00:00:00 2001
From: Daniel Riehm <daniel.riehm@kitware.com>
Date: Mon, 12 Jun 2023 10:59:17 -0500
Subject: [PATCH] Enforce dts validity when writing

---
 arrows/ffmpeg/ffmpeg_video_output.cxx | 28 +++++++++++++++++++++++++--
 doc/release-notes/master.txt          |  4 ++++
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/arrows/ffmpeg/ffmpeg_video_output.cxx b/arrows/ffmpeg/ffmpeg_video_output.cxx
index 10763377f7..a1d6ec240e 100644
--- a/arrows/ffmpeg/ffmpeg_video_output.cxx
+++ b/arrows/ffmpeg/ffmpeg_video_output.cxx
@@ -70,6 +70,7 @@ class ffmpeg_video_output::impl
     codec_context_uptr codec_context;
     AVCodec const* codec;
     sws_context_uptr image_conversion_context;
+    int64_t prev_video_dts;
   };
 
   impl();
@@ -413,7 +414,8 @@ ::open_video_state(
     metadata_stream{ nullptr },
     codec_context{ nullptr },
     codec{ nullptr },
-    image_conversion_context{ nullptr }
+    image_conversion_context{ nullptr },
+    prev_video_dts{ AV_NOPTS_VALUE }
 {
   // Allocate output format context
   {
@@ -740,8 +742,30 @@ ::add_image( kv::video_raw_image const& image )
     dynamic_cast< ffmpeg_video_raw_image const& >( image );
   for( auto const& packet : ffmpeg_image.packets )
   {
+    // Ensure this packet has sensible timestamps or FFmpeg will complain
+    if( packet->pts == AV_NOPTS_VALUE || packet->dts == AV_NOPTS_VALUE ||
+        packet->dts <= prev_video_dts || packet->dts > packet->pts )
+    {
+      LOG_ERROR(
+        parent->logger,
+        "Dropping video packet with invalid dts/pts "
+        << packet->dts << "/" << packet->pts << " "
+        << "with prev dts " << prev_video_dts );
+      continue;
+    }
+
+    // Copy the packet so we can switch the video stream index
+    packet_uptr tmp_packet{
+      throw_error_null(
+        av_packet_clone( packet.get() ), "Could not copy video packet" ) };
+    tmp_packet->stream_index = video_stream->index;
+
+    // Record this DTS for next time
+    prev_video_dts = packet->dts;
+
+    // Write the packet
     throw_error_code(
-      av_interleaved_write_frame( format_context.get(), packet.get() ),
+      av_interleaved_write_frame( format_context.get(), tmp_packet.get() ),
       "Could not write frame to file" );
   }
   ++frame_count;
diff --git a/doc/release-notes/master.txt b/doc/release-notes/master.txt
index 473387e8be..ff910ab11c 100644
--- a/doc/release-notes/master.txt
+++ b/doc/release-notes/master.txt
@@ -13,6 +13,10 @@ Arrows: Core
 
 * Made the transcode applet's failure to open a video result in a more graceful exit.
 
+Arrows: FFmpeg
+
+* Added check for incoming raw video packets' timestamps and stream indices.
+
 Arrows: KLV
 
 * Ensured that NaN comparisons happen consistently across all data structures.