From 49114173571201dd344dc1c5a75b4b8aa73e3025 Mon Sep 17 00:00:00 2001 From: Ayush7614 Date: Fri, 31 Jul 2026 17:11:20 +0530 Subject: [PATCH] Release held quota byte reservations on cancel, failure, and panic. reservingBytes was call-scoped only; multi-step writers now hold QuotaByteReservation headroom and release it when compression/write fails, capture is canceled, or panic wipe clears in-flight state. --- bitchat/Features/media/ImageUtils.swift | 45 +++++++--- .../Features/voice/VoiceCaptureSession.swift | 48 +++++++++-- bitchat/Features/voice/VoiceRecorder.swift | 23 ++++- .../Services/BLE/BLEIncomingFileStore.swift | 79 ++++++++++++++++-- .../ViewModels/ChatLiveVoiceCoordinator.swift | 35 +++++++- ...EIncomingFileStoreOutgoingQuotaTests.swift | 83 +++++++++++++++++++ 6 files changed, 284 insertions(+), 29 deletions(-) diff --git a/bitchat/Features/media/ImageUtils.swift b/bitchat/Features/media/ImageUtils.swift index 62a6de12..a28dfdb5 100644 --- a/bitchat/Features/media/ImageUtils.swift +++ b/bitchat/Features/media/ImageUtils.swift @@ -73,10 +73,24 @@ enum ImageUtils { } } - let outputURL = try makeOutputURL( - outputDirectory: outputDirectory, - reservingBytes: jpegData.count - ) + let outputURL = try makeOutputURL(outputDirectory: outputDirectory) + let reservation: BLEIncomingFileStore.QuotaByteReservation? + if outputDirectory == nil { + // Hold headroom only for the default Application Support tree; + // test / custom directories skip quota. Release even if the + // write throws so a failed encode→disk handoff cannot leak. + reservation = BLEIncomingFileStore.shared.reserveQuotaBytes( + jpegData.count, + scope: .outgoing + ) + } else { + reservation = nil + } + defer { + if let reservation { + BLEIncomingFileStore.shared.releaseQuotaReservation(reservation) + } + } try jpegData.write(to: outputURL, options: .atomic) return outputURL } @@ -154,10 +168,21 @@ enum ImageUtils { } } } - let outputURL = try makeOutputURL( - outputDirectory: outputDirectory, - reservingBytes: jpegData.count - ) + let outputURL = try makeOutputURL(outputDirectory: outputDirectory) + let reservation: BLEIncomingFileStore.QuotaByteReservation? + if outputDirectory == nil { + reservation = BLEIncomingFileStore.shared.reserveQuotaBytes( + jpegData.count, + scope: .outgoing + ) + } else { + reservation = nil + } + defer { + if let reservation { + BLEIncomingFileStore.shared.releaseQuotaReservation(reservation) + } + } try jpegData.write(to: outputURL, options: .atomic) return outputURL } @@ -201,7 +226,7 @@ enum ImageUtils { } #endif - private static func makeOutputURL(outputDirectory: URL? = nil, reservingBytes: Int = 0) throws -> URL { + private static func makeOutputURL(outputDirectory: URL? = nil) throws -> URL { let formatter = DateFormatter() formatter.dateFormat = "yyyyMMdd_HHmmss" let fileName = "img_\(formatter.string(from: Date()))_\(UUID().uuidString).jpg" @@ -212,7 +237,7 @@ enum ImageUtils { } else { // Default Application Support outgoing tree shares the process-wide // store with BLE deletion/delivery so quota exclusions stay live. - BLEIncomingFileStore.shared.enforceOutgoingQuota(reservingBytes: reservingBytes) + // Callers reserve/release quota bytes around the write itself. directory = try applicationFilesDirectory().appendingPathComponent("images/outgoing", isDirectory: true) } try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true, attributes: BLEIncomingFileStore.mediaProtectionAttributes) diff --git a/bitchat/Features/voice/VoiceCaptureSession.swift b/bitchat/Features/voice/VoiceCaptureSession.swift index a6a6d867..6d483767 100644 --- a/bitchat/Features/voice/VoiceCaptureSession.swift +++ b/bitchat/Features/voice/VoiceCaptureSession.swift @@ -100,6 +100,9 @@ final class PTTLiveVoiceSession: VoiceCaptureSession { private let stream: StreamState private var startDate: Date? private var completed = false + /// Held outgoing quota headroom for this capture; released on finish, + /// cancel, start failure, and panic cancel. + private var quotaReservation: BLEIncomingFileStore.QuotaByteReservation? var isLive: Bool { true } @@ -124,7 +127,8 @@ final class PTTLiveVoiceSession: VoiceCaptureSession { } func start() async throws { - let outputURL = try Self.makeOutputURL(burstID: burstID) + let (outputURL, reservation) = try Self.makeOutputURL(burstID: burstID) + quotaReservation = reservation let sendPacket = sendPacket let stream = stream capture.onFrames = { frames in @@ -158,6 +162,8 @@ final class PTTLiveVoiceSession: VoiceCaptureSession { // handed its token back — nothing to retry. A coordinator-side // interruption during handoff also cancels acquire, but that is // not a successful start and must propagate to the view model. + releaseQuotaReservation() + try? FileManager.default.removeItem(at: outputURL) guard completed else { throw CancellationError() } return } catch { @@ -171,9 +177,17 @@ final class PTTLiveVoiceSession: VoiceCaptureSession { // user let go, so bail instead of opening a hot mic. guard !completed else { capture.cancel() + releaseQuotaReservation() + try? FileManager.default.removeItem(at: outputURL) return } - try await capture.start(outputURL: outputURL) + do { + try await capture.start(outputURL: outputURL) + } catch { + releaseQuotaReservation() + try? FileManager.default.removeItem(at: outputURL) + throw error + } } startDate = now() SecureLogger.info("PTT: live burst \(burstID.hexEncodedString()) capture started", category: .session) @@ -182,6 +196,7 @@ final class PTTLiveVoiceSession: VoiceCaptureSession { func finish() async -> URL? { guard !completed else { return nil } completed = true + defer { releaseQuotaReservation() } let elapsed = startDate.map { now().timeIntervalSince($0) } ?? 0 let (url, encodedFrames) = capture.stop() @@ -217,6 +232,7 @@ final class PTTLiveVoiceSession: VoiceCaptureSession { // pause), and only capture.cancel() stops the mic and deactivates the // session. It is idempotent, so a redundant call is harmless. capture.cancel() + releaseQuotaReservation() if !alreadyCompleted { sendControlPacket(.canceled) } @@ -227,6 +243,7 @@ final class PTTLiveVoiceSession: VoiceCaptureSession { // conversation data racing the emergency transport reset. completed = true capture.cancel() + releaseQuotaReservation() } private func sendControlPacket(_ kind: VoiceBurstPacket.Kind) { @@ -234,20 +251,37 @@ final class PTTLiveVoiceSession: VoiceCaptureSession { sendPacket(packet.encode()) } - private static func makeOutputURL(burstID: Data) throws -> URL { + private func releaseQuotaReservation() { + guard let quotaReservation else { return } + BLEIncomingFileStore.shared.releaseQuotaReservation(quotaReservation) + self.quotaReservation = nil + } + + private static func makeOutputURL(burstID: Data) throws -> (URL, BLEIncomingFileStore.QuotaByteReservation) { let base = try FileManager.default.url( for: .applicationSupportDirectory, in: .userDomainMask, appropriateFor: nil, create: true ) - BLEIncomingFileStore.shared.enforceOutgoingQuota( - reservingBytes: FileTransferLimits.maxVoiceNoteBytes + let reservation = BLEIncomingFileStore.shared.reserveQuotaBytes( + FileTransferLimits.maxVoiceNoteBytes, + scope: .outgoing ) let directory = base .appendingPathComponent("files", isDirectory: true) .appendingPathComponent("voicenotes/outgoing", isDirectory: true) - try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true, attributes: BLEIncomingFileStore.mediaProtectionAttributes) - return directory.appendingPathComponent("voice_\(burstID.hexEncodedString()).m4a") + do { + try FileManager.default.createDirectory( + at: directory, + withIntermediateDirectories: true, + attributes: BLEIncomingFileStore.mediaProtectionAttributes + ) + } catch { + BLEIncomingFileStore.shared.releaseQuotaReservation(reservation) + throw error + } + let url = directory.appendingPathComponent("voice_\(burstID.hexEncodedString()).m4a") + return (url, reservation) } } diff --git a/bitchat/Features/voice/VoiceRecorder.swift b/bitchat/Features/voice/VoiceRecorder.swift index 049bea2e..f82f58a5 100644 --- a/bitchat/Features/voice/VoiceRecorder.swift +++ b/bitchat/Features/voice/VoiceRecorder.swift @@ -73,6 +73,9 @@ actor VoiceRecorder { /// True only while `startRecording()` is suspended in session acquire. /// A second start is rejected instead of superseding the first one. private var startInFlight = false + /// Held outgoing quota headroom for the in-flight capture; released on + /// stop, cancel, start failure, and panic cancel. + private var quotaReservation: BLEIncomingFileStore.QuotaByteReservation? init( sessionCoordinator: AudioSessionCoordinator = .shared, @@ -171,6 +174,7 @@ actor VoiceRecorder { return newURL } catch { releaseSessionToken() + releaseQuotaReservation() recorder = nil currentURL = nil activeOwner = nil @@ -189,12 +193,14 @@ actor VoiceRecorder { if startInFlight { activeOwner = nil startInFlight = false + releaseQuotaReservation() return nil } guard let activeRecorder = recorder else { let sessionURL = currentURL releaseSessionToken() + releaseQuotaReservation() currentURL = nil activeOwner = nil return sessionURL @@ -221,6 +227,7 @@ actor VoiceRecorder { activeRecorder.stop() } releaseSessionToken() + releaseQuotaReservation() self.recorder = nil currentURL = nil activeOwner = nil @@ -240,6 +247,7 @@ actor VoiceRecorder { recorder.stop() } releaseSessionToken() + releaseQuotaReservation() if let currentURL { try? FileManager.default.removeItem(at: currentURL) } @@ -303,8 +311,13 @@ actor VoiceRecorder { if let outputDirectory { baseDirectory = outputDirectory } else { - BLEIncomingFileStore.shared.enforceOutgoingQuota( - reservingBytes: FileTransferLimits.maxVoiceNoteBytes + // Reserve worst-case note size for the whole capture. Released on + // stop / cancel / start failure so a abandoned hold cannot keep + // the outgoing budget permanently tighter. + releaseQuotaReservation() + quotaReservation = BLEIncomingFileStore.shared.reserveQuotaBytes( + FileTransferLimits.maxVoiceNoteBytes, + scope: .outgoing ) baseDirectory = try applicationFilesDirectory() .appendingPathComponent("voicenotes/outgoing", isDirectory: true) @@ -317,6 +330,12 @@ actor VoiceRecorder { return baseDirectory.appendingPathComponent(fileName) } + private func releaseQuotaReservation() { + guard let quotaReservation else { return } + BLEIncomingFileStore.shared.releaseQuotaReservation(quotaReservation) + self.quotaReservation = nil + } + private func applicationFilesDirectory() throws -> URL { #if os(iOS) return try FileManager.default.url(for: .applicationSupportDirectory, in: .userDomainMask, appropriateFor: nil, create: true) diff --git a/bitchat/Services/BLE/BLEIncomingFileStore.swift b/bitchat/Services/BLE/BLEIncomingFileStore.swift index bcdd979b..e4c6ac45 100644 --- a/bitchat/Services/BLE/BLEIncomingFileStore.swift +++ b/bitchat/Services/BLE/BLEIncomingFileStore.swift @@ -107,10 +107,22 @@ struct BLEIncomingFileStore: @unchecked Sendable { fileprivate let id: UUID } + /// Held headroom for a multi-step write (image encode→disk, voice + /// capture, live PTT). Unlike the call-scoped `reservingBytes` argument + /// to `enforceQuota`, this stays charged against the quota until + /// `releaseQuotaReservation` — so cancel / encode failure / panic must + /// release it or later writers see a permanently tighter budget. + struct QuotaByteReservation: Sendable { + fileprivate let id: UUID + fileprivate let scope: MediaQuotaScope + fileprivate let bytes: Int64 + } + private final class PayloadCoordination: @unchecked Sendable { let lock = NSLock() var pendingDeliveryPaths: Set = [] var deletionReservations: [UUID: Set] = [:] + var byteReservations: [UUID: (scope: MediaQuotaScope, bytes: Int64)] = [:] } private static let defaultQuotaBytes: Int64 = 100 * 1024 * 1024 @@ -247,6 +259,9 @@ struct BLEIncomingFileStore: @unchecked Sendable { payloadCoordination.deletionReservations.removeAll( keepingCapacity: false ) + payloadCoordination.byteReservations.removeAll( + keepingCapacity: false + ) payloadCoordination.lock.unlock() } @@ -558,11 +573,17 @@ struct BLEIncomingFileStore: @unchecked Sendable { } /// Frees least-recently-modified files in `scope` until `reservingBytes` - /// fits under the quota. Files named `voice_live_*` (in-flight live - /// captures) are never evicted regardless of who triggers enforcement — - /// a finalized transfer can arrive at quota while a burst is still - /// streaming — but they still count toward usage. Paths reserved for an - /// in-flight delivery or private-media deletion are also skipped. + /// (plus any held `QuotaByteReservation`s for that scope) fits under the + /// quota. `reservingBytes` itself is call-scoped — it is not stored — + /// and is the right tool for a synchronous write that completes before + /// the next eviction. Multi-step writers should use + /// `reserveQuotaBytes` / `releaseQuotaReservation` instead. + /// + /// Files named `voice_live_*` (in-flight live captures) are never + /// evicted regardless of who triggers enforcement — a finalized + /// transfer can arrive at quota while a burst is still streaming — but + /// they still count toward usage. Paths reserved for an in-flight + /// delivery or private-media deletion are also skipped. func enforceQuota(reservingBytes: Int) { enforceQuota(reservingBytes: reservingBytes, scope: .incoming) } @@ -576,7 +597,49 @@ struct BLEIncomingFileStore: @unchecked Sendable { func enforceQuota(reservingBytes: Int, scope: MediaQuotaScope) { payloadCoordination.lock.lock() defer { payloadCoordination.lock.unlock() } + enforceQuotaLocked(reservingBytes: reservingBytes, scope: scope) + } + /// Evicts enough room for `bytes`, then holds that headroom against the + /// scope's quota until `releaseQuotaReservation`. Callers must release + /// on success, cancel, encode/write failure, and panic (panic also + /// clears every outstanding reservation). + @discardableResult + func reserveQuotaBytes(_ bytes: Int, scope: MediaQuotaScope) -> QuotaByteReservation { + payloadCoordination.lock.lock() + defer { payloadCoordination.lock.unlock() } + + let clamped = Int64(max(0, bytes)) + enforceQuotaLocked(reservingBytes: Int(clamped), scope: scope) + let reservation = QuotaByteReservation( + id: UUID(), + scope: scope, + bytes: clamped + ) + payloadCoordination.byteReservations[reservation.id] = ( + scope: scope, + bytes: clamped + ) + return reservation + } + + func releaseQuotaReservation(_ reservation: QuotaByteReservation) { + payloadCoordination.lock.lock() + defer { payloadCoordination.lock.unlock() } + payloadCoordination.byteReservations.removeValue(forKey: reservation.id) + } + + /// Test seam: bytes currently held by `reserveQuotaBytes` for `scope`. + func reservedQuotaBytes(for scope: MediaQuotaScope) -> Int64 { + payloadCoordination.lock.lock() + defer { payloadCoordination.lock.unlock() } + return payloadCoordination.byteReservations.values.reduce(into: Int64(0)) { + guard $1.scope == scope else { return } + $0 += $1.bytes + } + } + + private func enforceQuotaLocked(reservingBytes: Int, scope: MediaQuotaScope) { do { let base = try filesDirectory() let dirs = scope.subdirectories.map { @@ -599,8 +662,12 @@ struct BLEIncomingFileStore: @unchecked Sendable { } } + let heldBytes = payloadCoordination.byteReservations.values.reduce(into: Int64(0)) { + guard $1.scope == scope else { return } + $0 += $1.bytes + } let currentUsage = allFiles.reduce(0) { $0 + $1.size } - let targetUsage = quotaBytes - Int64(reservingBytes) + let targetUsage = quotaBytes - Int64(reservingBytes) - heldBytes guard currentUsage > targetUsage else { return } let needToFree = currentUsage - targetUsage diff --git a/bitchat/ViewModels/ChatLiveVoiceCoordinator.swift b/bitchat/ViewModels/ChatLiveVoiceCoordinator.swift index dd3bfbac..6f6e7ad8 100644 --- a/bitchat/ViewModels/ChatLiveVoiceCoordinator.swift +++ b/bitchat/ViewModels/ChatLiveVoiceCoordinator.swift @@ -113,10 +113,22 @@ final class ChatLiveVoiceCoordinator { var player: PTTBurstPlayer? var idleTimeout: Task? var gapRedrain: Task? + /// Incoming quota headroom held for this live capture; released on + /// finalize / cancel so a dropped burst cannot keep the budget tight. + var quotaReservation: BLEIncomingFileStore.QuotaByteReservation? var key: AssemblyKey { AssemblyKey(peerID: peerID, scope: scope, burstID: burstID) } - init(burstID: Data, peerID: PeerID, scope: VoiceBurstScope, nickname: String, message: BitchatMessage, fileURL: URL, fileHandle: FileHandle) { + init( + burstID: Data, + peerID: PeerID, + scope: VoiceBurstScope, + nickname: String, + message: BitchatMessage, + fileURL: URL, + fileHandle: FileHandle, + quotaReservation: BLEIncomingFileStore.QuotaByteReservation? + ) { self.burstID = burstID self.peerID = peerID self.scope = scope @@ -124,6 +136,7 @@ final class ChatLiveVoiceCoordinator { self.message = message self.fileURL = fileURL self.fileHandle = fileHandle + self.quotaReservation = quotaReservation self.firstPacketAt = Date() } } @@ -349,10 +362,14 @@ final class ChatLiveVoiceCoordinator { return nil } // BCH-01-002: live captures share the incoming-media quota with - // finalized transfers; reserve the burst's worst case up front. + // finalized transfers; hold the burst's worst case until finalize + // or cancel so a dropped assembly cannot leak headroom. // Eviction skips voice_live_* names, so partials still streaming in // are safe no matter which caller triggers enforcement. - fileStore.enforceQuota(reservingBytes: TransportConfig.pttMaxBurstBytes) + let reservation = fileStore.reserveQuotaBytes( + TransportConfig.pttMaxBurstBytes, + scope: .incoming + ) fileManager.createFile( atPath: fileURL.path, contents: nil, @@ -360,6 +377,7 @@ final class ChatLiveVoiceCoordinator { ) guard let handle = try? FileHandle(forWritingTo: fileURL) else { SecureLogger.error("PTT: cannot open capture file for burst \(burstID.hexEncodedString())", category: .session) + fileStore.releaseQuotaReservation(reservation) try? fileManager.removeItem(at: fileURL) return nil } @@ -383,7 +401,8 @@ final class ChatLiveVoiceCoordinator { nickname: nickname, message: message, fileURL: fileURL, - fileHandle: handle + fileHandle: handle, + quotaReservation: reservation ) // DM bubbles ride the full inbound pipeline (store append, unread, @@ -520,6 +539,7 @@ final class ChatLiveVoiceCoordinator { assembly.fileHandle = nil assemblies.removeValue(forKey: assembly.key) updatePublicTalkerIndicator() + releaseQuotaReservation(of: assembly) guard assembly.deliveredFrames > 0 else { // Nothing audible ever arrived — drop the empty bubble. @@ -628,12 +648,19 @@ final class ChatLiveVoiceCoordinator { assembly.fileHandle = nil assemblies.removeValue(forKey: assembly.key) updatePublicTalkerIndicator() + releaseQuotaReservation(of: assembly) removeBubble(of: assembly) WaveformCache.shared.purge(url: assembly.fileURL) try? fileManager.removeItem(at: assembly.fileURL) context.notifyUIChanged() } + private func releaseQuotaReservation(of assembly: Assembly) { + guard let reservation = assembly.quotaReservation else { return } + fileStore.releaseQuotaReservation(reservation) + assembly.quotaReservation = nil + } + // MARK: - Timers private func rescheduleIdleTimeout(for assembly: Assembly) { diff --git a/bitchatTests/Services/BLEIncomingFileStoreOutgoingQuotaTests.swift b/bitchatTests/Services/BLEIncomingFileStoreOutgoingQuotaTests.swift index e76e20ae..a6ae7481 100644 --- a/bitchatTests/Services/BLEIncomingFileStoreOutgoingQuotaTests.swift +++ b/bitchatTests/Services/BLEIncomingFileStoreOutgoingQuotaTests.swift @@ -168,4 +168,87 @@ struct BLEIncomingFileStoreOutgoingQuotaTests { #expect(!FileManager.default.fileExists(atPath: outgoingOld.path)) #expect(FileManager.default.fileExists(atPath: outgoingNew.path)) } + + @Test func quotaByteReservationReleasesOnExplicitRelease() throws { + let (store, _, cleanup) = try makeTempStore() + defer { cleanup() } + + let reservation = store.reserveQuotaBytes( + 12 * 1024 * 1024, + scope: .outgoing + ) + #expect(store.reservedQuotaBytes(for: .outgoing) == 12 * 1024 * 1024) + #expect(store.reservedQuotaBytes(for: .incoming) == 0) + + store.releaseQuotaReservation(reservation) + #expect(store.reservedQuotaBytes(for: .outgoing) == 0) + } + + @Test func quotaByteReservationReleasedAfterFailedWritePattern() throws { + // Mirrors ImageUtils: reserve → write throws → defer release must + // leave no held headroom behind. + let (store, _, cleanup) = try makeTempStore() + defer { cleanup() } + + do { + let reservation = store.reserveQuotaBytes(8 * 1024 * 1024, scope: .outgoing) + defer { store.releaseQuotaReservation(reservation) } + #expect(store.reservedQuotaBytes(for: .outgoing) == 8 * 1024 * 1024) + throw CocoaError(.fileWriteUnknown) + } catch { + #expect(store.reservedQuotaBytes(for: .outgoing) == 0) + } + } + + @Test func panicWipeClearsInFlightQuotaByteReservations() throws { + let (store, root, cleanup) = try makeTempStore() + defer { cleanup() } + + let outgoingReservation = store.reserveQuotaBytes( + 20 * 1024 * 1024, + scope: .outgoing + ) + let incomingReservation = store.reserveQuotaBytes( + 15 * 1024 * 1024, + scope: .incoming + ) + #expect(store.reservedQuotaBytes(for: .outgoing) == 20 * 1024 * 1024) + #expect(store.reservedQuotaBytes(for: .incoming) == 15 * 1024 * 1024) + + // Seed a file so panicWipe has a media tree to rebuild. + try writeBytes( + 1024, + to: root.appendingPathComponent("files/images/outgoing/seed.jpg"), + modified: Date() + ) + + try store.panicWipe() + + #expect(store.reservedQuotaBytes(for: .outgoing) == 0) + #expect(store.reservedQuotaBytes(for: .incoming) == 0) + // Stale tokens must not resurrect pre-panic headroom. + store.releaseQuotaReservation(outgoingReservation) + store.releaseQuotaReservation(incomingReservation) + #expect(store.reservedQuotaBytes(for: .outgoing) == 0) + #expect(store.reservedQuotaBytes(for: .incoming) == 0) + } + + @Test func heldQuotaByteReservationTightensEvictionTarget() throws { + let (store, root, cleanup) = try makeTempStore() + defer { cleanup() } + + let oldURL = root.appendingPathComponent("files/images/outgoing/old.jpg") + let newURL = root.appendingPathComponent("files/images/outgoing/new.jpg") + try writeBytes(60 * 1024 * 1024, to: oldURL, modified: Date(timeIntervalSinceNow: -3600)) + try writeBytes(45 * 1024 * 1024, to: newURL, modified: Date(timeIntervalSinceNow: -60)) + + // Hold 10 MB without an extra reservingBytes argument — eviction must + // still free the oldest file because usage + held exceeds quota. + let reservation = store.reserveQuotaBytes(10 * 1024 * 1024, scope: .outgoing) + defer { store.releaseQuotaReservation(reservation) } + + #expect(!FileManager.default.fileExists(atPath: oldURL.path)) + #expect(FileManager.default.fileExists(atPath: newURL.path)) + #expect(store.reservedQuotaBytes(for: .outgoing) == 10 * 1024 * 1024) + } }