344 lines
12 KiB
Swift
344 lines
12 KiB
Swift
import Foundation
|
||
import AVFoundation
|
||
import CoreMedia
|
||
|
||
/// “自编码”串流模式解码器：
|
||
/// 从 video_channel DataChannel 接收被控端 MediaCodec 编码的 H.264 裸流，
|
||
/// 重组分片 -> 解析单元 -> 转换为 AVCC 格式 -> 通过 AVSampleBufferDisplayLayer 硬解渲染。
|
||
///
|
||
/// 二进制协议（与被控端 SelfCodecEncoder 保持一致，均为大端序）：
|
||
/// 分片: MAGIC(0xAB,1B) + seq(int32) + total(int16) + idx(int16) + len(int32) + payload
|
||
/// 单元: type(1B: 1=CONFIG 2=FRAME) + pts(uint32,ms) + isKey(1B) + len(int32) + H.264 Annex-B 数据
|
||
final class SelfCodecDecoder {
|
||
|
||
// ---- 必须与被控端 SelfCodecEncoder 保持一致 ----
|
||
private static let magic: UInt8 = 0xAB
|
||
private static let unitTypeConfig: UInt8 = 1
|
||
private static let unitTypeFrame: UInt8 = 2
|
||
|
||
/// 渲染图层（由 SelfCodecDisplayView 提供）
|
||
weak var displayLayer: AVSampleBufferDisplayLayer?
|
||
/// 分辨率变化回调（主线程）
|
||
var onResolutionUpdate: ((Int, Int) -> Void)?
|
||
|
||
private(set) var videoWidth = 0
|
||
private(set) var videoHeight = 0
|
||
|
||
private let queue = DispatchQueue(label: "com.ttstd.selfcodec.decoder")
|
||
private var enabled = false
|
||
|
||
// 分片重组
|
||
private final class ChunkBuffer {
|
||
var total = 0
|
||
var received = 0
|
||
var chunks: [Data?] = []
|
||
}
|
||
private var assembling: [Int32: ChunkBuffer] = [:]
|
||
private var lastSeq: Int32 = -1
|
||
|
||
// H.264 参数集与格式
|
||
private var sps: Data?
|
||
private var pps: Data?
|
||
private var formatDescription: CMVideoFormatDescription?
|
||
|
||
// MARK: - 开关
|
||
|
||
func setEnabled(_ on: Bool) {
|
||
queue.async {
|
||
self.enabled = on
|
||
if !on {
|
||
self.resetLocked()
|
||
}
|
||
}
|
||
}
|
||
|
||
func release() {
|
||
queue.async {
|
||
self.enabled = false
|
||
self.resetLocked()
|
||
}
|
||
}
|
||
|
||
private func resetLocked() {
|
||
assembling.removeAll()
|
||
lastSeq = -1
|
||
sps = nil
|
||
pps = nil
|
||
formatDescription = nil
|
||
videoWidth = 0
|
||
videoHeight = 0
|
||
displayLayer?.flushAndRemoveImage()
|
||
}
|
||
|
||
// MARK: - 入口：DataChannel 二进制消息
|
||
|
||
func onBinaryMessage(_ raw: Data) {
|
||
queue.async {
|
||
self.handleChunk(raw)
|
||
}
|
||
}
|
||
|
||
private func handleChunk(_ raw: Data) {
|
||
guard enabled, raw.count >= 13 else { return }
|
||
var r = BigEndianReader(data: raw)
|
||
guard r.readByte() == Self.magic else { return }
|
||
guard let seq = r.readInt32() else { return }
|
||
|
||
// 过期检查：同一单元的所有分片 seq 相同，只丢弃明显更旧的单元
|
||
if lastSeq != -1 && seq < lastSeq && (lastSeq - seq) < 1000 {
|
||
return
|
||
}
|
||
guard let total = r.readInt16(), let idx = r.readInt16(),
|
||
let len = r.readInt32(), let chunk = r.readData(Int(len)),
|
||
total > 0 else { return }
|
||
|
||
let cb: ChunkBuffer
|
||
if let existing = assembling[seq] {
|
||
cb = existing
|
||
} else {
|
||
cb = ChunkBuffer()
|
||
cb.total = Int(total)
|
||
cb.chunks = Array(repeating: nil, count: Int(total))
|
||
assembling[seq] = cb
|
||
if assembling.count > 24, let minKey = assembling.keys.min() {
|
||
assembling.removeValue(forKey: minKey)
|
||
}
|
||
}
|
||
|
||
let i = Int(idx)
|
||
if i >= 0 && i < cb.total && cb.chunks[i] == nil {
|
||
cb.chunks[i] = chunk
|
||
cb.received += 1
|
||
}
|
||
if cb.received == cb.total {
|
||
assembling.removeValue(forKey: seq)
|
||
if seq > lastSeq { lastSeq = seq }
|
||
var unit = Data()
|
||
for c in cb.chunks where c != nil { unit.append(c!) }
|
||
handleUnit(unit)
|
||
}
|
||
}
|
||
|
||
// MARK: - 单元处理
|
||
|
||
private func handleUnit(_ unit: Data) {
|
||
var r = BigEndianReader(data: unit)
|
||
guard let type = r.readByte(),
|
||
let ptsRaw = r.readInt32(),
|
||
let isKey = r.readByte(),
|
||
let len = r.readInt32(),
|
||
let data = r.readData(Int(len)) else { return }
|
||
let ptsMs = Int64(UInt32(bitPattern: ptsRaw))
|
||
|
||
if type == Self.unitTypeConfig {
|
||
handleConfig(data)
|
||
} else if type == Self.unitTypeFrame {
|
||
handleFrame(data, ptsMs: ptsMs, isKey: isKey != 0)
|
||
}
|
||
}
|
||
|
||
/// CONFIG 单元：MediaCodec 的 csd-0(SPS)/csd-1(PPS)，Annex-B 格式
|
||
private func handleConfig(_ data: Data) {
|
||
for nal in Self.annexBNalUnits(in: data) {
|
||
guard let first = nal.first else { continue }
|
||
switch first & 0x1F {
|
||
case 7: if sps != nal { sps = nal; formatDescription = nil }
|
||
case 8: if pps != nal { pps = nal; formatDescription = nil }
|
||
default: break
|
||
}
|
||
}
|
||
rebuildFormatDescriptionIfNeeded()
|
||
}
|
||
|
||
private func rebuildFormatDescriptionIfNeeded() {
|
||
guard formatDescription == nil, let sps, let pps else { return }
|
||
var desc: CMVideoFormatDescription?
|
||
let status = sps.withUnsafeBytes { spsPtr -> OSStatus in
|
||
pps.withUnsafeBytes { ppsPtr -> OSStatus in
|
||
let paramSets: [UnsafePointer<UInt8>] = [
|
||
spsPtr.bindMemory(to: UInt8.self).baseAddress!,
|
||
ppsPtr.bindMemory(to: UInt8.self).baseAddress!
|
||
]
|
||
let sizes: [Int] = [sps.count, pps.count]
|
||
return CMVideoFormatDescriptionCreateFromH264ParameterSets(
|
||
allocator: kCFAllocatorDefault,
|
||
parameterSetCount: 2,
|
||
parameterSetPointers: paramSets,
|
||
parameterSetSizes: sizes,
|
||
nalUnitHeaderLength: 4,
|
||
formatDescriptionOut: &desc)
|
||
}
|
||
}
|
||
guard status == noErr, let desc else {
|
||
NSLog("[SelfCodec] create format description failed: \(status)")
|
||
return
|
||
}
|
||
formatDescription = desc
|
||
displayLayer?.flush()
|
||
|
||
let dims = CMVideoFormatDescriptionGetDimensions(desc)
|
||
let w = Int(dims.width), h = Int(dims.height)
|
||
if w != videoWidth || h != videoHeight {
|
||
videoWidth = w
|
||
videoHeight = h
|
||
DispatchQueue.main.async {
|
||
self.onResolutionUpdate?(w, h)
|
||
}
|
||
}
|
||
}
|
||
|
||
/// FRAME 单元：Annex-B -> AVCC（4 字节大端长度前缀），封装 CMSampleBuffer 后送显
|
||
private func handleFrame(_ data: Data, ptsMs: Int64, isKey: Bool) {
|
||
// 帧内可能携带 SPS/PPS（部分编码器关键帧前重发参数集）
|
||
var vclData = Data()
|
||
for nal in Self.annexBNalUnits(in: data) {
|
||
guard let first = nal.first else { continue }
|
||
let nalType = first & 0x1F
|
||
if nalType == 7 {
|
||
if sps != nal { sps = nal; formatDescription = nil }
|
||
continue
|
||
}
|
||
if nalType == 8 {
|
||
if pps != nal { pps = nal; formatDescription = nil }
|
||
continue
|
||
}
|
||
var lenBE = UInt32(nal.count).bigEndian
|
||
withUnsafeBytes(of: &lenBE) { vclData.append(contentsOf: $0) }
|
||
vclData.append(nal)
|
||
}
|
||
rebuildFormatDescriptionIfNeeded()
|
||
guard enabled, !vclData.isEmpty, let formatDescription, let layer = displayLayer else { return }
|
||
|
||
// 构造 CMBlockBuffer
|
||
var blockBuffer: CMBlockBuffer?
|
||
var status = CMBlockBufferCreateWithMemoryBlock(
|
||
allocator: kCFAllocatorDefault,
|
||
memoryBlock: nil,
|
||
blockLength: vclData.count,
|
||
blockAllocator: kCFAllocatorDefault,
|
||
customBlockSource: nil,
|
||
offsetToData: 0,
|
||
dataLength: vclData.count,
|
||
flags: 0,
|
||
blockBufferOut: &blockBuffer)
|
||
guard status == kCMBlockBufferNoErr, let blockBuffer else { return }
|
||
status = vclData.withUnsafeBytes { ptr in
|
||
CMBlockBufferReplaceDataBytes(
|
||
with: ptr.baseAddress!,
|
||
blockBuffer: blockBuffer,
|
||
offsetIntoDestination: 0,
|
||
dataLength: vclData.count)
|
||
}
|
||
guard status == kCMBlockBufferNoErr else { return }
|
||
|
||
// 构造 CMSampleBuffer（立即显示，不依赖时间轴）
|
||
var sampleBuffer: CMSampleBuffer?
|
||
var timing = CMSampleTimingInfo(
|
||
duration: .invalid,
|
||
presentationTimeStamp: CMTime(value: ptsMs, timescale: 1000),
|
||
decodeTimeStamp: .invalid)
|
||
var sampleSize = vclData.count
|
||
status = CMSampleBufferCreateReady(
|
||
allocator: kCFAllocatorDefault,
|
||
dataBuffer: blockBuffer,
|
||
formatDescription: formatDescription,
|
||
sampleCount: 1,
|
||
sampleTimingEntryCount: 1,
|
||
sampleTimingArray: &timing,
|
||
sampleSizeEntryCount: 1,
|
||
sampleSizeArray: &sampleSize,
|
||
sampleBufferOut: &sampleBuffer)
|
||
guard status == noErr, let sampleBuffer else { return }
|
||
|
||
if let attachments = CMSampleBufferGetSampleAttachmentsArray(sampleBuffer, createIfNecessary: true) as? [CFMutableDictionary],
|
||
let dict = attachments.first {
|
||
CFDictionarySetValue(
|
||
dict,
|
||
Unmanaged.passUnretained(kCMSampleAttachmentKey_DisplayImmediately).toOpaque(),
|
||
Unmanaged.passUnretained(kCFBooleanTrue).toOpaque())
|
||
if !isKey {
|
||
CFDictionarySetValue(
|
||
dict,
|
||
Unmanaged.passUnretained(kCMSampleAttachmentKey_NotSync).toOpaque(),
|
||
Unmanaged.passUnretained(kCFBooleanTrue).toOpaque())
|
||
}
|
||
}
|
||
|
||
if layer.status == .failed || layer.requiresFlushToResumeDecoding {
|
||
layer.flush()
|
||
}
|
||
layer.enqueue(sampleBuffer)
|
||
}
|
||
|
||
// MARK: - Annex-B 拆分
|
||
|
||
/// 拆分 Annex-B 码流（00 00 01 / 00 00 00 01 起始码）为 NAL 单元数组（不含起始码）
|
||
static func annexBNalUnits(in data: Data) -> [Data] {
|
||
var result: [Data] = []
|
||
let bytes = [UInt8](data)
|
||
let count = bytes.count
|
||
var starts: [Int] = []
|
||
var i = 0
|
||
while i + 2 < count {
|
||
if bytes[i] == 0 && bytes[i + 1] == 0 && bytes[i + 2] == 1 {
|
||
starts.append(i + 3)
|
||
i += 3
|
||
} else {
|
||
i += 1
|
||
}
|
||
}
|
||
guard !starts.isEmpty else { return data.isEmpty ? [] : [data] }
|
||
for (idx, start) in starts.enumerated() {
|
||
var end = count
|
||
if idx + 1 < starts.count {
|
||
end = starts[idx + 1] - 3
|
||
// 兼容 4 字节起始码（00 00 00 01）
|
||
if end > start && bytes[end - 1] == 0 { end -= 1 }
|
||
}
|
||
if end > start {
|
||
result.append(Data(bytes[start..<end]))
|
||
}
|
||
}
|
||
return result
|
||
}
|
||
}
|
||
|
||
// MARK: - 大端序读取器
|
||
|
||
private struct BigEndianReader {
|
||
let data: Data
|
||
private var offset: Int
|
||
|
||
init(data: Data) {
|
||
self.data = data
|
||
self.offset = data.startIndex
|
||
}
|
||
|
||
mutating func readByte() -> UInt8? {
|
||
guard offset < data.endIndex else { return nil }
|
||
defer { offset += 1 }
|
||
return data[offset]
|
||
}
|
||
|
||
mutating func readInt16() -> Int16? {
|
||
guard offset + 2 <= data.endIndex else { return nil }
|
||
let v = (UInt16(data[offset]) << 8) | UInt16(data[offset + 1])
|
||
offset += 2
|
||
return Int16(bitPattern: v)
|
||
}
|
||
|
||
mutating func readInt32() -> Int32? {
|
||
guard offset + 4 <= data.endIndex else { return nil }
|
||
var v: UInt32 = 0
|
||
for i in 0..<4 { v = (v << 8) | UInt32(data[offset + i]) }
|
||
offset += 4
|
||
return Int32(bitPattern: v)
|
||
}
|
||
|
||
mutating func readData(_ length: Int) -> Data? {
|
||
guard length >= 0, offset + length <= data.endIndex else { return nil }
|
||
defer { offset += length }
|
||
return data.subdata(in: offset..<(offset + length))
|
||
}
|
||
}
|