Support multipart form in URLSession helper (#699)

Author: crazytonyli

WordPressKit.xcodeproj/project.pbxproj

@@ -159,6 +159,7 @@
 		4A57A6832B54A326008D0660 /* WordPressAPIError+NSErrorBrdige.swift in Sources */ = {isa = PBXBuildFile; fileRef = 4A57A6822B54A326008D0660 /* WordPressAPIError+NSErrorBrdige.swift */; };
 		4A57A6872B54C68C008D0660 /* Constants.h in Headers */ = {isa = PBXBuildFile; fileRef = 4A57A6852B54C68C008D0660 /* Constants.h */; settings = {ATTRIBUTES = (Public, ); }; };
 		4A57A6882B54C68C008D0660 /* Constants.m in Sources */ = {isa = PBXBuildFile; fileRef = 4A57A6862B54C68C008D0660 /* Constants.m */; };
+		4A5BC1A82B59DE6600C7D037 /* Either.swift in Sources */ = {isa = PBXBuildFile; fileRef = 4A5BC1A72B59DE6600C7D037 /* Either.swift */; };
 		4A68E3CD29404181004AC3DC /* RemoteBlog.swift in Sources */ = {isa = PBXBuildFile; fileRef = 4A68E3CC29404181004AC3DC /* RemoteBlog.swift */; };
 		4A68E3CF29404289004AC3DC /* RemoteBlogOptionsHelper.swift in Sources */ = {isa = PBXBuildFile; fileRef = 4A68E3CE29404289004AC3DC /* RemoteBlogOptionsHelper.swift */; };
 		4A68E3D329406AA0004AC3DC /* RemoteMenu.swift in Sources */ = {isa = PBXBuildFile; fileRef = 4A68E3D029406AA0004AC3DC /* RemoteMenu.swift */; };
@@ -873,6 +874,7 @@
 		4A57A6822B54A326008D0660 /* WordPressAPIError+NSErrorBrdige.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "WordPressAPIError+NSErrorBrdige.swift"; sourceTree = "<group>"; };
 		4A57A6852B54C68C008D0660 /* Constants.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = Constants.h; sourceTree = "<group>"; };
 		4A57A6862B54C68C008D0660 /* Constants.m */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.objc; path = Constants.m; sourceTree = "<group>"; };
+		4A5BC1A72B59DE6600C7D037 /* Either.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Either.swift; sourceTree = "<group>"; };
 		4A68E3CC29404181004AC3DC /* RemoteBlog.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = RemoteBlog.swift; sourceTree = "<group>"; };
 		4A68E3CE29404289004AC3DC /* RemoteBlogOptionsHelper.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = RemoteBlogOptionsHelper.swift; sourceTree = "<group>"; };
 		4A68E3D029406AA0004AC3DC /* RemoteMenu.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = RemoteMenu.swift; sourceTree = "<group>"; };
@@ -2494,6 +2496,7 @@
 				4A11239F2B196821004690CF /* MultipartForm.swift */,
 				4AE278432B2FAF6200E4D9B1 /* HTTPProtocolHelpers.swift */,
 				3F391E192B50F3EB007975C4 /* Result+Callback.swift */,
+				4A5BC1A72B59DE6600C7D037 /* Either.swift */,
 				4A57A6852B54C68C008D0660 /* Constants.h */,
 				4A57A6862B54C68C008D0660 /* Constants.m */,
 			);
@@ -3483,6 +3486,7 @@
 				436D5641211B7F4400CEAA33 /* DomainContactInformation.swift in Sources */,
 				FA28A3D6259079960082C7B0 /* JetpackRestoreTypes.swift in Sources */,
 				32FC20CE255DCC6100CD0A7B /* JetpackScanThreat.swift in Sources */,
+				4A5BC1A82B59DE6600C7D037 /* Either.swift in Sources */,
 				FE50965F2A2E42A500DDD071 /* JetpackSocialServiceRemote.swift in Sources */,
 				3F3195AD266FF94B00397EE7 /* ZendeskMetadata.swift in Sources */,
 				40A71C6E220E1D8E002E3D25 /* StatsServiceRemoteV2.swift in Sources */,

WordPressKit/Either.swift

@@ -0,0 +1,15 @@
+import Foundation
+
+enum Either<L, R> {
+    case left(L)
+    case right(R)
+
+    func map<T>(left: (L) -> T, right: (R) -> T) -> T {
+        switch self {
+        case let .left(value):
+            return left(value)
+        case let .right(value):
+            return right(value)
+        }
+    }
+}

WordPressKit/HTTPClient.swift

@@ -52,15 +52,8 @@ extension URLSession {
             assert(parentProgress.cancellationHandler == nil, "The progress instance's cancellationHandler property must be nil")
         }
 
