100 lines
3 KiB
Swift
100 lines
3 KiB
Swift
//
|
|
// Copyright 2024 Signal Messenger, LLC
|
|
// SPDX-License-Identifier: AGPL-3.0-only
|
|
//
|
|
|
|
import CryptoKit
|
|
import Foundation
|
|
|
|
public class HmacStreamTransform: StreamTransform, FinalizableStreamTransform, BufferedStreamTransform {
|
|
public enum Error: Swift.Error {
|
|
case invalidHmac
|
|
case invalidFooter
|
|
}
|
|
|
|
public enum Constants {
|
|
static let FooterSize = 32
|
|
}
|
|
|
|
public enum Operation {
|
|
case generate
|
|
case validate
|
|
}
|
|
|
|
private var hmacState: HMAC<SHA256>
|
|
private let hmacKey: Data
|
|
private let operation: Operation
|
|
|
|
private var finalized = false
|
|
public var hasFinalized: Bool { finalized }
|
|
|
|
/// Because the HMAC footer is at the end of the input data, the input
|
|
/// needs to be buffered to always keep the latest 32 bytes around as
|
|
/// the provisional footer in case `finalize()` is called.
|
|
private var inputBuffer = Data()
|
|
|
|
/// If there is data in excess of the footer size, return true.
|
|
public var hasPendingBytes: Bool { inputBuffer.count > footerSize }
|
|
|
|
private let footerSize: Int
|
|
|
|
init(hmacKey: Data, operation: Operation) throws {
|
|
self.hmacKey = hmacKey
|
|
self.hmacState = HMAC(key: .init(data: hmacKey))
|
|
self.operation = operation
|
|
self.footerSize = {
|
|
switch operation {
|
|
case .generate: return 0
|
|
case .validate: return Constants.FooterSize
|
|
}
|
|
}()
|
|
}
|
|
|
|
public func readBufferedData() throws -> Data {
|
|
guard inputBuffer.count > self.footerSize else { return Data() }
|
|
|
|
// Get the data up to the point reserved for the footer, and preserve
|
|
// the provisional footer data in the input buffer.
|
|
let remainingDataLength = inputBuffer.count - self.footerSize
|
|
let remainingData = inputBuffer.subdata(in: 0..<remainingDataLength)
|
|
inputBuffer = inputBuffer.subdata(in: remainingDataLength..<inputBuffer.count)
|
|
return remainingData
|
|
}
|
|
|
|
public func transform(data: Data) throws -> Data {
|
|
inputBuffer.append(data)
|
|
let targetData = try readBufferedData()
|
|
if targetData.count > 0 {
|
|
// Update the hmac with the new block
|
|
hmacState.update(data: targetData)
|
|
}
|
|
return targetData
|
|
}
|
|
|
|
public func finalize() throws -> Data {
|
|
guard !finalized else { return Data() }
|
|
finalized = true
|
|
|
|
if inputBuffer.count < self.footerSize {
|
|
throw Error.invalidFooter
|
|
}
|
|
|
|
// Fetch the remaining non-footer data
|
|
var remainingData = try readBufferedData()
|
|
|
|
let hmac = Data(hmacState.finalize())
|
|
switch operation {
|
|
case .generate:
|
|
remainingData.append(hmac)
|
|
case .validate:
|
|
// footerdata remaining in inputBuffer
|
|
let footerData = inputBuffer
|
|
guard hmac.ows_constantTimeIsEqual(to: footerData) else {
|
|
throw Error.invalidHmac
|
|
}
|
|
inputBuffer = Data()
|
|
}
|
|
|
|
return remainingData
|
|
}
|
|
}
|