Cipher encrypt/decrypt returns optionals

CryptoSwift/ChaCha20.swift

@@ -10,30 +10,37 @@ import Foundation
 
 class ChaCha20 {
     private let stateSize = 16
-    private let blockSize = 16 * 4
-    private var context:Context = Context()
+    private var context:Context?
     
     private class Context {
         var input:[UInt32] = [UInt32](count: 16, repeatedValue: 0)
+        
+        deinit {
+            for (var i = 0; i < input.count; i++) {
+                input[i] = 0x00;
+            }
+        }
     }
     
     init(key:NSData, iv:NSData) {
-        context = contextSetup(iv: iv, key: key)
-    }
-
-    deinit
-    {
-        for (var i = 0; i < context.input.count; i++) {
-            context.input[i] = 0x00;
+        if let c = contextSetup(iv: iv, key: key) {
+            context = c
         }
     }
     
-    func encrypt(message:NSData) -> NSData {
-        let output = encryptBytes(message.arrayOfBytes())
-        return NSData(bytes: output, length: output.count)
+    func encrypt(message:NSData) -> NSData? {
+        if (context == nil) {
+            return nil
+        }
+        
+        if let output = encryptBytes(message.arrayOfBytes()) {
+            return NSData(bytes: output, length: output.count)
+        }
+        
+        return nil
     }
     
-    func decrypt(message:NSData) -> NSData {
+    func decrypt(message:NSData) -> NSData? {
         return encrypt(message)
     }
     
@@ -67,17 +74,22 @@ class ChaCha20 {
         return output;
     }
     
-    private func contextSetup(# iv:NSData, key:NSData, kbits:UInt32 = 256) -> Context {
-        return contextSetup(iv: iv.arrayOfBytes(), key: key.arrayOfBytes(), kbits: kbits)
+    private func contextSetup(# iv:NSData, key:NSData) -> Context? {
+        return contextSetup(iv: iv.arrayOfBytes(), key: key.arrayOfBytes())
     }
     
-    private func contextSetup(# iv:[Byte], key:[Byte], kbits:UInt32 = 256) -> Context {
-        var context = Context()
+    private func contextSetup(# iv:[Byte], key:[Byte]) -> Context? {
+        var ctx = Context()
+        let kbits = key.count * 8
+        
+        if (kbits != 128 && kbits != 256) {
+            return nil
+        }
         
         // 4 - 8
         for (var i = 0; i < 4; i++) {
             let start = i * 4
-            context.input[i + 4] = UInt32.withBytes(key[start..<(start + 4)]).bigEndian
+            ctx.input[i + 4] = UInt32.withBytes(key[start..<(start + 4)]).bigEndian
         }
         
         var addPos = 0;
@@ -85,62 +97,66 @@ class ChaCha20 {
         case 256:
             addPos += 16
             // sigma
-            context.input[0] = 0x61707865 //apxe
-            context.input[1] = 0x3320646e //3 dn
-            context.input[2] = 0x79622d32 //yb-2
-            context.input[3] = 0x6b206574 //k et
+            ctx.input[0] = 0x61707865 //apxe
+            ctx.input[1] = 0x3320646e //3 dn
+            ctx.input[2] = 0x79622d32 //yb-2
+            ctx.input[3] = 0x6b206574 //k et
         default:
             // tau
-            context.input[0] = 0x61707865 //apxe
-            context.input[1] = 0x3620646e //6 dn
-            context.input[2] = 0x79622d31 //yb-1
-            context.input[3] = 0x6b206574 //k et
+            ctx.input[0] = 0x61707865 //apxe
+            ctx.input[1] = 0x3620646e //6 dn
+            ctx.input[2] = 0x79622d31 //yb-1
+            ctx.input[3] = 0x6b206574 //k et
         break;
         }
         
         // 8 - 11
         for (var i = 0; i < 4; i++) {
             let start = addPos + (i*4)
-            context.input[i + 8] = UInt32.withBytes(key[start..<(start + 4)]).bigEndian
+            ctx.input[i + 8] = UInt32.withBytes(key[start..<(start + 4)]).bigEndian
         }
 
         // iv
-        context.input[12] = 0
-        context.input[13] = 0
-        context.input[14] = UInt32.withBytes(iv[0..<4]).bigEndian
-        context.input[15] = UInt32.withBytes(iv[4..<8]).bigEndian
+        ctx.input[12] = 0
+        ctx.input[13] = 0
+        ctx.input[14] = UInt32.withBytes(iv[0..<4]).bigEndian
+        ctx.input[15] = UInt32.withBytes(iv[4..<8]).bigEndian
         
-        return context
+        return ctx
     }
     
-    private func encryptBytes(message:[Byte]) -> [Byte] {
-        var cPos:Int = 0
-        var mPos:Int = 0
-        var bytes = message.count
-        
-        var c:[Byte] = [Byte](count: message.count, repeatedValue: 0)
-        
-        while (true) {
-            if let output = wordToByte(context.input) {
-                context.input[12] = context.input[12] &+ 1
-                if (context.input[12] == 0) {
-                    context.input[13] = context.input[13] &+ 1
-                    /* stopping at 2^70 bytes per nonce is user's responsibility */
-                }
-                if (bytes <= blockSize) {
-                    for (var i = 0; i < bytes; i++) {
+    private func encryptBytes(message:[Byte]) -> [Byte]? {
+        if let ctx = context {
+            var c:[Byte] = [Byte](count: message.count, repeatedValue: 0)
+            
+            let blockSize = 64
+            var cPos:Int = 0
+            var mPos:Int = 0
+            var bytes = message.count
+            
+            while (true) {
+                if let output = wordToByte(ctx.input) {
+                    ctx.input[12] = ctx.input[12] &+ 1
+                    if (ctx.input[12] == 0) {
+                        ctx.input[13] = ctx.input[13] &+ 1
+                        /* stopping at 2^70 bytes per nonce is user's responsibility */
+                    }
+                    if (bytes <= blockSize) {
+                        for (var i = 0; i < bytes; i++) {
+                            c[i + cPos] = message[i + mPos] ^ output[i]
+                        }
+                        return c
+                    }
+                    for (var i = 0; i < blockSize; i++) {
                         c[i + cPos] = message[i + mPos] ^ output[i]
                     }
-                    return c
-                }
-                for (var i = 0; i < blockSize; i++) {
-                    c[i + cPos] = message[i + mPos] ^ output[i]
+                    bytes -= blockSize
+                    cPos += blockSize
+                    mPos += blockSize
                 }
-                bytes -= blockSize
-                cPos += blockSize
-                mPos += blockSize
             }
         }
+        return nil;
     }
     
     private func quarterround(inout a:UInt32, inout _ b:UInt32, inout _ c:UInt32, inout _ d:UInt32) {

CryptoSwift/Cipher.swift

@@ -11,7 +11,7 @@ import Foundation
 public enum Cipher {
     case ChaCha20(key: NSData, iv: NSData)
     
-    public func encrypt(message: NSData) -> NSData {
+    public func encrypt(message: NSData) -> NSData? {
         switch (self) {
             case .ChaCha20(let key, let iv):
                 var chacha = CryptoSwift.ChaCha20(key: key, iv: iv);
@@ -19,7 +19,7 @@ public enum Cipher {
         }
     }
     
-    public func decrypt(message: NSData) -> NSData {
+    public func decrypt(message: NSData) -> NSData? {
         switch (self) {
         case .ChaCha20(let key, let iv):
             var chacha = CryptoSwift.ChaCha20(key: key, iv: iv);

CryptoSwift/NSDataExtension.swift

@@ -64,11 +64,11 @@ extension NSData {
         return CRC().crc32(self);
     }
 
-    public func encrypt(cipher: Cipher) -> NSData {
+    public func encrypt(cipher: Cipher) -> NSData? {
         return cipher.encrypt(self)
     }
 
-    public func decrypt(cipher: Cipher) -> NSData {
+    public func decrypt(cipher: Cipher) -> NSData? {
         return cipher.decrypt(self)
     }
 

CryptoSwift/StringExtension.swift

@@ -41,11 +41,11 @@ extension String {
     }
 
     public func encrypt(cipher: Cipher) -> String? {
-        return self.dataUsingEncoding(NSUTF8StringEncoding, allowLossyConversion: false)?.encrypt(cipher).toHexString()
+        return self.dataUsingEncoding(NSUTF8StringEncoding, allowLossyConversion: false)?.encrypt(cipher)?.toHexString()
     }
 
     public func decrypt(cipher: Cipher) -> String? {
-        return self.dataUsingEncoding(NSUTF8StringEncoding, allowLossyConversion: false)?.decrypt(cipher).toHexString()
+        return self.dataUsingEncoding(NSUTF8StringEncoding, allowLossyConversion: false)?.decrypt(cipher)?.toHexString()
     }
 
 }

CryptoSwiftTests/CipherTests.swift

@@ -54,14 +54,18 @@ class CipherTests: XCTestCase {
             let messageData = NSData(bytes: message, length: message.count);
             
             let setup = (key: keyData, iv: ivData)
-            let encrypted = Cipher.ChaCha20(setup).encrypt(messageData)
-            let decrypted = Cipher.ChaCha20(setup).decrypt(encrypted)
-            XCTAssertEqual(messageData, decrypted, "ChaCha20 decryption failed");
-            XCTAssertEqual(encrypted.hexString, expectedHex, "ChaCha20 failed");
+            var encrypted = Cipher.ChaCha20(setup).encrypt(messageData)
+            XCTAssertNotNil(encrypted, "")
+            var decrypted = Cipher.ChaCha20(setup).decrypt(encrypted!)
+            XCTAssertNotNil(decrypted, "")
+            
+            XCTAssertEqual(messageData, decrypted!, "ChaCha20 decryption failed");
+            XCTAssertEqual(encrypted!.hexString, expectedHex, "ChaCha20 failed");
 
             // check extension
             let encrypted2 = messageData.encrypt(Cipher.ChaCha20(setup))
-            XCTAssertEqual(encrypted, encrypted2, "ChaCha20 extension failed")
+            XCTAssertNotNil(encrypted2, "")
+            XCTAssertEqual(encrypted!, encrypted2!, "ChaCha20 extension failed")
             
         }
     }

README.md

@@ -61,27 +61,34 @@ Hashing a String and printing result
     
 Working with Cipher
 
-	// convenience setup tuplet
+	// convenience setup tuple
 	let setup = (key: keyData, iv: ivData)
 	
 	// encrypt
-	let encrypted = Cipher.ChaCha20(setup).encrypt(dataToEncrypt)
+	if let encrypted = Cipher.ChaCha20(setup).encrypt(dataToEncrypt) {
 	
-	// decrypt
-	let decrypted = Cipher.ChaCha20(setup).decrypt(encrypted)
-	
-	// validate result
-	if (encrypted.isEqual(decrypted)) {
-		print("Decryption failed!")
+	    // decrypt
+	    if let decrypted = Cipher.ChaCha20(setup).decrypt(encrypted) {
+	    
+	        // validate result
+	        if (encrypted.isEqual(decrypted)) {
+		        print("Decryption failed!")
+	        }
+	        
+	    }
 	}
+	
 
 with extensions
 	
-	// convenience setup tuplet
+	// convenience setup tuple
 	let setup = (key: keyData, iv: ivData)
 
-	let encrypted = dataToEncrypt.encrypt(Cipher.ChaCha20(setup))
-	let decrypted = encrypted.decrypt(Cipher.ChaCha20(setup))
+	if let encrypted = dataToEncrypt.encrypt(Cipher.ChaCha20(setup)) {
+		if let decrypted = encrypted.decrypt(Cipher.ChaCha20(setup)) {
+			println(decrypted)
+		}
+	}
 	
     
 ##Contact