refactor

CryptoSwift/Poly1305.swift

@@ -14,7 +14,7 @@ import Foundation
 
 public class Poly1305 {
     let blockSize = 16
-    private var ctx:Context
+    private var ctx:Context?
     
     private class Context {
         var r            = [Byte](count: 17, repeatedValue: 0)
@@ -77,19 +77,119 @@ public class Poly1305 {
         }
     }
     
-    class public func withKey(key: [Byte]) -> Poly1305 {
-        return Poly1305(key)
+    // MARK: - Public
+
+    /**
+    Calculate Message Authentication Code (MAC) for message.
+    Calculation context is discarder on instance deallocation.
+    
+    :param: key     256-bit key
+    :param: message Message
+    
+    :returns: Message Authentication Code
+    */
+    class public func authenticate(key: [Byte], message: [Byte]) -> [Byte]? {
+        var poly = Poly1305(key)
+        return poly.authenticate(message: message)
     }
     
+    // MARK: - Private
+    
     private init (_ key: [Byte]) {
         ctx = Context(key)
     }
     
-    public func auth(mac:[Byte], message:[Byte]) -> [Byte]? {
-        update(ctx, m: message)
-        return finish(ctx, mac: mac)
+    private func authenticate(# message:[Byte]) -> [Byte]? {
+        if let ctx = self.ctx {
+            update(ctx, message: message)
+            return finish(ctx)
+        }
+        return nil
     }
     
+    /**
+    Add message to be processed
+    
+    :param: context Context
+    :param: message message
+    :param: bytes   length of the message fragment to be processed
+    */
+    private func update(context:Context, message:[Byte], bytes:Int? = nil) {
+        var bytes = bytes ?? message.count
+        var mPos = 0
+        
+        /* handle leftover */
+        if (context.leftover > 0) {
+            var want = blockSize - context.leftover
+            if (want > bytes) {
+                want = bytes
+            }
+            
+            for i in 0..<want {
+                context.buffer[context.leftover + i] = message[mPos + i]
+            }
+            
+            bytes -= want
+            mPos += want
+            context.leftover += want
+            
+            if (context.leftover < blockSize) {
+                return
+            }
+            
+            blocks(context, m: context.buffer)
+            context.leftover = 0
+        }
+        
+        /* process full blocks */
+        if (bytes >= blockSize) {
+            var want = bytes & ~(blockSize - 1)
+            blocks(context, m: message, startPos: mPos)
+            mPos += want
+            bytes -= want;
+        }
+        
+        /* store leftover */
+        if (bytes > 0) {
+            for i in 0..<bytes {
+                context.buffer[context.leftover + i] = message[mPos + i]
+            }
+            
+            context.leftover += bytes
+        }
+    }
+    
+    private func finish(context:Context) -> [Byte]? {
+        var mac = [Byte](count: 16, repeatedValue: 0);
+        
+        /* process the remaining block */
+        if (context.leftover > 0) {
+            
+            var i = context.leftover
+            context.buffer[i++] = 1
+            for (; i < blockSize; i++) {
+                context.buffer[i] = 0
+            }
+            context.final = 1
+            
+            blocks(context, m: context.buffer)
+        }
+        
+        
+        /* fully reduce h */
+        freeze(context)
+        
+        /* h = (h + pad) % (1 << 128) */
+        add(context, c: context.pad)
+        for i in 0..<mac.count {
+            mac[i] = context.h[i]
+        }
+        
+        return mac
+    }
+    
+    // MARK: - Utils
+    
     private func add(context:Context, c:[Byte]) -> Bool {
         if (context.h.count != 17 && c.count != 17) {
             return false
@@ -108,9 +208,9 @@ public class Poly1305 {
         if (context.h.count != 17 && hr.count != 17) {
             return false
         }
-
+        
         var u:UInt32 = 0
-
+        
         for i in 0..<16 {
             u += hr[i];
             context.h[i] = Byte.withValue(u) // crash! h[i] = UInt8(u) & 0xff
@@ -176,9 +276,9 @@ public class Poly1305 {
                 c[i] = m[mPos + i]
             }
             c[16] = hibit
-
+            
             add(context, c: c)
-
+            
             /* h *= r */
             for i in 0..<17 {
                 u = 0
@@ -194,89 +294,10 @@ public class Poly1305 {
             }
             
             squeeze(context, hr: hr)
-
+            
             mPos += blockSize
             bytes -= blockSize
         }
         return mPos
     }
-    
-    private func finish(context:Context, mac:[Byte]) -> [Byte]? {
-        assert(mac.count == 16, "Invalid mac length")
-        if (mac.count != 16) {
-            return nil
-        }
-        
-        var resultMAC = mac;
-        
-        /* process the remaining block */
-        if (context.leftover > 0) {
-            
-            var i = context.leftover
-            context.buffer[i++] = 1
-            for (; i < blockSize; i++) {
-                context.buffer[i] = 0
-            }
-            context.final = 1
-            
-            blocks(context, m: context.buffer)
-        }
-        
-        
-        /* fully reduce h */
-        freeze(context)
-        
-        /* h = (h + pad) % (1 << 128) */
-        add(context, c: context.pad)
-        for i in 0..<16 {
-            resultMAC[i] = context.h[i]
-        }
-        
-        return resultMAC
-    }
-    
-    private func update(context:Context, m:[Byte]) {
-        var bytes = m.count
-        var mPos = 0
-        
-        /* handle leftover */
-        if (context.leftover > 0) {
-            var want = blockSize - context.leftover
-            if (want > bytes) {
-                want = bytes
-            }
-            
-            for i in 0..<want {
-                context.buffer[context.leftover + i] = m[mPos + i]
-            }
-            
-            bytes -= want
-            mPos += want
-            context.leftover += want
-            
-            if (context.leftover < blockSize) {
-                return
-            }
-            
-            blocks(context, m: context.buffer)
-            context.leftover = 0
-        }
-        
-        /* process full blocks */
-        if (bytes >= blockSize) {
-            var want = bytes & ~(blockSize - 1)
-            blocks(context, m: m, startPos: mPos)
-            mPos += want
-            bytes -= want;
-        }
-        
-        /* store leftover */
-        if (bytes > 0) {
-            for i in 0..<bytes {
-                context.buffer[context.leftover + i] = m[mPos + i]
-            }
-            
-            context.leftover += bytes
-        }
-    }
 }

CryptoSwiftTests/CipherTests.swift

@@ -23,12 +23,10 @@ class CipherTests: XCTestCase {
     func testPoly1305() {
         let key:[Byte] = [0xdd,0xde,0xdf,0xe0,0xe1,0xe2,0xe3,0xe4,0xe5,0xe6,0xe7,0xe8,0xe9,0xea,0xeb,0xec,0xed,0xee,0xef,0xf0,0xf1,0xf2,0xf3,0xf4,0xf5,0xf6,0xf7,0xf8,0xf9,0xfa,0xfb,0xfc]
         let msg:[Byte] = [0x79,0x7a,0x7b,0x7c,0x7d,0x7e,0x7f,0x80,0x81,0x82,0x83,0x84,0x85,0x86,0x87,0x88,0x89,0x8a,0x8b,0x8c,0x8d,0x8e,0x8f,0x90,0x91,0x92,0x93,0x94,0x95,0x96,0x97,0x98,0x99,0x9a,0x9b,0x9c,0x9d,0x9e,0x9f,0xa0,0xa1,0xa2,0xa3,0xa4,0xa5,0xa6,0xa7,0xa8,0xa9,0xaa,0xab,0xac,0xad,0xae,0xaf,0xb0,0xb1,0xb2,0xb3,0xb4,0xb5,0xb6,0xb7,0xb8,0xb9,0xba,0xbb,0xbc,0xbd,0xbe,0xbf,0xc0,0xc1]
-        let mac:[Byte] = [0xe0,0xa8,0x7d,0x58,0xff,0x7f,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00]
         let expectedMac:[Byte] = [0xdd,0xb9,0xda,0x7d,0xdd,0x5e,0x52,0x79,0x27,0x30,0xed,0x5c,0xda,0x5f,0x90,0xa4]
         
-        let poly = Poly1305.withKey(key)
-        var resultMac = poly.auth(mac, message: msg)
-        XCTAssertEqual(resultMac!, expectedMac, "Invalid auth mac")
+        let mac = Poly1305.authenticate(key, message: msg)
+        XCTAssertEqual(mac!, expectedMac, "Invalid authentication result")
     }
 
     func testChaCha20() {