Fix crash with VobSub headers and Directories

Signed-off-by: Ethan Dye <[email protected]>
ecdye · Oct 22, 2024 · b551948 · b551948
1 parent f8e223d
commit b551948
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 5 deletions.
diff --git a/Sources/macSubtitleOCR/FileHandler.swift b/Sources/macSubtitleOCR/FileHandler.swift
@@ -9,17 +9,25 @@
 import Foundation
 
 struct FileHandler {
-    let outputDirectory: String
+    private let outputDirectory: URL
 
     init(outputDirectory: String) {
-        self.outputDirectory = outputDirectory
+        self.outputDirectory = URL(fileURLWithPath: outputDirectory)
+        do {
+            try FileManager.default.createDirectory(
+                at: self.outputDirectory,
+                withIntermediateDirectories: true,
+                attributes: nil)
+        } catch {
+            fatalError("Failed to create output directory: \(error)")
+        }
     }
 
     func saveSRTFile(for result: macSubtitleOCRResult) throws {
         if result.srt.isEmpty {
             return
         }
-        let srtFilePath = URL(fileURLWithPath: outputDirectory).appendingPathComponent("track_\(result.trackNumber).srt")
+        let srtFilePath = outputDirectory.appendingPathComponent("track_\(result.trackNumber).srt")
         let srt = SRT(subtitles: result.srt.sorted { $0.index < $1.index })
         srt.write(toFileAt: srtFilePath)
     }
@@ -44,7 +52,7 @@ struct FileHandler {
 
         let jsonData = try JSONSerialization.data(withJSONObject: jsonResults, options: [.prettyPrinted, .sortedKeys])
         let jsonString = String(data: jsonData, encoding: .utf8) ?? "[]"
-        let jsonFilePath = URL(fileURLWithPath: outputDirectory).appendingPathComponent("track_\(result.trackNumber).json")
+        let jsonFilePath = outputDirectory.appendingPathComponent("track_\(result.trackNumber).json")
         try jsonString.write(to: jsonFilePath, atomically: true, encoding: .utf8)
     }
 }
diff --git a/Sources/macSubtitleOCR/Subtitles/VobSub/VobSubParser.swift b/Sources/macSubtitleOCR/Subtitles/VobSub/VobSubParser.swift
@@ -16,6 +16,7 @@ struct VobSubParser {
     let subtitle: Subtitle
     private let masterPalette: [UInt8]
     private let fps = 24.0 // TODO: Make this configurable / dynamic
+    private let minimumControlHeaderSize = 22
 
     // MARK: - Lifecycle
 
@@ -128,7 +129,10 @@ struct VobSubParser {
 
     private func parseCommandHeader(_ header: Data, offset: Int) {
         let relativeEndTimestamp = TimeInterval(Int(header.getUInt16BE()!)) * 1024 / 90000 / fps
-        let endOfControl = Int(header.getUInt16BE()!) - 4 - offset
+        let endOfControl = max(minimumControlHeaderSize, Int(header.getUInt16BE()!) - 4 - offset)
+        if endOfControl > header.count {
+            logger.warning("Control header is too short, \(header.count) bytes, trying to decode anyway, errors may occur")
+        }
         subtitle.endTimestamp = subtitle.startTimestamp! + relativeEndTimestamp
 
         var index = 2