Fix OAEP padding and update tests (#99)

TakeScoop · Dec 7, 2017 · bdb1d1c · bdb1d1c
1 parent 8a50bb3
commit bdb1d1c
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 5 deletions.
diff --git a/SwiftyRSA/ClearMessage.swift b/SwiftyRSA/ClearMessage.swift
@@ -56,7 +56,16 @@ public class ClearMessage: Message {
     public func encrypted(with key: PublicKey, padding: Padding) throws -> EncryptedMessage {
 
         let blockSize = SecKeyGetBlockSize(key.reference)
-        let maxChunkSize = (padding == []) ? blockSize : blockSize - 11
+
+        var maxChunkSize: Int
+        switch padding {
+        case []:
+            maxChunkSize = blockSize
+        case .OAEP:
+            maxChunkSize = blockSize - 42
+        default:
+            maxChunkSize = blockSize - 11
+        }
 
         var decryptedDataAsArray = [UInt8](repeating: 0, count: data.count)
         (data as NSData).getBytes(&decryptedDataAsArray, length: data.count)

diff --git a/SwiftyRSATests/EncryptDecryptTests.swift b/SwiftyRSATests/EncryptDecryptTests.swift
@@ -60,6 +60,16 @@ class EncryptDecryptTests: XCTestCase {
         XCTAssertEqual(decrypted.data, data)
     }
 
+    func test_OAEP() throws {
+        let data = TestUtils.randomData(count: 2048)
+        let clearMessage = ClearMessage(data: data)
+
+        let encrypted = try clearMessage.encrypted(with: publicKey, padding: .OAEP)
+        let decrypted = try encrypted.decrypted(with: privateKey, padding: .OAEP)
+
+        XCTAssertEqual(decrypted.data, data)
+    }
+
     func test_keyReferences() throws {
         let data = TestUtils.randomData(count: 2048)
         let clearMessage = ClearMessage(data: data)

diff --git a/SwiftyRSATests/TestUtils.swift b/SwiftyRSATests/TestUtils.swift
@@ -68,11 +68,12 @@ struct TestError: Error {
 
     @objc
     static public func randomData(count: Int) -> Data {
-        var data = Data(capacity: count)
-        data.withUnsafeMutableBytes { (bytes: UnsafeMutablePointer<UInt8>) -> Void in
-            _ = SecRandomCopyBytes(kSecRandomDefault, count, bytes)
+        var randomBytes = [UInt8](repeating: 0, count: count)
+        let status = SecRandomCopyBytes(kSecRandomDefault, count, &randomBytes)
+        if status != errSecSuccess {
+             XCTFail("Couldn't create random data")
         }
-        return data
+        return Data(bytes: randomBytes)
     }
 
     static func assertThrows(type: SwiftyRSAError, file: StaticString = #file, line: UInt = #line, block: () throws ->  Void) {