rastrillo / native Public

Reserve seal quota in blocks and cache account key state to avoid per-edit IPC

Paul Campbell pushed by paul@keymail.dev 1 month ago d91d10c190a995540818d3654bd91bcb1e5cee07 parent 4d40c6f

2 files changed, +51 −8

Sources/RastrilloApple/NativeSealedStore.swift +34 −7
Tests/RastrilloAppleTests/SealedStoreTests.swift +17 −1

diff --git a/Sources/RastrilloApple/NativeSealedStore.swift b/Sources/RastrilloApple/NativeSealedStore.swift

index 6d75f02..0b98ec4 100644

--- a/Sources/RastrilloApple/NativeSealedStore.swift

+++ b/Sources/RastrilloApple/NativeSealedStore.swift

@@ -92,6 +92,12 @@ public final class NativeSealedStore: BlobStore, @unchecked Sendable {

private let keys: KeyProvider

private let lock = NSLock()

private let fm = FileManager.default

+ // One store owns an account directory. Cached material remains protected on

+ // disk by completeFileProtection; cold opens still distinguish locked/lost.

+ private var cachedState: StoreKeyState?

+ private var cachedKey: SymmetricKey?

+ private var reservedUntil: UInt64 = 0

+ private var backupExcluded = false

var rotationLimit: UInt64 = 1 << 31

public init(root: URL, accountID: String, keys: KeyProvider) { self.root = root; self.accountID = accountID; self.keys = keys }

static let writeOptions: Data.WritingOptions = [.atomic, .completeFileProtection]

@@ -110,7 +116,12 @@ public final class NativeSealedStore: BlobStore, @unchecked Sendable {

return error.domain == NSCocoaErrorDomain && [NSFileReadNoSuchFileError, NSFileNoSuchFileError].contains(error.code)

}

private func io<T>(_ body: () throws -> T) throws -> T {

- do { return try body() } catch let error as BlobStoreError { throw error } catch { throw BlobStoreError.unavailable }

+ do { return try body() } catch {

+ // A failed write may have committed its reservation. Abandon its entire

+ // unused tail before retrying, including interrupted key rotation.

+ cachedState = nil; cachedKey = nil; reservedUntil = 0; backupExcluded = false

+ throw (error as? BlobStoreError) ?? .unavailable

+ }

}

private func hasFiles(_ directory: URL) throws -> Bool {

let children: [URL]

@@ -124,6 +135,7 @@ public final class NativeSealedStore: BlobStore, @unchecked Sendable {

return false

}

private func key(create: Bool) throws -> SymmetricKey? {

+ if let cachedKey { return cachedKey }

let (status, key) = keys.copyKey()

if status == errSecSuccess, let key { return key }

guard status == errSecItemNotFound else { throw BlobStoreError.unavailable }

@@ -143,18 +155,29 @@ public final class NativeSealedStore: BlobStore, @unchecked Sendable {

}

private func aad(_ path: String) -> Data { Data("\(accountID)|\(Self.formatVersion)|\(path)".utf8) }

private func readKeyState() throws -> StoreKeyState {

+ if let cachedState { return cachedState }

let (status, state) = keys.sealState()

guard status == errSecSuccess, let state else { throw BlobStoreError.unavailable }

+ cachedState = state; cachedKey = SymmetricKey(data: state.current)

+ reservedUntil = state.seals

return state

}

private func saveKeyState(_ state: StoreKeyState) throws {

- guard keys.saveSealState(state) == errSecSuccess else { throw BlobStoreError.unavailable }

+ var durable = state

+ durable.seals = max(state.seals, reservedUntil)

+ guard keys.saveSealState(durable) == errSecSuccess else { throw BlobStoreError.unavailable }

+ cachedState = state; cachedKey = SymmetricKey(data: state.current)

}

private func reserveSeal(_ state: inout StoreKeyState) throws {

guard state.seals < rotationLimit else { throw BlobStoreError.unavailable }

+ if state.seals == reservedUntil {

+ // Persist a high-water mark before using any slot. Relaunch starts at

+ // that mark, sacrificing the unused tail rather than reusing a quota.

+ reservedUntil = state.seals + min(1024, rotationLimit - state.seals)

+ try saveKeyState(state)

+ }

state.seals += 1

- // Count before sealing: a kill can overcount, but can never reuse a quota.

- try saveKeyState(state)

+ cachedState = state

}

private func open(_ data: Data, path: String, state: StoreKeyState) throws -> Data {

let box = try AES.GCM.SealedBox(combined: data)

@@ -176,6 +199,7 @@ public final class NativeSealedStore: BlobStore, @unchecked Sendable {

guard state.previous != nil || state.seals >= rotationLimit else { return }

if state.previous == nil {

state = StoreKeyState(current: SymmetricKey(size: .bits256).withUnsafeBytes { Data($0) }, previous: state.current)

+ reservedUntil = 0

try saveKeyState(state)

}

for file in try files(root) {

@@ -200,10 +224,12 @@ public final class NativeSealedStore: BlobStore, @unchecked Sendable {

var state = try readKeyState()

try rotateIfNeeded(&state)

try reserveSeal(&state)

- let key = SymmetricKey(data: state.current)

+ let key = cachedKey!

try fm.createDirectory(at: target.deletingLastPathComponent(), withIntermediateDirectories: true)

- var directory = root; var values = URLResourceValues(); values.isExcludedFromBackup = true

- try directory.setResourceValues(values)

+ if !backupExcluded {

+ var directory = root; var values = URLResourceValues(); values.isExcludedFromBackup = true

+ try directory.setResourceValues(values); backupExcluded = true

+ }

let sealed = try AES.GCM.seal(data, using: key, authenticating: aad(path))

guard let combined = sealed.combined else { throw BlobStoreError.unavailable }

try combined.write(to: target, options: Self.writeOptions)

@@ -244,6 +270,7 @@ public final class NativeSealedStore: BlobStore, @unchecked Sendable {

try lock.withLock { try io {

let status = keys.deleteKey()

guard status == errSecSuccess || status == errSecItemNotFound else { throw BlobStoreError.unavailable }

+ cachedState = nil; cachedKey = nil; reservedUntil = 0; backupExcluded = false

do { try fm.removeItem(at: root) } catch { if !missing(error) { throw error } }

} }

}

diff --git a/Tests/RastrilloAppleTests/SealedStoreTests.swift b/Tests/RastrilloAppleTests/SealedStoreTests.swift

index 63a0d15..bfbd049 100644

--- a/Tests/RastrilloAppleTests/SealedStoreTests.swift

+++ b/Tests/RastrilloAppleTests/SealedStoreTests.swift

@@ -172,6 +172,22 @@ final class SealedStoreTests: XCTestCase {

let lane = await relaunched.entries(lane: "doc:1")

XCTAssertEqual(lane.map(\.id), [entry.id])

}

+ func testSealBlocksCacheStateAndSkipUnusedQuotaAfterCrash() throws {

+ let keys = ScriptedKeyProvider()

+ var nonces = Set<Data>()

+ for launch in 1...2 {

+ let sealed = store(keys: keys)

+ for index in 0..<19 {

+ try sealed.write(Data([UInt8(index)]), to: "cache/a")

+ let bytes = try Data(contentsOf: file("account-a", "cache/a"))

+ XCTAssertTrue(nonces.insert(Data(bytes.prefix(12))).inserted, "A random 96-bit nonce must not repeat")

+ XCTAssertEqual(keys.sealState().1?.seals, UInt64(launch * 1024), "Persist the entire block before its first seal; a relaunch skips the unused tail")

+ XCTAssertEqual(try sealed.read("cache/a"), Data([UInt8(index)]))

+ }

+ XCTAssertEqual(keys.log.filter { $0 == "copy" }.count, launch, "Cache the key across reads and writes")

+ }

func testSealQuotaRotatesAndResealsInPlace() throws {

let keys = ScriptedKeyProvider()

let sealed = store(keys: keys)

@@ -183,7 +199,7 @@ final class SealedStoreTests: XCTestCase {

XCTAssertEqual(try sealed.read("cache/a"), Data("A".utf8))

XCTAssertEqual(try sealed.read("cache/b"), Data("B".utf8))

XCTAssertNil(keys.sealState().1?.previous)

- XCTAssertEqual(keys.sealState().1?.seals, 2, "One reseal and the requested write")

+ XCTAssertEqual(keys.sealState().1?.seals, 8, "Reserve a block, capped at this test key’s rotation limit")

let bytes = try Data(contentsOf: file("account-a", "cache/a"))

XCTAssertThrowsError(try AES.GCM.open(AES.GCM.SealedBox(combined: bytes), using: old, authenticating: Data("account-a|1|cache/a".utf8)))

}