From d4001c78b61c364231fbbdbff49fb71ee72debc8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marcin=20Krzyz=CC=87anowski?=
 <marcin.krzyzanowski@gmail.com>
Date: Tue, 5 Apr 2016 01:27:46 +0200
Subject: [PATCH] Allows direct use of HMAC and Poly1305

---
 Sources/CryptoSwift/Authenticator.swift |  4 +--
 Sources/CryptoSwift/HMAC.swift          | 12 ++-----
 Sources/CryptoSwift/PKCS5/PBKDF2.swift  |  6 ++--
 Sources/CryptoSwift/Poly1305.swift      | 44 +++++++++++--------------
 4 files changed, 27 insertions(+), 39 deletions(-)

diff --git a/Sources/CryptoSwift/Authenticator.swift b/Sources/CryptoSwift/Authenticator.swift
index 30987ac..cc15263 100644
--- a/Sources/CryptoSwift/Authenticator.swift
+++ b/Sources/CryptoSwift/Authenticator.swift
@@ -31,12 +31,12 @@ public enum Authenticator {
     public func authenticate(message: [UInt8]) throws -> [UInt8] {
         switch (self) {
         case .Poly1305(let key):
-            guard let auth = CryptoSwift.Poly1305.authenticate(key: key, message: message) else {
+            guard let auth = CryptoSwift.Poly1305(key: key)?.authenticate(message) else {
                 throw Error.AuthenticateError
             }
             return auth
         case .HMAC(let key, let variant):
-            guard let auth = CryptoSwift.HMAC.authenticate(key: key, message: message, variant: variant) else {
+            guard let auth = CryptoSwift.HMAC(key: key, variant: variant)?.authenticate(message) else {
                 throw Error.AuthenticateError
             }
             return auth
diff --git a/Sources/CryptoSwift/HMAC.swift b/Sources/CryptoSwift/HMAC.swift
index d546a72..a8f04e4 100644
--- a/Sources/CryptoSwift/HMAC.swift
+++ b/Sources/CryptoSwift/HMAC.swift
@@ -53,14 +53,8 @@ final public class HMAC {
     
     var key:[UInt8]
     let variant:Variant
-    
-    class internal func authenticate(key  key: [UInt8], message: [UInt8], variant:HMAC.Variant = .md5) -> [UInt8]? {
-        return HMAC(key, variant: variant)?.authenticate(message: message)
-    }
 
-    // MARK: - Private
-    
-    internal init? (_ key: [UInt8], variant:HMAC.Variant = .md5) {
+    public init? (key: [UInt8], variant:HMAC.Variant = .md5) {
         self.variant = variant
         self.key = key
 
@@ -74,8 +68,8 @@ final public class HMAC {
             self.key = key + [UInt8](count: variant.blockSize() - key.count, repeatedValue: 0)
         }
     }
-    
-    internal func authenticate(message  message:[UInt8]) -> [UInt8]? {
+
+    public func authenticate(message:[UInt8]) -> [UInt8]? {
         var opad = [UInt8](count: variant.blockSize(), repeatedValue: 0x5c)
         for (idx, _) in key.enumerate() {
             opad[idx] = key[idx] ^ opad[idx]
diff --git a/Sources/CryptoSwift/PKCS5/PBKDF2.swift b/Sources/CryptoSwift/PKCS5/PBKDF2.swift
index 46636a2..a8f21f9 100644
--- a/Sources/CryptoSwift/PKCS5/PBKDF2.swift
+++ b/Sources/CryptoSwift/PKCS5/PBKDF2.swift
@@ -21,7 +21,7 @@ public extension PKCS5 {
         private let prf: HMAC
 
         init(password: [UInt8], salt: [UInt8], iterations: Int /* c */, keyLength: Int /* dkLen */ , hashVariant: HMAC.Variant = .sha256) throws {
-            guard let prf = HMAC(password, variant: hashVariant) where (keyLength <= Int(pow(2,32) as Float - 1)) && (iterations > 0) && (password.count > 0) && (salt.count > 0) else {
+            guard let prf = HMAC(key: password, variant: hashVariant) where (keyLength <= Int(pow(2,32) as Float - 1)) && (iterations > 0) && (password.count > 0) && (salt.count > 0) else {
                 throw Error.InvalidInput
             }
 
@@ -57,7 +57,7 @@ private extension PKCS5.PBKDF2 {
     // F (P, S, c, i) = U_1 \xor U_2 \xor ... \xor U_c
     // U_1 = PRF (P, S || INT (i))
     private func calculateBlock(salt salt: [UInt8], blockNum: UInt) -> [UInt8]? {
-        guard let u1 = prf.authenticate(message: salt + INT(blockNum)) else {
+        guard let u1 = prf.authenticate(salt + INT(blockNum)) else {
             return nil
         }
 
@@ -67,7 +67,7 @@ private extension PKCS5.PBKDF2 {
             // U_2 = PRF (P, U_1) ,
             // U_c = PRF (P, U_{c-1}) .
             for _ in 2...self.iterations {
-                u = prf.authenticate(message: u)!
+                u = prf.authenticate(u)!
                 for x in 0..<ret.count {
                     ret[x] = ret[x] ^ u[x]
                 }
diff --git a/Sources/CryptoSwift/Poly1305.swift b/Sources/CryptoSwift/Poly1305.swift
index 64a5bf9..cf2869f 100644
--- a/Sources/CryptoSwift/Poly1305.swift
+++ b/Sources/CryptoSwift/Poly1305.swift
@@ -74,32 +74,17 @@ final public class Poly1305 {
             }
         }
     }
-    
-    // MARK: - Internal
-    
+
     /**
-    Calculate Message Authentication Code (MAC) for message.
-    Calculation context is discarder on instance deallocation.
-    
-    - parameter key:     256-bit key
-    - parameter message: Message
-    
-    - returns: Message Authentication Code
-    */
-    class internal func authenticate(key  key: [UInt8], message: [UInt8]) -> [UInt8]? {
-        return Poly1305(key)?.authenticate(message: message)
-    }
-    
-    // MARK: - Private
-    
-    private init? (_ key: [UInt8]) {
-        ctx = Context(key)
-        if (ctx == nil) {
-            return nil
-        }
-    }
-    
-    private func authenticate(message  message:[UInt8]) -> [UInt8]? {
+     Calculate Message Authentication Code (MAC) for message.
+     Calculation context is discarder on instance deallocation.
+
+     - parameter key:     256-bit key
+     - parameter message: Message
+
+     - returns: Message Authentication Code
+     */
+    public func authenticate(message:[UInt8]) -> [UInt8]? {
         if let ctx = self.ctx {
             update(ctx, message: message)
             return finish(ctx)
@@ -107,6 +92,15 @@ final public class Poly1305 {
         return nil
     }
     
+    public init? (key: [UInt8]) {
+        ctx = Context(key)
+        if (ctx == nil) {
+            return nil
+        }
+    }
+
+    // MARK: - Private
+
     /**
     Add message to be processed