fast-s3-swift

Upload.swift

/**
 * Upload.swift
 * fast-s3
 *
 * Created by Nick Spreen (spreen.co) on 3/3/25.
 *
 */

import Foundation
import Network

// MARK: Public Methods

/**
 * Upload a file to fast-s3, retrying if the upload fails. Waits for wifi connection.
 *
 * @param path The path to the file to upload.
 * @param apiKey The API key to use for the upload.
 * @param metadata Additional metadata to include with the upload.
 * @param retries The number of retries to attempt if the upload fails.
 * @return The URL of the uploaded file as string, or nil if the upload failed.
 */
public func uploadOnWifi(
	from path: String,
	for apiKey: String,
	with metadata: [String: String] = [:],
	retries: Int = 99
) async -> String? {
	for currentTry in 1...(retries + 1) {
		if currentTry > 1 {
			try? await Task.sleep(nanoseconds: UInt64(pow(1.5, Double(currentTry) - 2) * 1_000_000_000))
			if Task.isCancelled { return nil }
		}

		if await isCellular() { continue }

		return await upload(from: path, for: apiKey, with: metadata)
	}
	return nil
}

/**
 * Upload a file to fast-s3, retrying if the upload fails. Waits for wifi connection.
 *
 * @param path The path to the file to upload.
 * @param apiKey The API key to use for the upload.
 * @param metadata Additional metadata to include with the upload.
 * @param retries The number of retries to attempt if the upload fails.
 * @return The URL of the uploaded file as string, or nil if the upload failed.
 */
public func upload(from path: String, for apiKey: String, with metadata: [String: String] = [:]) async -> String?
{
	let fileURL = URL(fileURLWithPath: path)
	let fileName = fileURL.lastPathComponent
	var metadata = metadata
	metadata["filename"] = fileName

	guard let size = try? FileManager.default.attributesOfItem(atPath: path)[.size] as? Int else {
		return nil
	}

	if size < chunkSize {
		return await uploadSingle(fileURL: fileURL, apiKey: apiKey, metadata: metadata)
	}
	return await uploadMultipart(fileURL: fileURL, apiKey: apiKey, metadata: metadata, size: size)
}

// MARK: Internal Methods

func isCellular(completion: @escaping (Bool) -> Void) {
	let monitor = NWPathMonitor()
	let queue = DispatchQueue.global(qos: .background)

	monitor.pathUpdateHandler = { path in
		if path.usesInterfaceType(.cellular) {
			completion(true)
		} else {
			completion(false)
		}
		monitor.cancel() // Stop monitoring after first check
	}

	monitor.start(queue: queue)
}

func isCellular() async -> Bool {
	await withCheckedContinuation { continuation in
		isCellular {
			continuation.resume(returning: $0)
		}
	}
}

fileprivate func readBytes(from filePath: String, range: Range<Int>) -> Data? {
	guard let fileHandle = FileHandle(forReadingAtPath: filePath) else {
		print("[fast-s3] Failed to open file")
		return nil
	}

	defer {
		fileHandle.closeFile()
	}

	do {
		// Seek to the start position
		try fileHandle.seek(toOffset: UInt64(range.lowerBound))

		// Read required number of bytes
		let data = fileHandle.readData(ofLength: range.count)
		return data
	} catch {
		print("[fast-s3] Error reading file: \(error)")
		return nil
	}
}

fileprivate let baseURL = "https://fast-s3.spreen.co"
fileprivate let session = URLSession.shared
fileprivate let chunkSize = 3 * 1024 * 1024 // 3MiB
fileprivate let maxConcurrentTasks = 10

fileprivate func uploadSingle(
	fileURL: URL, apiKey: String, metadata: [String: String], retries: Int = 99
) async -> String? {
	guard let data = try? Data(contentsOf: fileURL) else {
		print("[fast-s3] Could not read file at \(fileURL)")
		return nil
	}

	var urlComponents = URLComponents(string: baseURL + "/upload")!
	urlComponents.queryItems = [URLQueryItem]()

	for (key, value) in metadata {
		urlComponents.queryItems!.append(URLQueryItem(name: key, value: value))
	}

	for currentTry in 1...(retries + 1) {
		if currentTry > 1 {
			try? await Task.sleep(nanoseconds: UInt64(pow(1.5, Double(currentTry) - 2) * 1_000_000_000))
			if Task.isCancelled { return nil }
		}

		var request = URLRequest(url: urlComponents.url!)
		request.httpMethod = "POST"
		request.httpBody = data
		request.setValue(apiKey, forHTTPHeaderField: "x-api-key")

		guard
			let (data, response) = try? await session.data(for: request),
			let httpResponse = response as? HTTPURLResponse
		else { continue }

		if
			(200...299).contains(httpResponse.statusCode),
			let json = try? JSONSerialization.jsonObject(with: data, options: []) as? [String: Any],
			let url = json["url"] as? String
		{
			print("[fast-s3] Upload succeded: \(url)")
			return url
		} else {
			print("[fast-s3] Upload failed with status code: \(httpResponse.statusCode)")
			print("[fast-s3] Response: \(String(data: data, encoding: .utf8) ?? "No response")")
		}
	}
	return nil
}

