Merge pull request #156 from kishikawakatsumi/Produkt-download-at-url…

…-with-progress

Improve download functions

Sources/SMBClient/FileReader.swift

@@ -42,7 +42,7 @@ public class FileReader {
     return buffer
   }
 
-  public func download() async throws -> Data {
+  public func download(progressHandler: (_ progress: Double) -> Void = { _ in }) async throws -> Data {
     let fileProxy = try await fileProxy()
 
     var offset: UInt64 = 0
@@ -57,11 +57,61 @@ public class FileReader {
 
       buffer.append(response.buffer)
       offset = UInt64(buffer.count)
+      progressHandler(Double(offset) / Double(fileProxy.size))
     } while NTStatus(response.header.status) != .endOfFile && buffer.count < fileProxy.size
 
+    progressHandler(1.0)
     return buffer
   }
 
+  public func download(to localPath: URL, overwrite: Bool = false, progressHandler: (_ progress: Double) -> Void = { _ in }) async throws {
+    let fileManger = FileManager()
+    let fileExists = fileManger.fileExists(atPath: localPath.path)
+
+    if fileExists {
+      if overwrite {
+        try Data().write(to: localPath)
+      } else {
+        throw CocoaError(.fileWriteFileExists)
+      }
+    }
+
+    if !fileExists {
+      try Data().write(to: localPath)
+    }
+    guard let fileHandle = FileHandle(forWritingAtPath: localPath.path) else {
+      throw URLError(.cannotWriteToFile)
+    }
+
+    defer {
+      fileHandle.closeFile()
+    }
+
+    try await download(fileHandle: fileHandle, progressHandler: progressHandler)
+  }
+
+  public func download(fileHandle: FileHandle, progressHandler: (_ progress: Double) -> Void = { _ in }) async throws {
+    let fileProxy = try await fileProxy()
+
+    var offset: UInt64 = 0
+    var response: Read.Response
+    repeat {
+      response = try await session.read(
+        fileId: fileProxy.id,
+        offset: offset
+      )
+
+      fileHandle.seekToEndOfFile()
+      fileHandle.write(response.buffer)
+      offset += UInt64(response.buffer.count)
+
+      let progress = Double(offset) / Double(fileProxy.size)
+      progressHandler(progress)
+    } while NTStatus(response.header.status) != .endOfFile && offset < fileProxy.size
+
+    progressHandler(1.0)
+  }
+
   public func close() async throws {
     if let createResponse {
       try await session.close(fileId: createResponse.fileId)

Sources/SMBClient/SMBClient.swift

@@ -117,14 +117,25 @@ public class SMBClient {
   }
 
   public func download(path: String) async throws -> Data {
+    return try await download(path: path, progressHandler: { _ in })
+  }
+
+  public func download(path: String, progressHandler: (_ progress: Double) -> Void) async throws -> Data {
     let fileReader = fileReader(path: Pathname.normalize(path))
 
-    let data = try await fileReader.download()
+    let data = try await fileReader.download(progressHandler: progressHandler)
     try await fileReader.close()
 
     return data
   }
 
+  public func download(path: String, localPath: URL, overwrite: Bool = false, progressHandler: (_ progress: Double) -> Void = { _ in }) async throws {
+    let fileReader = fileReader(path: Pathname.normalize(path))
+
+    try await fileReader.download(to: localPath, overwrite: overwrite, progressHandler: progressHandler)
+    try await fileReader.close()
+  }
+
   public func upload(content: Data, path: String) async throws {
     try await upload(content: content, path: Pathname.normalize(path), progressHandler: { _ in })
   }

Tests/SMBClientTests/SMBClientTests.swift

@@ -379,6 +379,58 @@ final class SMBClientTests: XCTestCase {
     XCTAssertEqual(data, try Data(contentsOf: fixtureURL.appending(component: "\(user.sharePath)/\(path)")))
   }
 
+  func testDownloadIntoFile01() async throws {
+    let user = bob
+    let client = SMBClient(host: "localhost", port: 4445)
+    try await client.login(username: user.username, password: user.password)
+    try await client.connectShare(user.share)
+
+    let path = "test_files/file_example_JPG_1MB.jpg"
+
+    let fileManager = FileManager()
+    let tempFolder = URL(fileURLWithPath: NSTemporaryDirectory(), isDirectory: true)
+
+    let destinationFile = tempFolder.appending(path: "downloadedfile.jpg", directoryHint: .notDirectory)
+
+    var progressWasUpdated = false
+    try await client.download(path: path, localPath: destinationFile, overwrite: true) { (progress) in
+      progressWasUpdated = true
+    }
+
+    XCTAssertTrue(fileManager.fileExists(atPath: destinationFile.path))
+    let data = try Data(contentsOf: destinationFile)
+    XCTAssertEqual(data, try Data(contentsOf: fixtureURL.appending(component: "\(user.sharePath)/\(path)")))
+    XCTAssertTrue(progressWasUpdated)
+  }
+
+  func testDownloadIntoFile02() async throws {
+    let user = bob
+    let client = SMBClient(host: "localhost", port: 4445)
+    try await client.login(username: user.username, password: user.password)
+    try await client.connectShare(user.share)
+
+    let path = "test_files/file_example_MP4_1920_18MG.mp4"
+
+    let fileManager = FileManager()
+    let tempFolder = URL(fileURLWithPath: NSTemporaryDirectory(), isDirectory: true)
+    print(tempFolder)
+    let destinationFile = tempFolder.appending(path: "downloadedfile.mp4", directoryHint: .notDirectory)
+
+    try Data().write(to: destinationFile)
+    let fileHandle = try FileHandle(forWritingTo: destinationFile)
+
+    var progressWasUpdated = false
+    let fileReader = client.fileReader(path: path)
+    try await fileReader.download(fileHandle: fileHandle) { (progress) in
+      progressWasUpdated = true
+    }
+
+    XCTAssertTrue(fileManager.fileExists(atPath: destinationFile.path))
+    let data = try Data(contentsOf: destinationFile)
+    XCTAssertEqual(data, try Data(contentsOf: fixtureURL.appending(component: "\(user.sharePath)/\(path)")))
+    XCTAssertTrue(progressWasUpdated)
+  }
+
   func testRandomRead01() async throws {
     let user = bob
     let client = SMBClient(host: "localhost", port: 4445)