AVAsset+Attachment.swift

//
// Copyright 2024 Signal Messenger, LLC
// SPDX-License-Identifier: AGPL-3.0-only
//

public import AVFoundation
import Foundation

extension AVAsset {

    public static func from(
        _ attachment: AttachmentStream,
    ) throws -> AVAsset {
        return try .fromEncryptedFile(
            at: attachment.fileURL,
            attachmentKey: AttachmentKey(combinedKey: attachment.attachment.encryptionKey),
            plaintextLength: attachment.info.unencryptedByteCount,
            mimeType: attachment.mimeType,
        )
    }

    static func fromEncryptedFile(
        at fileURL: URL,
        attachmentKey: AttachmentKey,
        plaintextLength: UInt32,
        mimeType: String,
    ) throws -> AVAsset {
        func createAsset(mimeTypeOverride: String? = nil) throws -> AVAsset {
            return try AVAsset._fromEncryptedFile(
                at: fileURL,
                attachmentKey: attachmentKey,
                plaintextLength: plaintextLength,
                mimeType: mimeTypeOverride ?? mimeType,
            )
        }

        guard let mimeTypeOverride = MimeTypeUtil.alternativeAudioMimeType(mimeType: mimeType) else {
            // If we have no override just return the first thing we get.
            return try createAsset()
        }

        if let asset = try? createAsset(), asset.isReadable {
            return asset
        }

        // Give it a second try with the overriden mimeType
        return try createAsset(mimeTypeOverride: mimeTypeOverride)
    }

    private static let videoDecryptionQueue = DispatchQueue(label: "Decrypt AVAsset")

    private static func _fromEncryptedFile(
        at fileURL: URL,
        attachmentKey: AttachmentKey,
        plaintextLength: UInt32,
        mimeType: String,
    ) throws -> AVAsset {
        let fileHandle = try Cryptography.encryptedAttachmentFileHandle(
            at: fileURL,
            plaintextLength: UInt64(safeCast: plaintextLength),
            attachmentKey: attachmentKey,
        )

        guard let utiType = MimeTypeUtil.utiTypeForMimeType(mimeType) else {
            throw OWSAssertionError("Invalid mime type")
        }

        let resourceLoader = EncryptedFileResourceLoader(
            utiType: utiType,
            fileHandle: fileHandle,
        )

        // Prioritize audio extensions; note these mappings differ from the generic
        // "fileExtensionForMimeType" because reasons.
        let mimeTypeExtensionOverride: [String: String] = [
            "audio/3gpp": "3gp",
            "audio/3gpp2": "3g2",
            "audio/aac": "m4a",
            "audio/mp3": "mp3",
            "audio/mp4": "mp4",
            "audio/x-m4a": "m4a",
            "audio/x-m4b": "m4b",
            "audio/x-m4p": "m4p",
            "audio/x-mp3": "mp3",
            "audio/x-mpeg": "mp3",
            "audio/x-mpeg3": "mp3",
        ]

        // AVAsset cares about the file extension. It shouldn't, but it does. If we
        // can map the mime type to a file extension, do so for the url we give the
        // AVAsset so it reads things correctly.
        let pathExtension = mimeTypeExtensionOverride[mimeType] ?? MimeTypeUtil.fileExtensionForMimeType(mimeType)
        let fileURLWithFakeExtension: URL
        if let pathExtension {
            fileURLWithFakeExtension = fileURL.appendingPathExtension(pathExtension)
        } else {
            fileURLWithFakeExtension = fileURL
        }
        guard let redirectURL = fileURLWithFakeExtension.convertToAVAssetRedirectURL(prefix: Self.customScheme) else {
            throw OWSAssertionError("Failed to prefix URL!")
        }
        let asset = AVURLAsset(url: redirectURL)
        asset.resourceLoader.preloadsEligibleContentKeys = true
        asset.resourceLoader.setDelegate(resourceLoader, queue: Self.videoDecryptionQueue)
        // The resource loader delegate is held via weak reference, but:
        // 1. it doesn't hold a reference to the AVAsset
        // 2. we dont want to impose on the caller to hold a strong reference to it
        // so we create a strong reference from the asset.
        ObjectRetainer.retainObject(resourceLoader, forLifetimeOf: asset)

        return asset
    }

    /// In order to get AVAsset to use the custom resource loader, we have to give it a URL scheme it doesn't
    /// understand how to load by itself. To do that, we prefix the url scheme with this string before handing
    /// it to AVAsset, and then strip the prefix in our own code.
    private static let customScheme = "signal"

    private class EncryptedFileResourceLoader: NSObject, AVAssetResourceLoaderDelegate {

        private let utiType: String
        private let fileHandle: EncryptedFileHandle

        init(utiType: String, fileHandle: EncryptedFileHandle) {
            self.utiType = utiType
            self.fileHandle = fileHandle
        }

        func resourceLoader(
            _ resourceLoader: AVAssetResourceLoader,
            shouldWaitForLoadingOfRequestedResource loadingRequest: AVAssetResourceLoadingRequest,
        ) -> Bool {
            if let _ = loadingRequest.contentInformationRequest {
                return handleContentInfoRequest(for: loadingRequest)
            } else if let _ = loadingRequest.dataRequest {
                return handleDataRequest(for: loadingRequest)
            } else {
                return false
            }
        }

        private func handleContentInfoRequest(for loadingRequest: AVAssetResourceLoadingRequest) -> Bool {
            guard let infoRequest = loadingRequest.contentInformationRequest else { return false }

            infoRequest.contentType = utiType
            infoRequest.contentLength = Int64(exactly: fileHandle.plaintextLength) ?? 0
            if #available(iOSApplicationExtension 16.0, *) {
                infoRequest.isEntireLengthAvailableOnDemand = true
            }
            infoRequest.isByteRangeAccessSupported = true
            loadingRequest.finishLoading()
            return true
        }

        private static let chunkSize = 4096

        private func handleDataRequest(for loadingRequest: AVAssetResourceLoadingRequest) -> Bool {
            guard
                let dataRequest = loadingRequest.dataRequest
            else {
                return false
            }

            let requestedOffset = UInt64(dataRequest.requestedOffset)
            var requestedLength = dataRequest.requestedLength
            if dataRequest.requestsAllDataToEndOfResource {
                requestedLength = Int(fileHandle.plaintextLength - requestedOffset)
            }

            do {
                if requestedOffset != fileHandle.offset() {
                    try fileHandle.seek(toOffset: requestedOffset)
                }
            } catch let error {
                loadingRequest.finishLoading(with: error)
                return true
            }

            var bytesReadSoFar = 0
            do {
                while bytesReadSoFar < requestedLength {
                    let lengthToRead = min(Self.chunkSize, requestedLength - bytesReadSoFar)
                    let data = try fileHandle.read(upToCount: lengthToRead)
                    bytesReadSoFar += data.count
                    dataRequest.respond(with: data)
                }
            } catch let error {
                loadingRequest.finishLoading(with: error)
                return true
            }

            loadingRequest.finishLoading()

            return true
        }
    }
}

private extension URL {
    func convertToAVAssetRedirectURL(prefix: String) -> URL? {
        guard
            var components = URLComponents(
                url: self,
                resolvingAgainstBaseURL: false,
            ),
            let scheme = components.scheme
        else {
            return nil
        }
        components.scheme = prefix + scheme
        return components.url
    }
}