1 files changed, 29 insertions, 15 deletions

diff --git a/CPP/7zip/Crypto/Pbkdf2HmacSha1.cpp b/CPP/7zip/Crypto/Pbkdf2HmacSha1.cpp
index cbbdec89..a7fcb728 100644
--- a/CPP/7zip/Crypto/Pbkdf2HmacSha1.cpp
+++ b/CPP/7zip/Crypto/Pbkdf2HmacSha1.cpp
@@ -2,26 +2,34 @@
 
 #include "StdAfx.h"
 
+#include "../../../C/CpuArch.h"
+
 #include "HmacSha1.h"
 
 namespace NCrypto {
 namespace NSha1 {
 
-void Pbkdf2Hmac(const Byte *pwd, size_t pwdSize, const Byte *salt, size_t saltSize,
-    UInt32 numIterations, Byte *key, size_t keySize)
+void Pbkdf2Hmac(const Byte *pwd, size_t pwdSize,
+    const Byte *salt, size_t saltSize,
+    UInt32 numIterations,
+    Byte *key, size_t keySize)
 {
   CHmac baseCtx;
   baseCtx.SetKey(pwd, pwdSize);
-  for (UInt32 i = 1; keySize > 0; i++)
+  
+  for (UInt32 i = 1; keySize != 0; i++)
   {
     CHmac ctx = baseCtx;
     ctx.Update(salt, saltSize);
-    Byte u[kDigestSize] = { (Byte)(i >> 24), (Byte)(i >> 16), (Byte)(i >> 8), (Byte)(i) };
-    const unsigned int curSize = (keySize < kDigestSize) ? (unsigned int)keySize : kDigestSize;
+  
+    Byte u[kDigestSize];
+    SetBe32(u, i);
+    
     ctx.Update(u, 4);
     ctx.Final(u, kDigestSize);
 
-    unsigned int s;
+    const unsigned curSize = (keySize < kDigestSize) ? (unsigned)keySize : kDigestSize;
+    unsigned s;
     for (s = 0; s < curSize; s++)
       key[s] = u[s];
     
@@ -39,26 +47,32 @@ void Pbkdf2Hmac(const Byte *pwd, size_t pwdSize, const Byte *salt, size_t saltSi
   }
 }
 
-void Pbkdf2Hmac32(const Byte *pwd, size_t pwdSize, const UInt32 *salt, size_t saltSize,
-    UInt32 numIterations, UInt32 *key, size_t keySize)
+void Pbkdf2Hmac32(const Byte *pwd, size_t pwdSize,
+    const UInt32 *salt, size_t saltSize,
+    UInt32 numIterations,
+    UInt32 *key, size_t keySize)
 {
   CHmac32 baseCtx;
   baseCtx.SetKey(pwd, pwdSize);
-  for (UInt32 i = 1; keySize > 0; i++)
+  
+  for (UInt32 i = 1; keySize != 0; i++)
   {
     CHmac32 ctx = baseCtx;
     ctx.Update(salt, saltSize);
-    UInt32 u[kDigestSizeInWords] = { i };
-    const unsigned int curSize = (keySize < kDigestSizeInWords) ? (unsigned int)keySize : kDigestSizeInWords;
+    
+    UInt32 u[kNumDigestWords];
+    u[0] = i;
+    
     ctx.Update(u, 1);
-    ctx.Final(u, kDigestSizeInWords);
+    ctx.Final(u, kNumDigestWords);
 
     // Speed-optimized code start
     ctx = baseCtx;
     ctx.GetLoopXorDigest(u, numIterations - 1);
     // Speed-optimized code end
     
-    unsigned int s;
+    const unsigned curSize = (keySize < kNumDigestWords) ? (unsigned)keySize : kNumDigestWords;
+    unsigned s;
     for (s = 0; s < curSize; s++)
       key[s] = u[s];
     
@@ -67,8 +81,8 @@ void Pbkdf2Hmac32(const Byte *pwd, size_t pwdSize, const UInt32 *salt, size_t sa
     for (UInt32 j = numIterations; j > 1; j--)
     {
       ctx = baseCtx;
-      ctx.Update(u, kDigestSizeInWords);
-      ctx.Final(u, kDigestSizeInWords);
+      ctx.Update(u, kNumDigestWords);
+      ctx.Final(u, kNumDigestWords);
       for (s = 0; s < curSize; s++)
         key[s] ^= u[s];
     }