add lru cache for _generate_salted_password for testing. staticmethod

asyncpg/protocol/scram.pxd

@@ -4,7 +4,6 @@
 # This module is part of asyncpg and is released under
 # the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0
 
-
 cdef class SCRAMAuthentication:
     cdef:
         readonly bytes authentication_method
@@ -24,8 +23,8 @@ cdef class SCRAMAuthentication:
     cdef create_client_final_message(self, str password)
     cdef parse_server_first_message(self, bytes server_response)
     cdef verify_server_final_message(self, bytes server_final_message)
-    cdef _bytes_xor(self, bytes a, bytes b)
+    @staticmethod
+    cdef _bytes_xor(bytes a, bytes b)
     cdef _generate_client_nonce(self, int num_bytes)
     cdef _generate_client_proof(self, str password)
-    cdef _generate_salted_password(self, str password, bytes salt, int iterations)
     cdef _normalize_password(self, str original_password)

asyncpg/protocol/scram.pyx

@@ -6,6 +6,7 @@
 
 
 import base64
+import functools
 import hashlib
 import hmac
 import re
@@ -14,6 +15,8 @@ import stringprep
 import unicodedata
 
 
+HMAC_CYCLES = 0
+
 @cython.final
 cdef class SCRAMAuthentication:
     """Contains the protocol for generating and a SCRAM hashed password.
@@ -168,6 +171,7 @@ cdef class SCRAMAuthentication:
 
     cdef verify_server_final_message(self, bytes server_final_message):
         """Verify the final message from the server"""
+        global HMAC_CYCLES
         cdef:
             bytes server_signature
 
@@ -179,11 +183,13 @@ cdef class SCRAMAuthentication:
 
         verify_server_signature = hmac.new(self.server_key.digest(),
             self.authorization_message, self.DIGEST)
+        HMAC_CYCLES += 1
         # validate the server signature against the verifier
         return server_signature == base64.b64encode(
             verify_server_signature.digest())
 
-    cdef _bytes_xor(self, bytes a, bytes b):
+    @staticmethod
+    cdef _bytes_xor(bytes a, bytes b):
         """XOR two bytestrings together"""
         return bytes(a_i ^ b_i for a_i, b_i in zip(a, b))
 
@@ -196,23 +202,30 @@ cdef class SCRAMAuthentication:
         return base64.b64encode(token)
 
     cdef _generate_client_proof(self, str password):
+        global HMAC_CYCLES
         """need to ensure a server response exists, i.e. """
         cdef:
             bytes salted_password
+            bytes tmp
+            str key
 
         if any([getattr(self, val) is None for val in
                 self.REQUIREMENTS_CLIENT_PROOF]):
             raise Exception(
                 "you need values from server to generate a client proof")
         # generate a salt password
-        salted_password = self._generate_salted_password(password,
+        salted_password = SCRAMAuthentication._generate_salted_password(password,
             self.password_salt, self.password_iterations)
+
         # client key is derived from the salted password
         client_key = hmac.new(salted_password, b"Client Key", self.DIGEST)
+
+        HMAC_CYCLES += 1
         # this allows us to compute the stored key that is residing on the server
         stored_key = self.DIGEST(client_key.digest())
         # as well as compute the server key
         self.server_key = hmac.new(salted_password, b"Server Key", self.DIGEST)
+        HMAC_CYCLES += 1
         # build the authorization message that will be used in the
         # client signature
         # the "c=" portion is for the channel binding, but this is not
@@ -224,11 +237,15 @@ cdef class SCRAMAuthentication:
         # sign!
         client_signature = hmac.new(stored_key.digest(),
             self.authorization_message, self.DIGEST)
+        HMAC_CYCLES += 1
         # and the proof
-        return self._bytes_xor(client_key.digest(), client_signature.digest())
+        return SCRAMAuthentication._bytes_xor(client_key.digest(), client_signature.digest())
 
-    cdef _generate_salted_password(self, str password, bytes salt, int iterations):
+    @functools.lru_cache
+    @staticmethod
+    def _generate_salted_password(str password, bytes salt, int iterations):
         """This follows the "Hi" algorithm specified in RFC5802"""
+        global HMAC_CYCLES
         cdef:
             bytes p
             bytes s
@@ -241,16 +258,18 @@ cdef class SCRAMAuthentication:
         s = base64.b64decode(salt)
         # the initial signature is the salt with a terminator of a 32-bit string
         # ending in 1
-        ui = hmac.new(p, s + b'\x00\x00\x00\x01', self.DIGEST)
+        ui = hmac.new(p, s + b'\x00\x00\x00\x01', SCRAMAuthentication.DIGEST)
+        HMAC_CYCLES += 1
         # grab the initial digest
         u = ui.digest()
         # for X number of iterations, recompute the HMAC signature against the
         # password and the latest iteration of the hash, and XOR it with the
         # previous version
         for x in range(iterations - 1):
             ui = hmac.new(p, ui.digest(), hashlib.sha256)
+            HMAC_CYCLES += 1
             # this is a fancy way of XORing two byte strings together
-            u = self._bytes_xor(u, ui.digest())
+            u = SCRAMAuthentication._bytes_xor(u, ui.digest())
         return u
 
     cdef _normalize_password(self, str original_password):