Linux Compatible URLSession Async (AsyncBytes)

FoundationNetworking+AsyncAwait.swift

//
//  FoundationNetworking+AsyncAwait.swift
//
//
//  Created by BadPirate on 7/10/24.
//

/// Currently, Foundation that is used for servers (Like Vapor) has split out URLSession into FoundationNetworking, while it has support for async / await
/// Generally, FoundationNetworking URLSession doesn't support the async/await helpers (like `data()` and `bytes()`) This extension will allow you to
/// Use these as normal if FoundationNetworking is present, by providing an implementation.
///
/// I could have tried to add this to FoundationNetworking, however, it seems that this will be deprecated in favor of Apple's 2023 open sourcing of
/// the overall Foundation class. So I am using this shim in the meantime.

import Foundation

#if os(Linux)
import FoundationNetworking

public class URLBytesSequence: NSObject, AsyncSequence, URLSessionDataDelegate {
    private let stream: AsyncStream<UInt8>
    private let continuation: AsyncStream<UInt8>.Continuation
    
    private var session: URLSession!
    private var dataTask: URLSessionDataTask!
    private var responseContinuation: CheckedContinuation<URLResponse,Error>!
    private var responseTask: Task<URLResponse,Error>!
    private var response: URLResponse?

    required init(for request: URLRequest) {
        let (stream, continuation) = AsyncStream.makeStream(of: UInt8.self)
        self.stream = stream
        self.continuation = continuation
        super.init()
        
        let session = URLSession(configuration: URLSessionConfiguration.default, delegate: self, delegateQueue: nil)
        self.session = session
        dataTask = session.dataTask(with: request)
        responseTask = Task {
            try await withCheckedThrowingContinuation { self.responseContinuation = $0 }
        }
    }
    
    public typealias Element = UInt8

    public func resume() { dataTask.resume() }

    // MARK: AsyncSequence

    public func makeAsyncIterator() -> AsyncStream<UInt8>.Iterator {
        return stream.makeAsyncIterator()
    }

    public typealias AsyncIterator = AsyncStream<UInt8>.Iterator

    // MARK: URLSessionDataDelegate

    public func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
        for byte in data { continuation.yield(byte) }
    }

    public func urlSession(
        _ session: URLSession, dataTask: URLSessionDataTask,
        didReceive response: URLResponse,
        completionHandler: @escaping (URLSession.ResponseDisposition) -> Void
    ) {
        responseContinuation.resume(returning: response)
        self.response = response
        completionHandler(.allow)
    }
    
    public func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
        if let error {
            if response == nil {
                responseContinuation.resume(throwing: error)
            }
        }
        continuation.finish()
    }

    public func bytes(
        for request: URLRequest,
        delegate: (any URLSessionTaskDelegate)? = nil) async throws -> (URLBytesSequence, URLResponse)
    {
        let session = URLSession(configuration: .default, delegate: self, delegateQueue: nil)
        let task = session.dataTask(with: request)
        task.resume()
        return try await (self, responseTask.value)
    }
}

public extension URLSession {
    func data(for request: URLRequest, delegate: (any URLSessionTaskDelegate)? = nil) async throws -> (Data, URLResponse) {
        try await withCheckedThrowingContinuation { continuation in
            let task = self.dataTask(with: request) { data, response, error in
                if let error {
                    continuation.resume(throwing: error)
                } else {
                    continuation.resume(returning: (data!, response!))
                }
            }
            task.resume()
        }
    }
    
    func bytes(for request: URLRequest, delegate: (any URLSessionTaskDelegate)? = nil) async throws -> (URLBytesSequence, URLResponse) {
        let sequence = URLBytesSequence(for: request)
        return try await sequence.bytes(for: request, delegate: delegate)
    }
}
#else
public typealias URLBytesSequence = URLSession.AsyncBytes
#endif