diff --git a/SignalServiceKit/src/Network/ProxiedContentDownloader.swift b/SignalServiceKit/src/Network/ProxiedContentDownloader.swift index f819355b8..de3c2c25f 100644 --- a/SignalServiceKit/src/Network/ProxiedContentDownloader.swift +++ b/SignalServiceKit/src/Network/ProxiedContentDownloader.swift @@ -162,8 +162,6 @@ public class ProxiedContentAssetRequest: NSObject { AssertIsOnMainThread() assert(oldValue == 0) assert(contentLength > 0) - - createSegments() } } public weak var contentLengthTask: URLSessionDataTask? @@ -204,7 +202,7 @@ public class ProxiedContentAssetRequest: NSObject { return contentLength } - private func createSegments() { + fileprivate func createSegments(withInitialData initialData: Data) { AssertIsOnMainThread() let segmentLength = segmentSize() @@ -213,8 +211,20 @@ public class ProxiedContentAssetRequest: NSObject { } let contentLength = UInt(self.contentLength) - var nextSegmentStart: UInt = 0 - var index: UInt = 0 + // Make the initial segment. + let assetSegment = ProxiedContentAssetSegment(index: 0, + segmentStart: 0, + segmentLength: UInt(initialData.count), + redundantLength: 0) + // "Download" the initial segment using the initialData. + assetSegment.state = .downloading + assetSegment.append(data: initialData) + // Mark the initial segment as complete. + assetSegment.state = .complete + segments.append(assetSegment) + + var nextSegmentStart = UInt(initialData.count) + var index: UInt = 1 while nextSegmentStart < contentLength { var segmentStart: UInt = nextSegmentStart var redundantLength: UInt = 0 @@ -534,31 +544,40 @@ open class ProxiedContentDownloader: NSObject, URLSessionTaskDelegate, URLSessio DispatchQueue.main.async { assetSegment.state = .complete - if assetRequest.areAllSegmentsComplete() { - // If the asset request has completed all of its segments, - // try to write the asset to file. - assetRequest.state = .complete - - // Move write off main thread. - DispatchQueue.global().async { - guard let downloadFolderPath = self.downloadFolderPath else { - owsFailDebug("Missing downloadFolderPath") - return - } - guard let asset = assetRequest.writeAssetToFile(downloadFolderPath: downloadFolderPath) else { - self.segmentRequestDidFail(assetRequest: assetRequest, assetSegment: assetSegment) - return - } - self.assetRequestDidSucceed(assetRequest: assetRequest, asset: asset) - } - } else { + if !self.tryToCompleteRequest(assetRequest: assetRequest) { self.processRequestQueueSync() } } } - private func assetRequestDidSucceed(assetRequest: ProxiedContentAssetRequest, asset: ProxiedContentAsset) { + // Returns true if the request is completed. + private func tryToCompleteRequest(assetRequest: ProxiedContentAssetRequest) -> Bool { + AssertIsOnMainThread() + + guard assetRequest.areAllSegmentsComplete() else { + return false + } + + // If the asset request has completed all of its segments, + // try to write the asset to file. + assetRequest.state = .complete + // Move write off main thread. + DispatchQueue.global().async { + guard let downloadFolderPath = self.downloadFolderPath else { + owsFailDebug("Missing downloadFolderPath") + return + } + guard let asset = assetRequest.writeAssetToFile(downloadFolderPath: downloadFolderPath) else { + self.segmentRequestDidFail(assetRequest: assetRequest) + return + } + self.assetRequestDidSucceed(assetRequest: assetRequest, asset: asset) + } + return true + } + + private func assetRequestDidSucceed(assetRequest: ProxiedContentAssetRequest, asset: ProxiedContentAsset) { DispatchQueue.main.async { self.assetMap.set(key: assetRequest.assetDescription.url, value: asset) self.removeAssetRequestFromQueue(assetRequest: assetRequest) @@ -566,11 +585,14 @@ open class ProxiedContentDownloader: NSObject, URLSessionTaskDelegate, URLSessio } } - // TODO: If we wanted to implement segment retry, we'll need to add - // a segmentRequestDidFail() method. - private func segmentRequestDidFail(assetRequest: ProxiedContentAssetRequest, assetSegment: ProxiedContentAssetSegment) { + private func segmentRequestDidFail(assetRequest: ProxiedContentAssetRequest, assetSegment: ProxiedContentAssetSegment? = nil) { DispatchQueue.main.async { - assetSegment.state = .failed + if let assetSegment = assetSegment { + assetSegment.state = .failed + + // TODO: If we wanted to implement segment retry, we'd do so here. + // For now, we just fail the entire asset request. + } assetRequest.state = .failed self.assetRequestDidFail(assetRequest: assetRequest) } @@ -637,21 +659,20 @@ open class ProxiedContentDownloader: NSObject, URLSessionTaskDelegate, URLSessio if assetRequest.state == .waiting { // If asset request hasn't yet determined the resource size, - // try to do so now. + // try to do so now, by requesting a small initial segment. assetRequest.state = .requestingSize + let segmentStart: UInt = 0 + // Vary the initial segment size to obscure the length of the response headers. + let segmentLength: UInt = 1024 + UInt(arc4random_uniform(1024)) var request = URLRequest(url: assetRequest.assetDescription.url as URL) - request.httpMethod = "HEAD" request.httpShouldUsePipelining = true - // Some services like Reddit will severely rate-limit requests without a user agent. - request.addValue("Signal iOS (+https://signal.org/download)", forHTTPHeaderField: "User-Agent") - + let rangeHeaderValue = "bytes=\(segmentStart)-\(segmentStart + segmentLength - 1)" + request.addValue(rangeHeaderValue, forHTTPHeaderField: "Range") let task = downloadSession.dataTask(with: request, completionHandler: { data, response, error -> Void in - if let data = data, data.count > 0 { - owsFailDebug("HEAD request has unexpected body: \(data.count).") - } - self.handleAssetSizeResponse(assetRequest: assetRequest, response: response, error: error) + self.handleAssetSizeResponse(assetRequest: assetRequest, data: data, response: response, error: error) }) + assetRequest.contentLengthTask = task task.resume() } else { @@ -678,25 +699,52 @@ open class ProxiedContentDownloader: NSObject, URLSessionTaskDelegate, URLSessio processRequestQueueSync() } - private func handleAssetSizeResponse(assetRequest: ProxiedContentAssetRequest, response: URLResponse?, error: Error?) { + private func handleAssetSizeResponse(assetRequest: ProxiedContentAssetRequest, data: Data?, response: URLResponse?, error: Error?) { guard error == nil else { assetRequest.state = .failed self.assetRequestDidFail(assetRequest: assetRequest) return } + guard let data = data, + data.count > 0 else { + owsFailDebug("Asset size response missing data.") + assetRequest.state = .failed + self.assetRequestDidFail(assetRequest: assetRequest) + return + } guard let httpResponse = response as? HTTPURLResponse else { owsFailDebug("Asset size response is invalid.") assetRequest.state = .failed self.assetRequestDidFail(assetRequest: assetRequest) return } - guard let contentLengthString = httpResponse.allHeaderFields["Content-Length"] as? String else { - owsFailDebug("Asset size response is missing content length.") + var firstContentRangeString: String? + for header in httpResponse.allHeaderFields.keys { + guard let headerString = header as? String else { + owsFailDebug("Invalid header: \(header)") + continue + } + if headerString.lowercased() == "content-range" { + firstContentRangeString = httpResponse.allHeaderFields[header] as? String + } + } + guard let contentRangeString = firstContentRangeString else { + owsFailDebug("Asset size response is missing content range.") assetRequest.state = .failed self.assetRequestDidFail(assetRequest: assetRequest) return } - guard let contentLength = Int(contentLengthString) else { + + // Example: content-range: bytes 0-1023/7630 + guard let contentLengthString = NSRegularExpression.parseFirstMatch(pattern: "^bytes \\d+\\-\\d+/(\\d+)$", + text: contentRangeString) else { + owsFailDebug("Asset size response has invalid content range.") + assetRequest.state = .failed + self.assetRequestDidFail(assetRequest: assetRequest) + return + } + guard contentLengthString.count > 0, + let contentLength = Int(contentLengthString) else { owsFailDebug("Asset size response has unparsable content length.") assetRequest.state = .failed self.assetRequestDidFail(assetRequest: assetRequest) @@ -711,8 +759,12 @@ open class ProxiedContentDownloader: NSObject, URLSessionTaskDelegate, URLSessio DispatchQueue.main.async { assetRequest.contentLength = contentLength + assetRequest.createSegments(withInitialData: data) assetRequest.state = .active - self.processRequestQueueSync() + + if !self.tryToCompleteRequest(assetRequest: assetRequest) { + self.processRequestQueueSync() + } } }