diff --git a/wolfcrypt/__init__.py b/wolfcrypt/__init__.py index 10893fb..4b286af 100644 --- a/wolfcrypt/__init__.py +++ b/wolfcrypt/__init__.py @@ -54,6 +54,8 @@ if top_level_py not in ["setup.py", "build_ffi.py"]: if TYPE_CHECKING: if _lib.CRYPTO_CB_ENABLED: from types import TracebackType + + from _cffi_backend import FFI from wolfcrypt.cryptocb import CryptoCallback # ty: ignore[possibly-missing-import] from wolfcrypt.exceptions import WolfCryptApiError @@ -71,10 +73,10 @@ if top_level_py not in ["setup.py", "build_ffi.py"]: if isinstance(exc_value, WolfCryptApiError): return exc_value.err_code log.error(exc_value) - return -1 + return _lib.WC_HW_E @_ffi.def_extern(onerror=crypto_cb_error_handler) - def py_wc_crypto_callback(device_id: int, info: _lib.wc_CryptoInfo, ctx: _ffi.CData) -> int: + def py_wc_crypto_callback(device_id: int, info: _lib.wc_CryptoInfo, ctx: FFI.CData) -> int: if ctx == _ffi.NULL: return _lib.CRYPTOCB_UNAVAILABLE crypto_cb: CryptoCallback = _ffi.from_handle(ctx) diff --git a/wolfcrypt/cryptocb.py b/wolfcrypt/cryptocb.py index b1c18d8..cf7b9bb 100644 --- a/wolfcrypt/cryptocb.py +++ b/wolfcrypt/cryptocb.py @@ -24,7 +24,7 @@ from __future__ import annotations import logging from collections import defaultdict -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, ClassVar, Final from typing_extensions import Self @@ -83,30 +83,38 @@ if _lib.CRYPTO_CB_ENABLED: class CryptoCallback: + # Keeps each registered instance (and its ctx handle) alive until it is + # unregistered, and records which instance owns a device id. + _registered: ClassVar[dict[int, CryptoCallback]] = {} + def __init__(self, device_id: int) -> None: self.device_id = device_id self.ctx = _ffi.new_handle(self) ret = _lib.wc_CryptoCb_RegisterDevice(device_id, _lib.py_wc_crypto_callback, self.ctx) if ret < 0: # pragma: no cover raise WolfCryptError(f"CryptoCb device registration error ({ret})") + CryptoCallback._registered[device_id] = self def __enter__(self) -> Self: return self def __exit__( self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None - ) -> bool: + ) -> None: self._unregister() - return False def __del__(self) -> None: - self._unregister() + try: + self._unregister() + except AttributeError: + # Can occur during interpreter shutdown + pass def callback(self, device_id: int, info: _lib.wc_CryptoInfo) -> int: log.debug("device_id=%d algo_type = %s", device_id, ALGO_TYPE_NAME[info.algo_type]) try: if info.algo_type == _lib.WC_ALGO_TYPE_HASH: - if info.hash.type not in DIGEST_SIZE: + if info.hash.type not in DIGEST_SIZE or not self._handles_hash(): return _lib.CRYPTOCB_UNAVAILABLE log.debug("hash_type = %s", HASH_TYPE_NAME[info.hash.type]) if info.hash.digest == _ffi.NULL: @@ -145,8 +153,19 @@ if _lib.CRYPTO_CB_ENABLED: def hash_finalize_callback(self, device_id: int, hash_type: int) -> bytes: raise NotImplementedError + def _handles_hash(self) -> bool: + # Update and finalize must both go to the callback, otherwise wolfCrypt + # would hash the data in software and finalize would ignore it. + cls = type(self) + return ( + cls.hash_update_callback is not CryptoCallback.hash_update_callback + and cls.hash_finalize_callback is not CryptoCallback.hash_finalize_callback + ) + def _unregister(self) -> None: - _lib.wc_CryptoCb_UnRegisterDevice(self.device_id) + if CryptoCallback._registered.get(self.device_id) is self: + del CryptoCallback._registered[self.device_id] + _lib.wc_CryptoCb_UnRegisterDevice(self.device_id) @classmethod def default_device_id(cls) -> int: diff --git a/wolfcrypt/hashes.py b/wolfcrypt/hashes.py index 3b952e7..8aa9a6a 100644 --- a/wolfcrypt/hashes.py +++ b/wolfcrypt/hashes.py @@ -73,7 +73,7 @@ class _Hash(ABC): @classmethod @abstractmethod - def new(cls, string: BytesOrStr | None) -> _Hash: ... + def new(cls, string: BytesOrStr | None, *, device_id: int = _lib.INVALID_DEVID) -> _Hash: ... def copy(self) -> _Hash: """ @@ -173,14 +173,14 @@ class _Hash(ABC): class _Sha(_Hash): @override @classmethod - def new(cls, string: BytesOrStr | None = None) -> _Hash: + def new(cls, string: BytesOrStr | None = None, *, device_id: int = _lib.INVALID_DEVID) -> _Hash: """ Creates a new hashing object and returns it. The optional **string** parameter, if supplied, will be immediately hashed into the object's starting state, as if obj.update(string) was called. """ - return cls(string) + return cls(string, device_id) if _lib.SHA_ENABLED: @@ -363,7 +363,7 @@ if _lib.SHA3_ENABLED: @override @classmethod - def new(cls, string: BytesOrStr | None = None, size: int = SHA3_384_DIGEST_SIZE, device_id: int = _lib.INVALID_DEVID) -> Sha3: + def new(cls, string: BytesOrStr | None = None, size: int = SHA3_384_DIGEST_SIZE, *, device_id: int = _lib.INVALID_DEVID) -> Sha3: return cls(string, size, device_id) @override @@ -488,7 +488,7 @@ if _lib.HMAC_ENABLED: @override @classmethod - def new(cls, key: BytesOrStr, string: BytesOrStr | None = None) -> _Hash: # pylint: disable=W0221 # ty: ignore[invalid-method-override] + def new(cls, key: BytesOrStr, string: BytesOrStr | None = None, *, device_id: int = _lib.INVALID_DEVID) -> _Hash: # pylint: disable=W0221 # ty: ignore[invalid-method-override] """ Creates a new hashing object and returns it. **key** is a required parameter containing a string giving the key @@ -496,7 +496,7 @@ if _lib.HMAC_ENABLED: will be immediately hashed into the object's starting state, as if obj.update(string) was called. """ - return cls(key, string) + return cls(key, string, device_id) @property