Fix OpenCL PBKDF2 buffer limits for long passwords

btcrecover/test/test_passwords.py

@@ -22,6 +22,8 @@
 
 
 import warnings, os, unittest, pickle, tempfile, shutil, multiprocessing, time, gc, filecmp, sys, hashlib, argparse
+
+from lib.opencl_brute import opencl
 if __name__ == '__main__':
     sys.path.append(os.path.join(os.path.dirname(__file__), "..", ".."))
 
@@ -3287,6 +3289,80 @@ def __init__(self):
         self.addTests(tl.loadTestsFromTestCase(Test08BIP39Passwords))
 
 
+@unittest.skipUnless(has_any_opencl_devices(), "requires OpenCL and a compatible device")
+class TestOpenCLPBKDF2BufferSizing(unittest.TestCase):
+    """Ensure OpenCL PBKDF2 buffer sizes track password limits."""
+
+    @classmethod
+    def setUpClass(cls):
+        super(TestOpenCLPBKDF2BufferSizing, cls).setUpClass()
+        if not has_any_opencl_devices():
+            raise unittest.SkipTest("requires OpenCL and a compatible device")
+
+        import pyopencl as cl
+
+        device = opencl_devices_list[0]
+        cls.platform_index = 0
+        cls.device_index = 0
+        for platform_index, platform in enumerate(cl.get_platforms()):
+            devices = platform.get_devices()
+            if device in devices:
+                cls.platform_index = platform_index
+                cls.device_index = devices.index(device)
+                break
+
+        cls.algos = opencl.opencl_algos(
+            cls.platform_index, 0, False, openclDevice=cls.device_index
+        )
+
+    def test_sha1_fast_kernel_limits_password_bytes(self):
+        password = b"a" * 32
+        salt = b"b" * 16
+        ctx = self.algos.cl_pbkdf2_init(
+            "sha1", saltlen=len(salt), dklen=16, max_password_bytes=len(password)
+        )
+        buf_structs = ctx[1]
+
+        self.assertEqual(buf_structs.pwdBufferSize_bytes, len(password))
+        self.assertEqual(buf_structs.inBufferSize_bytes, len(password))
+
+        result = self.algos.cl_pbkdf2(ctx, [password], salt, 1000, 16)
+        expected = [hashlib.pbkdf2_hmac("sha1", password, salt, 1000, 16)]
+        self.assertEqual(result, expected)
+
+    def test_sha1_falls_back_for_long_passwords(self):
+        password = b"l" * 200
+        salt = b"s" * 16
+        ctx = self.algos.cl_pbkdf2_init(
+            "sha1", saltlen=len(salt), dklen=20, max_password_bytes=len(password)
+        )
+        buf_structs = ctx[1]
+
+        self.assertEqual(buf_structs.pwdBufferSize_bytes, len(password))
+        self.assertEqual(buf_structs.inBufferSize_bytes, len(password))
+
+        result = self.algos.cl_pbkdf2(ctx, [password], salt, 2000, 20)
+        expected = [hashlib.pbkdf2_hmac("sha1", password, salt, 2000, 20)]
+        self.assertEqual(result, expected)
+
+    def test_saltlist_buffers_scale_with_password_length(self):
+        password = b"p" * 144
+        salts = [b"salt-one", b"salt-two-long"]
+        ctx = self.algos.cl_pbkdf2_saltlist_init(
+            "sha256", pwdlen=len(password), dklen=32
+        )
+        buf_structs = ctx[1]
+
+        self.assertEqual(buf_structs.pwdBufferSize_bytes, len(password))
+        self.assertEqual(buf_structs.inBufferSize_bytes, len(password))
+
+        result = self.algos.cl_pbkdf2_saltlist(ctx, password, salts, 4096, 32)
+        expected = [
+            hashlib.pbkdf2_hmac("sha256", password, salt, 4096, 32) for salt in salts
+        ]
+        self.assertEqual(result, expected)
+
+
 if __name__ == '__main__':
 
     import argparse