fileprivate func uploadMultipart(
	fileURL: URL,
	apiKey: String,
	metadata: [String: String],
	size: Int,
	retries: Int = 99
) async -> String? {
	let chunkCount = (size - 1) / chunkSize + 1

	guard
		let continuationId = await uploadChunk(
			chunk: 1, fileURL: fileURL, apiKey: apiKey, size: size, continuationId: nil,
			retries: retries)
	else {
		return nil
	}

	// Use task group for parallel uploads with max 10 concurrent tasks
	let group = await withTaskGroup(of: Bool.self) { group in
		var currentChunk = 2

		// Initially add up to maxConcurrentTasks
		for _ in 0..<min(maxConcurrentTasks, chunkCount - 1) {
			let chunk = currentChunk
			currentChunk += 1
			group.addTask {
				return await uploadChunk(
					chunk: chunk, fileURL: fileURL, apiKey: apiKey, size: size,
					continuationId: continuationId, retries: retries) != nil
			}
		}

		// Process remaining chunks as tasks complete
		while currentChunk <= chunkCount {
			// When a task completes, add another if there are more chunks
			if await group.next() == false {
				// If any chunk upload failed, cancel remaining work
				group.cancelAll()
				return false
			}

			let chunk = currentChunk
			currentChunk += 1
			group.addTask {
				return await uploadChunk(
					chunk: chunk, fileURL: fileURL, apiKey: apiKey, size: size,
					continuationId: continuationId, retries: retries) != nil
			}
		}

		// Wait for remaining tasks and check for failures
		for await result in group {
			if !result {
				return false
			}
		}
		return true
	}

	if !group { return nil }

	return await finishMultipartUpload(
		apiKey: apiKey, multiKey: continuationId, totalParts: chunkCount, metadata: metadata)
}

private func uploadChunk(
	chunk: Int,
	fileURL: URL,
	apiKey: String,
	size: Int,
	continuationId: String? = nil,
	retries: Int = 99
) async -> String? {
	let path = fileURL.path
	let isFirst = chunk == 1

	if continuationId == nil, !isFirst {
		return continuationId
	}

	// Calculate byte range for this chunk
	let offset = (chunk - 1) * chunkSize
	let endOffset = min(offset + chunkSize, size)

	guard let data = readBytes(from: path, range: offset..<endOffset) else {
		print("[fast-s3] Failed to read chunk \(chunk) from file")
		return nil
	}

	let uploadURL = "\(baseURL)/upload_multi"

	for currentTry in 1...(retries + 1) {
		if currentTry > 1 {
			try? await Task.sleep(nanoseconds: UInt64(pow(1.5, Double(currentTry) - 2) * 1_000_000_000))
			if Task.isCancelled { return nil }
		}

		var request = URLRequest(url: URL(string: uploadURL)!)
		request.httpMethod = "POST"
		request.httpBody = data
		request.setValue(apiKey, forHTTPHeaderField: "x-api-key")
		request.setValue("\(chunk)", forHTTPHeaderField: "x-part")
		if !isFirst, let continuationId = continuationId {
			request.setValue(continuationId, forHTTPHeaderField: "x-multi-key")
		}

		guard
			let (responseData, response) = try? await session.data(for: request),
			let httpResponse = response as? HTTPURLResponse,
			(200...299).contains(httpResponse.statusCode),
			let responseString = String(data: responseData, encoding: .utf8)
		else {
			continue
		}

		return extractMultiKey(from: responseString) ?? continuationId
	}

	return nil
}

private func extractMultiKey(from response: String) -> String? {
	guard
		let json = try? JSONSerialization.jsonObject(with: Data(response.utf8), options: []) as? [String: Any],
		let multiKey = json["multiKey"] as? String
	else {
		return nil
	}

	return multiKey
}

private func finishMultipartUpload(
	apiKey: String,
	multiKey: String,
	totalParts: Int,
	metadata: [String: String] = [:],
	retries: Int = 99
) async -> String? {
	let finishURL = "\(baseURL)/finish_multi"

	guard var urlComponents = URLComponents(string: finishURL) else {
		return nil
	}

	urlComponents.queryItems = [URLQueryItem]()
	for (key, value) in metadata {
		urlComponents.queryItems!.append(URLQueryItem(name: key, value: value))
	}

	guard let url = urlComponents.url else {
		return nil
	}

	for currentTry in 1...(retries + 1) {
		if currentTry > 1 {
			try? await Task.sleep(nanoseconds: UInt64(pow(1.5, Double(currentTry) - 2) * 1_000_000_000))
			if Task.isCancelled { return nil }
		}

		var request = URLRequest(url: url)
		request.httpMethod = "POST"
		request.setValue(apiKey, forHTTPHeaderField: "x-api-key")
		request.setValue(multiKey, forHTTPHeaderField: "x-multi-key")
		request.setValue("\(totalParts)", forHTTPHeaderField: "x-part")


		guard
			let (data, response) = try? await session.data(for: request),
			let httpResponse = response as? HTTPURLResponse
		else { continue }

		if
			(200...299).contains(httpResponse.statusCode),
			let json = try? JSONSerialization.jsonObject(with: data, options: []) as? [String: Any],
			let url = json["url"] as? String
		{
			print("[fast-s3] Upload succeded: \(url)")
			return url
		} else {
			print("[fast-s3] Upload failed with status code: \(httpResponse.statusCode)")
			print("[fast-s3] Response: \(String(data: data, encoding: .utf8) ?? "No response")")
		}
	}

	return nil
}