-        let request: URLRequest
-        do {
-            request = try builder.build()
-        } catch {
-            return .failure(.requestEncodingFailure(underlyingError: error))
-        }
-
         return await withCheckedContinuation { continuation in
-            let task = dataTask(with: request) { data, response, error in
+            let completion: @Sendable (Data?, URLResponse?, Error?) -> Void = { data, response, error in
                 let result: WordPressAPIResult<HTTPAPIResponse<Data>, E> = Self.parseResponse(
                     data: data,
                     response: response,
@@ -70,6 +63,16 @@ extension URLSession {
 
                 continuation.resume(returning: result)
             }
+
+            let task: URLSessionTask
+
+            do {
+                task = try self.task(for: builder, completion: completion)
+            } catch {
+                continuation.resume(returning: .failure(.requestEncodingFailure(underlyingError: error)))
+                return
+            }
+
             task.resume()
 
             if let parentProgress, parentProgress.totalUnitCount > parentProgress.completedUnitCount {
@@ -83,6 +86,32 @@ extension URLSession {
         }
     }
 
+    private func task(
+        for builder: HTTPRequestBuilder,
+        completion: @escaping @Sendable (Data?, URLResponse?, Error?) -> Void
+    ) throws -> URLSessionTask {
+        var request = try builder.build(encodeMultipartForm: false)
+
+        // Use special `URLSession.uploadTask` API for multipart POST requests.
+        if let multipart = builder.multipartForm, !multipart.isEmpty {
+            let isBackgroundSession = configuration.identifier != nil
+
+            return try builder
+                .encodeMultipartForm(request: &request, forceWriteToFile: isBackgroundSession)
+                .map(
+                    left: {
+                        uploadTask(with: request, from: $0, completionHandler: completion)
+                    },
+                    right: {
+                        uploadTask(with: request, fromFile: $0, completionHandler: completion)
+                    }
+                )
+        } else {
+            // Use `URLSession.dataTask` for all other request
+            return dataTask(with: request, completionHandler: completion)
+        }
+    }
+
     private static func parseResponse<E: LocalizedError>(
         data: Data?,
         response: URLResponse?,

WordPressKit/HTTPRequestBuilder.swift

@@ -17,6 +17,7 @@ final class HTTPRequestBuilder {
     private var method: Method = .get
     private var headers: [String: String] = [:]
     private var bodyBuilder: ((inout URLRequest) throws -> Void)?
+    private(set) var multipartForm: [MultipartFormField]?
 
     init(url: URL) {
         assert(url.scheme == "http" || url.scheme == "https")
@@ -83,6 +84,12 @@ final class HTTPRequestBuilder {
         return self
     }
 
+    func body(form: [MultipartFormField]) -> Self {
+        // Unlike other similar functions, multipart form encoding is handled by the `build` function.
+        multipartForm = form
+        return self
+    }
+
     func body(json: Encodable, jsonEncoder: JSONEncoder = JSONEncoder()) -> Self {
         body(json: {
             try jsonEncoder.encode(json)
@@ -111,7 +118,7 @@ final class HTTPRequestBuilder {
         return self
     }
 
-    func build() throws -> URLRequest {
+    func build(encodeMultipartForm: Bool = false) throws -> URLRequest {
         guard let url = urlComponents.url else {
             throw URLError(.badURL)
         }
@@ -123,13 +130,35 @@ final class HTTPRequestBuilder {
             request.addValue(value, forHTTPHeaderField: header)
         }
 
+        if encodeMultipartForm {
+            let encoded = try self.encodeMultipartForm(request: &request, forceWriteToFile: false)
+            switch encoded {
+            case let .left(data):
+                request.httpBody = data
+            case let .right(url):
+                request.httpBodyStream = InputStream(url: url)
+            }
+        }
+
         if let bodyBuilder {
             assert(method.allowsHTTPBody, "Can't include body in HTTP \(method.rawValue) requests")
             try bodyBuilder(&request)
         }
 
         return request
     }
+
+    func encodeMultipartForm(request: inout URLRequest, forceWriteToFile: Bool) throws -> Either<Data, URL> {
+        guard let multipartForm, !multipartForm.isEmpty else {
+            return .left(Data())
+        }
+
+        let boundery = String(format: "wordpresskit.%08x", Int.random(in: Int.min..<Int.max))
+        request.setValue("multipart/form-data; boundary=\(boundery)", forHTTPHeaderField: "Content-Type")
+        return try multipartForm
+            .multipartFormDataStream(boundary: boundery, forceWriteToFile: forceWriteToFile)
+
+    }
 }
 
 extension HTTPRequestBuilder {

WordPressKit/MultipartForm.swift

@@ -40,7 +40,7 @@ struct MultipartFormField {
 }
 
 extension Array where Element == MultipartFormField {
-    private func multipartFormDestination(forceWriteToFile: Bool = false) throws -> (outputStream: OutputStream, tempFilePath: String?) {
+    private func multipartFormDestination(forceWriteToFile: Bool) throws -> (outputStream: OutputStream, tempFilePath: String?) {
         let dest: OutputStream
         let tempFilePath: String?
 
@@ -62,12 +62,12 @@ extension Array where Element == MultipartFormField {
         return (dest, tempFilePath)
     }
 
-    func multipartFormDataStream(boundary: String = String(format: "wordpresskit.%08x", Int.random(in: Int.min..<Int.max))) throws -> InputStream {
+    func multipartFormDataStream(boundary: String, forceWriteToFile: Bool = false) throws -> Either<Data, URL> {
         guard !isEmpty else {
-            return InputStream(data: Data())
+            return .left(Data())
         }
 
-        let (dest, tempFilePath) = try multipartFormDestination()
+        let (dest, tempFilePath) = try multipartFormDestination(forceWriteToFile: forceWriteToFile)
 
         // Build the form content
         do {
@@ -79,14 +79,11 @@ extension Array where Element == MultipartFormField {
 
         // Return the result as `InputStream`
         if let tempFilePath {
-            guard let stream = InputStream(fileAtPath: tempFilePath) else {
-                throw MultipartFormError.inaccessbileFile(path: tempFilePath)
-            }
-            return stream
+            return .right(URL(fileURLWithPath: tempFilePath))
         }
 
         if let data = dest.property(forKey: .dataWrittenToMemoryStreamKey) as? Data {
-            return InputStream(data: data)
+            return .left(data)
         }
 
         throw MultipartFormError.impossible

WordPressKitTests/Utilities/HTTPRequestBuilderTests.swift

@@ -225,6 +225,32 @@ class HTTPRequestBuilderTests: XCTestCase {
         XCTAssertEqual(form, decodedForm)
     }
 
+    func testMultipartForm() throws {
+        XCTAssertNotNil(
+            try HTTPRequestBuilder(url: URL(string: "https://wordpress.org")!)
+                .method(.post)
+                .body(form: [MultipartFormField(text: "123456", name: "site")])
+                .build(encodeMultipartForm: true)
+                .httpBody
+        )
+
+        XCTAssertNil(
+            try HTTPRequestBuilder(url: URL(string: "https://wordpress.org")!)
+                .method(.post)
+                .body(form: [MultipartFormField(text: "123456", name: "site")])
+                .build(encodeMultipartForm: false)
+                .httpBody
+        )
+
+        XCTAssertNil(
+            try HTTPRequestBuilder(url: URL(string: "https://wordpress.org")!)
+                .method(.post)
+                .body(form: [MultipartFormField(text: "123456", name: "site")])
+                .build(encodeMultipartForm: false)
+                .httpBodyStream
+        )
+    }
+
 }
 
 private extension URLRequest {

WordPressKitTests/Utilities/MultipartFormTests.swift

@@ -81,7 +81,7 @@ class MutliparFormDataTests: XCTestCase {
     }
 
     func testEmptyForm() throws {
-        let formData = try [].multipartFormDataStream().readToEnd()
+        let formData = try [].multipartFormDataStream(boundary: "test").readToEnd()
         XCTAssertTrue(formData.isEmpty)
     }
 
@@ -141,7 +141,16 @@ class MutliparFormDataTests: XCTestCase {
 
 }
 
-private extension InputStream {
+extension Either<Data, URL> {
+    func readToEnd() -> Data {
+        map(
+            left: { $0 },
+            right: { InputStream(url: $0)?.readToEnd() ?? Data() }
+        )
+    }
+}
+
+extension InputStream {
     func readToEnd() -> Data {
         open()
         defer { close() }

WordPressKitTests/Utilities/URLSessionHelperTests.swift

@@ -177,6 +177,38 @@ class URLSessionHelperTests: XCTestCase {
             XCTFail("Unexpected result: \(result)")
         }
     }
+
+    func testMultipartForm() async throws {
+        var req: URLRequest?
+        stub(condition: isPath("/hello")) {
+            req = $0
+            return HTTPStubsResponse(data: "success".data(using: .utf8)!, statusCode: 200, headers: nil)
+        }
+
+        let builder = HTTPRequestBuilder(url: URL(string: "https://wordpress.org/hello")!)
+            .method(.post)
+            .body(form: [MultipartFormField(text: "value", name: "name", filename: nil)])
+
+        let _ = await URLSession.shared.perform(request: builder, errorType: TestError.self)
+
+        let request = try XCTUnwrap(req)
+        let boundary = try XCTUnwrap(
+            request
+                .value(forHTTPHeaderField: "Content-Type")?.split(separator: ";")
+                .map { $0.trimmingCharacters(in: .whitespaces) }
+                .reduce(into: [String: String]()) {
+                    let pair = $1.split(separator: "=")
+                    if pair.count == 2 {
+                        $0[String(pair[0])] = String(pair[1])
+                    }
+                }["boundary"]
+            )
+
+        let requestBody = try XCTUnwrap(request.httpBody ?? request.httpBodyStream?.readToEnd())
+
+        let expectedBody = "--\(boundary)\r\nContent-Disposition: form-data; name=\"name\"\r\n\r\nvalue\r\n--\(boundary)--\r\n"
+        XCTAssertEqual(String(data: requestBody, encoding: .utf8), expectedBody)
+    }
 }
 
 private enum TestError: LocalizedError, Equatable {