Skip to content

Commit 1294e7e

Browse files
Balamurugan-TestpressBalamurugan-Testpress
committed
Fix AES offline key handling
1 parent b89c85e commit 1294e7e

10 files changed

Lines changed: 151 additions & 149 deletions

File tree

‎Source/Constants/ContentProtectionConstants.swift‎

Lines changed: 0 additions & 24 deletions
This file was deleted.

‎Source/Database/TPStreamsDownloadManager.swift‎

Lines changed: 39 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,7 @@ public final class TPStreamsDownloadManager {
8585

8686
let localOfflineAsset = LocalOfflineAsset.create(
8787
assetId: asset.id,
88+
videoId: asset.video?.id,
8889
srcURL: asset.video!.playbackURL,
8990
title: asset.title,
9091
resolution:videoQuality.resolution,
@@ -102,42 +103,64 @@ public final class TPStreamsDownloadManager {
102103
tpStreamsDownloadDelegate?.onStart(offlineAsset: localOfflineAsset.asOfflineAsset())
103104
tpStreamsDownloadDelegate?.onStateChange(status: .inProgress, offlineAsset: localOfflineAsset.asOfflineAsset())
104105

105-
if (asset.video?.drmEncrypted == true){
106+
if asset.video?.drmEncrypted == true {
106107
M3U8Parser.extractContentID(url: URL(string: asset.video!.playbackURL)!) { result in
107108
switch result {
108109
case .success(let drmContentId):
109-
print("Extracted DRM content ID: \(drmContentId)")
110110
DispatchQueue.main.async {
111111
LocalOfflineAsset.manager.update(id: asset.id, with: ["drmContentId": drmContentId])
112112
self.requestPersistentKey(localOfflineAsset.assetId)
113113
}
114-
case .failure(let error):
115-
print("Error extracting content ID: \(error.localizedDescription)")
114+
case .failure:
115+
break
116116
}
117117
}
118118
}
119119
ToastHelper.show(message: DownloadMessages.started)
120120
}
121121

122122
private func preFetchEncryptionKey(for video: Video, assetId: String, accessToken: String?) {
123-
guard let videoId = video.id else {
124-
return
125-
}
126-
127-
ContentProtectionAPI.fetchEncryptionKey(videoId: videoId, accessToken: accessToken) { [weak self] keyData in
128-
if let keyData = keyData {
129-
self?.saveEncryptionKey(keyData, videoId: videoId)
123+
guard let videoId = video.id, !videoId.isEmpty else { return }
124+
guard let url = URL(string: video.playbackURL) else { return }
125+
parsePlaylist(url, videoId, accessToken)
126+
}
127+
128+
private func parsePlaylist(_ playlistURL: URL, _ videoId: String, _ accessToken: String?) {
129+
URLSession.shared.dataTask(with: playlistURL) { [weak self] responseData, _, _ in
130+
guard let responseData = responseData, let playlistString = String(data: responseData, encoding: .utf8) else {
131+
ContentProtectionAPI.fetchEncryptionKey(videoId: videoId, accessToken: accessToken) { [weak self] encryptionKeyData in
132+
if let encryptionKeyData = encryptionKeyData {
133+
self?.encryptionKeyRepository.save(encryptionKey: encryptionKeyData, for: videoId)
134+
}
135+
}
136+
return
130137
}
131-
}
138+
if let masterPlaylistPath = playlistString.components(separatedBy: .newlines).first(where: { $0.hasSuffix(".m3u8") }), let masterPlaylistURL = URL(string: masterPlaylistPath, relativeTo: playlistURL) {
139+
self?.parsePlaylist(masterPlaylistURL, videoId, accessToken)
140+
141+
} else {
142+
let encryptionKeyPattern = "#EXT-X-KEY:.*URI=\"([^\"]+)\"", regexMatch = try? NSRegularExpression(pattern: encryptionKeyPattern).firstMatch(in: playlistString, range: NSRange(playlistString.startIndex..., in: playlistString))
143+
if let regexMatch = regexMatch, let matchRange = Range(regexMatch.range(at: 1), in: playlistString), let keyURL = URL(string: String(playlistString[matchRange]), relativeTo: playlistURL) {
144+
ContentProtectionAPI.fetchEncryptionKey(url: keyURL, accessToken: accessToken) { [weak self] encryptionKeyData in
145+
if let encryptionKeyData = encryptionKeyData { [keyURL.absoluteString.components(separatedBy: "?").first, keyURL.pathComponents.last, videoId].compactMap { $0 }.forEach { self?.encryptionKeyRepository.save(encryptionKey: encryptionKeyData, for: $0) } }
146+
}
147+
} else {
148+
ContentProtectionAPI.fetchEncryptionKey(videoId: videoId, accessToken: accessToken) { [weak self] encryptionKeyData in
149+
if let encryptionKeyData = encryptionKeyData {
150+
self?.encryptionKeyRepository.save(encryptionKey: encryptionKeyData, for: videoId)
151+
}
152+
}
153+
}
154+
}
155+
}.resume()
132156
}
133157

134-
private func saveEncryptionKey(_ keyData: Data, videoId: String) {
135-
encryptionKeyRepository.save(encryptionKey: keyData, for: videoId)
158+
private func saveEncryptionKey(_ keyData: Data, identifier: String) {
159+
encryptionKeyRepository.save(encryptionKey: keyData, for: identifier)
136160
}
137161

138162
private func requestPersistentKey(_ assetID: String) {
139163
guard let localOfflineAsset = LocalOfflineAsset.manager.get(id: assetID) else {
140-
print("Asset with ID \(assetID) does not exist.")
141164
return
142165
}
143166
contentKeySession.processContentKeyRequest(
@@ -149,7 +172,6 @@ public final class TPStreamsDownloadManager {
149172

150173
public func pauseDownload(_ assetId: String) {
151174
guard let localOfflineAsset = LocalOfflineAsset.manager.get(id: assetId) else {
152-
print("Asset with ID \(assetId) does not exist.")
153175
return
154176
}
155177

@@ -164,7 +186,6 @@ public final class TPStreamsDownloadManager {
164186

165187
public func resumeDownload(_ assetId: String) {
166188
guard let localOfflineAsset = LocalOfflineAsset.manager.get(id: assetId) else {
167-
print("Asset with ID \(assetId) does not exist.")
168189
return
169190
}
170191

@@ -181,7 +202,6 @@ public final class TPStreamsDownloadManager {
181202

182203
public func cancelDownload(_ assetId: String) {
183204
guard let localOfflineAsset = LocalOfflineAsset.manager.get(id: assetId) else {
184-
print("Asset with ID \(assetId) does not exist.")
185205
return
186206
}
187207

@@ -262,12 +282,10 @@ public final class TPStreamsDownloadManager {
262282
if let videoId = localOfflineAsset.videoId {
263283
encryptionKeyRepository.delete(for: videoId)
264284
}
285+
encryptionKeyRepository.delete(for: localOfflineAsset.assetId)
265286
}
266287

267288
public func getAllOfflineAssets() -> [OfflineAsset] {
268-
// This method retrieves all offline assets from the local storage.
269-
// It filters out any assets that have a status of 'deleted',
270-
// ensuring that only available (non-deleted) video assets are returned.
271289
return LocalOfflineAsset.manager.getAll()
272290
.filter { $0.status != Status.deleted.rawValue }
273291
.map { $0.asOfflineAsset() }

‎Source/Managers/ContentKeyDelegate+Persistable.swift‎

Lines changed: 32 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -62,28 +62,33 @@ extension ContentKeyDelegate {
6262

6363
func retrieveAndStoreContentKey(_ session: AVContentKeySession, _ spcData: Data?, _ error: Error?, _ keyRequest: AVPersistableContentKeyRequest) {
6464
guard let spcData = spcData else { return }
65-
66-
self.requestCKC(spcData) { ckcData, error in
65+
self.requestCKC(spcData) { [weak self] ckcData, error in
66+
guard let self = self else { return }
6767
if let error = error {
6868
self.onError?(error)
6969
keyRequest.processContentKeyResponseError(error)
70+
self.requestingPersistentKey = false
71+
return
72+
}
73+
guard let ckcData = ckcData else {
74+
self.requestingPersistentKey = false
7075
return
7176
}
72-
guard let ckcData = ckcData else { return }
73-
77+
7478
do {
7579
if self.requestingPersistentKey {
7680
let persistentKey = try keyRequest.persistableContentKey(fromKeyVendorResponse: ckcData, options: nil)
7781
let expiryDate = Date().addingTimeInterval(self.licenseDurationSeconds ?? DEFAULT_LICENSE_EXPIRY_SECONDS)
7882
try self.storePersistentContentKey(contentKey: persistentKey, expiryDate: expiryDate)
7983
}
80-
84+
8185
let keyResponse = AVContentKeyResponse(fairPlayStreamingKeyResponseData: ckcData)
8286
keyRequest.processContentKeyResponse(keyResponse)
8387
} catch {
84-
print(error)
88+
print("Error processing content key: \(error)")
89+
self.onError?(error)
8590
}
86-
91+
8792
self.requestingPersistentKey = false
8893
}
8994
}
@@ -99,14 +104,22 @@ extension ContentKeyDelegate {
99104
}
100105

101106
func storePersistentContentKey(contentKey: Data, expiryDate: Date) throws {
102-
guard let fileURL = getPersistentContentKeyURL() else { return }
103-
104-
try contentKey.write(to: fileURL, options: Data.WritingOptions.atomicWrite)
107+
guard let fileURL = getPersistentContentKeyURL() else {
108+
throw NSError(domain: "ContentKeyDelegate", code: -1, userInfo: [NSLocalizedDescriptionKey: "Failed to resolve persistent content key URL."])
109+
}
110+
111+
let directoryURL = fileURL.deletingLastPathComponent()
112+
try FileManager.default.createDirectory(at: directoryURL, withIntermediateDirectories: true)
113+
114+
// Atomically write the key to disk
115+
try contentKey.write(to: fileURL, options: .atomic)
116+
117+
// Persist/refresh the expiry metadata via the download manager
105118
if let assetID = self.assetID {
106119
TPStreamsDownloadManager.shared.updateOfflineLicenseExpiry(assetID, expiryDate: expiryDate)
107120
}
108121
}
109-
122+
110123
func cleanupPersistentContentKey() {
111124
if let keyURL = getPersistentContentKeyURL() {
112125
try? FileManager.default.removeItem(at: keyURL)
@@ -116,10 +129,15 @@ extension ContentKeyDelegate {
116129
}
117130
}
118131

132+
private var persistedContentKeyDirectory: URL {
133+
let base = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask).first!
134+
return base.appendingPathComponent("ContentKeys", isDirectory: true)
135+
}
136+
119137
func getPersistentContentKeyURL() -> URL?{
120138
guard let contentID = self.contentID else { return nil }
121139

122-
return contentKeyDirectory.appendingPathComponent("\(contentID)-Key")
140+
return persistedContentKeyDirectory.appendingPathComponent("\(contentID)-Key")
123141
}
124142

125143
private func requestOfflineLicenseCredentials(completion: @escaping () -> Void) {
@@ -128,7 +146,7 @@ extension ContentKeyDelegate {
128146
return
129147
}
130148

131-
onRequestOfflineLicenseRenewal?(assetID) { [weak self] accessToken, licenseDuration in
149+
onRequestOfflineLicenseRenewal?(assetID) { [weak self] (accessToken: String?, licenseDuration: Double?) in
132150
if let accessToken = accessToken {
133151
self?.accessToken = accessToken
134152
}
@@ -138,4 +156,5 @@ extension ContentKeyDelegate {
138156
completion()
139157
}
140158
}
159+
141160
}

‎Source/Managers/ResourceLoaderDelegate.swift‎

Lines changed: 20 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -10,14 +10,16 @@ import AVFoundation
1010

1111
class ResourceLoaderDelegate: NSObject, AVAssetResourceLoaderDelegate {
1212
private let accessToken: String?
13+
private let assetId: String?
1314
private let isPlaybackOffline: Bool
1415
private let offlineAssetId: String?
1516
private let localOfflineAsset: LocalOfflineAsset?
1617

1718
private let encryptionKeyRepository = EncryptionKeyRepository.shared
1819

19-
init(accessToken: String?, isPlaybackOffline: Bool = false, offlineAssetId: String? = nil, localOfflineAsset: LocalOfflineAsset? = nil) {
20+
init(accessToken: String?, assetId: String? = nil, isPlaybackOffline: Bool = false, offlineAssetId: String? = nil, localOfflineAsset: LocalOfflineAsset? = nil) {
2021
self.accessToken = accessToken
22+
self.assetId = assetId
2123
self.isPlaybackOffline = isPlaybackOffline
2224
self.offlineAssetId = offlineAssetId
2325
self.localOfflineAsset = localOfflineAsset
@@ -27,35 +29,31 @@ class ResourceLoaderDelegate: NSObject, AVAssetResourceLoaderDelegate {
2729
func resourceLoader(_ resourceLoader: AVAssetResourceLoader, shouldWaitForLoadingOfRequestedResource loadingRequest: AVAssetResourceLoadingRequest) -> Bool {
2830
guard let url = loadingRequest.request.url else { return false }
2931

30-
guard url.path.contains("key") else { return false }
32+
guard url.path.contains("key") else {
33+
return false
34+
}
3135

3236
if isPlaybackOffline {
3337
serveOfflineEncryptionKey(for: loadingRequest, url: url)
3438
return true
3539
}
36-
37-
guard let videoId = VideoURLParser.extractVideoId(from: url) else { return false }
3840

39-
ContentProtectionAPI.fetchEncryptionKey(videoId: videoId, accessToken: accessToken) { [weak self] keyData in
41+
let completion: (Data?) -> Void = { [weak self] keyData in
4042
self?.setEncryptionKeyResponse(for: loadingRequest, data: keyData)
4143
}
44+
45+
ContentProtectionAPI.fetchEncryptionKey(url: url, accessToken: accessToken, completion: completion)
4246
return true
4347
}
4448

4549
private func serveOfflineEncryptionKey(for loadingRequest: AVAssetResourceLoadingRequest, url: URL) {
46-
let identifiers = [offlineAssetId, localOfflineAsset?.videoId, url.lastPathComponent]
47-
.compactMap { $0 }
48-
49-
let keyData = identifiers.lazy.compactMap { encryptionKeyRepository.get(for: $0) }.first
50-
51-
if let keyData = keyData {
52-
setEncryptionKeyResponse(for: loadingRequest, data: keyData)
50+
let encryptionKeyIdentifiers = [url.absoluteString.components(separatedBy: "?").first, url.pathComponents.last,
51+
url.pathComponents.firstIndex(where: { ["assets", "encryption_key", "aes_key"].contains($0) }).map { url.pathComponents[$0 + 1] },
52+
offlineAssetId, localOfflineAsset?.videoId, assetId].compactMap { $0 }
53+
if let encryptionKey = encryptionKeyIdentifiers.compactMap({ encryptionKeyRepository.get(for: $0) }).first {
54+
setEncryptionKeyResponse(for: loadingRequest, data: encryptionKey)
5355
} else {
54-
loadingRequest.finishLoading(with: NSError(
55-
domain: "TPStreamsSDK",
56-
code: -1,
57-
userInfo: [NSLocalizedDescriptionKey: "Missing offline AES encryption key"]
58-
))
56+
loadingRequest.finishLoading(with: NSError(domain: "TPStreamsSDK", code: -1, userInfo: [NSLocalizedDescriptionKey: "Missing offline key"]))
5957
}
6058
}
6159

@@ -83,6 +81,10 @@ class ResourceLoaderDelegate: NSObject, AVAssetResourceLoaderDelegate {
8381
if #available(iOS 11.2, *), allowedType == AVStreamingKeyDeliveryContentKeyType {
8482
return AVStreamingKeyDeliveryContentKeyType
8583
}
86-
return AVStreamingKeyDeliveryPersistentContentKeyType
84+
85+
// If the allowed type is set, use it. Otherwise, fallback to a sensible default.
86+
// For AES-128 keys, it might be nil or "public.data" in many cases.
87+
// We want to avoid FairPlay types if we are not handling FairPlay.
88+
return allowedType ?? "application/octet-stream"
8789
}
8890
}

0 commit comments

Comments
 (0)