Skip to content

Commit

Permalink
POC for improving CryptoInputStream and computing output length
Browse files Browse the repository at this point in the history
  • Loading branch information
@jguz-pubnub
jguz-pubnub committed Oct 5, 2023
1 parent d814f5d commit facc3c4
Showing 1 changed file with 76 additions and 93 deletions.
169 changes: 76 additions & 93 deletions Sources/PubNub/Helpers/Crypto/Miscellaneous/CryptoInputStream.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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?
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -342,102 +404,23 @@ public class CryptoInputStream: InputStream {
// MARK: - Encrypt

func encrypt(_ buffer: UnsafeMutablePointer<UInt8>, 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<UInt8>, 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!
)
}

Expand Down

0 comments on commit facc3c4

Please sign in to comment.