From 849cb1519e3f3c6dfbd3d1cdf31c93fa41c451fb Mon Sep 17 00:00:00 2001
From: Moritz Bunkus <moritz@bunkus.org>
Date: Sun, 25 Dec 2011 12:12:45 +0100
Subject: [PATCH] Don't buffer whole MPEG PS files if no packets are found for
 a track

---
 ChangeLog               |  8 ++++++++
 src/input/r_mpeg_ps.cpp | 13 +++++++++++--
 src/input/r_mpeg_ps.h   |  3 ++-
 3 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/ChangeLog b/ChangeLog
index f7957584a..bd10743a3 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,3 +1,11 @@
+2011-12-25  Moritz Bunkus  <moritz@bunkus.org>
+
+	* mkvmerge: bug fix: MPEG program streams in which a track
+	suddenly ends and others continue or in which a track has huge
+	gaps will no longer cause mkvmerge to try to read the whole file
+	at once. This could lead to excessive swapping and finally
+	mkvmerge crashing if no more memory was available.
+
 2011-12-24  Moritz Bunkus  <moritz@bunkus.org>
 
 	* mkvextract: bug fix: The track IDs used for extraction are
diff --git a/src/input/r_mpeg_ps.cpp b/src/input/r_mpeg_ps.cpp
index faa85ade6..8f7e5fab9 100644
--- a/src/input/r_mpeg_ps.cpp
+++ b/src/input/r_mpeg_ps.cpp
@@ -1204,6 +1204,8 @@ mpeg_ps_reader_c::create_packetizer(int64_t id) {
 
   if (-1 != track->timecode_offset)
     PTZR(track->ptzr)->m_ti.m_tcsync.displacement += track->timecode_offset;
+
+  m_ptzr_to_track_map[ PTZR(track->ptzr) ] = track;
 }
 
 void
@@ -1220,8 +1222,8 @@ mpeg_ps_reader_c::add_available_track_ids() {
 }
 
 file_status_e
-mpeg_ps_reader_c::read(generic_packetizer_c *,
-                       bool) {
+mpeg_ps_reader_c::read(generic_packetizer_c *requested_ptzr,
+                       bool force) {
   int64_t timecode, packet_pos;
   unsigned int length, full_length;
   unsigned char *buf;
@@ -1229,6 +1231,13 @@ mpeg_ps_reader_c::read(generic_packetizer_c *,
   if (file_done)
     return flush_packetizers();
 
+  int64_t num_queued_bytes = get_queued_bytes();
+  if (!force && (20 * 1024 * 1024 < num_queued_bytes)) {
+    mpeg_ps_track_ptr requested_ptzr_track = m_ptzr_to_track_map[requested_ptzr];
+    if (!requested_ptzr_track || (('a' != requested_ptzr_track->type) && ('v' != requested_ptzr_track->type)) || (64 * 1024 * 1024 < num_queued_bytes))
+      return FILE_STATUS_HOLDING;
+  }
+
   try {
     mpeg_ps_id_t new_id;
     while (find_next_packet(new_id)) {
diff --git a/src/input/r_mpeg_ps.h b/src/input/r_mpeg_ps.h
index 0cfc97279..5e2a17d6f 100644
--- a/src/input/r_mpeg_ps.h
+++ b/src/input/r_mpeg_ps.h
@@ -130,6 +130,7 @@ private:
   bool file_done;
 
   std::vector<mpeg_ps_track_ptr> tracks;
+  std::map<generic_packetizer_c *, mpeg_ps_track_ptr> m_ptzr_to_track_map;
 
 public:
   mpeg_ps_reader_c(const track_info_c &ti, const mm_io_cptr &in);
@@ -140,7 +141,7 @@ public:
   }
 
   virtual void read_headers();
-  virtual file_status_e read(generic_packetizer_c *ptzr, bool force = false);
+  virtual file_status_e read(generic_packetizer_c *requested_ptzr, bool force = false);
   virtual void identify();
   virtual void create_packetizer(int64_t id);
   virtual void create_packetizers();