From 2ac3bdf87068f318dbbe1a99f0b87afab44f7fe2 Mon Sep 17 00:00:00 2001
From: arch1t3cht <arch1t3cht@gmail.com>
Date: Wed, 25 Oct 2023 20:20:04 +0200
Subject: [PATCH] Allow reading compressed subtitles from mkv

The error handling inside of read_subtitles is a bit clumsy here,
but it's probably the best we can get without bigger refactors.
---
 src/mkv_wrap.cpp | 56 ++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 49 insertions(+), 7 deletions(-)

diff --git a/src/mkv_wrap.cpp b/src/mkv_wrap.cpp
index 2564d902c..672272a81 100644
--- a/src/mkv_wrap.cpp
+++ b/src/mkv_wrap.cpp
@@ -114,19 +114,49 @@ struct MkvStdIO final : InputStream {
 	}
 };
 
-static void read_subtitles(agi::ProgressSink *ps, MatroskaFile *file, MkvStdIO *input, bool srt, double totalTime, AssParser *parser) {
+static bool read_subtitles(agi::ProgressSink *ps, MatroskaFile *file, MkvStdIO *input, bool srt, double totalTime, AssParser *parser, CompressedStream *cs) {
 	std::vector<std::pair<int, std::string>> subList;
 
 	// Load blocks
 	uint64_t startTime, endTime, filePos;
 	unsigned int rt, frameSize, frameFlags;
 
+	std::vector<char> uncompBuf(cs ? 256 : 0);
+
 	while (mkv_ReadFrame(file, 0, &rt, &startTime, &endTime, &filePos, &frameSize, &frameFlags) == 0) {
-		if (ps->IsCancelled()) return;
+		if (ps->IsCancelled()) return true;
 		if (frameSize == 0) continue;
 
-		const auto readBuf = input->file.read(filePos, frameSize);
-		const auto readBufEnd = readBuf + frameSize;
+		const char *readBuf;
+		const char *readBufEnd;
+
+		if (cs) {
+			cs_NextFrame(cs, filePos, frameSize);
+			int bytesRead = 0;
+
+			int res;
+			do {
+				res = cs_ReadData(cs, &uncompBuf[bytesRead], uncompBuf.size() - bytesRead);
+				if (res == -1) {
+					const char *err = cs_GetLastError(cs);
+					if (!err) err = "Unknown error";
+					ps->Log("Failed to decompress subtitles: " + std::string(err));
+					ps->SetStayOpen(true);
+					return false;
+				}
+
+				bytesRead += res;
+
+				if (bytesRead >= uncompBuf.size())
+					uncompBuf.resize(2 * uncompBuf.size());
+			} while (res != 0);
+
+			readBuf = &uncompBuf[0];
+			readBufEnd = readBuf + bytesRead;
+		} else {
+			readBuf = input->file.read(filePos, frameSize);
+			readBufEnd = readBuf + frameSize;
+		}
 
 		// Get start and end times
 		int64_t timecodeScaleLow = 1000000;
@@ -170,6 +200,7 @@ static void read_subtitles(agi::ProgressSink *ps, MatroskaFile *file, MkvStdIO *
 	sort(begin(subList), end(subList));
 	for (auto order_value_pair : subList)
 		parser->AddLine(order_value_pair.second);
+	return true;
 }
 
 void MatroskaWrapper::GetSubtitles(agi::fs::path const& filename, AssFile *target) {
@@ -186,7 +217,7 @@ void MatroskaWrapper::GetSubtitles(agi::fs::path const& filename, AssFile *targe
 	// Find tracks
 	for (auto track : boost::irange(0u, tracks)) {
 		auto trackInfo = mkv_GetTrackInfo(file, track);
-		if (trackInfo->Type != 0x11 || trackInfo->CompEnabled) continue;
+		if (trackInfo->Type != 0x11) continue;
 
 		// Known subtitle format
 		std::string CodecID(trackInfo->CodecID);
@@ -242,6 +273,13 @@ void MatroskaWrapper::GetSubtitles(agi::fs::path const& filename, AssFile *targe
 
 	parser.AddLine("[Events]");
 
+	agi::scoped_holder<CompressedStream *, decltype(&cs_Destroy)> cs(nullptr, cs_Destroy);
+	if (trackInfo->CompEnabled) {
+		cs = cs_Create(file, trackToRead, err, sizeof(err));
+		if (!cs)
+			throw MatroskaException(err);
+	}
+
 	// Read timecode scale
 	auto segInfo = mkv_GetFileInfo(file);
 	int64_t timecodeScale = mkv_TruncFloat(trackInfo->TimecodeScale) * segInfo->TimecodeScale;
@@ -249,7 +287,11 @@ void MatroskaWrapper::GetSubtitles(agi::fs::path const& filename, AssFile *targe
 	// Progress bar
 	auto totalTime = double(segInfo->Duration) / timecodeScale;
 	DialogProgress progress(nullptr, _("Parsing Matroska"), _("Reading subtitles from Matroska file."));
-	progress.Run([&](agi::ProgressSink *ps) { read_subtitles(ps, file, &input, srt, totalTime, &parser); });
+	bool result;
+	progress.Run([&](agi::ProgressSink *ps) { result = read_subtitles(ps, file, &input, srt, totalTime, &parser, cs); });
+
+	if (!result)
+		throw MatroskaException("Failed to read subtitles");
 }
 
 bool MatroskaWrapper::HasSubtitles(agi::fs::path const& filename) {
@@ -264,7 +306,7 @@ bool MatroskaWrapper::HasSubtitles(agi::fs::path const& filename) {
 		for (auto track : boost::irange(0u, tracks)) {
 			auto trackInfo = mkv_GetTrackInfo(file, track);
 
-			if (trackInfo->Type == 0x11 && !trackInfo->CompEnabled) {
+			if (trackInfo->Type == 0x11) {
 				std::string CodecID(trackInfo->CodecID);
 				if (CodecID == "S_TEXT/SSA" || CodecID == "S_TEXT/ASS" || CodecID == "S_TEXT/UTF8")
 					return true;