From facc3c4c701e83a9a126743ce36044e2d899bf2f Mon Sep 17 00:00:00 2001 From: jguz-pubnub Date: Thu, 5 Oct 2023 14:39:15 +0200 Subject: [PATCH] POC for improving CryptoInputStream and computing output length --- .../Miscellaneous/CryptoInputStream.swift | 169 ++++++++---------- 1 file changed, 76 insertions(+), 93 deletions(-) diff --git a/Sources/PubNub/Helpers/Crypto/Miscellaneous/CryptoInputStream.swift b/Sources/PubNub/Helpers/Crypto/Miscellaneous/CryptoInputStream.swift index 5207e8ff..30e54c81 100644 --- a/Sources/PubNub/Helpers/Crypto/Miscellaneous/CryptoInputStream.swift +++ b/Sources/PubNub/Helpers/Crypto/Miscellaneous/CryptoInputStream.swift @@ -37,14 +37,12 @@ public class CryptoInputStream: InputStream { let options: CCOptions let cipher: Cipher } + public struct Cipher { let algorithm: CCAlgorithm let blockSize: Int - - func outputSize(from inputBytes: Int) -> Int { - inputBytes + (blockSize - (inputBytes % blockSize)) - } } + public enum Operation { case decrypt case encrypt @@ -57,7 +55,7 @@ public class CryptoInputStream: InputStream { private var cipherStream: InputStream private var rawDataBufferRead: Int = -1 private var rawDataBuffer: [UInt8]? - private var rawDataLength: Int + private var rawDataLength: Int = 0 private var rawDataRead: Int = 0 private var cryptoStream: CryptoStream? @@ -67,7 +65,6 @@ public class CryptoInputStream: InputStream { // Buffer for data which has been encrypted / decrypted from cipher stream. private var cryptedBuffer: [UInt8]? private var cryptedBufferRead: Int = -1 - private var cryptedDataLength: Int = -1 private var cryptedDataFinalysed: Bool = false private var cryptedBytes: Int = 0 private var totalCrypted: Int = 0 @@ -91,14 +88,79 @@ public class CryptoInputStream: InputStream { self.crypto = crypto self.includeInitializationVectorInContent = includeInitializationVectorInContent - if includeInitializationVectorInContent { - rawDataLength = contentLength - (operation == .decrypt ? crypto.cipher.blockSize : 0) - // The estimated content length is the IV length plus the crypted length - estimatedCryptoCount = crypto.cipher.blockSize + crypto.cipher.outputSize(from: rawDataLength) + if operation == .encrypt { + do { + cryptedBufferRead = 0 + let ivBuffer = crypto.iv.map { $0 } + + let cryptoStream = try CryptoStream( + operation: CCOperation(operation == .decrypt ? kCCDecrypt : kCCEncrypt), + algorithm: crypto.cipher.algorithm, + options: crypto.options, + keyBuffer: crypto.key.map { $0 }, + keyLength: crypto.key.count, + ivBuffer: ivBuffer + ) + if includeInitializationVectorInContent { + cryptedBuffer = ivBuffer + rawDataLength = contentLength + estimatedCryptoCount = crypto.cipher.blockSize + cryptoStream.getOutputLength(inputLength: rawDataLength, isFinal: true) + } else { + rawDataLength = contentLength + estimatedCryptoCount = cryptoStream.getOutputLength(inputLength: rawDataLength, isFinal: true) + } + self.cryptoStream = cryptoStream + + } catch { + _streamError = error + _streamStatus = .error + } } else { - rawDataLength = contentLength - estimatedCryptoCount = crypto.cipher.outputSize(from: rawDataLength) + + // Create a buffer to store the IV in that matches the cipher block size + var initializationVectorBuffer = [UInt8](repeating: 0, count: crypto.cipher.blockSize) + + if includeInitializationVectorInContent { + switch input.read(&initializationVectorBuffer, maxLength: crypto.cipher.blockSize) { + case let bytesRead where bytesRead < 0: + // -1 means that the operation failed; more information about the error can be obtained with `streamError`. + _streamStatus = .error + _streamError = input.streamError + default: + // 0 represents end of the current buffer + break + } + } else { + initializationVectorBuffer = crypto.iv.map { $0 } + } + + // Init the Crypto Stream + do { + let decryptStream = try CryptoStream( + operation: CCOperation(operation == .decrypt ? kCCDecrypt : kCCEncrypt), + algorithm: crypto.cipher.algorithm, + options: crypto.options, + keyBuffer: crypto.key.map { $0 }, + keyLength: crypto.key.count, + ivBuffer: initializationVectorBuffer + ) + + if includeInitializationVectorInContent { + rawDataLength = contentLength - crypto.cipher.blockSize + // The estimated content length is the IV length plus the crypted length + estimatedCryptoCount = crypto.cipher.blockSize + decryptStream.getOutputLength(inputLength: rawDataLength, isFinal: true) + } else { + rawDataLength = contentLength + estimatedCryptoCount = decryptStream.getOutputLength(inputLength: rawDataLength, isFinal: true) + } + + self.cryptoStream = decryptStream + } catch { + _streamError = error + _streamStatus = .error + } } + cipherStream = input // required because `init()` is not marked as a designated initializer @@ -342,102 +404,23 @@ public class CryptoInputStream: InputStream { // MARK: - Encrypt func encrypt(_ buffer: UnsafeMutablePointer, maxLength: Int) -> Int { - // Create Crypto if it does not exist - let encryptStream: CryptoStream - if let cryptoStream = cryptoStream { - encryptStream = cryptoStream - } else { - // Init the Crypto Stream - do { - let ivBuffer = crypto.iv.map { $0 } - - if includeInitializationVectorInContent { - cryptedBuffer = ivBuffer - } - cryptedBufferRead = 0 - - encryptStream = try CryptoStream( - operation: CCOperation(operation == .decrypt ? kCCDecrypt : kCCEncrypt), - algorithm: crypto.cipher.algorithm, - options: crypto.options, - keyBuffer: crypto.key.map { $0 }, - keyLength: crypto.key.count, - ivBuffer: ivBuffer - ) - - cryptedDataLength = encryptStream.getOutputLength( - inputLength: rawDataLength, isFinal: true - ) + (includeInitializationVectorInContent ? crypto.cipher.blockSize : 0) - } catch { - _streamError = error - _streamStatus = .error - return -1 - } - - // Assign the created stream to self so we retain it for future loops - cryptoStream = encryptStream - } - return crypt( outputBuffer: buffer, maxLength: maxLength, inputStream: cipherStream, readByteOffset: 0, - cryptoStream: encryptStream + cryptoStream: self.cryptoStream! ) } // MARK: - Decrypt func decrypt(_ buffer: UnsafeMutablePointer, maxLength: Int) -> Int { - // Create Crypto if it does not exist - let decryptStream: CryptoStream - if let cryptoStream = cryptoStream { - decryptStream = cryptoStream - } else { - // Create a buffer to store the IV in that matches the cipher block size - var initializationVectorBuffer = [UInt8](repeating: 0, count: crypto.cipher.blockSize) - - if includeInitializationVectorInContent { - switch cipherStream.read(&initializationVectorBuffer, maxLength: crypto.cipher.blockSize) { - case let bytesRead where bytesRead < 0: - // -1 means that the operation failed; more information about the error can be obtained with `streamError`. - _streamStatus = .error - _streamError = cipherStream.streamError - return bytesRead - default: - // 0 represents end of the current buffer - break - } - } else { - initializationVectorBuffer = crypto.iv.map { $0 } - } - - // Init the Crypto Stream - do { - decryptStream = try CryptoStream( - operation: CCOperation(operation == .decrypt ? kCCDecrypt : kCCEncrypt), - algorithm: crypto.cipher.algorithm, - options: crypto.options, - keyBuffer: crypto.key.map { $0 }, - keyLength: crypto.key.count, - ivBuffer: initializationVectorBuffer - ) - } catch { - _streamError = error - _streamStatus = .error - return -1 - } - - // Assign the created stream to self so we retain it for future loops - cryptoStream = decryptStream - } - return crypt( outputBuffer: buffer, maxLength: maxLength, inputStream: cipherStream, - cryptoStream: decryptStream + cryptoStream: self.cryptoStream! ) }