Refactor CMFormatDescription.configurationBox.

shogo4405 · May 4, 2024 · 9cf530c · 9cf530c
1 parent 1f14f35
commit 9cf530c
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 27 deletions.
diff --git a/Sources/Extension/CMVideoFormatDescription+Extension.swift b/Sources/Extension/CMVideoFormatDescription+Extension.swift
@@ -66,4 +66,43 @@ extension CMVideoFormatDescription {
             return true
         }
     }
+
+    var streamType: ESStreamType {
+        switch mediaSubType {
+        case .hevc:
+            return .h265
+        case .h264:
+            return .h264
+        default:
+            return .unspecific
+        }
+    }
+
+    var configurationBox: Data? {
+        guard let atoms = CMFormatDescriptionGetExtension(self, extensionKey: kCMFormatDescriptionExtension_SampleDescriptionExtensionAtoms) as? NSDictionary else {
+            return nil
+        }
+        switch mediaSubType {
+        case .h264:
+            return atoms["avcC"] as? Data
+        case .hevc:
+            return atoms["hvcC"] as? Data
+        default:
+            return nil
+        }
+    }
+
+    func makeDecodeConfigurtionRecord() -> (any DecoderConfigurationRecord)? {
+        guard let configurationBox else {
+            return nil
+        }
+        switch mediaSubType {
+        case .h264:
+            return AVCDecoderConfigurationRecord(data: configurationBox)
+        case .hevc:
+            return HEVCDecoderConfigurationRecord(data: configurationBox)
+        default:
+            return nil
+        }
+    }
 }
diff --git a/Sources/ISO/AVCDecoderConfigurationRecord.swift b/Sources/ISO/AVCDecoderConfigurationRecord.swift
@@ -10,16 +10,6 @@ protocol DecoderConfigurationRecord {
  - seealso: ISO/IEC 14496-15 2010
  */
 struct AVCDecoderConfigurationRecord: DecoderConfigurationRecord {
-    static func getData(_ formatDescription: CMFormatDescription?) -> Data? {
-        guard let formatDescription else {
-            return nil
-        }
-        if let atoms = CMFormatDescriptionGetExtension(formatDescription, extensionKey: "SampleDescriptionExtensionAtoms" as CFString) as? NSDictionary {
-            return atoms["avcC"] as? Data
-        }
-        return nil
-    }
-
     static let reserveLengthSizeMinusOne: UInt8 = 0x3F
     static let reserveNumOfSequenceParameterSets: UInt8 = 0xE0
     static let reserveChromaFormat: UInt8 = 0xFC

diff --git a/Sources/ISO/HEVCDecoderConfigurationRecord.swift b/Sources/ISO/HEVCDecoderConfigurationRecord.swift
@@ -3,16 +3,6 @@ import Foundation
 
 /// ISO/IEC 14496-15 8.3.3.1.2
 struct HEVCDecoderConfigurationRecord: DecoderConfigurationRecord {
-    static func getData(_ formatDescription: CMFormatDescription?) -> Data? {
-        guard let formatDescription else {
-            return nil
-        }
-        if let atoms = CMFormatDescriptionGetExtension(formatDescription, extensionKey: "SampleDescriptionExtensionAtoms" as CFString) as? NSDictionary {
-            return atoms["hvcC"] as? Data
-        }
-        return nil
-    }
-
     var configurationVersion: UInt8 = 1
     var generalProfileSpace: UInt8 = 0
     var generalTierFlag = false

diff --git a/Sources/ISO/TSWriter.swift b/Sources/ISO/TSWriter.swift
@@ -43,15 +43,15 @@ public final class TSWriter<T: TSWriterDelegate> {
         didSet {
             guard
                 let videoFormat,
-                let avcC = AVCDecoderConfigurationRecord.getData(videoFormat) else {
+                let config = videoFormat.makeDecodeConfigurtionRecord() as? AVCDecoderConfigurationRecord else {
                 return
             }
             var data = ESSpecificData()
-            data.streamType = .h264
+            data.streamType = videoFormat.streamType
             data.elementaryPID = kTSWriter_defaultVideoPID
             PMT.elementaryStreamSpecificData.append(data)
             videoContinuityCounter = 0
-            videoConfig = AVCDecoderConfigurationRecord(data: avcC)
+            videoConfig = config
         }
     }
 

diff --git a/Sources/RTMP/RTMPMuxer.swift b/Sources/RTMP/RTMPMuxer.swift
@@ -36,22 +36,22 @@ final class RTMPMuxer {
             case .publishing:
                 switch videoFormat?.mediaSubType {
                 case .h264?:
-                    guard let avcC = AVCDecoderConfigurationRecord.getData(videoFormat) else {
+                    guard let configurationBox = videoFormat?.configurationBox else {
                         return
                     }
                     var buffer = Data([FLVFrameType.key.rawValue << 4 | FLVVideoCodec.avc.rawValue, FLVAVCPacketType.seq.rawValue, 0, 0, 0])
-                    buffer.append(avcC)
+                    buffer.append(configurationBox)
                     stream?.doOutput(
                         .zero,
                         chunkStreamId: FLVTagType.video.streamId,
                         message: RTMPVideoMessage(streamId: 0, timestamp: 0, payload: buffer)
                     )
                 case .hevc?:
-                    guard let hvcC = HEVCDecoderConfigurationRecord.getData(videoFormat) else {
+                    guard let configurationBox = videoFormat?.configurationBox else {
                         return
                     }
                     var buffer = Data([0b10000000 | FLVFrameType.key.rawValue << 4 | FLVVideoPacketType.sequenceStart.rawValue, 0x68, 0x76, 0x63, 0x31])
-                    buffer.append(hvcC)
+                    buffer.append(configurationBox)
                     stream?.doOutput(
                         .zero,
                         chunkStreamId: FLVTagType.video.streamId,