From 6c18feb4edb21d7285d6e1f95e09f2e397b52bac Mon Sep 17 00:00:00 2001
From: uvjustin <46082645+uvjustin@users.noreply.github.com>
Date: Mon, 12 Oct 2020 00:52:51 +0800
Subject: [PATCH] Allow only one packet with no dts in stream init (#41129)

* Allow only one packet with no dts in stream init

* Call finalize_stream() on exit
---
 homeassistant/components/stream/worker.py | 24 ++++++++++++++++-------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/homeassistant/components/stream/worker.py b/homeassistant/components/stream/worker.py
index 22f67432a1b..9e2d2221602 100644
--- a/homeassistant/components/stream/worker.py
+++ b/homeassistant/components/stream/worker.py
@@ -103,6 +103,7 @@ def _stream_worker_internal(hass, stream, quit_event):
 
     def peek_first_pts():
         nonlocal first_pts, audio_stream
+        missing_dts = False
 
         def empty_stream_dict():
             return {
@@ -117,10 +118,13 @@ def _stream_worker_internal(hass, stream, quit_event):
             while first_packet[video_stream] is None:
                 packet = next(container.demux())
                 if (
-                    packet.stream == video_stream
-                    and packet.is_keyframe
-                    and packet.dts is not None
-                ):
+                    packet.dts is None
+                ):  # Allow single packet with no dts, raise error on second
+                    if missing_dts:
+                        raise av.AVError
+                    missing_dts = True
+                    continue
+                if packet.stream == video_stream and packet.is_keyframe:
                     first_packet[video_stream] = packet
                     initial_packets.append(packet)
             # Get first_pts from subsequent frame to first keyframe
@@ -128,8 +132,13 @@ def _stream_worker_internal(hass, stream, quit_event):
                 [pts is None for pts in {**first_packet, **first_pts}.values()]
             ) and (len(initial_packets) < PACKETS_TO_WAIT_FOR_AUDIO):
                 packet = next(container.demux((video_stream, audio_stream)))
-                if packet.dts is None:
-                    continue  # Discard packets with no dts
+                if (
+                    packet.dts is None
+                ):  # Allow single packet with no dts, raise error on second
+                    if missing_dts:
+                        raise av.AVError
+                    missing_dts = True
+                    continue
                 if (
                     first_packet[packet.stream] is None
                 ):  # actually video already found above so only for audio
@@ -158,6 +167,7 @@ def _stream_worker_internal(hass, stream, quit_event):
             _LOGGER.error(
                 "Error demuxing stream while finding first packet: %s", str(ex)
             )
+            finalize_stream()
             return False
         return True
 
@@ -203,7 +213,7 @@ def _stream_worker_internal(hass, stream, quit_event):
     def finalize_stream():
         if not stream.keepalive:
             # End of stream, clear listeners and stop thread
-            for fmt, _ in outputs.items():
+            for fmt in stream.outputs.keys():
                 hass.loop.call_soon_threadsafe(stream.outputs[fmt].put, None)
 
     if not peek_first_pts():