feat: Create JWK-RSA key with given bits

fix: Compression issue with large data
chore: RSA, Compression tests

Package.swift

@@ -20,7 +20,8 @@ let package = Package(
     products: [
         .library(
             name: "JWSETKit",
-            targets: ["JWSETKit"]),
+            targets: ["JWSETKit"]
+        ),
     ],
     dependencies: [
         .package(url: "https://github.com/Flight-School/AnyCodable", .upToNextMajor(from: "0.6.7")),

Sources/JWSETKit/Base/RawType.swift

@@ -18,7 +18,7 @@ extension StringRepresentable {
 
     public init(from decoder: any Decoder) throws {
         let container = try decoder.singleValueContainer()
-        self.init(rawValue: try container.decode(String.self))
+        try self.init(rawValue: container.decode(String.self))
     }
 
     public func encode(to encoder: any Encoder) throws {

Sources/JWSETKit/Cryptography/Compression/Apple.swift

@@ -25,56 +25,31 @@ extension JSONWebCompressionAlgorithm {
 /// Compressor contain compress and decompress implementation using `Compression` framework.
 public struct AppleCompressor<Codec>: JSONWebCompressor where Codec: CompressionCodec {
     public static func compress<D>(_ data: D) throws -> Data where D: DataProtocol {
-        var result = Data()
-        let outputFilter = try OutputFilter(.compress, using: Codec.algorithm.appleAlgorithm, writingTo: { data in
-            if let data = data {
-                result.append(data)
-            }
-        })
-
-        var index = 0
-        let bufferSize = data.count
-        var buffer = [UInt8](repeating: 0, count: Codec.pageSize)
-
-        while true {
-            let rangeLength = min(Codec.pageSize, bufferSize - index)
-            let startIndex = data.index(data.startIndex, offsetBy: index)
-            let endIndex = data.index(startIndex, offsetBy: rangeLength)
-            buffer.withUnsafeMutableBytes { bytes in
-                index += data.copyBytes(to: bytes, from: startIndex ..< endIndex)
-            }
-
-            try outputFilter.write(rangeLength == Codec.pageSize ? buffer : .init(buffer.prefix(rangeLength)))
-
-            if rangeLength == 0 {
-                break
-            }
+        var compressedData = Data()
+        let filter = try OutputFilter(.compress, using: Codec.algorithm.appleAlgorithm) {
+            compressedData.append($0 ?? .init())
         }
-        return result.prefix(index)
+
+        // Compress the data
+        try filter.write(data)
+        try filter.finalize()
+        return compressedData
     }
 
     public static func decompress<D>(_ data: D) throws -> Data where D: DataProtocol {
-        let data = Data(data)
-        var result = Data()
-        var index = 0
-        let bufferSize = data.count
-        var buffer = Data(repeating: 0, count: Codec.pageSize)
-
-        let inputFilter = try InputFilter(.decompress, using: Codec.algorithm.appleAlgorithm) { length -> Data in
-            let rangeLength = min(length, bufferSize - index)
-            let startIndex = data.index(data.startIndex, offsetBy: index)
-            let endIndex = data.index(startIndex, offsetBy: rangeLength)
-            buffer.withUnsafeMutableBytes {
-                index += data.copyBytes(to: $0, from: startIndex ..< endIndex)
-            }
-
-            return rangeLength == Codec.pageSize ? buffer : buffer.prefix(rangeLength)
+        var data = Data(data)
+        var decompressedData = Data()
+        let filter = try InputFilter(.decompress, using: Codec.algorithm.appleAlgorithm) { count in
+            defer { data = data.dropFirst(count) }
+            return data.prefix(count)
         }
 
-        while let page = try inputFilter.readData(ofLength: Codec.pageSize) {
-            result.append(page)
+        // Decompress the data
+        while let chunk = try filter.readData(ofLength: Codec.pageSize), !chunk.isEmpty {
+            decompressedData.append(chunk)
         }
-        return result
+
+        return decompressedData
     }
 }
 #endif

Sources/JWSETKit/Cryptography/RSA/JWK-RSA.swift

@@ -80,6 +80,28 @@ public struct JSONWebRSAPublicKey: MutableJSONWebKey, JSONWebValidatingKey, JSON
 
 /// JWK container for RSA private keys.
 public struct JSONWebRSAPrivateKey: MutableJSONWebKey, JSONWebSigningKey, JSONWebDecryptingKey, Sendable {
+    public struct KeySize {
+        public let bitCount: Int
+
+        /// RSA key size of 2048 bits
+        public static let bits2048 = KeySize(bitCount: 2048)
+
+        /// RSA key size of 3072 bits
+        public static let bits3072 = KeySize(bitCount: 3072)
+
+        /// RSA key size of 4096 bits
+        public static let bits4096 = KeySize(bitCount: 4096)
+
+        /// RSA key size with a custom number of bits.
+        ///
+        /// Params:
+        ///     - bitsCount: Positive integer that is a multiple of 8.
+        public init(bitCount: Int) {
+            precondition(bitCount % 8 == 0 && bitCount > 0)
+            self.bitCount = bitCount
+        }
+    }
+
     public var storage: JSONWebValueStorage
 
     public var publicKey: JSONWebRSAPublicKey {
@@ -94,7 +116,19 @@ public struct JSONWebRSAPrivateKey: MutableJSONWebKey, JSONWebSigningKey, JSONWe
     }
 
     public init(algorithm _: any JSONWebAlgorithm) throws {
-        self.storage = try _RSA.Signing.PrivateKey(keySize: .bits2048).storage
+        try self.init(bitCount: .bits2048)
+    }
+
+    public init(bitCount: KeySize) throws {
+#if canImport(CommonCrypto)
+        self.storage = try SecKey(rsaBitCounts: bitCount.bitCount).storage
+#elseif canImport(_CryptoExtras)
+        self.storage = try _RSA.Signing.PrivateKey(keySize: .init(bitCount: size)).storage
+#else
+        // This should never happen as CommonCrypto is available on Darwin platforms
+        // and CryptoExtras is used on non-Darwin platform.
+        fatalError("Unimplemented")
+#endif
     }
 
     public init(storage: JSONWebValueStorage) {

Sources/JWSETKit/Cryptography/RSA/SecKey.swift

@@ -200,6 +200,13 @@ extension JSONWebSigningKey where Self: SecKey {
         self = result
     }
 
+    public init(rsaBitCounts bits: Int) throws {
+        guard let result = try Self.createPairKey(type: .rsa, bits: bits) as? Self else {
+            throw JSONWebKeyError.operationNotAllowed
+        }
+        self = result
+    }
+
     public init(derRepresentation: Data, keyType: JSONWebKeyType) throws {
         var derRepresentation = derRepresentation
         let secKeyType: CFString

Sources/JWSETKit/Extensions/Localizing.swift

@@ -7,7 +7,7 @@
 
 import Foundation
 
-private class Decoy { }
+private class Decoy {}
 
 extension Bundle {
     func forLocale(_ locale: Locale) -> Bundle {
@@ -19,7 +19,7 @@ extension Bundle {
         return .module
     }
 
-    static let current: Bundle = Bundle(for: Decoy.self)
+    static let current: Bundle = .init(for: Decoy.self)
 }
 
 extension String {

Tests/JWSETKitTests/Cryptography/CompressionTests.swift

@@ -0,0 +1,43 @@
+//
+//  CompressionTests.swift
+//
+//
+//  Created by Amir Abbas Mousavian on 11/24/23.
+//
+
+import XCTest
+@testable import JWSETKit
+
+final class CompressionTests: XCTestCase {
+    let decompressed = Data("Data compression test. This text must be compressed.".utf8)
+    let deflateCompressed = Data(base64Encoded: "c0ksSVRIzs8tKEotLs7Mz1MoSS0u0VMIycgsBjIrShRyS4tLFJJS4WpSU/QA")!
+
+    var deflateCompressor: (any JSONWebCompressor.Type)? {
+        guard JSONWebCompressionAlgorithm.registeredAlgorithms.contains(.deflate) else { return nil }
+        return JSONWebCompressionAlgorithm.deflate.compressor
+    }
+
+    func testDeflateCompression() throws {
+        guard let deflateCompressor else { return }
+        let testCompressed = try deflateCompressor.compress(decompressed)
+        XCTAssertEqual(testCompressed, deflateCompressed)
+        XCTAssertLessThan(testCompressed.count, decompressed.count)
+    }
+
+    func testDeflateDecompression() throws {
+        guard let deflateCompressor else { return }
+        let testDecompressed = try deflateCompressor.decompress(deflateCompressed)
+        XCTAssertEqual(testDecompressed, decompressed)
+    }
+
+    func testCompressionDecompression() throws {
+        let length = Int.random(in: (1 << 17) ... (1 << 20)) // 128KB to 1MB
+        let decompressed = (0 ..< length).map { _ in UInt8.random(in: 0 ... 255) }
+        for algorithm in JSONWebCompressionAlgorithm.registeredAlgorithms {
+            guard let compressor = algorithm.compressor else { continue }
+            let testCompressed = try compressor.compress(decompressed)
+            let testDecompressed = try compressor.decompress(testCompressed)
+            XCTAssertEqual(Data(decompressed), testDecompressed)
+        }
+    }
+}

Tests/JWSETKitTests/Cryptography/RSATests.swift

@@ -50,15 +50,153 @@ final class RSATests: XCTestCase {
     AgMBAAE=
     """)!
 
-    func testPKCS8() throws {
+    let plaintext = Data("The quick brown fox jumps over the lazy dog.".utf8)
+
+    func testPKCS8Init() throws {
         XCTAssertNoThrow(try JSONWebRSAPrivateKey(derRepresentation: privateKeyDER))
         XCTAssertNoThrow(try JSONWebRSAPublicKey(derRepresentation: publicKeyDER))
     }
 
 #if canImport(CommonCrypto)
-    func testSecKey() throws {
+    func testSecKeyInit() throws {
         XCTAssertNoThrow(try SecKey(derRepresentation: publicKeyDER, keyType: .rsa))
         XCTAssertNoThrow(try SecKey(derRepresentation: privateKeyDER, keyType: .rsa))
     }
 #endif
+
+    func testEncrypt_RSA2048_OAEP_SHA1() throws {
+        let publicKey = try JSONWebRSAPublicKey(derRepresentation: publicKeyDER)
+        let privateKey = try JSONWebRSAPrivateKey(derRepresentation: privateKeyDER)
+
+        let ciphertext = try publicKey.encrypt(plaintext, using: .rsaEncryptionOAEP)
+        let decrypted = try privateKey.decrypt(ciphertext, using: .rsaEncryptionOAEP)
+
+        XCTAssertEqual(plaintext, decrypted)
+        XCTAssertNotEqual(plaintext, ciphertext)
+        XCTAssertEqual(ciphertext.count, 2048 / 8)
+    }
+
+    func testEncrypt_RSA2048_OAEP_SHA256() throws {
+        let publicKey = try JSONWebRSAPublicKey(derRepresentation: publicKeyDER)
+        let privateKey = try JSONWebRSAPrivateKey(derRepresentation: privateKeyDER)
+
+        let ciphertext = try publicKey.encrypt(plaintext, using: .rsaEncryptionOAEPSHA256)
+        let decrypted = try privateKey.decrypt(ciphertext, using: .rsaEncryptionOAEPSHA256)
+
+        XCTAssertEqual(plaintext, decrypted)
+        XCTAssertNotEqual(plaintext, ciphertext)
+        XCTAssertEqual(ciphertext.count, 2048 / 8)
+    }
+
+    func testEcrypt_RSA2048_OAEP_SHA384() throws {
+        let publicKey = try JSONWebRSAPublicKey(derRepresentation: publicKeyDER)
+        let privateKey = try JSONWebRSAPrivateKey(derRepresentation: privateKeyDER)
+
+        let ciphertext = try publicKey.encrypt(plaintext, using: .rsaEncryptionOAEPSHA384)
+        let decrypted = try privateKey.decrypt(ciphertext, using: .rsaEncryptionOAEPSHA384)
+
+        XCTAssertEqual(plaintext, decrypted)
+        XCTAssertNotEqual(plaintext, ciphertext)
+        XCTAssertEqual(ciphertext.count, 2048 / 8)
+    }
+
+    func testEncrypt_RSA2048_OAEP_SHA512() throws {
+        let publicKey = try JSONWebRSAPublicKey(derRepresentation: publicKeyDER)
+        let privateKey = try JSONWebRSAPrivateKey(derRepresentation: privateKeyDER)
+
+        let ciphertext = try publicKey.encrypt(plaintext, using: .rsaEncryptionOAEPSHA512)
+        let decrypted = try privateKey.decrypt(ciphertext, using: .rsaEncryptionOAEPSHA512)
+
+        XCTAssertEqual(plaintext, decrypted)
+        XCTAssertNotEqual(plaintext, ciphertext)
+        XCTAssertEqual(ciphertext.count, 2048 / 8)
+    }
+
+    func testSigning_RSA2048_PKCS1_SHA256() throws {
+        let publicKey = try JSONWebRSAPublicKey(derRepresentation: publicKeyDER)
+        let privateKey = try JSONWebRSAPrivateKey(derRepresentation: privateKeyDER)
+
+        let signature = try privateKey.signature(plaintext, using: .rsaSignaturePKCS1v15SHA256)
+        XCTAssertNoThrow(try publicKey.verifySignature(signature, for: plaintext, using: .rsaSignaturePKCS1v15SHA256))
+
+        XCTAssertNotEqual(plaintext, signature)
+        XCTAssertEqual(signature.count, 2048 / 8)
+    }
+
+    func testSigning_RSA2048_PKCS1_SHA384() throws {
+        let publicKey = try JSONWebRSAPublicKey(derRepresentation: publicKeyDER)
+        let privateKey = try JSONWebRSAPrivateKey(derRepresentation: privateKeyDER)
+
+        let signature = try privateKey.signature(plaintext, using: .rsaSignaturePKCS1v15SHA384)
+        XCTAssertNoThrow(try publicKey.verifySignature(signature, for: plaintext, using: .rsaSignaturePKCS1v15SHA384))
+
+        XCTAssertNotEqual(plaintext, signature)
+        XCTAssertEqual(signature.count, 2048 / 8)
+    }
+
+    func testSigning_RSA2048_PKCS1_SHA512() throws {
+        let publicKey = try JSONWebRSAPublicKey(derRepresentation: publicKeyDER)
+        let privateKey = try JSONWebRSAPrivateKey(derRepresentation: privateKeyDER)
+
+        let signature = try privateKey.signature(plaintext, using: .rsaSignaturePKCS1v15SHA512)
+        XCTAssertNoThrow(try publicKey.verifySignature(signature, for: plaintext, using: .rsaSignaturePKCS1v15SHA512))
+
+        XCTAssertNotEqual(plaintext, signature)
+        XCTAssertEqual(signature.count, 2048 / 8)
+    }
+
+    func testSigning_RSA2048_PSS_SHA256() throws {
+        let publicKey = try JSONWebRSAPublicKey(derRepresentation: publicKeyDER)
+        let privateKey = try JSONWebRSAPrivateKey(derRepresentation: privateKeyDER)
+
+        let signature = try privateKey.signature(plaintext, using: .rsaSignaturePSSSHA256)
+        XCTAssertNoThrow(try publicKey.verifySignature(signature, for: plaintext, using: .rsaSignaturePSSSHA256))
+
+        XCTAssertNotEqual(plaintext, signature)
+        XCTAssertEqual(signature.count, 2048 / 8)
+    }
+
+    func testSigning_RSA2048_PSS_SHA384() throws {
+        let publicKey = try JSONWebRSAPublicKey(derRepresentation: publicKeyDER)
+        let privateKey = try JSONWebRSAPrivateKey(derRepresentation: privateKeyDER)
+
+        let signature = try privateKey.signature(plaintext, using: .rsaSignaturePSSSHA384)
+        XCTAssertNoThrow(try publicKey.verifySignature(signature, for: plaintext, using: .rsaSignaturePSSSHA384))
+
+        XCTAssertNotEqual(plaintext, signature)
+        XCTAssertEqual(signature.count, 2048 / 8)
+    }
+
+    func testSigning_RSA2048_PSS_SHA512() throws {
+        let publicKey = try JSONWebRSAPublicKey(derRepresentation: publicKeyDER)
+        let privateKey = try JSONWebRSAPrivateKey(derRepresentation: privateKeyDER)
+
+        let signature = try privateKey.signature(plaintext, using: .rsaSignaturePSSSHA512)
+        XCTAssertNoThrow(try publicKey.verifySignature(signature, for: plaintext, using: .rsaSignaturePSSSHA512))
+
+        XCTAssertNotEqual(plaintext, signature)
+        XCTAssertEqual(signature.count, 2048 / 8)
+    }
+
+    func testSigning_RSA3072_PSS_SHA256() throws {
+        let privateKey = try JSONWebRSAPrivateKey(bitCount: .bits3072)
+        let publicKey = privateKey.publicKey
+
+        let signature = try privateKey.signature(plaintext, using: .rsaSignaturePSSSHA256)
+        XCTAssertNoThrow(try publicKey.verifySignature(signature, for: plaintext, using: .rsaSignaturePSSSHA256))
+
+        XCTAssertNotEqual(plaintext, signature)
+        XCTAssertEqual(signature.count, 3072 / 8)
+    }
+
+    func testSigning_RSA4096_PSS_SHA256() throws {
+        let privateKey = try JSONWebRSAPrivateKey(bitCount: .bits4096)
+        let publicKey = privateKey.publicKey
+
+        let signature = try privateKey.signature(plaintext, using: .rsaSignaturePSSSHA256)
+        XCTAssertNoThrow(try publicKey.verifySignature(signature, for: plaintext, using: .rsaSignaturePSSSHA256))
+
+        XCTAssertNotEqual(plaintext, signature)
+        XCTAssertEqual(signature.count, 4096 / 8)
+    }
 }

Tests/JWSETKitTests/ExampleKeys.swift

@@ -14,7 +14,7 @@ import Crypto
 #endif
 @testable import JWSETKit
 
-struct ExampleKeys {
+enum ExampleKeys {
     static let publicEC256 = try! JSONWebECPublicKey(jsonWebKeyData: Data(
         """
         {"kty":"EC",