97 lines
3.3 KiB
Swift
97 lines
3.3 KiB
Swift
//
|
|
// Copyright 2024 Signal Messenger, LLC
|
|
// SPDX-License-Identifier: AGPL-3.0-only
|
|
//
|
|
|
|
import Foundation
|
|
|
|
public class ChunkedInputStreamTransform: StreamTransform, BufferedStreamTransform {
|
|
|
|
private var buffer = Data()
|
|
|
|
public var hasPendingBytes: Bool { return buffer.count > 0 }
|
|
public func readBufferedData() throws -> Data { try getNextChunk() }
|
|
|
|
public func transform(data: Data) throws -> Data {
|
|
buffer.append(data)
|
|
return try getNextChunk()
|
|
}
|
|
|
|
/// Decode the next chunk of data, if enough data is present in the buffer.
|
|
private func getNextChunk() throws -> Data {
|
|
// decode the next variable length int
|
|
let (dataSize, intLength) = decodeVariableLengthInteger(buffer)
|
|
|
|
guard dataSize > 0 else {
|
|
// Don't have enough data to decode an int, so return for now
|
|
return Data()
|
|
}
|
|
|
|
guard let intDataSize = Int(exactly: dataSize) else {
|
|
// The decoded integer is to large to fit into an Int
|
|
// The Data operations all require Int Ranges, so
|
|
owsFailDebug("Decoded data size too large")
|
|
return Data()
|
|
}
|
|
|
|
// Only advance if there is enough data present to both
|
|
// decode the variable length integer and read the specified
|
|
// number of bytes.
|
|
let endOfBuffer = intDataSize + intLength
|
|
guard buffer.count >= endOfBuffer else {
|
|
return Data()
|
|
}
|
|
|
|
// Return a chunk of data from the buffer and advence the buffer.
|
|
let returnBuffer = buffer.subdata(in: intLength..<endOfBuffer)
|
|
if endOfBuffer >= buffer.count {
|
|
buffer = Data()
|
|
} else {
|
|
buffer = buffer.subdata(in: endOfBuffer..<buffer.count)
|
|
}
|
|
|
|
return returnBuffer
|
|
}
|
|
|
|
/// Inspect the incoming data and return the next variable length integer.
|
|
/// Because it's not guaranteed that there's enought data to either (a) decode the integer or
|
|
/// (b) read the amount of data specified by the returned integer, dont' remove the decoded
|
|
/// integer bytes from the buffer until we're certain there's enough data to fulfill reading
|
|
/// the specified amount of data.
|
|
private func decodeVariableLengthInteger(_ data: Data) -> (result: UInt64, length: Int) {
|
|
guard data.count > 0 else { return (result: 0, length: 0) }
|
|
|
|
return data.withUnsafeBytes {
|
|
guard let baseAddress = $0.bindMemory(to: UInt8.self).baseAddress else {
|
|
return (result: 0, length: 0)
|
|
}
|
|
var start = baseAddress
|
|
var length = data.count
|
|
var count = 1
|
|
var c = start[0]
|
|
start += 1
|
|
length -= 1
|
|
|
|
if c & 0x80 == 0 {
|
|
return (result: UInt64(c), length: count)
|
|
}
|
|
var value = UInt64(c & 0x7f)
|
|
var shift = UInt64(7)
|
|
while true {
|
|
if length < 1 || shift > 63 {
|
|
return (result: 0, length: 0)
|
|
}
|
|
c = start[0]
|
|
start += 1
|
|
length -= 1
|
|
count += 1
|
|
value |= UInt64(c & 0x7f) << shift
|
|
if c & 0x80 == 0 {
|
|
return (result: value, length: count)
|
|
}
|
|
shift += 7
|
|
}
|
|
}
|
|
}
|
|
}
|