lib/opencl_brute/opencl.py

@@ -8,7 +8,13 @@
 from collections import deque
 from itertools import chain, repeat, zip_longest
 import numpy as np
-import pyopencl as cl
+
+try:
+    import pyopencl as cl
+except ImportError as _pyopencl_import_error:  # pragma: no cover - exercised in environments without PyOpenCL
+    cl = None
+else:
+    _pyopencl_import_error = None
 
 # Minimum number of items to execute in a single OpenCL batch.  Some
 # OpenCL runtimes (such as PoCL) crash when a kernel is launched with a
@@ -17,6 +23,15 @@
 from lib.opencl_brute.buffer_structs import buffer_structs
 import os, sys, inspect
 
+
+def _require_pyopencl():
+    """Ensure PyOpenCL is available before executing OpenCL operations."""
+
+    if cl is None:
+        raise ImportError(
+            "pyopencl is required for OpenCL acceleration but is not installed"
+        ) from _pyopencl_import_error
+
 current_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
 parent_dir = os.path.dirname(current_dir)
 
@@ -67,6 +82,7 @@ def __init__(
         N_value=15,
         openclDevice=0,
     ):
+        _require_pyopencl()
         self.workgroupsize = 0
         self.computeunits = 0
         self.wordSize = None
@@ -921,30 +937,70 @@ def func(s, pwdim, pass_g, salt_g, result_g):
             result = [hexRes[:dklen] for hexRes in result]
         return result
 
-    def cl_pbkdf2_init(self, rtype, saltlen, dklen):
+    def cl_pbkdf2_init(self, rtype, saltlen, dklen, max_password_bytes=256):
         bufStructs = buffer_structs()
+        if max_password_bytes is None:
+            max_password_bytes = 256
+        assert max_password_bytes > 0, "max_password_bytes must be positive"
         if rtype == "md5":
-            self.max_out_bytes = bufStructs.specifyMD5(128, saltlen, dklen)
+            max_in_bytes = max(128, max_password_bytes)
+            self.max_out_bytes = bufStructs.specifyMD5(
+                max_in_bytes,
+                saltlen,
+                dklen,
+                max_password_bytes=max_password_bytes,
+            )
             # hmac is defined in with pbkdf2, as a kernel function
             prg = self.opencl_ctx.compile(bufStructs, "md5.cl", "pbkdf2.cl")
         elif rtype == "sha1":
-            if saltlen < 32 and dklen < 32:
+            use_fast_kernel = (
+                saltlen < 32 and dklen < 32 and max_password_bytes <= 32
+            )
+            if use_fast_kernel:
                 dklen = 32
-                self.max_out_bytes = bufStructs.specifySHA1(32, saltlen, dklen)
+                self.max_out_bytes = bufStructs.specifySHA1(
+                    32,
+                    saltlen,
+                    dklen,
+                    max_password_bytes=32,
+                )
                 prg = self.opencl_ctx.compile(bufStructs, "pbkdf2_sha1_32.cl", None)
             else:
-                self.max_out_bytes = bufStructs.specifySHA1(128, saltlen, dklen)
+                max_in_bytes = max(128, max_password_bytes)
+                self.max_out_bytes = bufStructs.specifySHA1(
+                    max_in_bytes,
+                    saltlen,
+                    dklen,
+                    max_password_bytes=max_password_bytes,
+                )
                 prg = self.opencl_ctx.compile(bufStructs, "sha1.cl", "pbkdf2.cl")
         elif rtype == "sha256":
-            if saltlen <= 64 and dklen <= 64:
+            use_fast_kernel = (
+                saltlen <= 64 and dklen <= 64 and max_password_bytes <= 32
+            )
+            if use_fast_kernel:
                 dklen = 64
