Stream file digest hashing

Sources/tart/OCI/Digest.swift

@@ -7,6 +7,8 @@ enum DigestError: Error {
 }
 
 class Digest {
+  private static let fileBufferSize = 4 * 1024 * 1024
+
   var hash: SHA256 = SHA256()
 
   func update(_ data: Data) {
@@ -22,7 +24,10 @@ class Digest {
   }
 
   static func hash(_ url: URL) throws -> String {
-    hash(try Data(contentsOf: url))
+    let file = try FileHandle(forReadingFrom: url)
+    defer { try? file.close() }
+
+    return try hashContents(from: file)
   }
 
   static func hash(_ url: URL, offset: UInt64, size: UInt64) throws -> String {
@@ -36,20 +41,53 @@ class Digest {
       throw DigestError.InvalidOffset
     }
 
-    if (offset + size) > fileSize {
+    if size > fileSize - offset {
       throw DigestError.InvalidSize
     }
 
-    // Read a chunk of size ``size`` at offset ``offset``
-    // and calculate it's digest
+    // Read the requested range incrementally and calculate its digest.
     let fh = try FileHandle(forReadingFrom: url)
-    defer { try! fh.close() }
+    defer { try? fh.close() }
 
     try fh.seek(toOffset: offset)
 
-    let data = try fh.read(upToCount: Int(size))!
+    return try hashContents(from: fh, size: size)
+  }
 
-    return hash(data)
+  /// Streams a file into SHA-256 while keeping Foundation's temporary read
+  /// buffers scoped to one chunk.
+  private static func hashContents(from file: FileHandle, size: UInt64? = nil) throws -> String {
+    let digest = Digest()
+    var remaining = size
+
+    while remaining.map({ $0 > 0 }) ?? true {
+      let didRead = try autoreleasepool { () throws -> Bool in
+        let count = remaining.map {
+          Int(min(UInt64(fileBufferSize), $0))
+        } ?? fileBufferSize
+
+        guard let data = try file.read(upToCount: count), !data.isEmpty else {
+          if remaining != nil {
+            throw DigestError.InvalidSize
+          }
+
+          return false
+        }
+
+        digest.update(data)
+        if let bytesRemaining = remaining {
+          remaining = bytesRemaining - UInt64(data.count)
+        }
+
+        return true
+      }
+
+      if !didRead {
+        break
+      }
+    }
+
+    return digest.finalize()
   }
 }
 

Tests/TartTests/DigestTests.swift

@@ -1,3 +1,4 @@
+import Foundation
 import XCTest
 @testable import tart
 
@@ -21,4 +22,34 @@ final class DigestTests: XCTestCase {
 
     XCTAssertEqual(Digest.hash(data), "sha256:d7a8fbb307d7809469ca9abcb0082e4f8d5651e46d3cdb762d02d0bf37c9e592")
   }
+
+  func testFileAndRangeHashingMatchDataHashing() throws {
+    let prefix = Data(repeating: 0x61, count: 4 * 1024 * 1024 + 17)
+    let range = Data("range".utf8)
+    let suffix = Data(repeating: 0x62, count: 23)
+    let data = prefix + range + suffix
+    let url = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
+    try data.write(to: url)
+    defer { try? FileManager.default.removeItem(at: url) }
+
+    XCTAssertEqual(try Digest.hash(url), Digest.hash(data))
+    XCTAssertEqual(try Digest.hash(url, offset: UInt64(prefix.count), size: UInt64(range.count)), Digest.hash(range))
+  }
+
+  func testRangeHashingRejectsOutOfBoundsRanges() throws {
+    let url = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
+    try Data("range".utf8).write(to: url)
+    defer { try? FileManager.default.removeItem(at: url) }
+
+    XCTAssertThrowsError(try Digest.hash(url, offset: 6, size: 0)) { error in
+      guard case DigestError.InvalidOffset = error else {
+        return XCTFail("unexpected error: \(error)")
+      }
+    }
+    XCTAssertThrowsError(try Digest.hash(url, offset: 1, size: UInt64.max)) { error in
+      guard case DigestError.InvalidSize = error else {
+        return XCTFail("unexpected error: \(error)")
+      }
+    }
+  }
 }