-            self.max_out_bytes = bufStructs.specifySHA2(256, 128, saltlen, dklen)
-            if saltlen <= 64 and dklen <= 64:
+            max_in_bytes = max(128, max_password_bytes)
+            self.max_out_bytes = bufStructs.specifySHA2(
+                256,
+                max_in_bytes,
+                saltlen,
+                dklen,
+                max_password_bytes=max_password_bytes,
+            )
+            if use_fast_kernel:
                 prg = self.opencl_ctx.compile(bufStructs, "pbkdf2_sha256_32.cl", None)
             else:
                 prg = self.opencl_ctx.compile(bufStructs, "sha256.cl", "pbkdf2.cl")
         elif rtype == "sha512":
-            self.max_out_bytes = bufStructs.specifySHA2(512, 256, saltlen, dklen)
+            max_in_bytes = max(256, max_password_bytes)
+            self.max_out_bytes = bufStructs.specifySHA2(
+                512,
+                max_in_bytes,
+                saltlen,
+                dklen,
+                max_password_bytes=max_password_bytes,
+            )
             prg = self.opencl_ctx.compile(bufStructs, "sha512.cl", "pbkdf2.cl")
         else:
             assert "Error on hash type, unknown !!!"
@@ -985,7 +1041,7 @@ def cl_pbkdf2_saltlist_init(self, type, pwdlen, dklen):
         bufStructs = buffer_structs()
         if type == "md5":
             self.max_out_bytes = bufStructs.specifyMD5(
-                max_in_bytes=128,
+                max_in_bytes=max(128, pwdlen),
                 max_salt_bytes=128,
                 dklen=dklen,
                 max_password_bytes=pwdlen,
@@ -994,7 +1050,7 @@ def cl_pbkdf2_saltlist_init(self, type, pwdlen, dklen):
             prg = self.opencl_ctx.compile(bufStructs, "md5.cl", "pbkdf2.cl")
         elif type == "sha1":
             self.max_out_bytes = bufStructs.specifySHA1(
-                max_in_bytes=128,
+                max_in_bytes=max(128, pwdlen),
                 max_salt_bytes=128,
                 dklen=dklen,
                 max_password_bytes=pwdlen,
@@ -1004,7 +1060,7 @@ def cl_pbkdf2_saltlist_init(self, type, pwdlen, dklen):
         elif type == "sha256":
             self.max_out_bytes = bufStructs.specifySHA2(
                 hashDigestSize_bits=256,
-                max_in_bytes=128,
+                max_in_bytes=max(128, pwdlen),
                 max_salt_bytes=128,
                 dklen=dklen,
                 max_password_bytes=pwdlen,
@@ -1013,7 +1069,7 @@ def cl_pbkdf2_saltlist_init(self, type, pwdlen, dklen):
         elif type == "sha512":
             self.max_out_bytes = bufStructs.specifySHA2(
                 hashDigestSize_bits=512,
-                max_in_bytes=256,
+                max_in_bytes=max(256, pwdlen),
                 max_salt_bytes=128,
                 dklen=dklen,
                 max_password_bytes=pwdlen,

lib/opencl_brute/opencl_information.py

@@ -10,17 +10,31 @@
     Refactored out of 'opencl.py'
 '''
 
-import pyopencl as cl
+try:
+    import pyopencl as cl
+except ImportError as _pyopencl_import_error:  # pragma: no cover - exercised when PyOpenCL missing
+    cl = None
+else:
+    _pyopencl_import_error = None
+
+
+def _require_pyopencl():
+    if cl is None:
+        raise ImportError(
+            "pyopencl is required for querying OpenCL information but is not installed"
+        ) from _pyopencl_import_error
 
 class opencl_information:
     def __init__(self):
         pass
 
     def printplatforms(self):
+        _require_pyopencl()
         for i,platformNum in enumerate(cl.get_platforms()):
             print('Platform %d - Name %s, Vendor %s' %(i,platformNum.name,platformNum.vendor))
 
     def printfullinfo(self):
+        _require_pyopencl()
         print('\n' + '=' * 60 + '\nOpenCL Platforms and Devices')
         for i,platformNum in enumerate(cl.get_platforms()):
             print('=' * 60)