From 75798de4ddd95ffdc67af52e7179b279c71c7379 Mon Sep 17 00:00:00 2001 From: libalpm64 <159192189+libalpm64@users.noreply.github.com> Date: Thu, 20 Aug 2026 06:33:28 -0400 Subject: [PATCH 1/4] ~ --- pixi.lock | 152 +++++++++++----------- pixi.toml | 6 +- src/thistle/__init__.mojo | 2 +- src/thistle/aes.mojo | 170 +++++++++++++++--------- src/thistle/aes_gpu.mojo | 67 +++++----- src/thistle/aes_ni.mojo | 185 ++++++++++++++------------- src/thistle/argon2.mojo | 169 ++++++++++++++---------- src/thistle/blake2b.mojo | 30 ++--- src/thistle/blake3.mojo | 6 +- src/thistle/camellia.mojo | 46 +++---- src/thistle/chacha20.mojo | 24 ++-- src/thistle/chacha20poly1305.mojo | 87 ++++++++----- src/thistle/curve25519.mojo | 36 ++---- src/thistle/ed25519.mojo | 124 +++++++++--------- src/thistle/ed25519_table.mojo | 4 +- src/thistle/kcipher2.mojo | 2 +- src/thistle/ml_dsa.mojo | 37 +++--- src/thistle/ml_kem.mojo | 76 ++++++----- src/thistle/p256.mojo | 79 +++++++----- src/thistle/p384.mojo | 79 +++++++----- src/thistle/pbkdf2.mojo | 65 ++++++---- src/thistle/poly1305.mojo | 20 ++- src/thistle/random.mojo | 26 ++-- src/thistle/rsa.mojo | 169 ++++++++++++------------ src/thistle/sha2.mojo | 81 +++++------- src/thistle/sha3.mojo | 54 +++++--- src/thistle/sha_ni.mojo | 21 +-- src/thistle/utils.mojo | 58 ++++++--- src/thistle/x25519.mojo | 23 +++- tests/benchmark.mojo | 169 ++++++++++++------------ tests/dudect.mojo | 23 ++-- tests/gen_ed25519_table.mojo | 2 +- tests/test_aes_gpu.mojo | 100 +++++---------- tests/test_random.mojo | 4 +- tests/test_security_boundaries.mojo | 119 +++++++++++++++++ tests/test_sha3_capacity_guard.mojo | 10 ++ tests/test_sha3_capacity_guard.sh | 18 +++ tests/test_signing.mojo | 14 +- tests/test_wycheproof_p256_ecdh.mojo | 4 +- tests/test_wycheproof_p384_ecdh.mojo | 4 +- tests/test_wycheproof_x25519.mojo | 8 +- tests/thistle_test_vectors.mojo | 81 ++++++------ 42 files changed, 1399 insertions(+), 1055 deletions(-) create mode 100644 tests/test_security_boundaries.mojo create mode 100644 tests/test_sha3_capacity_guard.mojo create mode 100644 tests/test_sha3_capacity_guard.sh diff --git a/pixi.lock b/pixi.lock index 801bbef..38824cb 100644 --- a/pixi.lock +++ b/pixi.lock @@ -42,13 +42,13 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/libxcrypt-4.4.36-hd590300_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/libzlib-1.3.2-h25fd6f3_2.conda - conda: https://conda.anaconda.org/conda-forge/noarch/markdown-it-py-4.2.0-pyhd8ed1ab_0.conda - - conda: https://conda.modular.com/max/linux-64/max-26.4.0-3.12release.conda - - conda: https://conda.modular.com/max/linux-64/max-core-26.4.0-release.conda - - conda: https://conda.modular.com/max/noarch/mblack-26.4.0-release.conda + - conda: https://conda.modular.com/max/linux-64/max-26.5.0-3.12release.conda + - conda: https://conda.modular.com/max/linux-64/max-core-26.5.0-release.conda + - conda: https://conda.modular.com/max/noarch/mblack-26.5.0-release.conda - conda: https://conda.anaconda.org/conda-forge/noarch/mdurl-0.1.2-pyhd8ed1ab_1.conda - - conda: https://conda.modular.com/max/linux-64/mojo-1.0.0b2-release.conda - - conda: https://conda.modular.com/max/linux-64/mojo-compiler-1.0.0b2-release.conda - - conda: https://conda.modular.com/max/noarch/mojo-python-1.0.0b2-release.conda + - conda: https://conda.modular.com/max/linux-64/mojo-1.0.0-release.conda + - conda: https://conda.modular.com/max/linux-64/mojo-compiler-1.0.0-release.conda + - conda: https://conda.modular.com/max/noarch/mojo-python-1.0.0-release.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/msgspec-0.21.1-py312h4c3975b_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/mypy_extensions-1.1.0-pyha770c72_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/ncurses-6.6-hdb14827_0.conda @@ -107,13 +107,13 @@ environments: - conda: https://conda.anaconda.org/conda-forge/osx-arm64/libzlib-1.3.2-h8088a28_2.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/llvm-openmp-22.1.8-hc7d1edf_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/markdown-it-py-4.2.0-pyhd8ed1ab_0.conda - - conda: https://conda.modular.com/max/osx-arm64/max-26.4.0-3.12release.conda - - conda: https://conda.modular.com/max/osx-arm64/max-core-26.4.0-release.conda - - conda: https://conda.modular.com/max/noarch/mblack-26.4.0-release.conda + - conda: https://conda.modular.com/max/osx-arm64/max-26.5.0-3.12release.conda + - conda: https://conda.modular.com/max/osx-arm64/max-core-26.5.0-release.conda + - conda: https://conda.modular.com/max/noarch/mblack-26.5.0-release.conda - conda: https://conda.anaconda.org/conda-forge/noarch/mdurl-0.1.2-pyhd8ed1ab_1.conda - - conda: https://conda.modular.com/max/osx-arm64/mojo-1.0.0b2-release.conda - - conda: https://conda.modular.com/max/osx-arm64/mojo-compiler-1.0.0b2-release.conda - - conda: https://conda.modular.com/max/noarch/mojo-python-1.0.0b2-release.conda + - conda: https://conda.modular.com/max/osx-arm64/mojo-1.0.0-release.conda + - conda: https://conda.modular.com/max/osx-arm64/mojo-compiler-1.0.0-release.conda + - conda: https://conda.modular.com/max/noarch/mojo-python-1.0.0-release.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/msgspec-0.21.1-py312h2bbb03f_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/mypy_extensions-1.1.0-pyha770c72_0.conda - conda: https://conda.anaconda.org/conda-forge/osx-arm64/ncurses-6.6-h1d4f5a5_0.conda @@ -745,58 +745,62 @@ packages: license_family: MIT size: 69017 timestamp: 1778169663339 -- conda: https://conda.modular.com/max/linux-64/max-26.4.0-3.12release.conda - sha256: 710bd1aa27a22f13e02defb0b38debe1a5ded2320f6cf74111e5f79e8ab5058c - md5: b0a88366853d518a7c27500922c1091d +- conda: https://conda.modular.com/max/linux-64/max-26.5.0-3.12release.conda + sha256: b027cc9d8d82e1ec07002e4022c3afecc5f9947aa55239a703d3bca7fac79179 + md5: a7785a6360c0ad216908199db53ca5cc depends: + - click >=8.0.0 + - exceptiongroup >=0.2.2 - msgspec >=0.19.0 - numpy >=1.18 - - psutil >=6.1.1 + - psutil >=7.0.0 - rich >=13.0.1 - taskgroup >=0.2.2 - typing-extensions >=4.12.2 - python 3.12.* - python-gil - - max-core ==26.4.0 + - max-core ==26.5.0 license: LicenseRef-Modular-Proprietary - size: 7817633 - timestamp: 1781237249041 -- conda: https://conda.modular.com/max/osx-arm64/max-26.4.0-3.12release.conda - sha256: 251834db6616953c3cbf73088671767e6837ed9bc6f8e29d4b65710a08fc9fba - md5: b403012fc7b0a2f399063f8f7ee0ad98 + size: 18560201 + timestamp: 1786292559995 +- conda: https://conda.modular.com/max/osx-arm64/max-26.5.0-3.12release.conda + sha256: 653a6a041d91da89c9d575fb049ed776c3c92ddf3ec4f03e4e7e02cd544d4e7b + md5: 4ab4be050e1068272eb3d28f47c3e602 depends: + - click >=8.0.0 + - exceptiongroup >=0.2.2 - msgspec >=0.19.0 - numpy >=1.18 - - psutil >=6.1.1 + - psutil >=7.0.0 - rich >=13.0.1 - taskgroup >=0.2.2 - typing-extensions >=4.12.2 - python 3.12.* - python-gil - - max-core ==26.4.0 + - max-core ==26.5.0 license: LicenseRef-Modular-Proprietary - size: 7072840 - timestamp: 1781237594992 -- conda: https://conda.modular.com/max/linux-64/max-core-26.4.0-release.conda - sha256: c036b1d9966ef6600a0c7e63bd72e358267bb33b7c83a53b0205e6a4edb31023 - md5: bc10308c187b5cd300578080aa748dbf + size: 16362991 + timestamp: 1786294957687 +- conda: https://conda.modular.com/max/linux-64/max-core-26.5.0-release.conda + sha256: a2aca8a7dc43cdac10f286ba31a811468794631e514f0df2687464cba15deeb0 + md5: 85c39f2cdbb36c2837c59150fad012d9 depends: - - mojo-compiler ==1.0.0b2 + - mojo-compiler ==1.0.0 license: LicenseRef-Modular-Proprietary - size: 88516478 - timestamp: 1781237250702 -- conda: https://conda.modular.com/max/osx-arm64/max-core-26.4.0-release.conda - sha256: 7fbd591ab7b9f5fd7633daf9012006bef92c683ceb31f62ac25c1f4bdbbd13cc - md5: 57a8a3e62959934dd4550ccd1ea5f617 + size: 117588665 + timestamp: 1786292557073 +- conda: https://conda.modular.com/max/osx-arm64/max-core-26.5.0-release.conda + sha256: 774f2b368d1473553debbfb61566e2d1091951990f003340eb836fd4baeb50ea + md5: 24bbf832b1dc7d486b3e8b051b7c001c depends: - - mojo-compiler ==1.0.0b2 + - mojo-compiler ==1.0.0 license: LicenseRef-Modular-Proprietary - size: 50764608 - timestamp: 1781237550544 -- conda: https://conda.modular.com/max/noarch/mblack-26.4.0-release.conda + size: 66719553 + timestamp: 1786294908386 +- conda: https://conda.modular.com/max/noarch/mblack-26.5.0-release.conda noarch: python - sha256: e2497cbefbf5d962cd0db8c1b893e7e48d8f3603cab84a9d8bad0030b5bb26b1 - md5: ec6bf2b9588f5ba4bd3472c0773c4892 + sha256: 49102b55366eab56b04620533302b8693b4fa4a6c5972b3cc978b3f3bad83ebe + md5: 0cbfc119215c063cd279f0a6beb04285 depends: - python >=3.10 - click >=8.0.0 @@ -806,8 +810,8 @@ packages: - platformdirs >=2 - tomli >=1.1.0 license: LicenseRef-Modular-Proprietary - size: 135301 - timestamp: 1781220755102 + size: 137187 + timestamp: 1786151006088 - conda: https://conda.anaconda.org/conda-forge/noarch/mdurl-0.1.2-pyhd8ed1ab_1.conda sha256: 78c1bbe1723449c52b7a9df1af2ee5f005209f67e40b6e1d3c7619127c43b1c7 md5: 592132998493b3ff25fd7479396e8351 @@ -817,53 +821,53 @@ packages: license_family: MIT size: 14465 timestamp: 1733255681319 -- conda: https://conda.modular.com/max/linux-64/mojo-1.0.0b2-release.conda - sha256: 7553366c138bae2fcf8617582f1bb6d6b6742087bd6a83ea83aa41c023499382 - md5: 8d4390edde2dcf29880c57b3088f15aa +- conda: https://conda.modular.com/max/linux-64/mojo-1.0.0-release.conda + sha256: 5778f999b69cf77bd6f07ccaf32d0bafc258fe75206afc8e88919269cebb6e75 + md5: afeb126acb57602a5f9b8153e75e1c3c depends: - python >=3.10 - - mojo-compiler ==1.0.0b2 - - mblack ==26.4.0 + - mojo-compiler ==1.0.0 + - mblack ==26.5.0 - jupyter_client >=8.6.2,<8.7 license: LicenseRef-Modular-Proprietary - size: 96475709 - timestamp: 1781237127405 -- conda: https://conda.modular.com/max/osx-arm64/mojo-1.0.0b2-release.conda - sha256: 39bb59a309af5e4c78c077263e5824e33bba8159a1e0fc39ff2786a1f9d5762b - md5: 1a6b6dd12763032e46b956407b39ab48 + size: 115641078 + timestamp: 1786292404342 +- conda: https://conda.modular.com/max/osx-arm64/mojo-1.0.0-release.conda + sha256: 200bd9ace6e06ad2a9b4f4cce9afa3e5f3c00b3d87a6d57206249df2a5568e38 + md5: 954433eb905b424cc3d4325c9c2242de depends: - python >=3.10 - - mojo-compiler ==1.0.0b2 - - mblack ==26.4.0 + - mojo-compiler ==1.0.0 + - mblack ==26.5.0 - jupyter_client >=8.6.2,<8.7 license: LicenseRef-Modular-Proprietary - size: 84573003 - timestamp: 1781237324882 -- conda: https://conda.modular.com/max/linux-64/mojo-compiler-1.0.0b2-release.conda - sha256: 87f8e9ef573f37072423aeb56a7c468bdb5f5889476bce4aa7c3396665c651fe - md5: 6be85d17872f5c2370015ca4ec07dd0d + size: 101155040 + timestamp: 1786294689438 +- conda: https://conda.modular.com/max/linux-64/mojo-compiler-1.0.0-release.conda + sha256: 4394c6146d47ec7794a9a3ed5775ae158f59f83f8e1aed59408b17c4909821b3 + md5: 75d2e2ca9d87fdc76513d7939202e7dc depends: - - mojo-python ==1.0.0b2 + - mojo-python ==1.0.0 license: LicenseRef-Modular-Proprietary - size: 85342385 - timestamp: 1781237127405 -- conda: https://conda.modular.com/max/osx-arm64/mojo-compiler-1.0.0b2-release.conda - sha256: 91c4d590a152ec2e26846955fcd7ec02796dfaffefa006a1c0c5790575be2051 - md5: cd8d3f4a22e5f14dd909d0d5570c4940 + size: 68511359 + timestamp: 1786292410712 +- conda: https://conda.modular.com/max/osx-arm64/mojo-compiler-1.0.0-release.conda + sha256: c52054bc444d851e5c38cc33e790fb011a1080470244aaa59351ac5056d08c59 + md5: b5103c9d9978173a1e97c5e5eed1aeba depends: - - mojo-python ==1.0.0b2 + - mojo-python ==1.0.0 license: LicenseRef-Modular-Proprietary - size: 62655166 - timestamp: 1781237320083 -- conda: https://conda.modular.com/max/noarch/mojo-python-1.0.0b2-release.conda + size: 61221085 + timestamp: 1786294690001 +- conda: https://conda.modular.com/max/noarch/mojo-python-1.0.0-release.conda noarch: python - sha256: 37bf44a74edcd2ebaec7fe8299ef183e006ad051bfcf21caa06becc0e0eeffd5 - md5: 8781bd673909835fa59ce818e3835839 + sha256: fc449a47c7707ed96a794768e29d5d2b0bf7afab373e2a9751f281ed7c4f822c + md5: 4a27e1596a5c2d27330a2377090b12c5 depends: - python >=3.10 license: LicenseRef-Modular-Proprietary - size: 24682 - timestamp: 1781220753246 + size: 24040 + timestamp: 1785891558838 - conda: https://conda.anaconda.org/conda-forge/linux-64/msgspec-0.21.1-py312h4c3975b_0.conda sha256: 25eb262c378a922eeed85c941ab7de2687ea842daed80521b861b7472b5a7f9a md5: 5e07dc45b4458c19fdc085bd6c1aa51f diff --git a/pixi.toml b/pixi.toml index c8f7269..c02d509 100644 --- a/pixi.toml +++ b/pixi.toml @@ -13,9 +13,11 @@ test-ml-dsa = "mojo run -I src/ tests/test_ml_dsa.mojo" test-p256-ecdh = "mojo run -I src/ tests/test_wycheproof_p256_ecdh.mojo" test-p384-ecdh = "mojo run -I src/ tests/test_wycheproof_p384_ecdh.mojo" test-signing = "mojo run -I src/ tests/test_signing.mojo" +test-security-boundaries = "mojo run -I src/ tests/test_security_boundaries.mojo" +test-sha3-guard = "zsh tests/test_sha3_capacity_guard.sh" bench = "mojo run -I src/ tests/benchmark.mojo" [dependencies] python = "3.12.*" -max = "==26.4.0" -mojo = ">=1.0.0b2,<2" +max = "==26.5.0" +mojo = "==1.0.0" diff --git a/src/thistle/__init__.mojo b/src/thistle/__init__.mojo index 41d801d..0c25d74 100644 --- a/src/thistle/__init__.mojo +++ b/src/thistle/__init__.mojo @@ -13,7 +13,7 @@ from .pbkdf2 import hmac_sha256, hmac_sha384, hmac_sha512 from .pbkdf2 import pbkdf2_hmac_sha256, pbkdf2_hmac_sha512 from .argon2 import Argon2id, argon2id_hash_string -from .aes import AESKey, expand_key_128, expand_key_192, expand_key_256 +from .aes import AESKey, AESExpandedKey, expand_key_128, expand_key_192, expand_key_256 from .aes_ni import has_aes_ni, aes_gcm_ctr_kernel, aes_gcm_encrypt, aes_gcm_decrypt, AESGCMContext from .aes_gpu import aes_gpu_kernel_ecb, aes_gpu_kernel_ctr, aes_gpu_kernel_gcm_ctr from .camellia import CamelliaCipher diff --git a/src/thistle/aes.mojo b/src/thistle/aes.mojo index 3a93c3f..e25d823 100644 --- a/src/thistle/aes.mojo +++ b/src/thistle/aes.mojo @@ -3,7 +3,7 @@ AES CPU implementation """ from std.bit import byte_swap -from std.memory import alloc, memset_zero +from std.memory import memset_zero from std.utils import StaticTuple from .utils import StackBuffer @@ -11,7 +11,7 @@ comptime ROUNDS_128: Int = 10 @always_inline def _ct_encrypt1( - block: UnsafePointer[UInt8, MutAnyOrigin], + block: UnsafePointer[mut=True, UInt8, _, address_space=_], skey: List[UInt64], rounds: Int, ) -> None: @@ -25,15 +25,15 @@ def _ct_encrypt1( @always_inline def cpu_aes_encrypt( - pt_bytes: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin], + pt_bytes: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], ) -> None: cpu_aes_encrypt(pt_bytes, round_keys, 10) @always_inline def cpu_aes_encrypt( - pt_bytes: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin], + pt_bytes: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], rounds: Int, ) -> None: var skey = cpu_aes_ct_skey(round_keys, rounds) @@ -41,9 +41,9 @@ def cpu_aes_encrypt( @always_inline def cpu_aes_ecb_kernel( - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, rounds: Int ) -> None: @@ -64,11 +64,11 @@ def cpu_aes_ecb_kernel( @always_inline def cpu_aes_cbc_kernel( - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - iv_ptr: UnsafePointer[UInt8, MutAnyOrigin], + iv_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: var skey = cpu_aes_ct_skey(round_keys, rounds) @@ -96,8 +96,8 @@ def cpu_aes_cbc_kernel( @always_inline def _ctr_write_block( - dst: UnsafePointer[UInt8, MutAnyOrigin], - nonce_ptr: UnsafePointer[UInt8, MutAnyOrigin], + dst: UnsafePointer[mut=True, UInt8, _, address_space=_], + nonce_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], offset: Int, ) -> None: for j in range(16): @@ -113,11 +113,11 @@ def _ctr_write_block( @always_inline def cpu_aes_ctr_kernel( - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - nonce_ptr: UnsafePointer[UInt8, MutAnyOrigin], + nonce_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: var skey = cpu_aes_ct_skey(round_keys, rounds) @@ -139,7 +139,7 @@ def cpu_aes_ctr_kernel( i += n @always_inline -def cpu_xts_mul_alpha_inplace(tweak_ptr: UnsafePointer[UInt8, MutAnyOrigin]) -> None: +def cpu_xts_mul_alpha_inplace(tweak_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_]) -> None: var carry = (tweak_ptr.load(15) & 0x80) != 0 for i in range(15, 0, -1): tweak_ptr.store(i, (tweak_ptr.load(i) << UInt8(1)) | (tweak_ptr.load(i - 1) >> UInt8(7))) @@ -150,12 +150,12 @@ def cpu_xts_mul_alpha_inplace(tweak_ptr: UnsafePointer[UInt8, MutAnyOrigin]) -> @always_inline def cpu_aes_xts_kernel( - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], - round_keys1: UnsafePointer[UInt32, MutAnyOrigin], - round_keys2: UnsafePointer[UInt32, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys1: UnsafePointer[mut=True, UInt32, _, address_space=_], + round_keys2: UnsafePointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - tweak_ptr: UnsafePointer[UInt8, MutAnyOrigin], + tweak_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: var skey1 = cpu_aes_ct_skey(round_keys1, rounds) @@ -197,10 +197,10 @@ def sub_word(w: UInt32) -> UInt32: var pair = _ct_interleave_in[1](_ct_le32(bp, 0), 0, 0, 0) q[0] = pair[0] q[4] = pair[1] - _ct_ortho(q) - _ct_sbox(q) - _ct_ortho(q) - var ws = _ct_interleave_out(q[0], q[4]) + _ct_ortho[1](q) + _ct_sbox[1](q) + _ct_ortho[1](q) + var ws = _ct_interleave_out[1](q[0], q[4]) var v = UInt64(ws[0]) return ( (UInt32(v & 0xFF) << 24) | (UInt32((v >> 8) & 0xFF) << 16) @@ -212,8 +212,8 @@ comptime RCON: StaticTuple[UInt8, 11] = StaticTuple[UInt8, 11]( ) def expand_key_128_into( - key_bytes: UnsafePointer[UInt8, MutAnyOrigin], - w: UnsafePointer[UInt32, MutAnyOrigin], + key_bytes: UnsafePointer[mut=False, UInt8, _, address_space=_], + w: UnsafePointer[mut=True, UInt32, _, address_space=_], ) raises -> None: for i in range(4): var key_val: UInt32 = 0 @@ -229,8 +229,8 @@ def expand_key_128_into( w.store(i, w.load(i - 4) ^ temp) def expand_key_192_into( - key_bytes: UnsafePointer[UInt8, MutAnyOrigin], - w: UnsafePointer[UInt32, MutAnyOrigin], + key_bytes: UnsafePointer[mut=False, UInt8, _, address_space=_], + w: UnsafePointer[mut=True, UInt32, _, address_space=_], ) raises -> None: for i in range(6): var key_val: UInt32 = 0 @@ -246,8 +246,8 @@ def expand_key_192_into( w.store(i, w.load(i - 6) ^ temp) def expand_key_256_into( - key_bytes: UnsafePointer[UInt8, MutAnyOrigin], - w: UnsafePointer[UInt32, MutAnyOrigin], + key_bytes: UnsafePointer[mut=False, UInt8, _, address_space=_], + w: UnsafePointer[mut=True, UInt32, _, address_space=_], ) raises -> None: for i in range(8): var key_val: UInt32 = 0 @@ -264,20 +264,59 @@ def expand_key_256_into( temp = sub_word(temp) w.store(i, w.load(i - 8) ^ temp) -def expand_key_128(key_bytes: UnsafePointer[UInt8, MutAnyOrigin]) raises -> UnsafePointer[UInt32, MutAnyOrigin]: - var w = alloc[UInt32](44) - expand_key_128_into(key_bytes, w) - return w +struct AESExpandedKey(Movable): + """Owned AES-128/192/256 round-key schedule, wiped on destruction.""" + + var _round_keys: StackBuffer[UInt32, 60] + var rounds: Int + + def __init__(out self, key: Span[UInt8, ...]) raises: + self._round_keys = StackBuffer[UInt32, 60](fill=0) + if len(key) == 16: + self.rounds = 10 + expand_key_128_into(key.unsafe_ptr(), self._round_keys.ptr()) + elif len(key) == 24: + self.rounds = 12 + expand_key_192_into(key.unsafe_ptr(), self._round_keys.ptr()) + elif len(key) == 32: + self.rounds = 14 + expand_key_256_into(key.unsafe_ptr(), self._round_keys.ptr()) + else: + raise Error("AES keys must contain exactly 16, 24, or 32 bytes") + + def __deinit__(deinit self): + memset_zero(self._round_keys.ptr(), 60) + + @always_inline + def ptr[ + origin: Origin, address_space: AddressSpace, // + ](ref[origin, address_space] self) -> UnsafePointer[ + UInt32, origin, address_space=address_space + ]: + return ( + self._round_keys.ptr() + .unsafe_mut_cast[origin.mut]() + .unsafe_origin_cast[origin]() + .address_space_cast[address_space]() + ) + -def expand_key_192(key_bytes: UnsafePointer[UInt8, MutAnyOrigin]) raises -> UnsafePointer[UInt32, MutAnyOrigin]: - var w = alloc[UInt32](52) - expand_key_192_into(key_bytes, w) - return w +def expand_key_128(key: Span[UInt8, ...]) raises -> AESExpandedKey: + if len(key) != 16: + raise Error("AES-128 keys must contain exactly 16 bytes") + return AESExpandedKey(key) -def expand_key_256(key_bytes: UnsafePointer[UInt8, MutAnyOrigin]) raises -> UnsafePointer[UInt32, MutAnyOrigin]: - var w = alloc[UInt32](60) - expand_key_256_into(key_bytes, w) - return w + +def expand_key_192(key: Span[UInt8, ...]) raises -> AESExpandedKey: + if len(key) != 24: + raise Error("AES-192 keys must contain exactly 24 bytes") + return AESExpandedKey(key) + + +def expand_key_256(key: Span[UInt8, ...]) raises -> AESExpandedKey: + if len(key) != 32: + raise Error("AES-256 keys must contain exactly 32 bytes") + return AESExpandedKey(key) struct AESKey: var _data: StackBuffer[UInt8, 16] @@ -290,12 +329,21 @@ struct AESKey: self._round_keys = StackBuffer[UInt32, 44]() expand_key_128_into(self._data.ptr(), self._round_keys.ptr()) - def __del__(deinit self): + def __deinit__(deinit self): memset_zero(self._data.ptr(), 16) memset_zero(self._round_keys.ptr(), 44) - def round_keys(mut self) -> UnsafePointer[UInt32, MutAnyOrigin]: - return self._round_keys.ptr() + def round_keys[ + origin: Origin, address_space: AddressSpace, // + ](ref[origin, address_space] self) -> UnsafePointer[ + UInt32, origin, address_space=address_space + ]: + return ( + self._round_keys.ptr() + .unsafe_mut_cast[origin.mut]() + .unsafe_origin_cast[origin]() + .address_space_cast[address_space]() + ) @always_inline @@ -580,7 +628,7 @@ def _ct_mix_columns[W: Int](mut q: InlineArray[SIMD[DType.uint64, W], 8]): def cpu_aes_ct_skey( - round_keys: UnsafePointer[UInt32, MutAnyOrigin], rounds: Int + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], rounds: Int ) -> List[UInt64]: var skey = List[UInt64](capacity=(rounds + 1) * 8) for r in range(rounds + 1): @@ -590,17 +638,17 @@ def cpu_aes_ct_skey( var w3 = SIMD[DType.uint64, 1](UInt64(byte_swap(round_keys.load(r * 4 + 3)))) var q = InlineArray[SIMD[DType.uint64, 1], 8](fill=0) for i in range(4): - var pair = _ct_interleave_in(w0, w1, w2, w3) + var pair = _ct_interleave_in[1](w0, w1, w2, w3) q[i] = pair[0] q[i + 4] = pair[1] - _ct_ortho(q) + _ct_ortho[1](q) for i in range(8): skey.append(UInt64(q[i])) return skey^ @always_inline -def _ct_le32(p: UnsafePointer[UInt8, MutAnyOrigin], off: Int) -> UInt64: +def _ct_le32(p: UnsafePointer[mut=True, UInt8, _, address_space=_], off: Int) -> UInt64: return ( UInt64(p.load(off)) | (UInt64(p.load(off + 1)) << 8) @@ -610,7 +658,7 @@ def _ct_le32(p: UnsafePointer[UInt8, MutAnyOrigin], off: Int) -> UInt64: @always_inline -def _ct_store_le32(p: UnsafePointer[UInt8, MutAnyOrigin], off: Int, w: UInt64): +def _ct_store_le32(p: UnsafePointer[mut=True, UInt8, _, address_space=_], off: Int, w: UInt64): p.store(off, UInt8(w & 0xFF)) p.store(off + 1, UInt8((w >> 8) & 0xFF)) p.store(off + 2, UInt8((w >> 16) & 0xFF)) @@ -619,8 +667,8 @@ def _ct_store_le32(p: UnsafePointer[UInt8, MutAnyOrigin], off: Int, w: UInt64): @always_inline def _ct_encrypt_blocks[W: Int]( - blocks: UnsafePointer[UInt8, MutAnyOrigin], - skp: UnsafePointer[UInt64, ImmutAnyOrigin], + blocks: UnsafePointer[mut=True, UInt8, _, address_space=_], + skp: UnsafePointer[mut=False, UInt64, _, address_space=_], rounds: Int, ) -> None: var q = InlineArray[SIMD[DType.uint64, W], 8](fill=0) @@ -665,7 +713,7 @@ def _ct_encrypt_blocks[W: Int]( def cpu_aes_ct_encrypt4( - blocks: UnsafePointer[UInt8, MutAnyOrigin], + blocks: UnsafePointer[mut=True, UInt8, _, address_space=_], skey: List[UInt64], rounds: Int, ) -> None: @@ -673,7 +721,7 @@ def cpu_aes_ct_encrypt4( def cpu_aes_ct_encrypt16( - blocks: UnsafePointer[UInt8, MutAnyOrigin], + blocks: UnsafePointer[mut=True, UInt8, _, address_space=_], skey: List[UInt64], rounds: Int, ) -> None: @@ -681,8 +729,8 @@ def cpu_aes_ct_encrypt16( def cpu_aes_ct_encrypt( - pt_bytes: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin], + pt_bytes: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], rounds: Int = 10, ) -> None: var skey = cpu_aes_ct_skey(round_keys, rounds) diff --git a/src/thistle/aes_gpu.mojo b/src/thistle/aes_gpu.mojo index 005ddf0..b0b87c3 100644 --- a/src/thistle/aes_gpu.mojo +++ b/src/thistle/aes_gpu.mojo @@ -8,7 +8,7 @@ from std.memory.unsafe_pointer import UnsafePointer from .aes import _ct_encrypt_blocks @always_inline -def add_counter_offset(counter: UnsafePointer[UInt8, MutAnyOrigin], offset: Int) -> None: +def add_counter_offset(counter: UnsafePointer[mut=True, UInt8, _, address_space=_], offset: Int) -> None: var carry = offset for i in range(15, -1, -1): if carry == 0: @@ -23,9 +23,9 @@ def add_counter_offset(counter: UnsafePointer[UInt8, MutAnyOrigin], offset: Int) @always_inline def _gcm_counter_from_j0( - j0: UnsafePointer[UInt8, MutAnyOrigin], + j0: UnsafePointer[mut=True, UInt8, _, address_space=_], block_index: Int, - counter: UnsafePointer[UInt8, MutAnyOrigin], + counter: UnsafePointer[mut=True, UInt8, _, address_space=_], ) -> None: for i in range(12): counter[i] = j0[i] @@ -41,56 +41,58 @@ def _gcm_counter_from_j0( @always_inline def aes_gpu_kernel_ecb( - input_data: UnsafePointer[UInt8, MutAnyOrigin], - output_data: UnsafePointer[UInt8, MutAnyOrigin], - skey: UnsafePointer[UInt64, MutAnyOrigin], - n: Int, - rounds: Int, + input_data: UnsafePointer[mut=True, UInt8, MutUntrackedOrigin], + output_data: UnsafePointer[mut=True, UInt8, MutUntrackedOrigin], + skey: UnsafePointer[mut=True, UInt64, MutUntrackedOrigin], + n: Int32, + rounds: Int32, ) -> None: var tid = global_idx.x var base_block = Int(tid) * 4 - if base_block >= n: + var num_blocks = Int(n) + if base_block >= num_blocks: return var buf = stack_allocation[64, UInt8]() for k in range(4): var b = base_block + k - if b >= n: + if b >= num_blocks: b = base_block for j in range(16): buf[k * 16 + j] = input_data[b * 16 + j] - _ct_encrypt_blocks[1](buf, skey, rounds) + _ct_encrypt_blocks[1](buf, skey, Int(rounds)) for k in range(4): var blk = base_block + k - if blk < n: + if blk < num_blocks: for j in range(16): output_data[blk * 16 + j] = buf[k * 16 + j] @always_inline def aes_gpu_kernel_ctr( - input_data: UnsafePointer[UInt8, MutAnyOrigin], - output_data: UnsafePointer[UInt8, MutAnyOrigin], - skey: UnsafePointer[UInt64, MutAnyOrigin], - n: Int, - nonce: UnsafePointer[UInt8, MutAnyOrigin], - rounds: Int, + input_data: UnsafePointer[mut=True, UInt8, MutUntrackedOrigin], + output_data: UnsafePointer[mut=True, UInt8, MutUntrackedOrigin], + skey: UnsafePointer[mut=True, UInt64, MutUntrackedOrigin], + n: Int32, + nonce: UnsafePointer[mut=True, UInt8, MutUntrackedOrigin], + rounds: Int32, ) -> None: var tid = global_idx.x var base_block = Int(tid) * 4 - if base_block >= n: + var num_blocks = Int(n) + if base_block >= num_blocks: return var buf = stack_allocation[64, UInt8]() for k in range(4): var b = base_block + k - if b >= n: + if b >= num_blocks: b = base_block var dst = buf + k * 16 for j in range(16): dst[j] = nonce[j] add_counter_offset(dst, b) - _ct_encrypt_blocks[1](buf, skey, rounds) + _ct_encrypt_blocks[1](buf, skey, Int(rounds)) for k in range(4): var blk = base_block + k - if blk < n: + if blk < num_blocks: var bp = input_data + blk * 16 var op = output_data + blk * 16 for j in range(16): @@ -98,27 +100,28 @@ def aes_gpu_kernel_ctr( @always_inline def aes_gpu_kernel_gcm_ctr( - input_data: UnsafePointer[UInt8, MutAnyOrigin], - output_data: UnsafePointer[UInt8, MutAnyOrigin], - skey: UnsafePointer[UInt64, MutAnyOrigin], - n: Int, - j0: UnsafePointer[UInt8, MutAnyOrigin], - rounds: Int, + input_data: UnsafePointer[mut=True, UInt8, MutUntrackedOrigin], + output_data: UnsafePointer[mut=True, UInt8, MutUntrackedOrigin], + skey: UnsafePointer[mut=True, UInt64, MutUntrackedOrigin], + n: Int32, + j0: UnsafePointer[mut=True, UInt8, MutUntrackedOrigin], + rounds: Int32, ) -> None: var tid = global_idx.x var base_block = Int(tid) * 4 - if base_block >= n: + var num_blocks = Int(n) + if base_block >= num_blocks: return var buf = stack_allocation[64, UInt8]() for k in range(4): var b = base_block + k - if b >= n: + if b >= num_blocks: b = base_block _gcm_counter_from_j0(j0, b, buf + k * 16) - _ct_encrypt_blocks[1](buf, skey, rounds) + _ct_encrypt_blocks[1](buf, skey, Int(rounds)) for k in range(4): var blk = base_block + k - if blk < n: + if blk < num_blocks: var bp = input_data + blk * 16 var op = output_data + blk * 16 for j in range(16): diff --git a/src/thistle/aes_ni.mojo b/src/thistle/aes_ni.mojo index 5162ab2..32c480b 100644 --- a/src/thistle/aes_ni.mojo +++ b/src/thistle/aes_ni.mojo @@ -4,7 +4,7 @@ AES-NI implementation from std.collections import List, InlineArray from std.sys import llvm_intrinsic, CompilationTarget -from std.memory import bitcast, memset_zero, memcpy, UnsafePointer, Span +from std.memory import bitcast, memset_zero, memcpy, UnsafePointer from std.utils import StaticTuple from .aes import cpu_aes_encrypt, cpu_aes_ct_encrypt, cpu_aes_ct_encrypt16, cpu_aes_ct_skey, expand_key_128_into, expand_key_192_into, expand_key_256_into from .utils import StackBuffer, load_64be, store_64be @@ -59,18 +59,18 @@ def _mm_aesenclast_si128(lhs: SIMD128, rhs: SIMD128) -> SIMD128: return SIMD128(0) @always_inline -def _mm_loadu_si128(ptr: UnsafePointer[UInt8, MutAnyOrigin]) -> SIMD128: +def _mm_loadu_si128(ptr: UnsafePointer[mut=True, UInt8, _, address_space=_]) -> SIMD128: return ptr.bitcast[UInt64]().load[width=2, alignment=1]() @always_inline -def _mm_storeu_si128(ptr: UnsafePointer[UInt8, MutAnyOrigin], data: SIMD128) -> None: +def _mm_storeu_si128(ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], data: SIMD128) -> None: var bytes: SIMD[DType.uint8, 16] = bitcast[DType.uint8, 16](data) ptr.store[width=16, alignment=1](0, bytes) @always_inline def _write_gcm_counter( - counter_ptr: UnsafePointer[UInt8, MutAnyOrigin], - j0_ptr: UnsafePointer[UInt8, MutAnyOrigin], + counter_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + j0_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], block_index: Int, ) -> None: for j in range(12): @@ -87,7 +87,7 @@ def _write_gcm_counter( counter_ptr.store(15, UInt8(ctr & 0xFF)) @always_inline -def _load_round_key(idx: Int, round_keys: UnsafePointer[UInt32, MutAnyOrigin]) -> SIMD128: +def _load_round_key(idx: Int, round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_]) -> SIMD128: var w0 = round_keys.load(idx * 4) var w1 = round_keys.load(idx * 4 + 1) var w2 = round_keys.load(idx * 4 + 2) @@ -102,8 +102,8 @@ def _load_round_key(idx: Int, round_keys: UnsafePointer[UInt32, MutAnyOrigin]) - def x86_aes_encrypt_128( - pt: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin] + pt: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_] ) -> None: var state = _mm_loadu_si128(pt) x86_aes_encrypt_128_direct(state, round_keys) @@ -111,8 +111,8 @@ def x86_aes_encrypt_128( def x86_aes_encrypt_192( - pt: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin] + pt: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_] ) -> None: var state = _mm_loadu_si128(pt) x86_aes_encrypt_192_direct(state, round_keys) @@ -120,8 +120,8 @@ def x86_aes_encrypt_192( def x86_aes_encrypt_256( - pt: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin] + pt: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_] ) -> None: var state = _mm_loadu_si128(pt) x86_aes_encrypt_256_direct(state, round_keys) @@ -130,7 +130,7 @@ def x86_aes_encrypt_256( @always_inline def x86_aes_encrypt_128_direct( mut state: SIMD128, - round_keys: UnsafePointer[UInt32, MutAnyOrigin] + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_] ) -> None: var keys = StaticTuple[SIMD128, 11]() comptime for i in range(11): @@ -144,7 +144,7 @@ def x86_aes_encrypt_128_direct( @always_inline def x86_aes_encrypt_192_direct( mut state: SIMD128, - round_keys: UnsafePointer[UInt32, MutAnyOrigin] + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_] ) -> None: var keys = StaticTuple[SIMD128, 13]() comptime for i in range(13): @@ -158,7 +158,7 @@ def x86_aes_encrypt_192_direct( @always_inline def x86_aes_encrypt_256_direct( mut state: SIMD128, - round_keys: UnsafePointer[UInt32, MutAnyOrigin] + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_] ) -> None: var keys = StaticTuple[SIMD128, 15]() comptime for i in range(15): @@ -172,7 +172,7 @@ def x86_aes_encrypt_256_direct( @always_inline def _arm_load_keys[N: Int]( - round_keys: UnsafePointer[UInt32, MutAnyOrigin] + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_] ) -> InlineArray[SIMD16, N]: var keys = InlineArray[SIMD16, N](uninitialized=True) comptime for i in range(N): @@ -182,7 +182,7 @@ def _arm_load_keys[N: Int]( keys[i] = raw.shuffle[ 3, 2, 1, 0, 7, 6, 5, 4, 11, 10, 9, 8, 15, 14, 13, 12 ]() - return keys + return keys^ @always_inline @@ -197,8 +197,8 @@ def _arm_enc_block[NR: Int]( def arm_aes_encrypt_128( - pt: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin] + pt: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_] ) -> None: var keys = _arm_load_keys[11](round_keys) var x = pt.load[width=16, alignment=1](0) @@ -206,8 +206,8 @@ def arm_aes_encrypt_128( def arm_aes_encrypt_192( - pt: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin] + pt: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_] ) -> None: var keys = _arm_load_keys[13](round_keys) var x = pt.load[width=16, alignment=1](0) @@ -215,8 +215,8 @@ def arm_aes_encrypt_192( def arm_aes_encrypt_256( - pt: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin] + pt: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_] ) -> None: var keys = _arm_load_keys[15](round_keys) var x = pt.load[width=16, alignment=1](0) @@ -224,9 +224,9 @@ def arm_aes_encrypt_256( @always_inline def _arm_ecb_loop[NR: Int]( - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, ) -> None: var keys = _arm_load_keys[NR + 1](round_keys) @@ -260,9 +260,9 @@ def _arm_ecb_loop[NR: Int]( @always_inline def arm_aes_ecb_kernel( - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, rounds: Int ) -> None: @@ -276,11 +276,11 @@ def arm_aes_ecb_kernel( @always_inline def _arm_cbc_loop[NR: Int]( - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - iv_ptr: UnsafePointer[UInt8, MutAnyOrigin], + iv_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], ) -> None: var keys = _arm_load_keys[NR + 1](round_keys) var prev = iv_ptr.load[width=16, alignment=1](0) @@ -298,11 +298,11 @@ def _arm_cbc_loop[NR: Int]( @always_inline def arm_aes_cbc_kernel( - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - iv_ptr: UnsafePointer[UInt8, MutAnyOrigin], + iv_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: if rounds == 10: @@ -314,12 +314,12 @@ def arm_aes_cbc_kernel( @always_inline def arm_aes_xts_kernel( - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], - round_keys1: UnsafePointer[UInt32, MutAnyOrigin], - round_keys2: UnsafePointer[UInt32, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys1: UnsafePointer[mut=True, UInt32, _, address_space=_], + round_keys2: UnsafePointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - tweak_ptr: UnsafePointer[UInt8, MutAnyOrigin], + tweak_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: var tweak = tweak_ptr.load[width=16, alignment=1](0) @@ -352,9 +352,9 @@ def arm_aes_xts_kernel( @always_inline def x86_aes_ecb_kernel( - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, rounds: Int ) -> None: @@ -374,11 +374,11 @@ def x86_aes_ecb_kernel( @always_inline def x86_aes_cbc_kernel( - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - iv_ptr: UnsafePointer[UInt8, MutAnyOrigin], + iv_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: var prev_block = _mm_loadu_si128(iv_ptr) @@ -409,12 +409,12 @@ def _gf_mul2_xts_simd(val: SIMD128) -> SIMD128: @always_inline def x86_aes_xts_kernel( - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], - round_keys1: UnsafePointer[UInt32, MutAnyOrigin], - round_keys2: UnsafePointer[UInt32, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys1: UnsafePointer[mut=True, UInt32, _, address_space=_], + round_keys2: UnsafePointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - tweak_ptr: UnsafePointer[UInt8, MutAnyOrigin], + tweak_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: var tweak = _mm_loadu_si128(tweak_ptr) @@ -442,8 +442,8 @@ def x86_aes_xts_kernel( @always_inline def aes_encrypt( - pt: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin], + pt: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], rounds: Int = 10 ) -> None: """Unified AES encryption function. @@ -471,11 +471,11 @@ def aes_encrypt( @always_inline def aes_gcm_ctr_kernel( - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - j0_ptr: UnsafePointer[UInt8, MutAnyOrigin], + j0_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: comptime if CompilationTarget._has_feature["sse"]() and CompilationTarget._has_feature["aes"](): @@ -489,11 +489,11 @@ def aes_gcm_ctr_kernel( @always_inline def _hw_gcm_ctr_kernel( - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - j0_ptr: UnsafePointer[UInt8, MutAnyOrigin], + j0_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: comptime if has_arm_crypto(): @@ -529,11 +529,11 @@ def _hw_gcm_ctr_kernel( @always_inline def _arm_gcm_ctr_loop[NR: Int]( - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - j0_ptr: UnsafePointer[UInt8, MutAnyOrigin], + j0_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], ) -> None: var keys = _arm_load_keys[NR + 1](round_keys) var ctr = StackBuffer[UInt8, 64]() @@ -575,11 +575,11 @@ def _arm_gcm_ctr_loop[NR: Int]( @always_inline def _arm_gcm_fused_loop[NR: Int]( - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - j0_ptr: UnsafePointer[UInt8, MutAnyOrigin], + j0_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], mut gh: _GHash, ghash_ciphertext: Bool, ) -> None: @@ -677,11 +677,11 @@ def _arm_gcm_fused_loop[NR: Int]( @always_inline def _soft_gcm_ctr_kernel( - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], - round_keys: UnsafePointer[UInt32, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - j0_ptr: UnsafePointer[UInt8, MutAnyOrigin], + j0_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: var skey = cpu_aes_ct_skey(round_keys, rounds) @@ -734,7 +734,7 @@ def _rev128(v: SIMD16) -> SIMD128: @always_inline("nodebug") -def _load_nat128(p: UnsafePointer[UInt8, MutAnyOrigin]) -> SIMD128: +def _load_nat128(p: UnsafePointer[mut=True, UInt8, _, address_space=_]) -> SIMD128: return _rev128(p.load[width=16, alignment=1](0)) @@ -843,7 +843,7 @@ struct _GHash(Copyable, Movable): self.y_lo = z_lo @always_inline - def update(mut self, data: UnsafePointer[UInt8, MutAnyOrigin], length: Int): + def update(mut self, data: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int): comptime if CompilationTarget.has_neon() and CompilationTarget._has_feature["aes"]() and not CompilationTarget.is_x86(): self._update_pmull(data, length) return @@ -851,7 +851,7 @@ struct _GHash(Copyable, Movable): self._update_soft(data, length) @always_inline - def _update_pmull(mut self, data: UnsafePointer[UInt8, MutAnyOrigin], length: Int): + def _update_pmull(mut self, data: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int): var y = SIMD128(_bitrev64(self.y_hi), _bitrev64(self.y_lo)) var off = 0 @@ -871,21 +871,23 @@ struct _GHash(Copyable, Movable): if off < length: var block = InlineArray[UInt8, 16](fill=0) - memcpy(dest=block.unsafe_ptr(), src=data + off, count=length - off) + for i in range(length - off): + block[i] = data[off + i] y = _gf_mul_nat(_load_nat128(block.unsafe_ptr()) ^ y, self.hn) self.y_hi = _bitrev64(y[0]) self.y_lo = _bitrev64(y[1]) @always_inline - def _update_soft(mut self, data: UnsafePointer[UInt8, MutAnyOrigin], length: Int): + def _update_soft(mut self, data: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int): var off = 0 while off < length: var block = InlineArray[UInt8, 16](fill=0) var n = length - off if n > 16: n = 16 - memcpy(dest=block.unsafe_ptr(), src=data + off, count=n) + for i in range(n): + block[i] = data[off + i] self.y_hi ^= load_64be(block.unsafe_ptr(), 0) self.y_lo ^= load_64be(block.unsafe_ptr(), 8) self._mul_y_by_h_soft() @@ -900,8 +902,8 @@ struct _GHash(Copyable, Movable): @always_inline def _encrypt_block( - rk: UnsafePointer[UInt32, MutAnyOrigin], rounds: Int, - src: UnsafePointer[UInt8, MutAnyOrigin], dst: UnsafePointer[UInt8, MutAnyOrigin] + rk: UnsafePointer[mut=True, UInt32, _, address_space=_], rounds: Int, + src: UnsafePointer[mut=True, UInt8, _, address_space=_], dst: UnsafePointer[mut=True, UInt8, _, address_space=_] ): for i in range(16): dst[i] = src[i] @@ -928,7 +930,8 @@ def _derive_j0( return var gh = _GHash(h_hi, h_lo) var iv_buf = List[UInt8](capacity=len(iv)) - iv_buf.extend(iv) + for i in range(len(iv)): + iv_buf.append(iv[i]) gh.update(iv_buf.unsafe_ptr(), len(iv)) gh.update_lengths(UInt64(0), UInt64(len(iv)) * 8) store_64be(j0.unsafe_ptr(), 0, gh.y_hi) @@ -936,10 +939,10 @@ def _derive_j0( def _gctr_and_ghash( - rk: UnsafePointer[UInt32, MutAnyOrigin], rounds: Int, + rk: UnsafePointer[mut=True, UInt32, _, address_space=_], rounds: Int, j0: InlineArray[UInt8, 16], - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int, mut gh: _GHash, ghash_ciphertext: Bool, @@ -994,11 +997,11 @@ def _gctr_and_ghash( def _gcm_core_keyed( - rk: UnsafePointer[UInt32, MutAnyOrigin], rounds: Int, + rk: UnsafePointer[mut=True, UInt32, _, address_space=_], rounds: Int, mut gh: _GHash, iv: Span[UInt8, ...], aad: Span[UInt8, ...], - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int, mut tag: InlineArray[UInt8, 16], ghash_ciphertext: Bool, diff --git a/src/thistle/argon2.mojo b/src/thistle/argon2.mojo index ff0b0be..cd6afbd 100644 --- a/src/thistle/argon2.mojo +++ b/src/thistle/argon2.mojo @@ -5,27 +5,27 @@ RFC 9106 from std.collections import List from std.memory import alloc, UnsafePointer, memcpy, memset_zero -from std.algorithm import parallelize +from max.algorithm import parallelize from std.bit import rotate_bits_left from .blake2b import Blake2b comptime MASK32 = 0xFFFFFFFF @always_inline -def zero_buffer(ptr: UnsafePointer[UInt8, MutAnyOrigin], len: Int): +def zero_buffer(ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], len: Int): memset_zero(ptr, len) @always_inline -def zero_buffer_u64(ptr: UnsafePointer[UInt64, MutAnyOrigin], len: Int): +def zero_buffer_u64(ptr: UnsafePointer[mut=True, UInt64, _, address_space=_], len: Int): memset_zero(ptr, len) @always_inline -def zero_and_free(ptr: UnsafePointer[UInt8, MutAnyOrigin], len: Int): +def zero_and_free(ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], len: Int): zero_buffer(ptr, len) ptr.free() @always_inline -def zero_and_free_u64(ptr: UnsafePointer[UInt64, MutAnyOrigin], len: Int): +def zero_and_free_u64(ptr: UnsafePointer[mut=True, UInt64, _, address_space=_], len: Int): zero_buffer_u64(ptr, len) ptr.free() @@ -46,7 +46,7 @@ def gb(a: UInt64, b: UInt64, c: UInt64, d: UInt64) -> Tuple[UInt64, UInt64, UInt return (a_new, b_new, c_new, d_new) @always_inline -def _p_column(base: Int, v: UnsafePointer[UInt64, MutAnyOrigin]): +def _p_column(base: Int, v: UnsafePointer[mut=True, UInt64, _, address_space=_]): var v0, v4, v8, v12 = gb(v[base + 0], v[base + 4], v[base + 8], v[base + 12]) var v1, v5, v9, v13 = gb(v[base + 1], v[base + 5], v[base + 9], v[base + 13]) var v2, v6, v10, v14 = gb(v[base + 2], v[base + 6], v[base + 10], v[base + 14]) @@ -69,7 +69,7 @@ def _p_column(base: Int, v: UnsafePointer[UInt64, MutAnyOrigin]): v[base + 15] = v15 @always_inline -def _p_diagonal(base: Int, v: UnsafePointer[UInt64, MutAnyOrigin]): +def _p_diagonal(base: Int, v: UnsafePointer[mut=True, UInt64, _, address_space=_]): var v0, v5, v10, v15 = gb(v[base + 0], v[base + 5], v[base + 10], v[base + 15]) var v1, v6, v11, v12 = gb(v[base + 1], v[base + 6], v[base + 11], v[base + 12]) var v2, v7, v8, v13 = gb(v[base + 2], v[base + 7], v[base + 8], v[base + 13]) @@ -92,8 +92,8 @@ def _p_diagonal(base: Int, v: UnsafePointer[UInt64, MutAnyOrigin]): v[base + 14] = v14 struct MemoryPool: - var block_buffer: UnsafePointer[UInt64, MutAnyOrigin] - var temp_buffer: UnsafePointer[UInt64, MutAnyOrigin] + var block_buffer: UnsafePointer[UInt64, MutUntrackedOrigin] + var temp_buffer: UnsafePointer[UInt64, MutUntrackedOrigin] var buffer_size: Int def __init__(out self, size: Int): @@ -106,18 +106,18 @@ struct MemoryPool: zero_and_free_u64(self.temp_buffer, self.buffer_size) @always_inline - def get_block(self) -> UnsafePointer[UInt64, MutAnyOrigin]: + def get_block(self) -> UnsafePointer[UInt64, MutUntrackedOrigin]: return self.block_buffer @always_inline - def get_temp(self) -> UnsafePointer[UInt64, MutAnyOrigin]: + def get_temp(self) -> UnsafePointer[UInt64, MutUntrackedOrigin]: return self.temp_buffer @always_inline def compression_g_with_pool( - out_ptr: UnsafePointer[UInt64, MutAnyOrigin], - x_ptr: UnsafePointer[UInt64, ImmutAnyOrigin], - y_ptr: UnsafePointer[UInt64, ImmutAnyOrigin], + out_ptr: UnsafePointer[mut=True, UInt64, _, address_space=_], + x_ptr: UnsafePointer[mut=False, UInt64, _, address_space=_], + y_ptr: UnsafePointer[mut=False, UInt64, _, address_space=_], with_xor: Bool, pool: MemoryPool, ): @@ -186,66 +186,89 @@ def compression_g_with_pool( out_ptr[i] = block[i] ^ block_xy[i] @always_inline -def store_le32(ptr: UnsafePointer[UInt8, MutAnyOrigin], offset: Int, val: Int): +def store_le32(ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], offset: Int, val: Int): ptr[offset + 0] = UInt8(val & 0xFF) ptr[offset + 1] = UInt8((val >> 8) & 0xFF) ptr[offset + 2] = UInt8((val >> 16) & 0xFF) ptr[offset + 3] = UInt8((val >> 24) & 0xFF) -def variable_length_hash_to_ptr(t_len: Int, input: Span[UInt8, ...], out_ptr: UnsafePointer[UInt8, MutAnyOrigin]) raises: - var le_buf = alloc[UInt8](4) +def variable_length_hash_into( + t_len: Int, input: Span[UInt8, ...], output: Span[mut=True, UInt8, ...] +) raises: + if t_len < 1: + raise Error("Argon2 variable-length hash output must not be empty") + if t_len > len(output): + raise Error("Argon2 variable-length hash output exceeds destination") + if t_len > Int.MAX - 31: + raise Error("Argon2 variable-length hash output is too large") + + var out_ptr = output.unsafe_ptr() if t_len <= 64: - var ctx = Blake2b(t_len) - store_le32(le_buf, 0, t_len) - ctx.update(Span[UInt8, ...](ptr=le_buf, length=4)) - ctx.update(input) - ctx.finalize_into(out_ptr) - zero_and_free(le_buf, 4) + var le_buf = alloc[UInt8](4) + try: + var ctx = Blake2b(t_len) + store_le32(le_buf, 0, t_len) + ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf, length=4)) + ctx.update(input) + ctx.finalize_into(out_ptr) + finally: + zero_and_free(le_buf, 4) return + var le_buf = alloc[UInt8](4) var r = (t_len + 31) // 32 - 2 var v_buf = alloc[UInt8](64) - - var ctx1 = Blake2b(64) - store_le32(le_buf, 0, t_len) - ctx1.update(Span[UInt8, ...](ptr=le_buf, length=4)) - ctx1.update(input) - ctx1.finalize_into(v_buf) - - var out_offset = 0 - for _ in range(r - 1): - memcpy(dest=out_ptr + out_offset, src=v_buf, count=32) + try: + var ctx1 = Blake2b(64) + store_le32(le_buf, 0, t_len) + ctx1.update(Span[UInt8, ...](unsafe_ptr=le_buf, length=4)) + ctx1.update(input) + ctx1.finalize_into(v_buf) + + var out_offset = 0 + for _ in range(r - 1): + for j in range(32): + out_ptr[out_offset + j] = v_buf[j] + out_offset += 32 + + var ctx = Blake2b(64) + ctx.update(Span[UInt8, ...](unsafe_ptr=v_buf, length=64)) + ctx.finalize_into(v_buf) + + for j in range(32): + out_ptr[out_offset + j] = v_buf[j] out_offset += 32 - var ctx = Blake2b(64) - ctx.update(Span[UInt8, ...](ptr=v_buf, length=64)) - ctx.finalize_into(v_buf) - - memcpy(dest=out_ptr + out_offset, src=v_buf, count=32) - out_offset += 32 - - var last_len = t_len - 32 * r - var ctx_last = Blake2b(last_len) - ctx_last.update(Span[UInt8, ...](ptr=v_buf, length=64)) - ctx_last.finalize_into(out_ptr + out_offset) - - zero_and_free(v_buf, 64) - zero_and_free(le_buf, 4) + var last_len = t_len - 32 * r + var ctx_last = Blake2b(last_len) + ctx_last.update(Span[UInt8, ...](unsafe_ptr=v_buf, length=64)) + ctx_last.finalize_into(out_ptr + out_offset) + finally: + zero_and_free(v_buf, 64) + zero_and_free(le_buf, 4) def variable_length_hash(t_len: Int, input: Span[UInt8, ...]) raises -> List[UInt8]: + if t_len < 1: + raise Error("Argon2 variable-length hash output must not be empty") var out_buf = alloc[UInt8](t_len) - variable_length_hash_to_ptr(t_len, input, out_buf) - var result = List[UInt8](capacity=t_len) - for i in range(t_len): - result.append(out_buf[i]) - zero_and_free(out_buf, t_len) - return result^ + try: + variable_length_hash_into( + t_len, + input, + Span[mut=True, UInt8, ...](unsafe_ptr=out_buf, length=t_len), + ) + var result = List[UInt8](capacity=t_len) + for i in range(t_len): + result.append(out_buf[i]) + return result^ + finally: + zero_and_free(out_buf, t_len) @always_inline def _argon2_process_lane( - memory: UnsafePointer[UInt64, MutAnyOrigin], + memory: UnsafePointer[mut=True, UInt64, _, address_space=_], lane: Int, t: Int, slice_idx: Int, @@ -348,7 +371,13 @@ def _argon2_process_lane( var r_ptr = memory + (ref_lane * q * 128 + ref_index * 128) var c_ptr = memory + (lane * q * 128 + index * 128) - compression_g_with_pool(c_ptr, p_ptr, r_ptr, t > 0, pool) + compression_g_with_pool( + c_ptr, + p_ptr.unsafe_origin_cast[MutAnyOrigin](), + r_ptr.unsafe_origin_cast[MutAnyOrigin](), + t > 0, + pool, + ) zero_and_free_u64(addressing_block, 128) zero_and_free_u64(z_u64, 128) @@ -439,28 +468,28 @@ struct Argon2id: var h0_ctx = Blake2b(64) var le_buf = alloc[UInt8](4) store_le32(le_buf, 0, self.parallelism) - h0_ctx.update(Span[UInt8, ...](ptr=le_buf, length=4)) + h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf, length=4)) store_le32(le_buf, 0, self.tag_length) - h0_ctx.update(Span[UInt8, ...](ptr=le_buf, length=4)) + h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf, length=4)) store_le32(le_buf, 0, self.memory_size_kb) - h0_ctx.update(Span[UInt8, ...](ptr=le_buf, length=4)) + h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf, length=4)) store_le32(le_buf, 0, self.iterations) - h0_ctx.update(Span[UInt8, ...](ptr=le_buf, length=4)) + h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf, length=4)) store_le32(le_buf, 0, self.version) - h0_ctx.update(Span[UInt8, ...](ptr=le_buf, length=4)) + h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf, length=4)) store_le32(le_buf, 0, self.type_code) - h0_ctx.update(Span[UInt8, ...](ptr=le_buf, length=4)) + h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf, length=4)) store_le32(le_buf, 0, len(password)) - h0_ctx.update(Span[UInt8, ...](ptr=le_buf, length=4)) + h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf, length=4)) h0_ctx.update(password) store_le32(le_buf, 0, len(self.salt)) - h0_ctx.update(Span[UInt8, ...](ptr=le_buf, length=4)) + h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf, length=4)) h0_ctx.update(Span[UInt8, ...](self.salt)) store_le32(le_buf, 0, len(self.secret)) - h0_ctx.update(Span[UInt8, ...](ptr=le_buf, length=4)) + h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf, length=4)) h0_ctx.update(Span[UInt8, ...](self.secret)) store_le32(le_buf, 0, len(self.ad)) - h0_ctx.update(Span[UInt8, ...](ptr=le_buf, length=4)) + h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf, length=4)) h0_ctx.update(Span[UInt8, ...](self.ad)) zero_and_free(le_buf, 4) @@ -488,7 +517,13 @@ struct Argon2id: store_le32(h0_input, 68, i) var b_bytes = alloc[UInt8](1024) - variable_length_hash_to_ptr(1024, Span[UInt8, ...](ptr=h0_input, length=72), b_bytes) + variable_length_hash_into( + 1024, + Span[UInt8, ...](unsafe_ptr=h0_input, length=72), + Span[mut=True, UInt8, ...]( + unsafe_ptr=b_bytes, length=1024 + ), + ) for k in range(128): var word = (b_bytes + k * 8).bitcast[UInt64]().load[width=1, alignment=1]() @@ -536,7 +571,7 @@ struct Argon2id: zero_and_free_u64(c_block, 128) zero_and_free_u64(memory, m_prime_blocks * 128) - var result = variable_length_hash(self.tag_length, Span[UInt8, ...](ptr=c_bytes, length=1024)) + var result = variable_length_hash(self.tag_length, Span[UInt8, ...](unsafe_ptr=c_bytes, length=1024)) zero_and_free(c_bytes, 1024) return result^ diff --git a/src/thistle/blake2b.mojo b/src/thistle/blake2b.mojo index 8245d36..2e63193 100644 --- a/src/thistle/blake2b.mojo +++ b/src/thistle/blake2b.mojo @@ -53,7 +53,7 @@ def g(a: UInt64, b: UInt64, c: UInt64, d: UInt64, x: UInt64, y: UInt64) -> Tuple @always_inline -def _mload(m: UnsafePointer[UInt8, ImmutAnyOrigin], i: Int) -> UInt64: +def _mload(m: UnsafePointer[mut=False, UInt8, _, address_space=_], i: Int) -> UInt64: return (m + i * 8).bitcast[UInt64]().load[width=1, alignment=1]() @@ -63,7 +63,7 @@ def round_fn[r: Int]( mut v4: UInt64, mut v5: UInt64, mut v6: UInt64, mut v7: UInt64, mut v8: UInt64, mut v9: UInt64, mut v10: UInt64, mut v11: UInt64, mut v12: UInt64, mut v13: UInt64, mut v14: UInt64, mut v15: UInt64, - m: UnsafePointer[UInt8, ImmutAnyOrigin], + m: UnsafePointer[mut=False, UInt8, _, address_space=_], ) -> Tuple[UInt64, UInt64, UInt64, UInt64, UInt64, UInt64, UInt64, UInt64, UInt64, UInt64, UInt64, UInt64, UInt64, UInt64, UInt64, UInt64]: comptime s = SIGMA[r] @@ -157,7 +157,7 @@ struct Blake2b(Movable): if self.t_low < 128: self.t_high += 1 - def compress(mut self, m: UnsafePointer[UInt8, ImmutAnyOrigin], is_last: Bool): + def compress(mut self, m: UnsafePointer[mut=False, UInt8, _, address_space=_], is_last: Bool): var v0 = self.h[0] var v1 = self.h[1] var v2 = self.h[2] @@ -214,11 +214,8 @@ struct Blake2b(Movable): var to_copy = 128 - self.buffer_len if total < to_copy: to_copy = total - memcpy( - dest=self._buf_ptr() + self.buffer_len, - src=data.unsafe_ptr(), - count=to_copy, - ) + for j in range(to_copy): + self._buf_ptr()[self.buffer_len + j] = data[j] self.buffer_len += to_copy i += to_copy if i == total: @@ -232,14 +229,11 @@ struct Blake2b(Movable): self.compress(data.unsafe_ptr() + i, False) i += 128 - memcpy( - dest=self._buf_ptr(), - src=data.unsafe_ptr() + i, - count=total - i, - ) + for j in range(total - i): + self._buf_ptr()[j] = data[i + j] self.buffer_len = total - i - def finalize_into(mut self, output: UnsafePointer[UInt8, MutAnyOrigin]): + def finalize_into(mut self, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): var old_low = self.t_low self.t_low += UInt64(self.buffer_len) if self.t_low < old_low: @@ -252,11 +246,9 @@ struct Blake2b(Movable): self.compress(self._buf_ptr(), True) var h_copy = self.h - memcpy( - dest=output, - src=UnsafePointer(to=h_copy).bitcast[UInt8](), - count=self.out_len, - ) + var h_bytes = UnsafePointer(to=h_copy).bitcast[UInt8]() + for i in range(self.out_len): + output[i] = h_bytes[i] def finalize(mut self) -> List[UInt8]: var output = List[UInt8](capacity=self.out_len) diff --git a/src/thistle/blake3.mojo b/src/thistle/blake3.mojo index 31f6f00..2152422 100644 --- a/src/thistle/blake3.mojo +++ b/src/thistle/blake3.mojo @@ -2,7 +2,7 @@ BLAKE3 cryptographic hash function """ -from std.algorithm import parallelize +from max.algorithm import parallelize from std.collections import List from std.memory import UnsafePointer, bitcast from std.bit import count_trailing_zeros @@ -112,7 +112,7 @@ def compress_internal[ counter: UInt64, blen: UInt8, flags: UInt8, - out_ptr: UnsafePointer[SIMD[DType.uint32, w], MutAnyOrigin], + out_ptr: UnsafePointer[mut=True, SIMD[DType.uint32, w], _, address_space=_], ): """BLAKE3 compression: 7 rounds of G with message permutation.""" # fmt: off @@ -191,7 +191,7 @@ def compress_internal_16way( base_counter: UInt64, blen: UInt8, flags: UInt8, - out_ptr: UnsafePointer[SIMD[DType.uint32, 16], MutAnyOrigin], + out_ptr: UnsafePointer[mut=True, SIMD[DType.uint32, 16], _, address_space=_], ): """16-way SIMD compression with per-lane sequential counters.""" var counters_low = SIMD[DType.uint32, 16]( diff --git a/src/thistle/camellia.mojo b/src/thistle/camellia.mojo index ffe63fd..fe6f5a0 100644 --- a/src/thistle/camellia.mojo +++ b/src/thistle/camellia.mojo @@ -150,7 +150,7 @@ def _slice_subkey(k: UInt64) -> InlineArray[UInt64, 8]: for i in range(8): p |= (((lanes >> UInt64(8 * i + kk)) & 1) * UInt64(0xFF)) << UInt64(8 * i) out[kk] = p - return out + return out^ @always_inline @@ -260,7 +260,7 @@ comptime _ONE_VALUE_LANES: UInt64 = 0xFF @always_inline -def _wipe_u64(ptr: UnsafePointer[UInt64, MutAnyOrigin], count: Int): +def _wipe_u64(ptr: UnsafePointer[mut=True, UInt64, _, address_space=_], count: Int): for i in range(count): ptr.store[volatile=True](i, UInt64(0)) @@ -327,7 +327,7 @@ def _f_scalar(f_in: UInt64, ke: UInt64) -> UInt64: a[k] = acc a[0] = UInt64(a[0]) ^ _ONE_VALUE_LANES - _ct_sbox(a) + _ct_sbox[1](a) var b = InlineArray[SIMD[DType.uint64, 1], 8](fill=0) comptime for k in range(8): @@ -703,7 +703,7 @@ struct CamelliaCipher: @always_inline def _load_half[W: Int]( - buf: UnsafePointer[UInt8, MutAnyOrigin], off: Int, kw: UInt64 + buf: UnsafePointer[mut=True, UInt8, _, address_space=_], off: Int, kw: UInt64 ) -> InlineArray[SIMD[DType.uint64, W], 8]: var q = InlineArray[SIMD[DType.uint64, W], 8](fill=0) var kwl = byte_swap(kw) @@ -716,13 +716,13 @@ def _load_half[W: Int]( )[0] ^ kwl ) - _ct_ortho(q) - return q + _ct_ortho[W](q) + return q^ @always_inline def _store_half[W: Int]( - buf: UnsafePointer[UInt8, MutAnyOrigin], + buf: UnsafePointer[mut=True, UInt8, _, address_space=_], off: Int, mut q: InlineArray[SIMD[DType.uint64, W], 8], kw: UInt64, @@ -757,7 +757,7 @@ def _six_rounds[forward: Bool, W: Int]( @always_inline def _encrypt_batch[W: Int]( - cipher: CamelliaCipher, buf: UnsafePointer[UInt8, MutAnyOrigin] + cipher: CamelliaCipher, buf: UnsafePointer[mut=True, UInt8, _, address_space=_] ): var a = _load_half[W](buf, 0, cipher.kw[0]) var b = _load_half[W](buf, 8, cipher.kw[1]) @@ -780,7 +780,7 @@ def _encrypt_batch[W: Int]( @always_inline def _decrypt_batch[W: Int]( - cipher: CamelliaCipher, buf: UnsafePointer[UInt8, MutAnyOrigin] + cipher: CamelliaCipher, buf: UnsafePointer[mut=True, UInt8, _, address_space=_] ): var b = _load_half[W](buf, 0, cipher.kw[2]) var a = _load_half[W](buf, 8, cipher.kw[3]) @@ -803,7 +803,7 @@ def _decrypt_batch[W: Int]( @always_inline def _batch[encrypt: Bool, W: Int]( - cipher: CamelliaCipher, buf: UnsafePointer[UInt8, MutAnyOrigin] + cipher: CamelliaCipher, buf: UnsafePointer[mut=True, UInt8, _, address_space=_] ): comptime if encrypt: _encrypt_batch[W](cipher, buf) @@ -931,7 +931,7 @@ def _six_rounds_bs[forward: Bool]( @always_inline def _batch16_hw[encrypt: Bool]( - cipher: CamelliaCipher, buf: UnsafePointer[UInt8, MutAnyOrigin] + cipher: CamelliaCipher, buf: UnsafePointer[mut=True, UInt8, _, address_space=_] ): var m = InlineArray[_U8x16, 16](fill=_U8x16(0)) comptime for i in range(16): @@ -1094,7 +1094,7 @@ def _camellia_block[encrypt: Bool]( def _camellia_blocks[encrypt: Bool]( cipher: CamelliaCipher, - data: UnsafePointer[UInt8, MutAnyOrigin], + data: UnsafePointer[mut=True, UInt8, _, address_space=_], num_blocks: Int, ): comptime if _has_hw_sbox(): @@ -1142,7 +1142,7 @@ def camellia_decrypt_block( def camellia_encrypt_blocks( cipher: CamelliaCipher, - data: UnsafePointer[UInt8, MutAnyOrigin], + data: UnsafePointer[mut=True, UInt8, _, address_space=_], num_blocks: Int, ): _camellia_blocks[True](cipher, data, num_blocks) @@ -1150,18 +1150,18 @@ def camellia_encrypt_blocks( def camellia_decrypt_blocks( cipher: CamelliaCipher, - data: UnsafePointer[UInt8, MutAnyOrigin], + data: UnsafePointer[mut=True, UInt8, _, address_space=_], num_blocks: Int, ): _camellia_blocks[False](cipher, data, num_blocks) def camellia_cbc_encrypt_kernel( - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], cipher: CamelliaCipher, num_blocks: Int, - iv_ptr: UnsafePointer[UInt8, MutAnyOrigin], + iv_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], ): var prev = iv_ptr.load[width=16, alignment=1](0) for i in range(num_blocks): @@ -1171,11 +1171,11 @@ def camellia_cbc_encrypt_kernel( def camellia_cbc_decrypt_kernel( - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], cipher: CamelliaCipher, num_blocks: Int, - iv_ptr: UnsafePointer[UInt8, MutAnyOrigin], + iv_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], ): var ct = InlineArray[UInt8, 1024](fill=0) var pt = InlineArray[UInt8, 1024](fill=0) @@ -1200,11 +1200,11 @@ def camellia_cbc_decrypt_kernel( def camellia_ctr_kernel( - input_ptr: UnsafePointer[UInt8, MutAnyOrigin], - output_ptr: UnsafePointer[UInt8, MutAnyOrigin], + input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], cipher: CamelliaCipher, num_blocks: Int, - nonce_ptr: UnsafePointer[UInt8, MutAnyOrigin], + nonce_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], ): var ks = InlineArray[UInt8, 512](fill=0) var kp = ks.unsafe_ptr() diff --git a/src/thistle/chacha20.mojo b/src/thistle/chacha20.mojo index 80e06f1..655a263 100644 --- a/src/thistle/chacha20.mojo +++ b/src/thistle/chacha20.mojo @@ -227,7 +227,7 @@ comptime _CTR_INC4 = SIMD[DType.uint32, 16]( def _quad_rows_init( key: SIMD[DType.uint32, 8], counter: UInt32, - nonce: SIMD[DType.uint32, 3], + nonce: SIMD[DType.uint32, 4], ) -> Tuple[ SIMD[DType.uint32, 16], SIMD[DType.uint32, 16], @@ -456,7 +456,7 @@ def _qr_scalar(mut a: UInt32, mut b: UInt32, mut c: UInt32, mut d: UInt32): def chacha20_block_core( key: SIMD[DType.uint32, 8], counter: UInt32, - nonce: SIMD[DType.uint32, 3], + nonce: SIMD[DType.uint32, 4], ) -> SIMD[DType.uint32, 16]: var row0 = CHACHA_CONSTANTS @@ -490,7 +490,7 @@ def chacha20_block_core( def _chacha20_block_scalar( key: SIMD[DType.uint32, 8], counter: UInt32, - nonce: SIMD[DType.uint32, 3], + nonce: SIMD[DType.uint32, 4], ) -> SIMD[DType.uint32, 16]: var x0: UInt32 = 0x61707865 var x1: UInt32 = 0x3320646E @@ -529,19 +529,19 @@ def _chacha20_block_scalar( def chacha20_block( - key: SIMD[DType.uint8, 32], counter: UInt32, nonce: SIMD[DType.uint8, 12] + key: SIMD[DType.uint8, 32], counter: UInt32, nonce: SIMD[DType.uint8, 16] ) -> SIMD[DType.uint8, 64]: var key_words = bitcast[DType.uint32, 8](key) - var nonce_words = bitcast[DType.uint32, 3](nonce) + var nonce_words = bitcast[DType.uint32, 4](nonce) var state = chacha20_block_core(key_words, counter, nonce_words) return bitcast[DType.uint8, 64](state) @always_inline def _xor_block64( - src: UnsafePointer[UInt8, MutAnyOrigin], - dst: UnsafePointer[UInt8, MutAnyOrigin], + src: UnsafePointer[mut=True, UInt8, _, address_space=_], + dst: UnsafePointer[mut=True, UInt8, _, address_space=_], keystream: SIMD[DType.uint32, 16], offset: Int, ): @@ -552,17 +552,17 @@ def _xor_block64( struct ChaCha20: var key: SIMD[DType.uint32, 8] - var nonce: SIMD[DType.uint32, 3] + var nonce: SIMD[DType.uint32, 4] var counter: UInt32 def __init__( out self, key_bytes: SIMD[DType.uint8, 32], - nonce_bytes: SIMD[DType.uint8, 12], + nonce_bytes: SIMD[DType.uint8, 16], counter: UInt32 = 1, ): self.key = bitcast[DType.uint32, 8](key_bytes) - self.nonce = bitcast[DType.uint32, 3](nonce_bytes) + self.nonce = bitcast[DType.uint32, 4](nonce_bytes) self.counter = counter def _check_counter_space(self, data_len: Int) raises: @@ -574,8 +574,8 @@ struct ChaCha20: @always_inline def _stream_xor( mut self, - src: UnsafePointer[UInt8, MutAnyOrigin], - dst: UnsafePointer[UInt8, MutAnyOrigin], + src: UnsafePointer[mut=True, UInt8, _, address_space=_], + dst: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int, ) raises: self._check_counter_space(length) diff --git a/src/thistle/chacha20poly1305.mojo b/src/thistle/chacha20poly1305.mojo index 97826f3..3e4d268 100644 --- a/src/thistle/chacha20poly1305.mojo +++ b/src/thistle/chacha20poly1305.mojo @@ -8,12 +8,16 @@ from std.collections import InlineArray from .chacha20 import ChaCha20, chacha20_block_core, simd_double_round, CHACHA_CONSTANTS from .poly1305 import Poly1305 -comptime _ZEROS16 = InlineArray[UInt8, 16](fill=0) - - -def hchacha20(key: Span[UInt8, ...], input16: Span[UInt8, ...], output: UnsafePointer[UInt8, MutAnyOrigin]) raises: +def hchacha20( + key: Span[UInt8, ...], + input16: Span[UInt8, ...], + output: Span[mut=True, UInt8, ...], +) raises: if len(key) < 32 or len(input16) < 16: raise Error("HChaCha20 needs a 32-byte key and 16-byte input") + if len(output) < 32: + raise Error("HChaCha20 output needs at least 32 writable bytes") + var out_ptr = output.unsafe_ptr() var kw = key.unsafe_ptr().bitcast[UInt32]().load[width=8, alignment=1]() var iw = input16.unsafe_ptr().bitcast[UInt32]().load[width=4, alignment=1]() @@ -26,28 +30,29 @@ def hchacha20(key: Span[UInt8, ...], input16: Span[UInt8, ...], output: UnsafePo var dr = simd_double_round(row0, row1, row2, row3) row0 = dr[0]; row1 = dr[1]; row2 = dr[2]; row3 = dr[3] - output.bitcast[UInt32]().store[alignment=1](0, row0) - (output + 16).bitcast[UInt32]().store[alignment=1](0, row3) + out_ptr.bitcast[UInt32]().store[alignment=1](0, row0) + (out_ptr + 16).bitcast[UInt32]().store[alignment=1](0, row3) def _aead_tag( poly_key: Span[UInt8, ...], aad: Span[UInt8, ...], ciphertext: Span[UInt8, ...], - output: UnsafePointer[UInt8, MutAnyOrigin], + output: UnsafePointer[mut=True, UInt8, _, address_space=_], ) raises: var p = Poly1305(poly_key) - var zp = _ZEROS16.unsafe_ptr() + var zeros16 = InlineArray[UInt8, 16](fill=0) + var zp = zeros16.unsafe_ptr() p.update(aad) if len(aad) % 16 != 0: - p.update(Span[UInt8, ...](ptr=zp, length=16 - len(aad) % 16)) + p.update(Span[UInt8, ...](unsafe_ptr=zp, length=16 - len(aad) % 16)) p.update(ciphertext) if len(ciphertext) % 16 != 0: - p.update(Span[UInt8, ...](ptr=zp, length=16 - len(ciphertext) % 16)) + p.update(Span[UInt8, ...](unsafe_ptr=zp, length=16 - len(ciphertext) % 16)) var lens = InlineArray[UInt8, 16](uninitialized=True) lens.unsafe_ptr().bitcast[UInt64]().store[alignment=1](0, UInt64(len(aad))) (lens.unsafe_ptr() + 8).bitcast[UInt64]().store[alignment=1](0, UInt64(len(ciphertext))) - p.update(Span[UInt8, ...](ptr=lens.unsafe_ptr(), length=16)) + p.update(Span[UInt8, ...](unsafe_ptr=lens.unsafe_ptr(), length=16)) p.finalize_into(output) @@ -56,20 +61,22 @@ def _aead_core[encrypt: Bool]( nonce: Span[UInt8, ...], aad: Span[UInt8, ...], input: Span[UInt8, ...], - output: UnsafePointer[UInt8, MutAnyOrigin], - tag: UnsafePointer[UInt8, MutAnyOrigin], + output: UnsafePointer[mut=True, UInt8, _, address_space=_], + tag: UnsafePointer[mut=True, UInt8, _, address_space=_], ) raises: var key_bytes = key.unsafe_ptr().load[width=32, alignment=1](0) - var nonce_bytes = nonce.unsafe_ptr().load[width=12, alignment=1](0) + var nonce_bytes = SIMD[DType.uint8, 16](0) + for i in range(12): + nonce_bytes[i] = nonce[i] var kw = bitcast[DType.uint32, 8](key_bytes) - var nw = bitcast[DType.uint32, 3](nonce_bytes) + var nw = bitcast[DType.uint32, 4](nonce_bytes) var block0 = chacha20_block_core(kw, 0, nw) var poly_key = InlineArray[UInt8, 32](uninitialized=True) poly_key.unsafe_ptr().store[alignment=1]( 0, bitcast[DType.uint8, 64](block0).slice[32]() ) - var poly_key_span = Span[UInt8, ...](ptr=poly_key.unsafe_ptr(), length=32) + var poly_key_span = Span[UInt8, ...](unsafe_ptr=poly_key.unsafe_ptr(), length=32) var cipher = ChaCha20(key_bytes, nonce_bytes, counter=1) var src = input.unsafe_ptr().unsafe_mut_cast[True]().unsafe_origin_cast[MutAnyOrigin]() @@ -78,7 +85,7 @@ def _aead_core[encrypt: Bool]( comptime if encrypt: _aead_tag( poly_key_span, aad, - Span[UInt8, ...](ptr=output, length=len(input)), tag, + Span[UInt8, ...](unsafe_ptr=output, length=len(input)), tag, ) else: _aead_tag(poly_key_span, aad, input, tag) @@ -92,14 +99,20 @@ def chacha20_poly1305_encrypt( nonce: Span[UInt8, ...], aad: Span[UInt8, ...], plaintext: Span[UInt8, ...], - ciphertext: UnsafePointer[UInt8, MutAnyOrigin], - tag: UnsafePointer[UInt8, MutAnyOrigin], + ciphertext: Span[mut=True, UInt8, ...], + tag: Span[mut=True, UInt8, ...], ) raises: if len(key) != 32: raise Error("ChaCha20-Poly1305 key must be 32 bytes") if len(nonce) != 12: raise Error("ChaCha20-Poly1305 nonce must be 12 bytes") - _aead_core[True](key, nonce, aad, plaintext, ciphertext, tag) + if len(ciphertext) < len(plaintext): + raise Error("ChaCha20-Poly1305 ciphertext output is too small") + if len(tag) < 16: + raise Error("ChaCha20-Poly1305 tag output is too small") + _aead_core[True]( + key, nonce, aad, plaintext, ciphertext.unsafe_ptr(), tag.unsafe_ptr() + ) def chacha20_poly1305_decrypt( @@ -108,7 +121,7 @@ def chacha20_poly1305_decrypt( aad: Span[UInt8, ...], ciphertext: Span[UInt8, ...], tag: Span[UInt8, ...], - plaintext: UnsafePointer[UInt8, MutAnyOrigin], + plaintext: Span[mut=True, UInt8, ...], ) raises -> Bool: if len(key) != 32: raise Error("ChaCha20-Poly1305 key must be 32 bytes") @@ -116,11 +129,15 @@ def chacha20_poly1305_decrypt( raise Error("ChaCha20-Poly1305 nonce must be 12 bytes") if len(tag) != 16: return False + if len(plaintext) < len(ciphertext): + raise Error("ChaCha20-Poly1305 plaintext output is too small") var key_bytes = key.unsafe_ptr().load[width=32, alignment=1](0) - var nonce_bytes = nonce.unsafe_ptr().load[width=12, alignment=1](0) + var nonce_bytes = SIMD[DType.uint8, 16](0) + for i in range(12): + nonce_bytes[i] = nonce[i] var kw = bitcast[DType.uint32, 8](key_bytes) - var nw = bitcast[DType.uint32, 3](nonce_bytes) + var nw = bitcast[DType.uint32, 4](nonce_bytes) var block0 = chacha20_block_core(kw, 0, nw) var poly_key = InlineArray[UInt8, 32](uninitialized=True) poly_key.unsafe_ptr().store[alignment=1]( @@ -129,7 +146,7 @@ def chacha20_poly1305_decrypt( var expected = InlineArray[UInt8, 16](uninitialized=True) _aead_tag( - Span[UInt8, ...](ptr=poly_key.unsafe_ptr(), length=32), + Span[UInt8, ...](unsafe_ptr=poly_key.unsafe_ptr(), length=32), aad, ciphertext, expected.unsafe_ptr(), ) var diff: UInt8 = 0 @@ -143,7 +160,7 @@ def chacha20_poly1305_decrypt( var cipher = ChaCha20(key_bytes, nonce_bytes, counter=1) var src = ciphertext.unsafe_ptr().unsafe_mut_cast[True]().unsafe_origin_cast[MutAnyOrigin]() - cipher._stream_xor(src, plaintext, len(ciphertext)) + cipher._stream_xor(src, plaintext.unsafe_ptr(), len(ciphertext)) var poly_key_ptr = poly_key.unsafe_ptr() for i in range(32): poly_key_ptr.store[volatile=True](i, UInt8(0)) @@ -155,8 +172,8 @@ def xchacha20_poly1305_encrypt( nonce: Span[UInt8, ...], aad: Span[UInt8, ...], plaintext: Span[UInt8, ...], - ciphertext: UnsafePointer[UInt8, MutAnyOrigin], - tag: UnsafePointer[UInt8, MutAnyOrigin], + ciphertext: Span[mut=True, UInt8, ...], + tag: Span[mut=True, UInt8, ...], ) raises: if len(key) != 32: raise Error("XChaCha20-Poly1305 key must be 32 bytes") @@ -165,8 +182,8 @@ def xchacha20_poly1305_encrypt( var sub = _xchacha_subkey_nonce(key, nonce) var sp = sub.unsafe_ptr().unsafe_origin_cast[MutAnyOrigin]() chacha20_poly1305_encrypt( - Span[UInt8, ...](ptr=sp, length=32), - Span[UInt8, ...](ptr=sp + 32, length=12), + Span[UInt8, ...](unsafe_ptr=sp, length=32), + Span[UInt8, ...](unsafe_ptr=sp + 32, length=12), aad, plaintext, ciphertext, tag, ) for i in range(44): @@ -179,7 +196,7 @@ def xchacha20_poly1305_decrypt( aad: Span[UInt8, ...], ciphertext: Span[UInt8, ...], tag: Span[UInt8, ...], - plaintext: UnsafePointer[UInt8, MutAnyOrigin], + plaintext: Span[mut=True, UInt8, ...], ) raises -> Bool: if len(key) != 32: raise Error("XChaCha20-Poly1305 key must be 32 bytes") @@ -188,8 +205,8 @@ def xchacha20_poly1305_decrypt( var sub = _xchacha_subkey_nonce(key, nonce) var sp = sub.unsafe_ptr().unsafe_origin_cast[MutAnyOrigin]() var ok = chacha20_poly1305_decrypt( - Span[UInt8, ...](ptr=sp, length=32), - Span[UInt8, ...](ptr=sp + 32, length=12), + Span[UInt8, ...](unsafe_ptr=sp, length=32), + Span[UInt8, ...](unsafe_ptr=sp + 32, length=12), aad, ciphertext, tag, plaintext, ) for i in range(44): @@ -201,9 +218,9 @@ def _xchacha_subkey_nonce(key: Span[UInt8, ...], nonce: Span[UInt8, ...]) raises var out = InlineArray[UInt8, 44](fill=0) hchacha20( key, - Span[UInt8, ...](ptr=nonce.unsafe_ptr(), length=16), - out.unsafe_ptr(), + Span[UInt8, ...](unsafe_ptr=nonce.unsafe_ptr(), length=16), + Span[mut=True, UInt8, ...](unsafe_ptr=out.unsafe_ptr(), length=32), ) for i in range(8): out[36 + i] = nonce[16 + i] - return out + return out^ diff --git a/src/thistle/curve25519.mojo b/src/thistle/curve25519.mojo index 2e7cbf0..ef4bac2 100644 --- a/src/thistle/curve25519.mojo +++ b/src/thistle/curve25519.mojo @@ -1,5 +1,4 @@ from std.builtin.dtype import DType -from std.collections import InlineArray @always_inline def _u128_shr[shift: Int](x: UInt128) -> UInt128: @@ -22,39 +21,20 @@ def _u128_shr[shift: Int](x: UInt128) -> UInt128: return UInt128(hi >> UInt64(shift - 64)) struct FieldElement51(Movable, Copyable, ImplicitlyCopyable): - var limbs: InlineArray[UInt64, 5] + var limbs: SIMD[DType.uint64, 8] @always_inline def __init__(out self): - self.limbs = InlineArray[UInt64, 5](uninitialized=True) - self.limbs[0] = 0 - self.limbs[1] = 0 - self.limbs[2] = 0 - self.limbs[3] = 0 - self.limbs[4] = 0 + self.limbs = SIMD[DType.uint64, 8](0) @always_inline def __init__(out self, l0: UInt64, l1: UInt64, l2: UInt64, l3: UInt64, l4: UInt64): - self.limbs = InlineArray[UInt64, 5](uninitialized=True) - self.limbs[0] = l0 - self.limbs[1] = l1 - self.limbs[2] = l2 - self.limbs[3] = l3 - self.limbs[4] = l4 + self.limbs = SIMD[DType.uint64, 8](l0, l1, l2, l3, l4, 0, 0, 0) @always_inline - def __init__(out self, limbs: InlineArray[UInt64, 5]): + def __init__(out self, limbs: SIMD[DType.uint64, 8]): self.limbs = limbs - @always_inline - def __init__(out self, limbs: SIMD[DType.uint64, 5]): - self.limbs = InlineArray[UInt64, 5](uninitialized=True) - self.limbs[0] = limbs[0] - self.limbs[1] = limbs[1] - self.limbs[2] = limbs[2] - self.limbs[3] = limbs[3] - self.limbs[4] = limbs[4] - @always_inline def __copyinit__(out self, copy: Self): self.limbs = copy.limbs @@ -146,7 +126,7 @@ struct FieldElement51(Movable, Copyable, ImplicitlyCopyable): @always_inline def __sub__(self, other: FieldElement51) -> FieldElement51: - var l = InlineArray[UInt64, 5](uninitialized=True) + var l = SIMD[DType.uint64, 8](0) l[0] = (self.limbs[0] + 0x7FFFFFFFFFFED0) - other.limbs[0] l[1] = (self.limbs[1] + 0x7FFFFFFFFFFFF0) - other.limbs[1] l[2] = (self.limbs[2] + 0x7FFFFFFFFFFFF0) - other.limbs[2] @@ -293,7 +273,7 @@ struct FieldElement51(Movable, Copyable, ImplicitlyCopyable): return FieldElement51(l0_2, l1_2, l2, l3, l4) @always_inline - def _reduce(self, limbs: InlineArray[UInt64, 5]) -> FieldElement51: + def _reduce(self, limbs: SIMD[DType.uint64, 8]) -> FieldElement51: var l = limbs var MASK = UInt64(0x7FFFFFFFFFFFF) for _ in range(5): @@ -309,7 +289,7 @@ struct FieldElement51(Movable, Copyable, ImplicitlyCopyable): if len(bytes) < 32: raise Error("FieldElement51 input must be at least 32 bytes") @always_inline - def load8(ptr: UnsafePointer[UInt8, _]) -> UInt64: + def load8(ptr: UnsafePointer[mut=False, UInt8, _, address_space=_]) -> UInt64: return ptr.bitcast[UInt64]().load[width=1, alignment=1]() var ptr = bytes.unsafe_ptr() @@ -323,7 +303,7 @@ struct FieldElement51(Movable, Copyable, ImplicitlyCopyable): return FieldElement51(l0, l1, l2, l3, l4) - def to_bytes_into(self, output: UnsafePointer[UInt8, MutAnyOrigin]): + def to_bytes_into(self, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): var res = self._reduce(self.limbs) var limbs = res.limbs diff --git a/src/thistle/ed25519.mojo b/src/thistle/ed25519.mojo index d2e0ad2..278936e 100644 --- a/src/thistle/ed25519.mojo +++ b/src/thistle/ed25519.mojo @@ -9,37 +9,41 @@ from .ed25519_table import ed25519_base_table, ed25519_b_odd_table from .sha2 import SHA512Context, sha512_update, sha512_final_to_buffer -comptime L_LIMBS = SIMD[DType.uint64, 5]( +comptime L_LIMBS = SIMD[DType.uint64, 8]( 0x0002631a5cf5d3ed, 0x000dea2f79cd6581, 0x000000000014def9, 0x0000000000000000, 0x0000100000000000, + 0, 0, 0, ) comptime LFACTOR: UInt64 = 0x51da312547e1b -comptime RR_LIMBS = SIMD[DType.uint64, 5]( +comptime RR_LIMBS = SIMD[DType.uint64, 8]( 0x0009d265e952d13b, 0x000d63c715bea69f, 0x0005be65cb687604, 0x0003dceec73d217f, 0x000009411b7c309a, + 0, 0, 0, ) -comptime ED25519_D_LIMBS = SIMD[DType.uint64, 5]( +comptime ED25519_D_LIMBS = SIMD[DType.uint64, 8]( 0x34dca135978a3, 0x001a8283b156ebd, 0x005e7a26001c029, 0x00739c663a03cbb, 0x0052036cee2b6ff, + 0, 0, 0, ) -comptime POW2_256_LIMBS = SIMD[DType.uint64, 5]( +comptime POW2_256_LIMBS = SIMD[DType.uint64, 8]( 0x0009f4e532df7449, 0x000da9f725df7382, 0x000f5be65cc244cc, 0x000a3dceec73d217, 0x0000099411b7c309, + 0, 0, 0, ) comptime L_BYTES = SIMD[DType.uint8, 32]( @@ -83,7 +87,7 @@ def _encoded_y_lt_p(y: Span[UInt8, ...]) -> Bool: gt = 1 return lt == 1 -def _pack_limbs_into(limbs: SIMD[DType.uint64, 5], output: UnsafePointer[UInt8, MutAnyOrigin]): +def _pack_limbs_into(limbs: SIMD[DType.uint64, 8], output: UnsafePointer[mut=True, UInt8, _, address_space=_]): var words = SIMD[DType.uint64, 4](0, 0, 0, 0) words[0] = limbs[0] | (limbs[1] << 52) words[1] = (limbs[1] >> 12) | (limbs[2] << 40) @@ -93,12 +97,12 @@ def _pack_limbs_into(limbs: SIMD[DType.uint64, 5], output: UnsafePointer[UInt8, for i in range(32): output[i] = bytes[i] -def _unpack_limbs(bytes: Span[UInt8, ...]) -> SIMD[DType.uint64, 5]: +def _unpack_limbs(bytes: Span[UInt8, ...]) -> SIMD[DType.uint64, 8]: # Input may be byte-aligned; use alignment=1 for the UInt64 wide load. var words = bytes.unsafe_ptr().bitcast[UInt64]().load[width=4, alignment=1]() comptime MASK = (UInt64(1) << 52) - 1 comptime TOP_MASK = (UInt64(1) << 48) - 1 - var s = SIMD[DType.uint64, 5](0, 0, 0, 0, 0) + var s = SIMD[DType.uint64, 8](0) s[0] = words[0] & MASK s[1] = ((words[0] >> UInt64(52)) | (words[1] << UInt64(12))) & MASK s[2] = ((words[1] >> UInt64(40)) | (words[2] << UInt64(24))) & MASK @@ -106,11 +110,11 @@ def _unpack_limbs(bytes: Span[UInt8, ...]) -> SIMD[DType.uint64, 5]: s[4] = (words[3] >> UInt64(16)) & TOP_MASK return s -def _from_512_raw(bytes: Span[UInt8, ...]) -> SIMD[DType.uint64, 5]: +def _from_512_raw(bytes: Span[UInt8, ...]) -> SIMD[DType.uint64, 8]: # RFC 8032 5.1.6: reduce 64-byte SHA-512 output modulo L. var ptr = bytes.unsafe_ptr() - var lo_span = Span[UInt8, ...](ptr=ptr, length=32) - var hi_span = Span[UInt8, ...](ptr=ptr + 32, length=32) + var lo_span = Span[UInt8, ...](unsafe_ptr=ptr, length=32) + var hi_span = Span[UInt8, ...](unsafe_ptr=ptr + 32, length=32) var lo = Scalar.from_bytes(lo_span) var hi = Scalar.from_bytes(hi_span) var pow2_256 = Scalar(POW2_256_LIMBS) @@ -138,14 +142,14 @@ def ed25519_base_point() -> EdwardsPoint: struct Scalar(Movable, Copyable, ImplicitlyCopyable): - var limbs: SIMD[DType.uint64, 5] + var limbs: SIMD[DType.uint64, 8] @always_inline def __init__(out self): - self.limbs = SIMD[DType.uint64, 5](0, 0, 0, 0, 0) + self.limbs = SIMD[DType.uint64, 8](0) @always_inline - def __init__(out self, limbs: SIMD[DType.uint64, 5]): + def __init__(out self, limbs: SIMD[DType.uint64, 8]): self.limbs = limbs @always_inline @@ -163,8 +167,10 @@ struct Scalar(Movable, Copyable, ImplicitlyCopyable): var raw = _unpack_limbs(bytes) return Scalar(raw)._montgomery_mul(Scalar(RR_LIMBS)) - def to_bytes_into(self, output: UnsafePointer[UInt8, MutAnyOrigin]): - var raw = self._montgomery_mul(Scalar(SIMD[DType.uint64, 5](1, 0, 0, 0, 0))) + def to_bytes_into(self, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): + var raw = self._montgomery_mul( + Scalar(SIMD[DType.uint64, 8](1, 0, 0, 0, 0, 0, 0, 0)) + ) _pack_limbs_into(raw.limbs, output) @staticmethod @@ -181,11 +187,11 @@ struct Scalar(Movable, Copyable, ImplicitlyCopyable): s[0] &= 0xF8 s[31] &= 0x7F s[31] |= 0x40 - return Scalar.from_bytes(Span[UInt8, ...](ptr=s.unsafe_ptr(), length=32)) + return Scalar.from_bytes(Span[UInt8, ...](unsafe_ptr=s.unsafe_ptr(), length=32)) def __add__(self, other: Scalar) -> Scalar: comptime MASK = (UInt64(1) << 52) - 1 - var sum = SIMD[DType.uint64, 5](0, 0, 0, 0, 0) + var sum = SIMD[DType.uint64, 8](0) var carry: UInt64 = 0 for i in range(5): carry = self.limbs[i] + other.limbs[i] + (carry >> 52) @@ -202,7 +208,7 @@ struct Scalar(Movable, Copyable, ImplicitlyCopyable): return self._montgomery_mul(other) @staticmethod - def _montgomery_mul_raw(a: SIMD[DType.uint64, 5], b: SIMD[DType.uint64, 5]) -> SIMD[DType.uint64, 5]: + def _montgomery_mul_raw(a: SIMD[DType.uint64, 8], b: SIMD[DType.uint64, 8]) -> SIMD[DType.uint64, 8]: var z = InlineArray[UInt128, 9](uninitialized=True) for i in range(9): z[i] = 0 for i in range(5): @@ -217,7 +223,7 @@ struct Scalar(Movable, Copyable, ImplicitlyCopyable): var p = (UInt64(sum.cast[DType.uint64]() * LFACTOR)) & ((UInt64(1) << 52) - 1) n[i] = p carry = (sum + UInt128(p) * UInt128(L_LIMBS[0])) >> 52 - var r = SIMD[DType.uint64, 5](0, 0, 0, 0, 0) + var r = SIMD[DType.uint64, 8](0) for i in range(4): var sum = carry + z[5 + i] for j in range(i + 1, 5): @@ -234,13 +240,13 @@ struct Scalar(Movable, Copyable, ImplicitlyCopyable): def wipe(mut self): UnsafePointer(to=self.limbs).bitcast[UInt64]().store[volatile=True]( - 0, SIMD[DType.uint64, 5](0, 0, 0, 0, 0) + 0, SIMD[DType.uint64, 8](0) ) def _sub(self, other: Scalar) -> Scalar: # Branchless subtract modulo L; used on secret scalar paths. comptime MASK = (UInt64(1) << 52) - 1 - var diff = SIMD[DType.uint64, 5](0, 0, 0, 0, 0) + var diff = SIMD[DType.uint64, 8](0) var borrow: UInt64 = 0 for i in range(5): var x = other.limbs[i] + borrow @@ -309,7 +315,7 @@ def edwards_negate(p: EdwardsPoint) -> EdwardsPoint: def _ct_select_fe(a: FieldElement51, b: FieldElement51, choice: UInt8) -> FieldElement51: # constant-time select via mask var mask = UInt64(0) - UInt64(choice) - var limbs = SIMD[DType.uint64, 5](0, 0, 0, 0, 0) + var limbs = SIMD[DType.uint64, 8](0) for i in range(5): limbs[i] = a.limbs[i] ^ (mask & (a.limbs[i] ^ b.limbs[i])) return FieldElement51(limbs) @@ -346,7 +352,7 @@ def _edwards_double_standalone(p: EdwardsPoint) -> EdwardsPoint: @no_inline def fe_from_bytes(bytes: Span[UInt8, ...]) -> FieldElement51: # Decode 255-bit little-endian field element; caller clears x-parity bit. - def load8(ptr: UnsafePointer[UInt8, _]) -> UInt64: + def load8(ptr: UnsafePointer[mut=False, UInt8, _, address_space=_]) -> UInt64: var v: UInt64 = 0 for j in range(8): v |= UInt64(ptr[j]) << UInt64(j * 8) @@ -361,12 +367,12 @@ def fe_from_bytes(bytes: Span[UInt8, ...]) -> FieldElement51: return FieldElement51(l0, l1, l2, l3, l4) @no_inline -def edwards_encode_into(p: EdwardsPoint, output: UnsafePointer[UInt8, MutAnyOrigin]): +def edwards_encode_into(p: EdwardsPoint, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): # RFC 8032 5.1.2: encode y and store x parity in bit 255. _edwards_encode_with_zinv(p, p.Z.invert(), output) @no_inline -def _edwards_encode_with_zinv(p: EdwardsPoint, z_inv: FieldElement51, output: UnsafePointer[UInt8, MutAnyOrigin]): +def _edwards_encode_with_zinv(p: EdwardsPoint, z_inv: FieldElement51, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): var x = p.X * z_inv var y = p.Y * z_inv y.to_bytes_into(output) @@ -384,9 +390,9 @@ def edwards_decode(data: Span[UInt8, ...], strict: Bool = True) -> DecodeResult: y_bytes[i] = data[i] var sign = (y_bytes[31] >> 7) & 1 y_bytes[31] &= 0x7F - if strict and not _encoded_y_lt_p(Span[UInt8, ...](ptr=y_bytes.unsafe_ptr(), length=32)): + if strict and not _encoded_y_lt_p(Span[UInt8, ...](unsafe_ptr=y_bytes.unsafe_ptr(), length=32)): return DecodeResult(False, EdwardsPoint()) - var y = fe_from_bytes(Span[UInt8, ...](ptr=y_bytes.unsafe_ptr(), length=32)) + var y = fe_from_bytes(Span[UInt8, ...](unsafe_ptr=y_bytes.unsafe_ptr(), length=32)) var y2 = y.square() var u = y2 - FieldElement51.ONE() @@ -596,7 +602,7 @@ def _radix16_digits(scalar: Span[UInt8, ...]) -> InlineArray[Int, 64]: var carry = (digits[i] + 8) >> 4 digits[i] -= carry << 4 digits[i + 1] += carry - return digits + return digits^ @always_inline def _base_table_lookup(ptr: UnsafePointer[UInt64, _], j: Int, digit: Int) -> AffineNielsPoint: @@ -613,12 +619,12 @@ def _base_table_lookup(ptr: UnsafePointer[UInt64, _], j: Int, digit: Int) -> Aff var m = UInt64(0) - ((diff - 1) >> 63) acc = acc ^ ((acc ^ cand) & SIMD[DType.uint64, 16](m)) var sm = sign.cast[DType.uint64]() - var yp = SIMD[DType.uint64, 5](acc[0], acc[1], acc[2], acc[3], acc[4]) - var ym = SIMD[DType.uint64, 5](acc[5], acc[6], acc[7], acc[8], acc[9]) - var swap = (yp ^ ym) & SIMD[DType.uint64, 5](sm) + var yp = SIMD[DType.uint64, 8](acc[0], acc[1], acc[2], acc[3], acc[4], 0, 0, 0) + var ym = SIMD[DType.uint64, 8](acc[5], acc[6], acc[7], acc[8], acc[9], 0, 0, 0) + var swap = (yp ^ ym) & SIMD[DType.uint64, 8](sm) yp ^= swap ym ^= swap - var xy = FieldElement51(SIMD[DType.uint64, 5](acc[10], acc[11], acc[12], acc[13], acc[14])) + var xy = FieldElement51(acc[10], acc[11], acc[12], acc[13], acc[14]) var xy_neg = FieldElement51.ZERO() - xy var xy_sel = _ct_select_fe(xy, xy_neg, UInt8(sm & 1)) return AffineNielsPoint(FieldElement51(yp), FieldElement51(ym), xy_sel) @@ -668,7 +674,7 @@ def _naf5(scalar: Span[UInt8, ...]) -> InlineArray[Int, 256]: carry = 1 naf[pos] = Int(window) - 32 pos += 5 - return naf + return naf^ @always_inline def _b_odd_entry(ptr: UnsafePointer[UInt64, _], k: Int) -> AffineNielsPoint: @@ -712,7 +718,7 @@ def _double_scalar_mult_vartime(a: Span[UInt8, ...], A: EdwardsPoint, b: Span[UI return Q @no_inline -def ed25519_generate_public_key(private_key: Span[UInt8, ...], output: UnsafePointer[UInt8, MutAnyOrigin]) raises: +def ed25519_generate_public_key(private_key: Span[UInt8, ...], output: UnsafePointer[mut=True, UInt8, _, address_space=_]) raises: # RFC 8032 5.1.5: public key A = [pruned SHA512(secret)]B. if len(private_key) != 32: raise Error("Ed25519 private key must be 32 bytes") @@ -720,10 +726,10 @@ def ed25519_generate_public_key(private_key: Span[UInt8, ...], output: UnsafePoi var ctx = SHA512Context() sha512_update(ctx, private_key) sha512_final_to_buffer(ctx, hash.unsafe_ptr()) - var s = Scalar.from_bytes_clamped(Span[UInt8, ...](ptr=hash.unsafe_ptr(), length=32)) + var s = Scalar.from_bytes_clamped(Span[UInt8, ...](unsafe_ptr=hash.unsafe_ptr(), length=32)) var s_bytes = InlineArray[UInt8, 32](uninitialized=True) s.to_bytes_into(s_bytes.unsafe_ptr()) - var pub_point = _mul_base_ct(Span[UInt8, ...](ptr=s_bytes.unsafe_ptr(), length=32)) + var pub_point = _mul_base_ct(Span[UInt8, ...](unsafe_ptr=s_bytes.unsafe_ptr(), length=32)) edwards_encode_into(pub_point, output) ctx.wipe() s.wipe() @@ -735,7 +741,7 @@ def ed25519_generate_public_key(private_key: Span[UInt8, ...], output: UnsafePoi s_ptr.store[volatile=True](i, UInt8(0)) @no_inline -def ed25519_sign(private_key: Span[UInt8, ...], message: Span[UInt8, ...], output: UnsafePointer[UInt8, MutAnyOrigin]) raises: +def ed25519_sign(private_key: Span[UInt8, ...], message: Span[UInt8, ...], output: UnsafePointer[mut=True, UInt8, _, address_space=_]) raises: # RFC 8032 5.1.6 pure Ed25519: # r = SHA512(prefix || M), R = [r]B, # k = SHA512(R || A || M), S = r + k*s mod L. @@ -745,22 +751,22 @@ def ed25519_sign(private_key: Span[UInt8, ...], message: Span[UInt8, ...], outpu var ctx = SHA512Context() sha512_update(ctx, private_key) sha512_final_to_buffer(ctx, hash.unsafe_ptr()) - var s_scalar = Scalar.from_bytes_clamped(Span[UInt8, ...](ptr=hash.unsafe_ptr(), length=32)) + var s_scalar = Scalar.from_bytes_clamped(Span[UInt8, ...](unsafe_ptr=hash.unsafe_ptr(), length=32)) var s_bytes = InlineArray[UInt8, 32](uninitialized=True) s_scalar.to_bytes_into(s_bytes.unsafe_ptr()) - var A_point = _mul_base_ct(Span[UInt8, ...](ptr=s_bytes.unsafe_ptr(), length=32)) + var A_point = _mul_base_ct(Span[UInt8, ...](unsafe_ptr=s_bytes.unsafe_ptr(), length=32)) var r_hash = InlineArray[UInt8, 64](uninitialized=True) var r_ctx = SHA512Context() - sha512_update(r_ctx, Span[UInt8, ...](ptr=hash.unsafe_ptr() + 32, length=32)) + sha512_update(r_ctx, Span[UInt8, ...](unsafe_ptr=hash.unsafe_ptr() + 32, length=32)) sha512_update(r_ctx, message) sha512_final_to_buffer(r_ctx, r_hash.unsafe_ptr()) - var r_scalar = Scalar.from_bytes_wide(Span[UInt8, ...](ptr=r_hash.unsafe_ptr(), length=64)) + var r_scalar = Scalar.from_bytes_wide(Span[UInt8, ...](unsafe_ptr=r_hash.unsafe_ptr(), length=64)) var r_bytes = InlineArray[UInt8, 32](uninitialized=True) r_scalar.to_bytes_into(r_bytes.unsafe_ptr()) - var R_point = _mul_base_ct(Span[UInt8, ...](ptr=r_bytes.unsafe_ptr(), length=32)) + var R_point = _mul_base_ct(Span[UInt8, ...](unsafe_ptr=r_bytes.unsafe_ptr(), length=32)) var zz_inv = (A_point.Z * R_point.Z).invert() var A_enc = InlineArray[UInt8, 32](uninitialized=True) @@ -770,12 +776,12 @@ def ed25519_sign(private_key: Span[UInt8, ...], message: Span[UInt8, ...], outpu var k_hash = InlineArray[UInt8, 64](uninitialized=True) var k_ctx = SHA512Context() - sha512_update(k_ctx, Span[UInt8, ...](ptr=R_enc.unsafe_ptr(), length=32)) - sha512_update(k_ctx, Span[UInt8, ...](ptr=A_enc.unsafe_ptr(), length=32)) + sha512_update(k_ctx, Span[UInt8, ...](unsafe_ptr=R_enc.unsafe_ptr(), length=32)) + sha512_update(k_ctx, Span[UInt8, ...](unsafe_ptr=A_enc.unsafe_ptr(), length=32)) sha512_update(k_ctx, message) sha512_final_to_buffer(k_ctx, k_hash.unsafe_ptr()) - var k_scalar = Scalar.from_bytes_wide(Span[UInt8, ...](ptr=k_hash.unsafe_ptr(), length=64)) + var k_scalar = Scalar.from_bytes_wide(Span[UInt8, ...](unsafe_ptr=k_hash.unsafe_ptr(), length=64)) var S_scalar = r_scalar + k_scalar * s_scalar var S_bytes = InlineArray[UInt8, 32](uninitialized=True) S_scalar.to_bytes_into(S_bytes.unsafe_ptr()) @@ -813,7 +819,7 @@ struct Ed25519SigningKey(Copyable, Movable): var ctx = SHA512Context() sha512_update(ctx, private_key) sha512_final_to_buffer(ctx, hash.unsafe_ptr()) - self._s = Scalar.from_bytes_clamped(Span[UInt8, ...](ptr=hash.unsafe_ptr(), length=32)) + self._s = Scalar.from_bytes_clamped(Span[UInt8, ...](unsafe_ptr=hash.unsafe_ptr(), length=32)) self._prefix = InlineArray[UInt8, 32](uninitialized=True) for i in range(32): @@ -821,7 +827,7 @@ struct Ed25519SigningKey(Copyable, Movable): var s_bytes = InlineArray[UInt8, 32](uninitialized=True) self._s.to_bytes_into(s_bytes.unsafe_ptr()) - var A_point = _mul_base_ct(Span[UInt8, ...](ptr=s_bytes.unsafe_ptr(), length=32)) + var A_point = _mul_base_ct(Span[UInt8, ...](unsafe_ptr=s_bytes.unsafe_ptr(), length=32)) self._a_enc = InlineArray[UInt8, 32](uninitialized=True) edwards_encode_into(A_point, self._a_enc.unsafe_ptr()) @@ -839,34 +845,34 @@ struct Ed25519SigningKey(Copyable, Movable): for i in range(32): p.store[volatile=True](i, UInt8(0)) - def public_key_into(self, output: UnsafePointer[UInt8, MutAnyOrigin]): + def public_key_into(self, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): for i in range(32): output[i] = self._a_enc[i] @no_inline - def sign(self, message: Span[UInt8, ...], output: UnsafePointer[UInt8, MutAnyOrigin]) raises: + def sign(self, message: Span[UInt8, ...], output: UnsafePointer[mut=True, UInt8, _, address_space=_]) raises: var r_hash = InlineArray[UInt8, 64](uninitialized=True) var r_ctx = SHA512Context() - sha512_update(r_ctx, Span[UInt8, ...](ptr=self._prefix.unsafe_ptr(), length=32)) + sha512_update(r_ctx, Span[UInt8, ...](unsafe_ptr=self._prefix.unsafe_ptr(), length=32)) sha512_update(r_ctx, message) sha512_final_to_buffer(r_ctx, r_hash.unsafe_ptr()) - var r_scalar = Scalar.from_bytes_wide(Span[UInt8, ...](ptr=r_hash.unsafe_ptr(), length=64)) + var r_scalar = Scalar.from_bytes_wide(Span[UInt8, ...](unsafe_ptr=r_hash.unsafe_ptr(), length=64)) var r_bytes = InlineArray[UInt8, 32](uninitialized=True) r_scalar.to_bytes_into(r_bytes.unsafe_ptr()) - var R_point = _mul_base_ct(Span[UInt8, ...](ptr=r_bytes.unsafe_ptr(), length=32)) + var R_point = _mul_base_ct(Span[UInt8, ...](unsafe_ptr=r_bytes.unsafe_ptr(), length=32)) var R_enc = InlineArray[UInt8, 32](uninitialized=True) edwards_encode_into(R_point, R_enc.unsafe_ptr()) var k_hash = InlineArray[UInt8, 64](uninitialized=True) var k_ctx = SHA512Context() - sha512_update(k_ctx, Span[UInt8, ...](ptr=R_enc.unsafe_ptr(), length=32)) - sha512_update(k_ctx, Span[UInt8, ...](ptr=self._a_enc.unsafe_ptr(), length=32)) + sha512_update(k_ctx, Span[UInt8, ...](unsafe_ptr=R_enc.unsafe_ptr(), length=32)) + sha512_update(k_ctx, Span[UInt8, ...](unsafe_ptr=self._a_enc.unsafe_ptr(), length=32)) sha512_update(k_ctx, message) sha512_final_to_buffer(k_ctx, k_hash.unsafe_ptr()) - var k_scalar = Scalar.from_bytes_wide(Span[UInt8, ...](ptr=k_hash.unsafe_ptr(), length=64)) + var k_scalar = Scalar.from_bytes_wide(Span[UInt8, ...](unsafe_ptr=k_hash.unsafe_ptr(), length=64)) var S_scalar = r_scalar + k_scalar * self._s var S_bytes = InlineArray[UInt8, 32](uninitialized=True) S_scalar.to_bytes_into(S_bytes.unsafe_ptr()) @@ -908,23 +914,23 @@ def ed25519_verify(public_key: Span[UInt8, ...], message: Span[UInt8, ...], sign var S_bytes = InlineArray[UInt8, 32](uninitialized=True) for i in range(32): S_bytes[i] = signature[32 + i] - var S_bytes_span = Span[UInt8, ...](ptr=S_bytes.unsafe_ptr(), length=32) + var S_bytes_span = Span[UInt8, ...](unsafe_ptr=S_bytes.unsafe_ptr(), length=32) if not _s_lt_l(S_bytes_span): return False var k_hash = InlineArray[UInt8, 64](uninitialized=True) var k_ctx = SHA512Context() - sha512_update(k_ctx, Span[UInt8, ...](ptr=R_enc.unsafe_ptr(), length=32)) + sha512_update(k_ctx, Span[UInt8, ...](unsafe_ptr=R_enc.unsafe_ptr(), length=32)) sha512_update(k_ctx, public_key) sha512_update(k_ctx, message) sha512_final_to_buffer(k_ctx, k_hash.unsafe_ptr()) - var k_hash_span = Span[UInt8, ...](ptr=k_hash.unsafe_ptr(), length=64) + var k_hash_span = Span[UInt8, ...](unsafe_ptr=k_hash.unsafe_ptr(), length=64) var k_scalar = Scalar.from_bytes_wide(k_hash_span) var k_bytes = InlineArray[UInt8, 32](uninitialized=True) k_scalar.to_bytes_into(k_bytes.unsafe_ptr()) - var k_bytes_span = Span[UInt8, ...](ptr=k_bytes.unsafe_ptr(), length=32) + var k_bytes_span = Span[UInt8, ...](unsafe_ptr=k_bytes.unsafe_ptr(), length=32) var Q = _double_scalar_mult_vartime(k_bytes_span, edwards_negate(A), S_bytes_span) var Q_enc = InlineArray[UInt8, 32](uninitialized=True) diff --git a/src/thistle/ed25519_table.mojo b/src/thistle/ed25519_table.mojo index d27d317..fad4e16 100644 --- a/src/thistle/ed25519_table.mojo +++ b/src/thistle/ed25519_table.mojo @@ -401,10 +401,10 @@ def ed25519_base_table() -> InlineArray[UInt64, 4096]: p.store[alignment=8](3712, _ED25519_BT29) p.store[alignment=8](3840, _ED25519_BT30) p.store[alignment=8](3968, _ED25519_BT31) - return t + return t^ @no_inline def ed25519_b_odd_table() -> InlineArray[UInt64, 128]: var t = InlineArray[UInt64, 128](uninitialized=True) t.unsafe_ptr().store[alignment=8](0, _ED25519_B_ODD) - return t + return t^ diff --git a/src/thistle/kcipher2.mojo b/src/thistle/kcipher2.mojo index 7961d89..8c9648f 100644 --- a/src/thistle/kcipher2.mojo +++ b/src/thistle/kcipher2.mojo @@ -395,7 +395,7 @@ struct KCipher2: ik[10] = ik[6] ^ ik[9] ik[11] = ik[7] ^ ik[10] - return ik + return ik^ @always_inline def _select_u32(self, a: UInt32, b: UInt32, choice: UInt32) -> UInt32: diff --git a/src/thistle/ml_dsa.mojo b/src/thistle/ml_dsa.mojo index cc88bbb..01729d5 100644 --- a/src/thistle/ml_dsa.mojo +++ b/src/thistle/ml_dsa.mojo @@ -110,8 +110,8 @@ struct DSAPolyVec[ROWS: Int](Movable): return self.data[i] @always_inline - def __setitem__(mut self, i: Int, value: DSAPoly): - self.data[i] = value + def __setitem__(mut self, i: Int, var value: DSAPoly): + self.data[i] = value^ @fieldwise_init @@ -228,7 +228,8 @@ def _zero_poly_vec(count: Int) -> List[List[UInt32]]: def _copy_bytes(src: Span[UInt8, ...]) -> List[UInt8]: var out = List[UInt8](capacity=len(src)) - out.extend(src) + for i in range(len(src)): + out.append(src[i]) return out^ @@ -237,7 +238,8 @@ def _slice_bytes(src: Span[UInt8, ...], start: Int, count: Int) -> List[UInt8]: def _append_bytes(mut out: List[UInt8], src: Span[UInt8, ...]): - out.extend(src) + for i in range(len(src)): + out.append(src[i]) def _append_bytes_stack(mut out: StackBuffer[UInt8, ...], src: Span[UInt8, ...]): @@ -393,7 +395,8 @@ def _dsa_poly_sub_into(mut r: DSAPoly, b: DSAPoly): def _dsa_copy_poly_into(mut r: DSAPoly, a: DSAPoly): - r = a + for i in range(N): + r[i] = a[i] def _ntt_mul_into(mut r: List[UInt32], a: List[UInt32], b: List[UInt32]): @@ -403,9 +406,9 @@ def _ntt_mul_into(mut r: List[UInt32], a: List[UInt32], b: List[UInt32]): @always_inline def _ntt_mul_ptrs( - r: UnsafePointer[UInt32, MutAnyOrigin], - a: UnsafePointer[UInt32, ImmutAnyOrigin], - b: UnsafePointer[UInt32, ImmutAnyOrigin], + r: UnsafePointer[mut=True, UInt32, _, address_space=_], + a: UnsafePointer[mut=False, UInt32, _, address_space=_], + b: UnsafePointer[mut=False, UInt32, _, address_space=_], ): var i = 0 while i < N: @@ -557,7 +560,7 @@ def _sample_ntt(rho: Span[UInt8, ...], s: UInt8, r: UInt8) raises -> List[UInt32 var out_len = 840 while True: - var buf = _shake128_expand(Span[UInt8, ...](ptr=input.ptr(), length=input.len()), out_len) + var buf = _shake128_expand(Span[UInt8, ...](unsafe_ptr=input.ptr(), length=input.len()), out_len) var a = List[UInt32](unsafe_uninit_length=N) var ap = a.unsafe_ptr() var bp = buf.unsafe_ptr() @@ -601,7 +604,7 @@ def _sample_bounded_poly(rho: Span[UInt8, ...], r: UInt8, p: MLDSAParams) -> Lis var out_len = 272 while True: - var buf = _shake256_expand(Span[UInt8, ...](ptr=input.ptr(), length=input.len()), out_len) + var buf = _shake256_expand(Span[UInt8, ...](unsafe_ptr=input.ptr(), length=input.len()), out_len) var a = List[UInt32](capacity=N) for i in range(len(buf)): var z0 = buf[i] & 0x0F @@ -676,7 +679,7 @@ def _sample_in_ball_dsa(rho: Span[UInt8, ...], p: MLDSAParams) -> DSAPoly: break _zero_list_u8(buf) if ok: - return c + return c^ out_len *= 2 @@ -1118,7 +1121,7 @@ def _compute_message_hash(tr: Span[UInt8, ...], msg: Span[UInt8, ...], context: var prefix = StackBuffer[UInt8, 2]() prefix.push_unchecked(UInt8(0)) prefix.push_unchecked(UInt8(len(context))) - sha3_update(ctx, Span[UInt8, ...](ptr=prefix.ptr(), length=prefix.len())) + sha3_update(ctx, Span[UInt8, ...](unsafe_ptr=prefix.ptr(), length=prefix.len())) zero_stack_u8(prefix) sha3_update(ctx, context) @@ -1137,7 +1140,7 @@ def mldsa_private_key_from_seed(seed: Span[UInt8, ...], p: MLDSAParams) raises - _append_bytes_stack(xi, seed) xi.push_unchecked(UInt8(p.k)) xi.push_unchecked(UInt8(p.l)) - var expanded = shake256(Span[UInt8, ...](ptr=xi.ptr(), length=xi.len()), 128) + var expanded = shake256(Span[UInt8, ...](unsafe_ptr=xi.ptr(), length=xi.len()), 128) zero_stack_u8(xi) var rho = _slice_bytes(Span[UInt8, ...](expanded), 0, 32) var rho_s = _slice_bytes(Span[UInt8, ...](expanded), 32, 64) @@ -1282,7 +1285,7 @@ def mldsa_sign_external_mu(priv: MLDSAPrivateKey, mu: Span[UInt8, ...], random: _append_bytes_stack(h_input, Span[UInt8, ...](priv.k_seed)) _append_bytes_stack(h_input, random) _append_bytes_stack(h_input, mu) - var nonce = shake256(Span[UInt8, ...](ptr=h_input.ptr(), length=h_input.len()), 64) + var nonce = shake256(Span[UInt8, ...](unsafe_ptr=h_input.ptr(), length=h_input.len()), 64) zero_stack_u8(h_input) # signing scratch allocated once and reused across attempts @@ -1305,7 +1308,7 @@ def mldsa_sign_external_mu(priv: MLDSAPrivateKey, mu: Span[UInt8, ...], random: seed.push_unchecked(UInt8(kappa & 0xFF)) seed.push_unchecked(UInt8((kappa >> 8) & 0xFF)) kappa += 1 - var v = shake256(Span[UInt8, ...](ptr=seed.ptr(), length=seed.len()), (p.gamma1_log + 1) * N // 8) + var v = shake256(Span[UInt8, ...](unsafe_ptr=seed.ptr(), length=seed.len()), (p.gamma1_log + 1) * N // 8) zero_stack_u8(seed) _dsa_bit_unpack_into(y[_r], Span[UInt8, ...](v), p) _zero_list_u8(v) @@ -1325,7 +1328,7 @@ def mldsa_sign_external_mu(priv: MLDSAPrivateKey, mu: Span[UInt8, ...], random: _append_bytes_stack(ch_input, mu) for i in range(p.k): _dsa_append_w1_encoded_stack(ch_input, w[i], p) - var ch_list = shake256(Span[UInt8, ...](ptr=ch_input.ptr(), length=ch_input.len()), p.lambda_bits // 4) + var ch_list = shake256(Span[UInt8, ...](unsafe_ptr=ch_input.ptr(), length=ch_input.len()), p.lambda_bits // 4) zero_stack_u8(ch_input) var ch = InlineArray[UInt8, MLDSA_CRHBYTES](fill=0) for i in range(p.lambda_bits // 4): @@ -1482,7 +1485,7 @@ def mldsa_verify_external_mu(pub: MLDSAPublicKey, mu: Span[UInt8, ...], sig: Spa _append_bytes_stack(ch_input, mu) for i in range(p.k): _append_use_hint_encoded_stack(ch_input, w[i], h[i], p) - var computed = shake256(Span[UInt8, ...](ptr=ch_input.ptr(), length=ch_input.len()), p.lambda_bits // 4) + var computed = shake256(Span[UInt8, ...](unsafe_ptr=ch_input.ptr(), length=ch_input.len()), p.lambda_bits // 4) zero_stack_u8(ch_input) var ok = _bytes_equal(Span[UInt8, ...](ch), Span[UInt8, ...](computed)) _zero_list_u8(ch) diff --git a/src/thistle/ml_kem.mojo b/src/thistle/ml_kem.mojo index 78f622b..a15a3d3 100644 --- a/src/thistle/ml_kem.mojo +++ b/src/thistle/ml_kem.mojo @@ -343,7 +343,7 @@ def base_case_multiply(a0: Int16, a1: Int16, b0: Int16, b1: Int16, zeta: Int16) def poly_reduce(mut r: Poly): comptime W = simd_width_of[DType.int16]() - var ptr = r.coeffs.unsafe_ptr() + var ptr = r.coeffs.unsafe_ptr().unsafe_origin_cast[MutAnyOrigin]() def reduce_chunk[w: Int](i: Int) {ptr}: var a = ptr.load[width=w](i) @@ -379,7 +379,7 @@ def poly_basemul_montgomery(mut r: Poly, ref a: Poly, ref b: Poly): def poly_tomont(mut r: Poly): comptime F: Int16 = Int16((1 << 32) % Q) comptime W = simd_width_of[DType.int16]() - var ptr = r.coeffs.unsafe_ptr() + var ptr = r.coeffs.unsafe_ptr().unsafe_origin_cast[MutAnyOrigin]() def tomont_chunk[w: Int](i: Int) {ptr}: var a = ptr.load[width=w](i) @@ -391,8 +391,12 @@ def poly_tomont(mut r: Poly): def poly_add_inplace(mut r: Poly, ref b: Poly): comptime W = simd_width_of[DType.int16]() - var rp = r.coeffs.unsafe_ptr() - var bp = b.coeffs.unsafe_ptr() + var rp = r.coeffs.unsafe_ptr().unsafe_origin_cast[MutAnyOrigin]() + var bp = ( + b.coeffs.unsafe_ptr() + .unsafe_mut_cast[False]() + .unsafe_origin_cast[ImmutAnyOrigin]() + ) def add_inplace_chunk[w: Int](i: Int) {rp, bp}: rp.store[width=w](i, rp.load[width=w](i) + bp.load[width=w](i)) @@ -402,8 +406,12 @@ def poly_add_inplace(mut r: Poly, ref b: Poly): def poly_sub_from(mut r: Poly, ref a: Poly): comptime W = simd_width_of[DType.int16]() - var rp = r.coeffs.unsafe_ptr() - var ap = a.coeffs.unsafe_ptr() + var rp = r.coeffs.unsafe_ptr().unsafe_origin_cast[MutAnyOrigin]() + var ap = ( + a.coeffs.unsafe_ptr() + .unsafe_mut_cast[False]() + .unsafe_origin_cast[ImmutAnyOrigin]() + ) def sub_from_chunk[w: Int](i: Int) {rp, ap}: rp.store[width=w](i, ap.load[width=w](i) - rp.load[width=w](i)) @@ -823,7 +831,7 @@ def prf(key: Span[UInt8, ...], iv: UInt8, out_len: Int) raises -> List[UInt8]: for i in range(SYMBYTES): input.push_unchecked(key[i]) input.push_unchecked(iv) - return shake256(Span[UInt8, ...](ptr=input.ptr(), length=input.len()), out_len) + return shake256(Span[UInt8, ...](unsafe_ptr=input.ptr(), length=input.len()), out_len) def prf_into(mut out: StackBuffer[UInt8, ...], key: Span[UInt8, ...], iv: UInt8, out_len: Int) raises: @@ -833,7 +841,7 @@ def prf_into(mut out: StackBuffer[UInt8, ...], key: Span[UInt8, ...], iv: UInt8, for i in range(SYMBYTES): input.push_unchecked(key[i]) input.push_unchecked(iv) - shake256_into(out, Span[UInt8, ...](ptr=input.ptr(), length=input.len()), out_len) + shake256_into(out, Span[UInt8, ...](unsafe_ptr=input.ptr(), length=input.len()), out_len) def rkprf(key: Span[UInt8, ...], input: Span[UInt8, ...]) raises -> List[UInt8]: @@ -846,7 +854,7 @@ def rkprf(key: Span[UInt8, ...], input: Span[UInt8, ...]) raises -> List[UInt8]: buf.push_unchecked(key[i]) for i in range(len(input)): buf.push_unchecked(input[i]) - return shake256(Span[UInt8, ...](ptr=buf.ptr(), length=buf.len()), SYMBYTES) + return shake256(Span[UInt8, ...](unsafe_ptr=buf.ptr(), length=buf.len()), SYMBYTES) def rkprf_into(mut out: StackBuffer[UInt8, SYMBYTES], key: Span[UInt8, ...], input: Span[UInt8, ...]) raises: @@ -859,7 +867,7 @@ def rkprf_into(mut out: StackBuffer[UInt8, SYMBYTES], key: Span[UInt8, ...], inp buf.push_unchecked(key[i]) for i in range(len(input)): buf.push_unchecked(input[i]) - shake256_into(out, Span[UInt8, ...](ptr=buf.ptr(), length=buf.len()), SYMBYTES) + shake256_into(out, Span[UInt8, ...](unsafe_ptr=buf.ptr(), length=buf.len()), SYMBYTES) def hash_h_into(mut out: StackBuffer[UInt8, SYMBYTES], input: Span[UInt8, ...]): @@ -878,7 +886,7 @@ def xof(seed: Span[UInt8, ...], x: UInt8, y: UInt8, out_len: Int) raises -> List extseed.push_unchecked(seed[i]) extseed.push_unchecked(x) extseed.push_unchecked(y) - return shake128(Span[UInt8, ...](ptr=extseed.ptr(), length=extseed.len()), out_len) + return shake128(Span[UInt8, ...](unsafe_ptr=extseed.ptr(), length=extseed.len()), out_len) def rej_uniform(mut r: Poly, start: Int, buf: Span[UInt8, ...]) -> Int: @@ -907,14 +915,14 @@ def sample_ntt_into(mut out: Poly, seed: Span[UInt8, ...], x: UInt8, y: UInt8) r extseed.push_unchecked(y) var ctx = SHA3Context(1344) - sha3_update(ctx, Span[UInt8, ...](ptr=extseed.ptr(), length=extseed.len())) + sha3_update(ctx, Span[UInt8, ...](unsafe_ptr=extseed.ptr(), length=extseed.len())) shake_finalize(ctx) var block = StackBuffer[UInt8, 504]() var ctr = 0 while ctr < N: shake_squeeze_prefix_into(ctx, block, 504) - ctr += rej_uniform(out, ctr, Span[UInt8, ...](ptr=block.ptr(), length=block.len())) + ctr += rej_uniform(out, ctr, Span[UInt8, ...](unsafe_ptr=block.ptr(), length=block.len())) if ctr < N: shake_advance(ctx) @@ -922,19 +930,19 @@ def sample_ntt_into(mut out: Poly, seed: Span[UInt8, ...], x: UInt8, y: UInt8) r def poly_getnoise_eta1_512(mut r: Poly, seed: Span[UInt8, ...], iv: UInt8) raises: var buf = StackBuffer[UInt8, ETA1_512 * N // 4]() prf_into(buf, seed, iv, ETA1_512 * N // 4) - poly_cbd_eta1_512(r, Span[UInt8, ...](ptr=buf.ptr(), length=buf.len())) + poly_cbd_eta1_512(r, Span[UInt8, ...](unsafe_ptr=buf.ptr(), length=buf.len())) def poly_getnoise_eta1(mut r: Poly, seed: Span[UInt8, ...], iv: UInt8) raises: var buf = StackBuffer[UInt8, ETA1 * N // 4]() prf_into(buf, seed, iv, ETA1 * N // 4) - poly_cbd_eta1(r, Span[UInt8, ...](ptr=buf.ptr(), length=buf.len())) + poly_cbd_eta1(r, Span[UInt8, ...](unsafe_ptr=buf.ptr(), length=buf.len())) def poly_getnoise_eta2(mut r: Poly, seed: Span[UInt8, ...], iv: UInt8) raises: var buf = StackBuffer[UInt8, ETA2 * N // 4]() prf_into(buf, seed, iv, ETA2 * N // 4) - poly_cbd_eta2(r, Span[UInt8, ...](ptr=buf.ptr(), length=buf.len())) + poly_cbd_eta2(r, Span[UInt8, ...](unsafe_ptr=buf.ptr(), length=buf.len())) def gen_matrix(mut a: InlineArray[Polyvec, K_MAX], seed: Span[UInt8, ...], transposed: Bool, k: Int) raises: @@ -1057,9 +1065,9 @@ def k_pke_keygen(mut ek: KPKEEncryptionKey, mut dk: KPKEDecapsulationKey, d: Spa g_input.push_unchecked(d[i]) g_input.push_unchecked(UInt8(k)) var g = StackBuffer[UInt8, 2 * SYMBYTES]() - hash_g_into(g, Span[UInt8, ...](ptr=g_input.ptr(), length=g_input.len())) - var rho = Span[UInt8, ...](ptr=g.ptr(), length=SYMBYTES) - var sigma = Span[UInt8, ...](ptr=g.ptr() + SYMBYTES, length=SYMBYTES) + hash_g_into(g, Span[UInt8, ...](unsafe_ptr=g_input.ptr(), length=g_input.len())) + var rho = Span[UInt8, ...](unsafe_ptr=g.ptr(), length=SYMBYTES) + var sigma = Span[UInt8, ...](unsafe_ptr=g.ptr() + SYMBYTES, length=SYMBYTES) var a = InlineArray[Polyvec, K_MAX](fill=Polyvec()) gen_matrix(a, Span[UInt8, ...](rho), False, k) @@ -1108,9 +1116,9 @@ def k_pke_keygen_k[k: Int](mut ek: KPKEEncryptionKey, mut dk: KPKEDecapsulationK g_input.push_unchecked(d[i]) g_input.push_unchecked(UInt8(k)) var g = StackBuffer[UInt8, 2 * SYMBYTES]() - hash_g_into(g, Span[UInt8, ...](ptr=g_input.ptr(), length=g_input.len())) - var rho = Span[UInt8, ...](ptr=g.ptr(), length=SYMBYTES) - var sigma = Span[UInt8, ...](ptr=g.ptr() + SYMBYTES, length=SYMBYTES) + hash_g_into(g, Span[UInt8, ...](unsafe_ptr=g_input.ptr(), length=g_input.len())) + var rho = Span[UInt8, ...](unsafe_ptr=g.ptr(), length=SYMBYTES) + var sigma = Span[UInt8, ...](unsafe_ptr=g.ptr() + SYMBYTES, length=SYMBYTES) var a = InlineArray[Polyvec, K_MAX](fill=Polyvec()) gen_matrix_k_static[k, False](a, Span[UInt8, ...](rho)) @@ -1305,7 +1313,7 @@ def kem_keygen_internal(seed: Span[UInt8, ...], k: Int) raises -> DecapsulationK for i in range(ek_bytes.len()): dk.ek.raw_bytes[i] = ek_bytes[i] var h = StackBuffer[UInt8, SYMBYTES]() - hash_h_into(h, Span[UInt8, ...](ptr=ek_bytes.ptr(), length=ek_bytes.len())) + hash_h_into(h, Span[UInt8, ...](unsafe_ptr=ek_bytes.ptr(), length=ek_bytes.len())) for i in range(SYMBYTES): dk.ek.h[i] = h[i] dk.z[i] = z[i] @@ -1327,7 +1335,7 @@ def kem_keygen_internal_k[k: Int](seed: Span[UInt8, ...]) raises -> Decapsulatio for i in range(ek_bytes.len()): dk.ek.raw_bytes[i] = ek_bytes[i] var h = StackBuffer[UInt8, SYMBYTES]() - hash_h_into(h, Span[UInt8, ...](ptr=ek_bytes.ptr(), length=ek_bytes.len())) + hash_h_into(h, Span[UInt8, ...](unsafe_ptr=ek_bytes.ptr(), length=ek_bytes.len())) for i in range(SYMBYTES): dk.ek.h[i] = h[i] dk.z[i] = z[i] @@ -1635,11 +1643,11 @@ def mlkem_encaps_seed_into(mut ciphertext_out: StackBuffer[UInt8, CIPHERTEXTBYTE for i in range(SYMBYTES): g_input.push_unchecked(h[i]) var g = StackBuffer[UInt8, 2 * SYMBYTES]() - hash_g_into(g, Span[UInt8, ...](ptr=g_input.ptr(), length=g_input.len())) + hash_g_into(g, Span[UInt8, ...](unsafe_ptr=g_input.ptr(), length=g_input.len())) zero_stack_u8(g_input) for i in range(SYMBYTES): shared_out.push_unchecked(g[i]) - if not k_pke_encrypt_into(ciphertext_out, ek, m, Span[UInt8, ...](ptr=g.ptr() + SYMBYTES, length=SYMBYTES)): + if not k_pke_encrypt_into(ciphertext_out, ek, m, Span[UInt8, ...](unsafe_ptr=g.ptr() + SYMBYTES, length=SYMBYTES)): zero_stack_u8(g) zero_stack_u8(shared_out) zero_stack_u8(ciphertext_out) @@ -1667,11 +1675,11 @@ def mlkem_encaps_seed_into_k[k: Int](mut ciphertext_out: StackBuffer[UInt8, CIPH for i in range(SYMBYTES): g_input.push_unchecked(h[i]) var g = StackBuffer[UInt8, 2 * SYMBYTES]() - hash_g_into(g, Span[UInt8, ...](ptr=g_input.ptr(), length=g_input.len())) + hash_g_into(g, Span[UInt8, ...](unsafe_ptr=g_input.ptr(), length=g_input.len())) zero_stack_u8(g_input) for i in range(SYMBYTES): shared_out.push_unchecked(g[i]) - if not k_pke_encrypt_into_k[k](ciphertext_out, ek, m, Span[UInt8, ...](ptr=g.ptr() + SYMBYTES, length=SYMBYTES)): + if not k_pke_encrypt_into_k[k](ciphertext_out, ek, m, Span[UInt8, ...](unsafe_ptr=g.ptr() + SYMBYTES, length=SYMBYTES)): zero_stack_u8(g) zero_stack_u8(shared_out) zero_stack_u8(ciphertext_out) @@ -1732,11 +1740,11 @@ def mlkem_decaps_into(mut shared_out: StackBuffer[UInt8, SYMBYTES], dk_bytes: Sp for i in range(SYMBYTES): g_input.push_unchecked(dk.ek.h[i]) var g = StackBuffer[UInt8, 2 * SYMBYTES]() - hash_g_into(g, Span[UInt8, ...](ptr=g_input.ptr(), length=g_input.len())) + hash_g_into(g, Span[UInt8, ...](unsafe_ptr=g_input.ptr(), length=g_input.len())) zero_stack_u8(g_input) var ct_check = StackBuffer[UInt8, CIPHERTEXTBYTES_MAX]() - if not k_pke_encrypt_into(ct_check, dk.ek.pke_ek, Span[UInt8, ...](ptr=m.ptr(), length=m.len()), Span[UInt8, ...](ptr=g.ptr() + SYMBYTES, length=SYMBYTES)): + if not k_pke_encrypt_into(ct_check, dk.ek.pke_ek, Span[UInt8, ...](unsafe_ptr=m.ptr(), length=m.len()), Span[UInt8, ...](unsafe_ptr=g.ptr() + SYMBYTES, length=SYMBYTES)): zero_stack_u8(m) zero_stack_u8(g) zero_stack_u8(ct_check) @@ -1744,7 +1752,7 @@ def mlkem_decaps_into(mut shared_out: StackBuffer[UInt8, SYMBYTES], dk_bytes: Sp var rejection = StackBuffer[UInt8, SYMBYTES]() rkprf_into(rejection, Span[UInt8, ...](dk.z), ciphertext) - var equal = _ct_is_zero_u8(_bytes_diff(Span[UInt8, ...](ptr=ct_check.ptr(), length=ct_check.len()), ciphertext)) + var equal = _ct_is_zero_u8(_bytes_diff(Span[UInt8, ...](unsafe_ptr=ct_check.ptr(), length=ct_check.len()), ciphertext)) for i in range(SYMBYTES): shared_out.push_unchecked(_ct_select_u8(rejection[i], g[i], equal)) @@ -1776,11 +1784,11 @@ def mlkem_decaps_into_k[k: Int](mut shared_out: StackBuffer[UInt8, SYMBYTES], dk for i in range(SYMBYTES): g_input.push_unchecked(dk.ek.h[i]) var g = StackBuffer[UInt8, 2 * SYMBYTES]() - hash_g_into(g, Span[UInt8, ...](ptr=g_input.ptr(), length=g_input.len())) + hash_g_into(g, Span[UInt8, ...](unsafe_ptr=g_input.ptr(), length=g_input.len())) zero_stack_u8(g_input) var ct_check = StackBuffer[UInt8, CIPHERTEXTBYTES_MAX]() - if not k_pke_encrypt_into_k[k](ct_check, dk.ek.pke_ek, Span[UInt8, ...](ptr=m.ptr(), length=m.len()), Span[UInt8, ...](ptr=g.ptr() + SYMBYTES, length=SYMBYTES)): + if not k_pke_encrypt_into_k[k](ct_check, dk.ek.pke_ek, Span[UInt8, ...](unsafe_ptr=m.ptr(), length=m.len()), Span[UInt8, ...](unsafe_ptr=g.ptr() + SYMBYTES, length=SYMBYTES)): zero_stack_u8(m) zero_stack_u8(g) zero_stack_u8(ct_check) @@ -1788,7 +1796,7 @@ def mlkem_decaps_into_k[k: Int](mut shared_out: StackBuffer[UInt8, SYMBYTES], dk var rejection = StackBuffer[UInt8, SYMBYTES]() rkprf_into(rejection, Span[UInt8, ...](dk.z), ciphertext) - var equal = _ct_is_zero_u8(_bytes_diff(Span[UInt8, ...](ptr=ct_check.ptr(), length=ct_check.len()), ciphertext)) + var equal = _ct_is_zero_u8(_bytes_diff(Span[UInt8, ...](unsafe_ptr=ct_check.ptr(), length=ct_check.len()), ciphertext)) for i in range(SYMBYTES): shared_out.push_unchecked(_ct_select_u8(rejection[i], g[i], equal)) diff --git a/src/thistle/p256.mojo b/src/thistle/p256.mojo index a683922..899f090 100644 --- a/src/thistle/p256.mojo +++ b/src/thistle/p256.mojo @@ -6,6 +6,7 @@ from .p256_table import p256_base_table from .utils import u64_nonzero_choice, u64_zero_choice from .sha2 import sha256_hash from .pbkdf2 import hmac_sha256 +from std.utils import StaticTuple comptime P256_SIZE = 32 comptime P256_POINT_SIZE = 65 @@ -14,13 +15,15 @@ comptime _MASK64 = UInt128(0xFFFFFFFFFFFFFFFF) struct U256(Copyable, ImplicitlyCopyable, Movable): - var limbs: InlineArray[UInt64, 4] + var limbs: StaticTuple[UInt64, 4] def __init__(out self): - self.limbs = InlineArray[UInt64, 4](fill=0) + self.limbs = StaticTuple[UInt64, 4]() + comptime for i in range(4): + self.limbs[i] = 0 def __init__(out self, l0: UInt64, l1: UInt64, l2: UInt64, l3: UInt64): - self.limbs = InlineArray[UInt64, 4](uninitialized=True) + self.limbs = StaticTuple[UInt64, 4]() self.limbs[0] = l0 self.limbs[1] = l1 self.limbs[2] = l2 @@ -451,7 +454,7 @@ def _from_be(bytes: Span[UInt8, ...]) -> U256: return out -def _to_be(x: U256, output: UnsafePointer[UInt8, MutAnyOrigin]): +def _to_be(x: U256, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): for i in range(4): var limb = x.limbs[3 - i] for k in range(8): @@ -750,7 +753,7 @@ def _scalar_mult_base(k: U256) -> P256Point: def p256_decode_uncompressed(point: Span[UInt8, ...]) -> P256Point: if len(point) == 33 and (point[0] == 0x02 or point[0] == 0x03): var x = _from_be( - Span[UInt8, ...](ptr=point.unsafe_ptr() + 1, length=32) + Span[UInt8, ...](unsafe_ptr=point.unsafe_ptr() + 1, length=32) ) if _cmp(x, _p()) >= 0: return P256Point() @@ -769,8 +772,8 @@ def p256_decode_uncompressed(point: Span[UInt8, ...]) -> P256Point: return p if len(point) != P256_POINT_SIZE or point[0] != 0x04: return P256Point() - var x = _from_be(Span[UInt8, ...](ptr=point.unsafe_ptr() + 1, length=32)) - var y = _from_be(Span[UInt8, ...](ptr=point.unsafe_ptr() + 33, length=32)) + var x = _from_be(Span[UInt8, ...](unsafe_ptr=point.unsafe_ptr() + 1, length=32)) + var y = _from_be(Span[UInt8, ...](unsafe_ptr=point.unsafe_ptr() + 33, length=32)) var p = P256Point(x, y, False) if not _is_on_curve(p): return P256Point() @@ -778,31 +781,32 @@ def p256_decode_uncompressed(point: Span[UInt8, ...]) -> P256Point: def p256_encode_uncompressed( - point: P256Point, output: UnsafePointer[UInt8, MutAnyOrigin] + point: P256Point, output: Span[mut=True, UInt8, ...] ) -> Bool: - if point.infinity or not _is_on_curve(point): + if len(output) < P256_POINT_SIZE or point.infinity or not _is_on_curve(point): return False - output[0] = 0x04 - _to_be(point.x, output + 1) - _to_be(point.y, output + 33) + var out_ptr = output.unsafe_ptr() + out_ptr[0] = 0x04 + _to_be(point.x, out_ptr + 1) + _to_be(point.y, out_ptr + 33) return True @no_inline def p256_public_key( - private_key: Span[UInt8, ...], output: UnsafePointer[UInt8, MutAnyOrigin] + private_key: Span[UInt8, ...], output: Span[mut=True, UInt8, ...] ) -> Bool: - if len(private_key) != 32: + if len(private_key) != 32 or len(output) < P256_POINT_SIZE: return False var d = _from_be(private_key) if d.is_zero() or _cmp(d, _n()) >= 0: - var dp = d.limbs.unsafe_ptr() + var dp = UnsafePointer(to=d.limbs[0]).unsafe_mut_cast[True]() for i in range(4): dp.store[volatile=True](i, UInt64(0)) return False var q = _scalar_mult_base(d) var ok = p256_encode_uncompressed(q, output) - var dp = d.limbs.unsafe_ptr() + var dp = UnsafePointer(to=d.limbs[0]).unsafe_mut_cast[True]() for i in range(4): dp.store[volatile=True](i, UInt64(0)) return ok @@ -812,29 +816,29 @@ def p256_public_key( def p256_ecdh( private_key: Span[UInt8, ...], public_key: Span[UInt8, ...], - output: UnsafePointer[UInt8, MutAnyOrigin], + output: Span[mut=True, UInt8, ...], ) -> Bool: - if len(private_key) != 32: + if len(private_key) != 32 or len(output) < P256_SIZE: return False var d = _from_be(private_key) if d.is_zero() or _cmp(d, _n()) >= 0: - var dp = d.limbs.unsafe_ptr() + var dp = UnsafePointer(to=d.limbs[0]).unsafe_mut_cast[True]() for i in range(4): dp.store[volatile=True](i, UInt64(0)) return False var q = p256_decode_uncompressed(public_key) if q.infinity: - var dp = d.limbs.unsafe_ptr() + var dp = UnsafePointer(to=d.limbs[0]).unsafe_mut_cast[True]() for i in range(4): dp.store[volatile=True](i, UInt64(0)) return False var shared = _scalar_mult(d, q) - var dp = d.limbs.unsafe_ptr() + var dp = UnsafePointer(to=d.limbs[0]).unsafe_mut_cast[True]() for i in range(4): dp.store[volatile=True](i, UInt64(0)) if shared.infinity: return False - _to_be(shared.x, output) + _to_be(shared.x, output.unsafe_ptr()) return True @@ -925,7 +929,7 @@ def _reduce_n(x: U256) -> U256: def _wipe_u256(mut x: U256): - var ptr = x.limbs.unsafe_ptr() + var ptr = UnsafePointer(to=x.limbs[0]).unsafe_mut_cast[True]() for i in range(4): ptr.store[volatile=True](i, UInt64(0)) @@ -1008,10 +1012,14 @@ def _rfc6979_p256(private_key: Span[UInt8, ...], digest: Span[UInt8, ...], skip: def p256_ecdsa_sign_digest( private_key: Span[UInt8, ...], digest: Span[UInt8, ...], - signature: UnsafePointer[UInt8, MutAnyOrigin], + signature: Span[mut=True, UInt8, ...], ) -> Bool: - if len(private_key) != 32 or len(digest) != 32: + if ( + len(private_key) != 32 or len(digest) != 32 + or len(signature) < P256_SIGNATURE_SIZE + ): return False + var signature_ptr = signature.unsafe_ptr() var d = _from_be(private_key) if d.is_zero() or _cmp(d, _n()) >= 0: _wipe_u256(d) @@ -1037,8 +1045,8 @@ def p256_ecdsa_sign_digest( _wipe_u256(total) retry += 1 continue - _to_be(r, signature) - _to_be(s, signature + 32) + _to_be(r, signature_ptr) + _to_be(s, signature_ptr + 32) _wipe_u256(d) _wipe_u256(z) _wipe_u256(k) @@ -1051,10 +1059,14 @@ def p256_ecdsa_sign_digest( def p256_ecdsa_sign( private_key: Span[UInt8, ...], message: Span[UInt8, ...], - signature: UnsafePointer[UInt8, MutAnyOrigin], + signature: Span[mut=True, UInt8, ...], ) -> Bool: + if len(signature) < P256_SIGNATURE_SIZE: + return False var digest = sha256_hash(message) - var ok = p256_ecdsa_sign_digest(private_key, Span[UInt8, ...](digest), signature) + var ok = p256_ecdsa_sign_digest( + private_key, Span[UInt8, ...](digest), signature + ) _wipe_list_u8(digest) return ok @@ -1116,7 +1128,9 @@ def p256_ecdsa_sign_der( ) raises -> List[UInt8]: from .ecdsa_der import ecdsa_der_encode var raw = List[UInt8](unsafe_uninit_length=P256_SIGNATURE_SIZE) - if not p256_ecdsa_sign(private_key, message, raw.unsafe_ptr()): + if not p256_ecdsa_sign( + private_key, message, Span[mut=True, UInt8, ...](raw) + ): raise Error("P-256 ECDSA signing failed") return ecdsa_der_encode(Span[UInt8, ...](raw), P256_SIZE) @@ -1139,7 +1153,10 @@ def p256_keygen() raises -> Tuple[List[UInt8], List[UInt8]]: var d = _from_be(Span[UInt8, ...](private_key)) if not d.is_zero() and _cmp(d, _n()) < 0: var public_key = List[UInt8](unsafe_uninit_length=65) - if p256_public_key(Span[UInt8, ...](private_key), public_key.unsafe_ptr()): + if p256_public_key( + Span[UInt8, ...](private_key), + Span[mut=True, UInt8, ...](public_key), + ): _wipe_u256(d) return (private_key^, public_key^) _wipe_u256(d) diff --git a/src/thistle/p384.mojo b/src/thistle/p384.mojo index a99c710..22fe48d 100644 --- a/src/thistle/p384.mojo +++ b/src/thistle/p384.mojo @@ -6,6 +6,7 @@ from .p384_table import p384_base_table from .utils import u64_nonzero_choice, u64_zero_choice from .sha2 import sha384_hash from .pbkdf2 import hmac_sha384 +from std.utils import StaticTuple comptime P384_SIZE = 48 comptime P384_POINT_SIZE = 97 @@ -15,10 +16,12 @@ comptime _N0 = UInt64(0x0000000100000001) # -p^-1 mod 2^64 struct U384(Copyable, ImplicitlyCopyable, Movable): - var limbs: InlineArray[UInt64, 6] + var limbs: StaticTuple[UInt64, 6] def __init__(out self): - self.limbs = InlineArray[UInt64, 6](fill=0) + self.limbs = StaticTuple[UInt64, 6]() + comptime for i in range(6): + self.limbs[i] = 0 def __init__( out self, @@ -29,7 +32,7 @@ struct U384(Copyable, ImplicitlyCopyable, Movable): l4: UInt64, l5: UInt64, ): - self.limbs = InlineArray[UInt64, 6](uninitialized=True) + self.limbs = StaticTuple[UInt64, 6]() self.limbs[0] = l0 self.limbs[1] = l1 self.limbs[2] = l2 @@ -396,7 +399,7 @@ def _from_be(bytes: Span[UInt8, ...]) -> U384: return out -def _to_be(x: U384, output: UnsafePointer[UInt8, MutAnyOrigin]): +def _to_be(x: U384, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): for i in range(6): var limb = x.limbs[5 - i] for k in range(8): @@ -694,7 +697,7 @@ def _scalar_mult_base(k: U384) -> P384Point: def p384_decode_uncompressed(point: Span[UInt8, ...]) -> P384Point: if len(point) == 49 and (point[0] == 0x02 or point[0] == 0x03): var x = _from_be( - Span[UInt8, ...](ptr=point.unsafe_ptr() + 1, length=48) + Span[UInt8, ...](unsafe_ptr=point.unsafe_ptr() + 1, length=48) ) if _cmp(x, _p()) >= 0: return P384Point() @@ -713,8 +716,8 @@ def p384_decode_uncompressed(point: Span[UInt8, ...]) -> P384Point: return p if len(point) != P384_POINT_SIZE or point[0] != 0x04: return P384Point() - var x = _from_be(Span[UInt8, ...](ptr=point.unsafe_ptr() + 1, length=48)) - var y = _from_be(Span[UInt8, ...](ptr=point.unsafe_ptr() + 49, length=48)) + var x = _from_be(Span[UInt8, ...](unsafe_ptr=point.unsafe_ptr() + 1, length=48)) + var y = _from_be(Span[UInt8, ...](unsafe_ptr=point.unsafe_ptr() + 49, length=48)) var p = P384Point(x, y, False) if not _is_on_curve(p): return P384Point() @@ -722,31 +725,32 @@ def p384_decode_uncompressed(point: Span[UInt8, ...]) -> P384Point: def p384_encode_uncompressed( - point: P384Point, output: UnsafePointer[UInt8, MutAnyOrigin] + point: P384Point, output: Span[mut=True, UInt8, ...] ) -> Bool: - if point.infinity or not _is_on_curve(point): + if len(output) < P384_POINT_SIZE or point.infinity or not _is_on_curve(point): return False - output[0] = 0x04 - _to_be(point.x, output + 1) - _to_be(point.y, output + 49) + var out_ptr = output.unsafe_ptr() + out_ptr[0] = 0x04 + _to_be(point.x, out_ptr + 1) + _to_be(point.y, out_ptr + 49) return True @no_inline def p384_public_key( - private_key: Span[UInt8, ...], output: UnsafePointer[UInt8, MutAnyOrigin] + private_key: Span[UInt8, ...], output: Span[mut=True, UInt8, ...] ) -> Bool: - if len(private_key) != 48: + if len(private_key) != 48 or len(output) < P384_POINT_SIZE: return False var d = _from_be(private_key) if d.is_zero() or _cmp(d, _n()) >= 0: - var dp = d.limbs.unsafe_ptr() + var dp = UnsafePointer(to=d.limbs[0]).unsafe_mut_cast[True]() for i in range(6): dp.store[volatile=True](i, UInt64(0)) return False var q = _scalar_mult_base(d) var ok = p384_encode_uncompressed(q, output) - var dp = d.limbs.unsafe_ptr() + var dp = UnsafePointer(to=d.limbs[0]).unsafe_mut_cast[True]() for i in range(6): dp.store[volatile=True](i, UInt64(0)) return ok @@ -756,29 +760,29 @@ def p384_public_key( def p384_ecdh( private_key: Span[UInt8, ...], public_key: Span[UInt8, ...], - output: UnsafePointer[UInt8, MutAnyOrigin], + output: Span[mut=True, UInt8, ...], ) -> Bool: - if len(private_key) != 48: + if len(private_key) != 48 or len(output) < P384_SIZE: return False var d = _from_be(private_key) if d.is_zero() or _cmp(d, _n()) >= 0: - var dp = d.limbs.unsafe_ptr() + var dp = UnsafePointer(to=d.limbs[0]).unsafe_mut_cast[True]() for i in range(6): dp.store[volatile=True](i, UInt64(0)) return False var q = p384_decode_uncompressed(public_key) if q.infinity: - var dp = d.limbs.unsafe_ptr() + var dp = UnsafePointer(to=d.limbs[0]).unsafe_mut_cast[True]() for i in range(6): dp.store[volatile=True](i, UInt64(0)) return False var shared = _scalar_mult(d, q) - var dp = d.limbs.unsafe_ptr() + var dp = UnsafePointer(to=d.limbs[0]).unsafe_mut_cast[True]() for i in range(6): dp.store[volatile=True](i, UInt64(0)) if shared.infinity: return False - _to_be(shared.x, output) + _to_be(shared.x, output.unsafe_ptr()) return True @@ -875,7 +879,7 @@ def _reduce_n(x: U384) -> U384: def _wipe_u384(mut x: U384): - var ptr = x.limbs.unsafe_ptr() + var ptr = UnsafePointer(to=x.limbs[0]).unsafe_mut_cast[True]() for i in range(6): ptr.store[volatile=True](i, UInt64(0)) @@ -958,10 +962,14 @@ def _rfc6979_p384(private_key: Span[UInt8, ...], digest: Span[UInt8, ...], skip: def p384_ecdsa_sign_digest( private_key: Span[UInt8, ...], digest: Span[UInt8, ...], - signature: UnsafePointer[UInt8, MutAnyOrigin], + signature: Span[mut=True, UInt8, ...], ) -> Bool: - if len(private_key) != 48 or len(digest) != 48: + if ( + len(private_key) != 48 or len(digest) != 48 + or len(signature) < P384_SIGNATURE_SIZE + ): return False + var signature_ptr = signature.unsafe_ptr() var d = _from_be(private_key) if d.is_zero() or _cmp(d, _n()) >= 0: _wipe_u384(d) @@ -987,8 +995,8 @@ def p384_ecdsa_sign_digest( _wipe_u384(total) retry += 1 continue - _to_be(r, signature) - _to_be(s, signature + 48) + _to_be(r, signature_ptr) + _to_be(s, signature_ptr + 48) _wipe_u384(d) _wipe_u384(z) _wipe_u384(k) @@ -1001,10 +1009,14 @@ def p384_ecdsa_sign_digest( def p384_ecdsa_sign( private_key: Span[UInt8, ...], message: Span[UInt8, ...], - signature: UnsafePointer[UInt8, MutAnyOrigin], + signature: Span[mut=True, UInt8, ...], ) -> Bool: + if len(signature) < P384_SIGNATURE_SIZE: + return False var digest = sha384_hash(message) - var ok = p384_ecdsa_sign_digest(private_key, Span[UInt8, ...](digest), signature) + var ok = p384_ecdsa_sign_digest( + private_key, Span[UInt8, ...](digest), signature + ) _wipe_list_u8(digest) return ok @@ -1066,7 +1078,9 @@ def p384_ecdsa_sign_der( ) raises -> List[UInt8]: from .ecdsa_der import ecdsa_der_encode var raw = List[UInt8](unsafe_uninit_length=P384_SIGNATURE_SIZE) - if not p384_ecdsa_sign(private_key, message, raw.unsafe_ptr()): + if not p384_ecdsa_sign( + private_key, message, Span[mut=True, UInt8, ...](raw) + ): raise Error("P-384 ECDSA signing failed") return ecdsa_der_encode(Span[UInt8, ...](raw), P384_SIZE) @@ -1089,7 +1103,10 @@ def p384_keygen() raises -> Tuple[List[UInt8], List[UInt8]]: var d = _from_be(Span[UInt8, ...](private_key)) if not d.is_zero() and _cmp(d, _n()) < 0: var public_key = List[UInt8](unsafe_uninit_length=97) - if p384_public_key(Span[UInt8, ...](private_key), public_key.unsafe_ptr()): + if p384_public_key( + Span[UInt8, ...](private_key), + Span[mut=True, UInt8, ...](public_key), + ): _wipe_u384(d) return (private_key^, public_key^) _wipe_u384(d) diff --git a/src/thistle/pbkdf2.mojo b/src/thistle/pbkdf2.mojo index 713c103..f76a3b9 100644 --- a/src/thistle/pbkdf2.mojo +++ b/src/thistle/pbkdf2.mojo @@ -18,13 +18,13 @@ from .sha2 import ( ) @always_inline -def _xor_block_32(dst: UnsafePointer[UInt8, MutAnyOrigin], src: UnsafePointer[UInt8, MutAnyOrigin]): +def _xor_block_32(dst: UnsafePointer[mut=True, UInt8, _, address_space=_], src: UnsafePointer[mut=True, UInt8, _, address_space=_]): var d = dst.bitcast[UInt64]().load[width=4, alignment=1]() var s = src.bitcast[UInt64]().load[width=4, alignment=1]() dst.bitcast[UInt64]().store[width=4, alignment=1](0, d ^ s) @always_inline -def _xor_block_64(dst: UnsafePointer[UInt8, MutAnyOrigin], src: UnsafePointer[UInt8, MutAnyOrigin]): +def _xor_block_64(dst: UnsafePointer[mut=True, UInt8, _, address_space=_], src: UnsafePointer[mut=True, UInt8, _, address_space=_]): var d = dst.bitcast[UInt64]().load[width=8, alignment=1]() var s = src.bitcast[UInt64]().load[width=8, alignment=1]() dst.bitcast[UInt64]().store[width=8, alignment=1](0, d ^ s) @@ -54,7 +54,8 @@ struct PBKDF2SHA256(Movable): sha256_update(ctx, password) sha256_final_to_buffer(ctx, k.ptr()) else: - memcpy(dest=k.ptr(), src=password.unsafe_ptr(), count=len(password)) + for i in range(len(password)): + k[i] = password[i] for i in range(64): self.ipad[i] = k[i] ^ 0x36 @@ -71,13 +72,13 @@ struct PBKDF2SHA256(Movable): @always_inline def hmac(mut self, data: Span[UInt8, ...]): self.inner_ctx.reset() - sha256_update(self.inner_ctx, Span[UInt8, ...](ptr=self.ipad.ptr(), length=64)) + sha256_update(self.inner_ctx, Span[UInt8, ...](unsafe_ptr=self.ipad.ptr(), length=64)) sha256_update(self.inner_ctx, data) sha256_final_to_buffer(self.inner_ctx, self.inner_hash.ptr()) self.outer_ctx.reset() - sha256_update(self.outer_ctx, Span[UInt8, ...](ptr=self.opad.ptr(), length=64)) - sha256_update(self.outer_ctx, Span[UInt8, ...](ptr=self.inner_hash.ptr(), length=32)) + sha256_update(self.outer_ctx, Span[UInt8, ...](unsafe_ptr=self.opad.ptr(), length=64)) + sha256_update(self.outer_ctx, Span[UInt8, ...](unsafe_ptr=self.inner_hash.ptr(), length=32)) sha256_final_to_buffer(self.outer_ctx, self.u_block.ptr()) @always_inline @@ -88,14 +89,14 @@ struct PBKDF2SHA256(Movable): self.counter_bytes[3] = UInt8(counter & 0xFF) self.inner_ctx.reset() - sha256_update(self.inner_ctx, Span[UInt8, ...](ptr=self.ipad.ptr(), length=64)) + sha256_update(self.inner_ctx, Span[UInt8, ...](unsafe_ptr=self.ipad.ptr(), length=64)) sha256_update(self.inner_ctx, data) - sha256_update(self.inner_ctx, Span[UInt8, ...](ptr=self.counter_bytes.ptr(), length=4)) + sha256_update(self.inner_ctx, Span[UInt8, ...](unsafe_ptr=self.counter_bytes.ptr(), length=4)) sha256_final_to_buffer(self.inner_ctx, self.inner_hash.ptr()) self.outer_ctx.reset() - sha256_update(self.outer_ctx, Span[UInt8, ...](ptr=self.opad.ptr(), length=64)) - sha256_update(self.outer_ctx, Span[UInt8, ...](ptr=self.inner_hash.ptr(), length=32)) + sha256_update(self.outer_ctx, Span[UInt8, ...](unsafe_ptr=self.opad.ptr(), length=64)) + sha256_update(self.outer_ctx, Span[UInt8, ...](unsafe_ptr=self.inner_hash.ptr(), length=32)) sha256_final_to_buffer(self.outer_ctx, self.u_block.ptr()) @always_inline @@ -105,13 +106,21 @@ struct PBKDF2SHA256(Movable): var derived_key = List[UInt8](capacity=dklen) var t_block = StackBuffer[UInt8, 32](fill=0) + var input_block = StackBuffer[UInt8, 32](fill=0) for block_idx in range(1, num_blocks + 1): self.hmac_with_counter(salt, UInt32(block_idx)) memcpy(dest=t_block.ptr(), src=self.u_block.ptr(), count=32) for _ in range(1, iterations): - self.hmac(Span[UInt8, ...](ptr=self.u_block.ptr(), length=32)) + memcpy( + dest=input_block.ptr(), + src=self.u_block.ptr(), + count=32, + ) + self.hmac( + Span[UInt8, ...](unsafe_ptr=input_block.ptr(), length=32) + ) _xor_block_32(t_block.ptr(), self.u_block.ptr()) var remaining = dklen - len(derived_key) @@ -156,7 +165,8 @@ struct PBKDF2SHA512(Movable): sha512_update(ctx, password) sha512_final_to_buffer(ctx, k.ptr()) else: - memcpy(dest=k.ptr(), src=password.unsafe_ptr(), count=len(password)) + for i in range(len(password)): + k[i] = password[i] for i in range(128): self.ipad[i] = k[i] ^ 0x36 @@ -173,13 +183,13 @@ struct PBKDF2SHA512(Movable): @always_inline def hmac(mut self, data: Span[UInt8, ...]): self.inner_ctx.reset() - sha512_update(self.inner_ctx, Span[UInt8, ...](ptr=self.ipad.ptr(), length=128)) + sha512_update(self.inner_ctx, Span[UInt8, ...](unsafe_ptr=self.ipad.ptr(), length=128)) sha512_update(self.inner_ctx, data) sha512_final_to_buffer(self.inner_ctx, self.inner_hash.ptr()) self.outer_ctx.reset() - sha512_update(self.outer_ctx, Span[UInt8, ...](ptr=self.opad.ptr(), length=128)) - sha512_update(self.outer_ctx, Span[UInt8, ...](ptr=self.inner_hash.ptr(), length=64)) + sha512_update(self.outer_ctx, Span[UInt8, ...](unsafe_ptr=self.opad.ptr(), length=128)) + sha512_update(self.outer_ctx, Span[UInt8, ...](unsafe_ptr=self.inner_hash.ptr(), length=64)) sha512_final_to_buffer(self.outer_ctx, self.u_block.ptr()) @always_inline @@ -190,14 +200,14 @@ struct PBKDF2SHA512(Movable): self.counter_bytes[3] = UInt8(counter & 0xFF) self.inner_ctx.reset() - sha512_update(self.inner_ctx, Span[UInt8, ...](ptr=self.ipad.ptr(), length=128)) + sha512_update(self.inner_ctx, Span[UInt8, ...](unsafe_ptr=self.ipad.ptr(), length=128)) sha512_update(self.inner_ctx, data) - sha512_update(self.inner_ctx, Span[UInt8, ...](ptr=self.counter_bytes.ptr(), length=4)) + sha512_update(self.inner_ctx, Span[UInt8, ...](unsafe_ptr=self.counter_bytes.ptr(), length=4)) sha512_final_to_buffer(self.inner_ctx, self.inner_hash.ptr()) self.outer_ctx.reset() - sha512_update(self.outer_ctx, Span[UInt8, ...](ptr=self.opad.ptr(), length=128)) - sha512_update(self.outer_ctx, Span[UInt8, ...](ptr=self.inner_hash.ptr(), length=64)) + sha512_update(self.outer_ctx, Span[UInt8, ...](unsafe_ptr=self.opad.ptr(), length=128)) + sha512_update(self.outer_ctx, Span[UInt8, ...](unsafe_ptr=self.inner_hash.ptr(), length=64)) sha512_final_to_buffer(self.outer_ctx, self.u_block.ptr()) @always_inline @@ -207,13 +217,21 @@ struct PBKDF2SHA512(Movable): var derived_key = List[UInt8](capacity=dklen) var t_block = StackBuffer[UInt8, 64](fill=0) + var input_block = StackBuffer[UInt8, 64](fill=0) for block_idx in range(1, num_blocks + 1): self.hmac_with_counter(salt, UInt32(block_idx)) memcpy(dest=t_block.ptr(), src=self.u_block.ptr(), count=64) for _ in range(1, iterations): - self.hmac(Span[UInt8, ...](ptr=self.u_block.ptr(), length=64)) + memcpy( + dest=input_block.ptr(), + src=self.u_block.ptr(), + count=64, + ) + self.hmac( + Span[UInt8, ...](unsafe_ptr=input_block.ptr(), length=64) + ) _xor_block_64(t_block.ptr(), self.u_block.ptr()) var remaining = dklen - len(derived_key) @@ -259,7 +277,8 @@ def hmac_sha384(key: Span[UInt8, ...], data: Span[UInt8, ...]) -> List[UInt8]: for i in range(48): khp.store[volatile=True](i, UInt8(0)) else: - memcpy(dest=k.ptr(), src=key.unsafe_ptr(), count=len(key)) + for i in range(len(key)): + k[i] = key[i] var ipad = StackBuffer[UInt8, 128](fill=0) var opad = StackBuffer[UInt8, 128](fill=0) @@ -268,12 +287,12 @@ def hmac_sha384(key: Span[UInt8, ...], data: Span[UInt8, ...]) -> List[UInt8]: opad[i] = k[i] ^ 0x5C var inner = SHA512Context(SHA384_IV) - sha512_update(inner, Span[UInt8, ...](ptr=ipad.ptr(), length=128)) + sha512_update(inner, Span[UInt8, ...](unsafe_ptr=ipad.ptr(), length=128)) sha512_update(inner, data) var inner_hash = sha512_final_with_len(inner, 48) var outer = SHA512Context(SHA384_IV) - sha512_update(outer, Span[UInt8, ...](ptr=opad.ptr(), length=128)) + sha512_update(outer, Span[UInt8, ...](unsafe_ptr=opad.ptr(), length=128)) sha512_update(outer, Span[UInt8, ...](inner_hash)) var result = sha512_final_with_len(outer, 48) diff --git a/src/thistle/poly1305.mojo b/src/thistle/poly1305.mojo index 900655a..cd206cf 100644 --- a/src/thistle/poly1305.mojo +++ b/src/thistle/poly1305.mojo @@ -10,7 +10,7 @@ comptime _M42: UInt64 = 0x3FFFFFFFFFF @always_inline -def _le64(ptr: UnsafePointer[UInt8, _], offset: Int) -> UInt64: +def _le64(ptr: UnsafePointer[mut=False, UInt8, _, address_space=_], offset: Int) -> UInt64: return (ptr + offset).bitcast[UInt64]().load[width=1, alignment=1]() @@ -67,7 +67,7 @@ def _reduce(mut h0: UInt64, mut h1: UInt64, mut h2: UInt64, d0: UInt128, d1: UIn @always_inline -def _limbs_at(ptr: UnsafePointer[UInt8, _], offset: Int, hibit: UInt64) -> SIMD[DType.uint64, 4]: +def _limbs_at(ptr: UnsafePointer[mut=False, UInt8, _, address_space=_], offset: Int, hibit: UInt64) -> SIMD[DType.uint64, 4]: var t0 = _le64(ptr, offset) var t1 = _le64(ptr, offset + 8) return SIMD[DType.uint64, 4]( @@ -175,7 +175,7 @@ struct Poly1305: _reduce(self.h0, self.h1, self.h2, d0, d1, d2) @no_inline - def _blocks8(mut self, ptr: UnsafePointer[UInt8, _], count8: Int): + def _blocks8(mut self, ptr: UnsafePointer[mut=False, UInt8, _, address_space=_], count8: Int): var h0 = self.h0 var h1 = self.h1 var h2 = self.h2 @@ -207,7 +207,7 @@ struct Poly1305: self.h2 = h2 @no_inline - def _blocks4(mut self, ptr: UnsafePointer[UInt8, _], count4: Int): + def _blocks4(mut self, ptr: UnsafePointer[mut=False, UInt8, _, address_space=_], count4: Int): var h0 = self.h0 var h1 = self.h1 var h2 = self.h2 @@ -270,7 +270,7 @@ struct Poly1305: self.buf_len += 1 i += 1 - def finalize_into(mut self, output: UnsafePointer[UInt8, MutAnyOrigin]): + def finalize_into(mut self, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): if self.buf_len > 0: self.buf[self.buf_len] = 1 for j in range(self.buf_len + 1, 16): @@ -337,7 +337,13 @@ struct Poly1305: UnsafePointer(to=self.pad1).store[volatile=True](0, UInt64(0)) -def poly1305_mac(key: Span[UInt8, ...], message: Span[UInt8, ...], output: UnsafePointer[UInt8, MutAnyOrigin]) raises: +def poly1305_mac( + key: Span[UInt8, ...], + message: Span[UInt8, ...], + output: Span[mut=True, UInt8, ...], +) raises: + if len(output) < 16: + raise Error("Poly1305 output needs at least 16 writable bytes") var p = Poly1305(key) p.update(message) - p.finalize_into(output) + p.finalize_into(output.unsafe_ptr()) diff --git a/src/thistle/random.mojo b/src/thistle/random.mojo index 2a41adc..4707e17 100644 --- a/src/thistle/random.mojo +++ b/src/thistle/random.mojo @@ -4,7 +4,7 @@ from std.sys import CompilationTarget, inlined_assembly @always_inline -def _getrandom_linux_x86(buf: UnsafePointer[UInt8, MutAnyOrigin], length: Int) -> Int: +def _getrandom_linux_x86(buf: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int) -> Int: # Linux x86-64: getrandom(buf, len, flags=0), syscall 318. # rax is both syscall-number input and return-value output. return Int( @@ -18,7 +18,7 @@ def _getrandom_linux_x86(buf: UnsafePointer[UInt8, MutAnyOrigin], length: Int) - @always_inline -def _getrandom_linux_arm(buf: UnsafePointer[UInt8, MutAnyOrigin], length: Int) -> Int: +def _getrandom_linux_arm(buf: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int) -> Int: # Linux aarch64: getrandom(buf, len, flags=0), syscall 278. # inputs go via scratch registers to avoid asm constraint conflicts return Int( @@ -38,7 +38,7 @@ def _getrandom_linux_arm(buf: UnsafePointer[UInt8, MutAnyOrigin], length: Int) - @always_inline -def _getentropy_macos_arm(buf: UnsafePointer[UInt8, MutAnyOrigin], length: Int) -> Int: +def _getentropy_macos_arm(buf: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int) -> Int: # Darwin arm64: getentropy(buf, size), syscall 500. # x16 = syscall number, x0 = buffer, x1 = size. # Errors set carry and return errno in x0; convert that to -errno. @@ -61,7 +61,7 @@ def _getentropy_macos_arm(buf: UnsafePointer[UInt8, MutAnyOrigin], length: Int) ) -def _fill_linux_x86(buf: UnsafePointer[UInt8, MutAnyOrigin], length: Int) raises: +def _fill_linux_x86(buf: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int) raises: var offset = 0 while offset < length: var ret = _getrandom_linux_x86(buf + offset, length - offset) @@ -74,7 +74,7 @@ def _fill_linux_x86(buf: UnsafePointer[UInt8, MutAnyOrigin], length: Int) raises offset += ret -def _fill_linux_arm(buf: UnsafePointer[UInt8, MutAnyOrigin], length: Int) raises: +def _fill_linux_arm(buf: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int) raises: var offset = 0 while offset < length: var ret = _getrandom_linux_arm(buf + offset, length - offset) @@ -87,7 +87,7 @@ def _fill_linux_arm(buf: UnsafePointer[UInt8, MutAnyOrigin], length: Int) raises offset += ret -def _fill_macos_arm(buf: UnsafePointer[UInt8, MutAnyOrigin], length: Int) raises: +def _fill_macos_arm(buf: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int) raises: var offset = 0 while offset < length: var chunk = min(256, length - offset) @@ -99,20 +99,20 @@ def _fill_macos_arm(buf: UnsafePointer[UInt8, MutAnyOrigin], length: Int) raises offset += chunk -def random_fill(buf: UnsafePointer[UInt8, MutAnyOrigin], length: Int) raises: - if length < 0: - raise Error("random_fill length must be non-negative") +def random_fill(buf: Span[mut=True, UInt8, ...]) raises: + var length = len(buf) if length == 0: return + var ptr = buf.unsafe_ptr() comptime if CompilationTarget.is_linux() and CompilationTarget.is_x86(): - _fill_linux_x86(buf, length) + _fill_linux_x86(ptr, length) elif CompilationTarget.is_linux() and CompilationTarget.has_neon(): - _fill_linux_arm(buf, length) + _fill_linux_arm(ptr, length) elif CompilationTarget.is_macos() and CompilationTarget.is_apple_silicon(): - _fill_macos_arm(buf, length) + _fill_macos_arm(ptr, length) else: CompilationTarget.unsupported_target_error[operation="random_fill"]() @@ -122,5 +122,5 @@ def random_bytes(n: Int) raises -> List[UInt8]: if n < 0: raise Error("random_bytes length must be non-negative") var result = List[UInt8](length=n, fill=0) - random_fill(result.unsafe_ptr(), n) + random_fill(Span[mut=True, UInt8, ...](result)) return result^ diff --git a/src/thistle/rsa.mojo b/src/thistle/rsa.mojo index b4f70e2..2b9d136 100644 --- a/src/thistle/rsa.mojo +++ b/src/thistle/rsa.mojo @@ -5,6 +5,7 @@ RSASSA-PSS signature verification per RFC 8017 from std.collections import List, InlineArray from std.memory import UnsafePointer from std.bit import rotate_bits_left, count_leading_zeros +from std.utils import StaticTuple from .random import random_bytes from .sha2 import ( sha224_hash, sha256_hash, sha384_hash, sha512_hash, @@ -20,6 +21,14 @@ comptime SHA384: Int = 384 comptime SHA512: Int = 512 +@always_inline +def _bn_zero() -> StaticTuple[UInt64, _NL]: + var out = StaticTuple[UInt64, _NL]() + comptime for i in range(_NL): + out[i] = 0 + return out + + def _sha1(data: Span[UInt8, ...]) -> InlineArray[UInt8, 20]: var h0: UInt32 = 0x67452301 var h1: UInt32 = 0xEFCDAB89 @@ -95,7 +104,7 @@ def _sha1(data: Span[UInt8, ...]) -> InlineArray[UInt8, 20]: out[4 * i + 1] = UInt8((hs[i] >> 16) & 0xFF) out[4 * i + 2] = UInt8((hs[i] >> 8) & 0xFF) out[4 * i + 3] = UInt8(hs[i] & 0xFF) - return out + return out^ def _hash_len(alg: Int) raises -> Int: @@ -112,7 +121,7 @@ def _hash_len(alg: Int) raises -> Int: raise Error("unsupported hash for RSA-PSS") -def _hash_into(alg: Int, data: Span[UInt8, ...], output: UnsafePointer[UInt8, MutAnyOrigin]) raises -> Int: +def _hash_into(alg: Int, data: Span[UInt8, ...], output: UnsafePointer[mut=True, UInt8, _, address_space=_]) raises -> Int: if alg == SHA256: var ctx = SHA256Context() sha256_update(ctx, data) @@ -149,7 +158,7 @@ def _mgf1( seed: UnsafePointer[UInt8, _], seed_len: Int, mask_len: Int, - output: UnsafePointer[UInt8, MutAnyOrigin], + output: UnsafePointer[mut=True, UInt8, _, address_space=_], ) raises: var h_len = _hash_len(alg) var block = InlineArray[UInt8, 128](fill=0) @@ -165,7 +174,7 @@ def _mgf1( block[seed_len + 3] = UInt8(counter & 0xFF) _ = _hash_into( alg, - Span[UInt8, ...](ptr=block.unsafe_ptr(), length=seed_len + 4), + Span[UInt8, ...](unsafe_ptr=block.unsafe_ptr(), length=seed_len + 4), digest.unsafe_ptr(), ) var take = mask_len - done @@ -183,7 +192,7 @@ def _emsa_pss_encode( sha: Int, mgf_sha: Int, em_bits: Int, - output: UnsafePointer[UInt8, MutAnyOrigin], + output: UnsafePointer[mut=True, UInt8, _, address_space=_], ) raises -> Bool: var h_len = _hash_len(sha) _ = _hash_len(mgf_sha) @@ -201,7 +210,7 @@ def _emsa_pss_encode( var h = InlineArray[UInt8, 64](uninitialized=True) _ = _hash_into( sha, - Span[UInt8, ...](ptr=mprime.unsafe_ptr(), length=8 + h_len + len(salt)), + Span[UInt8, ...](unsafe_ptr=mprime.unsafe_ptr(), length=8 + h_len + len(salt)), h.unsafe_ptr(), ) @@ -235,7 +244,7 @@ def _digest_info_prefix_len(alg: Int) raises -> Int: def _digest_info_prefix( - alg: Int, output: UnsafePointer[UInt8, MutAnyOrigin] + alg: Int, output: UnsafePointer[mut=True, UInt8, _, address_space=_] ) raises: var hex: String if alg == SHA1: @@ -260,7 +269,7 @@ def _digest_info_prefix( @always_inline -def _bn_ge(a: InlineArray[UInt64, _NL], b: InlineArray[UInt64, _NL], k: Int) -> Bool: +def _bn_ge(a: StaticTuple[UInt64, _NL], b: StaticTuple[UInt64, _NL], k: Int) -> Bool: for i in range(k - 1, -1, -1): if a[i] > b[i]: return True @@ -271,8 +280,8 @@ def _bn_ge(a: InlineArray[UInt64, _NL], b: InlineArray[UInt64, _NL], k: Int) -> @always_inline def _bn_ge_ct( - a: InlineArray[UInt64, _NL], - b: InlineArray[UInt64, _NL], + a: StaticTuple[UInt64, _NL], + b: StaticTuple[UInt64, _NL], k: Int, ) -> Bool: var borrow: UInt64 = 0 @@ -288,7 +297,7 @@ def _bn_ge_ct( @always_inline -def _bn_sub(mut a: InlineArray[UInt64, _NL], b: InlineArray[UInt64, _NL], k: Int): +def _bn_sub(mut a: StaticTuple[UInt64, _NL], b: StaticTuple[UInt64, _NL], k: Int): var borrow: UInt64 = 0 for i in range(k): var d = (UInt128(1) << 64) + UInt128(a[i]) - UInt128(b[i]) - UInt128(borrow) @@ -298,8 +307,8 @@ def _bn_sub(mut a: InlineArray[UInt64, _NL], b: InlineArray[UInt64, _NL], k: Int @always_inline def _bn_sub_copy( - a: InlineArray[UInt64, _NL], b: InlineArray[UInt64, _NL], k: Int -) -> Tuple[InlineArray[UInt64, _NL], UInt64]: + a: StaticTuple[UInt64, _NL], b: StaticTuple[UInt64, _NL], k: Int +) -> Tuple[StaticTuple[UInt64, _NL], UInt64]: var out = a var borrow: UInt64 = 0 for i in range(k): @@ -311,11 +320,11 @@ def _bn_sub_copy( @always_inline def _bn_select( - a: InlineArray[UInt64, _NL], - b: InlineArray[UInt64, _NL], + a: StaticTuple[UInt64, _NL], + b: StaticTuple[UInt64, _NL], choice: UInt64, k: Int, -) -> InlineArray[UInt64, _NL]: +) -> StaticTuple[UInt64, _NL]: var out = a var mask = UInt64(0) - (choice & UInt64(1)) for i in range(k): @@ -329,7 +338,7 @@ def _nonzero_choice(x: UInt64) -> UInt64: @always_inline -def _bn_dbl_mod(mut a: InlineArray[UInt64, _NL], n: InlineArray[UInt64, _NL], k: Int): +def _bn_dbl_mod(mut a: StaticTuple[UInt64, _NL], n: StaticTuple[UInt64, _NL], k: Int): var carry: UInt64 = 0 for i in range(k): var v = a[i] @@ -343,12 +352,12 @@ def _bn_dbl_mod(mut a: InlineArray[UInt64, _NL], n: InlineArray[UInt64, _NL], k: @always_inline def _mont_mul_k[K: Int]( - a: InlineArray[UInt64, _NL], - b: InlineArray[UInt64, _NL], - n: InlineArray[UInt64, _NL], + a: StaticTuple[UInt64, _NL], + b: StaticTuple[UInt64, _NL], + n: StaticTuple[UInt64, _NL], n0: UInt64, -) -> InlineArray[UInt64, _NL]: - var t = InlineArray[UInt64, _NL](uninitialized=True) +) -> StaticTuple[UInt64, _NL]: + var t = StaticTuple[UInt64, _NL]() comptime for z in range(K): t[z] = 0 var t_hi: UInt64 = 0 @@ -403,13 +412,13 @@ def _mont_mul_k[K: Int]( def _mont_mul_any( - a: InlineArray[UInt64, _NL], - b: InlineArray[UInt64, _NL], - n: InlineArray[UInt64, _NL], + a: StaticTuple[UInt64, _NL], + b: StaticTuple[UInt64, _NL], + n: StaticTuple[UInt64, _NL], n0: UInt64, k: Int, -) -> InlineArray[UInt64, _NL]: - var t = InlineArray[UInt64, _NL](fill=0) +) -> StaticTuple[UInt64, _NL]: + var t = _bn_zero() var t_hi: UInt64 = 0 for i in range(k): var ai = a[i] @@ -434,10 +443,10 @@ def _mont_mul_any( @always_inline def _mont_sqr_k[K: Int]( - a: InlineArray[UInt64, _NL], - n: InlineArray[UInt64, _NL], + a: StaticTuple[UInt64, _NL], + n: StaticTuple[UInt64, _NL], n0: UInt64, -) -> InlineArray[UInt64, _NL]: +) -> StaticTuple[UInt64, _NL]: var t = InlineArray[UInt64, 2 * _NL + 2](uninitialized=True) comptime for z in range(2 * K + 1): t[z] = 0 @@ -479,7 +488,7 @@ def _mont_sqr_k[K: Int]( t[i + K] = s2.cast[DType.uint64]() ehold = (s2 >> 64).cast[DType.uint64]() - var out = InlineArray[UInt64, _NL](uninitialized=True) + var out = StaticTuple[UInt64, _NL]() comptime for i in range(K): out[i] = t[K + i] var reduced, borrow = _bn_sub_copy(out, n, K) @@ -489,11 +498,11 @@ def _mont_sqr_k[K: Int]( @always_inline def _mont_sqr( - a: InlineArray[UInt64, _NL], - n: InlineArray[UInt64, _NL], + a: StaticTuple[UInt64, _NL], + n: StaticTuple[UInt64, _NL], n0: UInt64, k: Int, -) -> InlineArray[UInt64, _NL]: +) -> StaticTuple[UInt64, _NL]: if k == 16: return _mont_sqr_k[16](a, n, n0) if k == 32: @@ -507,12 +516,12 @@ def _mont_sqr( @always_inline def _mont_mul( - a: InlineArray[UInt64, _NL], - b: InlineArray[UInt64, _NL], - n: InlineArray[UInt64, _NL], + a: StaticTuple[UInt64, _NL], + b: StaticTuple[UInt64, _NL], + n: StaticTuple[UInt64, _NL], n0: UInt64, k: Int, -) -> InlineArray[UInt64, _NL]: +) -> StaticTuple[UInt64, _NL]: if k == 16: return _mont_mul_k[16](a, b, n, n0) if k == 32: @@ -525,13 +534,13 @@ def _mont_mul( struct RsaPublicKey: - var n: InlineArray[UInt64, _NL] + var n: StaticTuple[UInt64, _NL] var k: Int var nb: Int var mod_bits: Int var n0: UInt64 - var rmod: InlineArray[UInt64, _NL] - var r2: InlineArray[UInt64, _NL] + var rmod: StaticTuple[UInt64, _NL] + var r2: StaticTuple[UInt64, _NL] var e: List[UInt8] def __init__(out self, modulus: Span[UInt8, ...], exponent: Span[UInt8, ...]) raises: @@ -569,7 +578,7 @@ struct RsaPublicKey: self.k = (nb + 7) // 8 var k = self.k - self.n = InlineArray[UInt64, _NL](fill=0) + self.n = _bn_zero() for i in range(nb): var byte = UInt64(modulus[lead + nb - 1 - i]) self.n[i >> 3] |= byte << UInt64(8 * (i & 7)) @@ -589,7 +598,7 @@ struct RsaPublicKey: inv = inv * (2 - n0 * inv) self.n0 = UInt64(0) - inv - self.rmod = InlineArray[UInt64, _NL](fill=0) + self.rmod = _bn_zero() if (self.n[k - 1] >> 63) != 0: _bn_sub(self.rmod, self.n, k) else: @@ -612,13 +621,13 @@ struct RsaPublicKey: def _public_op( self, sig: Span[UInt8, ...], - output: UnsafePointer[UInt8, MutAnyOrigin], + output: UnsafePointer[mut=True, UInt8, _, address_space=_], ) raises -> Bool: var nb = self.nb var k = self.k if len(sig) != nb: return False - var s = InlineArray[UInt64, _NL](fill=0) + var s = _bn_zero() for i in range(nb): var byte = UInt64(sig[nb - 1 - i]) s[i >> 3] |= byte << UInt64(8 * (i & 7)) @@ -642,7 +651,7 @@ struct RsaPublicKey: if not started: return False - var one = InlineArray[UInt64, _NL](fill=0) + var one = _bn_zero() one[0] = 1 var m = _mont_mul(acc, one, self.n, self.n0, k) for i in range(nb): @@ -714,7 +723,7 @@ struct RsaPublicKey: var h2 = InlineArray[UInt8, 64](uninitialized=True) _ = _hash_into( sha, - Span[UInt8, ...](ptr=mprime.unsafe_ptr(), length=8 + h_len + salt_len), + Span[UInt8, ...](unsafe_ptr=mprime.unsafe_ptr(), length=8 + h_len + salt_len), h2.unsafe_ptr(), ) @@ -724,8 +733,8 @@ struct RsaPublicKey: return diff == 0 -def _wipe_bn(mut value: InlineArray[UInt64, _NL], k: Int): - var ptr = value.unsafe_ptr() +def _wipe_bn(mut value: StaticTuple[UInt64, _NL], k: Int): + var ptr = UnsafePointer(to=value[0]).unsafe_mut_cast[True]() for i in range(k): ptr.store[volatile=True](i, UInt64(0)) @@ -750,7 +759,7 @@ struct RsaPrivateKey: var d_len = len(private_exponent) - lead if d_len == 0 or d_len > self.public.nb: raise Error("invalid RSA private exponent") - var dbn = InlineArray[UInt64, _NL](fill=0) + var dbn = _bn_zero() for i in range(d_len): var byte = UInt64(private_exponent[lead + d_len - 1 - i]) dbn[i >> 3] |= byte << UInt64(8 * (i & 7)) @@ -769,13 +778,13 @@ struct RsaPrivateKey: def _private_op( self, encoded: Span[UInt8, ...], - signature: UnsafePointer[UInt8, MutAnyOrigin], + signature: UnsafePointer[mut=True, UInt8, _, address_space=_], ) raises -> Bool: var nb = self.public.nb var k = self.public.k if len(encoded) != nb: return False - var input = InlineArray[UInt64, _NL](fill=0) + var input = _bn_zero() for i in range(nb): var byte = UInt64(encoded[nb - 1 - i]) input[i >> 3] |= byte << UInt64(8 * (i & 7)) @@ -784,7 +793,7 @@ struct RsaPrivateKey: return False var base = _mont_mul(input, self.public.r2, self.public.n, self.public.n0, k) - var table = InlineArray[InlineArray[UInt64, _NL], 16](uninitialized=True) + var table = StaticTuple[StaticTuple[UInt64, _NL], 16]() table[0] = self.public.rmod table[1] = base for i in range(2, 16): @@ -805,7 +814,7 @@ struct RsaPrivateKey: selected = _bn_select(selected, table[i], hit, k) acc = _mont_mul(acc, selected, self.public.n, self.public.n0, k) - var one = InlineArray[UInt64, _NL](fill=0) + var one = _bn_zero() one[0] = 1 var result = _mont_mul(acc, one, self.public.n, self.public.n0, k) for i in range(nb): @@ -813,7 +822,7 @@ struct RsaPrivateKey: signature[i] = UInt8((limb >> UInt64(8 * ((nb - 1 - i) & 7))) & 0xFF) var recovered = InlineArray[UInt8, 528](uninitialized=True) - var sig_span = Span[UInt8, ...](ptr=signature, length=nb) + var sig_span = Span[UInt8, ...](unsafe_ptr=signature, length=nb) var valid = self.public._public_op(sig_span, recovered.unsafe_ptr()) var diff = UInt8(0) for i in range(nb): @@ -837,7 +846,7 @@ struct RsaPrivateKey: salt: Span[UInt8, ...], sha: Int, mgf_sha: Int, - signature: UnsafePointer[UInt8, MutAnyOrigin], + signature: UnsafePointer[mut=True, UInt8, _, address_space=_], ) raises -> Bool: var em_bits = self.public.mod_bits - 1 var em_len = (em_bits + 7) // 8 @@ -848,7 +857,7 @@ struct RsaPrivateKey: ): return False var ok = self._private_op( - Span[UInt8, ...](ptr=encoded.unsafe_ptr(), length=self.public.nb), + Span[UInt8, ...](unsafe_ptr=encoded.unsafe_ptr(), length=self.public.nb), signature, ) var ep = encoded.unsafe_ptr() @@ -885,8 +894,8 @@ struct RsaPrivateKey: def _bn_reduce_bytes( data: Span[UInt8, ...], key: RsaPublicKey -) -> InlineArray[UInt64, _NL]: - var value = InlineArray[UInt64, _NL](fill=0) +) -> StaticTuple[UInt64, _NL]: + var value = _bn_zero() for bi in range(len(data)): var byte = data[bi] for bit in range(7, -1, -1): @@ -904,10 +913,10 @@ def _bn_reduce_bytes( def _private_pow( key: RsaPublicKey, exponent: InlineArray[UInt8, 528], - input: InlineArray[UInt64, _NL], -) -> InlineArray[UInt64, _NL]: + input: StaticTuple[UInt64, _NL], +) -> StaticTuple[UInt64, _NL]: var base = _mont_mul(input, key.r2, key.n, key.n0, key.k) - var table = InlineArray[InlineArray[UInt64, _NL], 16](uninitialized=True) + var table = StaticTuple[StaticTuple[UInt64, _NL], 16]() table[0] = key.rmod table[1] = base for i in range(2, 16): @@ -926,7 +935,7 @@ def _private_pow( var hit = _nonzero_choice(digit ^ UInt64(i)) ^ UInt64(1) selected = _bn_select(selected, table[i], hit, key.k) acc = _mont_mul(acc, selected, key.n, key.n0, key.k) - var one = InlineArray[UInt64, _NL](fill=0) + var one = _bn_zero() one[0] = 1 var result = _mont_mul(acc, one, key.n, key.n0, key.k) _wipe_bn(base, key.k) @@ -937,10 +946,10 @@ def _private_pow( def _bn_mul_parts( - a: InlineArray[UInt64, _NL], a_len: Int, - b: InlineArray[UInt64, _NL], b_len: Int, -) -> InlineArray[UInt64, _NL]: - var out = InlineArray[UInt64, _NL](fill=0) + a: StaticTuple[UInt64, _NL], a_len: Int, + b: StaticTuple[UInt64, _NL], b_len: Int, +) -> StaticTuple[UInt64, _NL]: + var out = _bn_zero() for i in range(a_len): var carry = UInt64(0) for j in range(b_len): @@ -952,7 +961,7 @@ def _bn_mul_parts( def _bn_equal( - a: InlineArray[UInt64, _NL], b: InlineArray[UInt64, _NL], k: Int + a: StaticTuple[UInt64, _NL], b: StaticTuple[UInt64, _NL], k: Int ) -> Bool: var diff = UInt64(0) for i in range(k): @@ -966,7 +975,7 @@ struct RsaCrtPrivateKey: var q: RsaPublicKey var dp: InlineArray[UInt8, 528] var dq: InlineArray[UInt8, 528] - var qinv: InlineArray[UInt64, _NL] + var qinv: StaticTuple[UInt64, _NL] def __init__( out self, @@ -979,11 +988,11 @@ struct RsaCrtPrivateKey: if self.public.mod_bits < 2048: raise Error("RSA signing requires a modulus of at least 2048 bits") var three = InlineArray[UInt8, 1](fill=3) - self.p = RsaPublicKey(prime1, Span[UInt8, ...](ptr=three.unsafe_ptr(), length=1)) - self.q = RsaPublicKey(prime2, Span[UInt8, ...](ptr=three.unsafe_ptr(), length=1)) + self.p = RsaPublicKey(prime1, Span[UInt8, ...](unsafe_ptr=three.unsafe_ptr(), length=1)) + self.q = RsaPublicKey(prime2, Span[UInt8, ...](unsafe_ptr=three.unsafe_ptr(), length=1)) self.dp = InlineArray[UInt8, 528](fill=0) self.dq = InlineArray[UInt8, 528](fill=0) - self.qinv = InlineArray[UInt64, _NL](fill=0) + self.qinv = _bn_zero() if self.p.k + self.q.k > _NL: raise Error("RSA CRT factors are too large") @@ -1006,8 +1015,8 @@ struct RsaCrtPrivateKey: self.dp[self.p.nb - dp_len + i] = exponent1[dp_lead + i] for i in range(dq_len): self.dq[self.q.nb - dq_len + i] = exponent2[dq_lead + i] - var dp_bn = InlineArray[UInt64, _NL](fill=0) - var dq_bn = InlineArray[UInt64, _NL](fill=0) + var dp_bn = _bn_zero() + var dq_bn = _bn_zero() for i in range(self.p.nb): dp_bn[i >> 3] |= UInt64(self.dp[self.p.nb - 1 - i]) << UInt64(8 * (i & 7)) for i in range(self.q.nb): @@ -1038,7 +1047,7 @@ struct RsaCrtPrivateKey: var limb = self.q.n[(self.q.nb - 1 - i) >> 3] q_bytes[i] = UInt8((limb >> UInt64(8 * ((self.q.nb - 1 - i) & 7))) & 0xFF) var q_mod_p = _bn_reduce_bytes( - Span[UInt8, ...](ptr=q_bytes.unsafe_ptr(), length=self.q.nb), self.p + Span[UInt8, ...](unsafe_ptr=q_bytes.unsafe_ptr(), length=self.q.nb), self.p ) var check = _mont_mul( _mont_mul(q_mod_p, self.p.r2, self.p.n, self.p.n0, self.p.k), @@ -1074,7 +1083,7 @@ struct RsaCrtPrivateKey: def _private_op( self, encoded: Span[UInt8, ...], - signature: UnsafePointer[UInt8, MutAnyOrigin], + signature: UnsafePointer[mut=True, UInt8, _, address_space=_], ) raises -> Bool: if len(encoded) != self.public.nb: return False @@ -1088,7 +1097,7 @@ struct RsaCrtPrivateKey: var limb = m2[(self.q.nb - 1 - i) >> 3] m2_bytes[i] = UInt8((limb >> UInt64(8 * ((self.q.nb - 1 - i) & 7))) & 0xFF) var m2_mod_p = _bn_reduce_bytes( - Span[UInt8, ...](ptr=m2_bytes.unsafe_ptr(), length=self.q.nb), self.p + Span[UInt8, ...](unsafe_ptr=m2_bytes.unsafe_ptr(), length=self.q.nb), self.p ) var h, borrow = _bn_sub_copy(m1, m2_mod_p, self.p.k) var h_plus_p = h @@ -1118,7 +1127,7 @@ struct RsaCrtPrivateKey: var recovered = InlineArray[UInt8, 528](uninitialized=True) var valid = self.public._public_op( - Span[UInt8, ...](ptr=signature, length=self.public.nb), recovered.unsafe_ptr() + Span[UInt8, ...](unsafe_ptr=signature, length=self.public.nb), recovered.unsafe_ptr() ) var diff = UInt8(0) for i in range(self.public.nb): @@ -1144,7 +1153,7 @@ struct RsaCrtPrivateKey: def pss_sign_with_salt( self, message: Span[UInt8, ...], salt: Span[UInt8, ...], sha: Int, mgf_sha: Int, - signature: UnsafePointer[UInt8, MutAnyOrigin], + signature: UnsafePointer[mut=True, UInt8, _, address_space=_], ) raises -> Bool: var em_bits = self.public.mod_bits - 1 var em_len = (em_bits + 7) // 8 @@ -1155,7 +1164,7 @@ struct RsaCrtPrivateKey: ): return False var ok = self._private_op( - Span[UInt8, ...](ptr=encoded.unsafe_ptr(), length=self.public.nb), signature + Span[UInt8, ...](unsafe_ptr=encoded.unsafe_ptr(), length=self.public.nb), signature ) var ep = encoded.unsafe_ptr() for i in range(self.public.nb): diff --git a/src/thistle/sha2.mojo b/src/thistle/sha2.mojo index 380abb4..1d8c06d 100644 --- a/src/thistle/sha2.mojo +++ b/src/thistle/sha2.mojo @@ -9,6 +9,7 @@ from std.bit import rotate_bits_right, byte_swap from std.builtin.simd import SIMD from std.builtin.dtype import DType from std.sys import CompilationTarget +from std.utils import StaticTuple from .sha_ni import sha512ni_transform_blocks, sha256ni_transform_blocks from .utils import bytes_to_hex, string_to_bytes, load_32be, load_64be @@ -28,7 +29,7 @@ comptime SHA256_K = SIMD[DType.uint32, 64]( 0x90BEFFFA, 0xA4506CEB, 0xBEF9A3F7, 0xC67178F2, ) -comptime SHA512_K = SIMD[DType.uint64, 80]( +comptime SHA512_K = StaticTuple[UInt64, 80]( 0x428A2F98D728AE22, 0x7137449123EF65CD, 0xB5C0FBCFEC4D3B2F, 0xE9B5DBA58189DBBC, 0x3956C25BF348B538, 0x59F111F1B605D019, 0x923F82A4AF194F9B, 0xAB1C5ED5DA6D8118, 0xD807AA98A3030242, 0x12835B0145706FBE, 0x243185BE4EE4B28C, 0x550C7DC3D5FFB4E2, 0x72BE5D74F27B896F, 0x80DEB1FE3B1696B1, 0x9BDC06A725C71235, @@ -193,7 +194,7 @@ struct SHA256Context(Movable): @always_inline def sha256_transform_blocks( mut state: SIMD[DType.uint32, 8], - data: UnsafePointer[UInt8, ImmutAnyOrigin], + data: UnsafePointer[mut=False, UInt8, _, address_space=_], nblocks: Int, ): comptime if (CompilationTarget.has_neon() and CompilationTarget._has_feature["sha2"]() and not CompilationTarget.is_x86()) or (CompilationTarget.is_x86() and CompilationTarget._has_feature["sse"]() and CompilationTarget._has_feature["sha"]()): @@ -284,23 +285,17 @@ def sha256_update(mut ctx: SHA256Context, data: Span[UInt8, ...]): if ctx.buffer_len > 0: var available = 64 - ctx.buffer_len if total_len >= available: - memcpy( - dest=buf_ptr + ctx.buffer_len, - src=data.unsafe_ptr(), - count=available, - ) + for j in range(available): + buf_ptr[ctx.buffer_len + j] = data[j] ctx.state = sha256_transform( - ctx.state, Span[UInt8, ...](ptr=buf_ptr, length=64) + ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=64) ) ctx.count += 512 i += available ctx.buffer_len = 0 else: - memcpy( - dest=buf_ptr + ctx.buffer_len, - src=data.unsafe_ptr(), - count=total_len, - ) + for j in range(total_len): + buf_ptr[ctx.buffer_len + j] = data[j] ctx.buffer_len += total_len return @@ -312,11 +307,8 @@ def sha256_update(mut ctx: SHA256Context, data: Span[UInt8, ...]): if i < total_len: var remaining = total_len - i - memcpy( - dest=buf_ptr + ctx.buffer_len, - src=data.unsafe_ptr() + i, - count=remaining, - ) + for j in range(remaining): + buf_ptr[ctx.buffer_len + j] = data[i + j] ctx.buffer_len += remaining @@ -333,7 +325,7 @@ def sha256_hash(data: Span[UInt8, ...]) -> List[UInt8]: return sha256_final(ctx) -def sha256_final_to_buffer(mut ctx: SHA256Context, output: UnsafePointer[UInt8, MutAnyOrigin]): +def sha256_final_to_buffer(mut ctx: SHA256Context, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): var buf_ptr = ctx.buffer.unsafe_ptr() var bit_count = ctx.count + UInt64(ctx.buffer_len) * 8 @@ -343,7 +335,7 @@ def sha256_final_to_buffer(mut ctx: SHA256Context, output: UnsafePointer[UInt8, if ctx.buffer_len > 56: memset_zero(buf_ptr + ctx.buffer_len, 64 - ctx.buffer_len) ctx.state = sha256_transform( - ctx.state, Span[UInt8, ...](ptr=buf_ptr, length=64) + ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=64) ) ctx.buffer_len = 0 @@ -354,7 +346,7 @@ def sha256_final_to_buffer(mut ctx: SHA256Context, output: UnsafePointer[UInt8, ctx.buffer[56 + i] = UInt8(UInt64(bit_count >> UInt64(56 - i * 8)) & 0xFF) ctx.state = sha256_transform( - ctx.state, Span[UInt8, ...](ptr=buf_ptr, length=64) + ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=64) ) for i in range(8): @@ -422,7 +414,7 @@ struct SHA512Context(Movable): @always_inline def sha512_transform_blocks( mut state: SIMD[DType.uint64, 8], - data: UnsafePointer[UInt8, ImmutAnyOrigin], + data: UnsafePointer[mut=False, UInt8, _, address_space=_], nblocks: Int, ): comptime if CompilationTarget.has_neon() and CompilationTarget._has_feature["sha3"]() and not CompilationTarget.is_x86(): @@ -513,13 +505,10 @@ def sha512_update(mut ctx: SHA512Context, data: Span[UInt8, ...]): if ctx.buffer_len > 0: var available = 128 - ctx.buffer_len if total_len >= available: - memcpy( - dest=buf_ptr + ctx.buffer_len, - src=data.unsafe_ptr(), - count=available, - ) + for j in range(available): + buf_ptr[ctx.buffer_len + j] = data[j] ctx.state = sha512_transform( - ctx.state, Span[UInt8, ...](ptr=buf_ptr, length=128) + ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=128) ) var old_low = ctx.count_low @@ -530,11 +519,8 @@ def sha512_update(mut ctx: SHA512Context, data: Span[UInt8, ...]): i += available ctx.buffer_len = 0 else: - memcpy( - dest=buf_ptr + ctx.buffer_len, - src=data.unsafe_ptr(), - count=total_len, - ) + for j in range(total_len): + buf_ptr[ctx.buffer_len + j] = data[j] ctx.buffer_len += total_len return @@ -551,11 +537,8 @@ def sha512_update(mut ctx: SHA512Context, data: Span[UInt8, ...]): if i < total_len: var remaining = total_len - i - memcpy( - dest=buf_ptr + ctx.buffer_len, - src=data.unsafe_ptr() + i, - count=remaining, - ) + for j in range(remaining): + buf_ptr[ctx.buffer_len + j] = data[i + j] ctx.buffer_len += remaining @@ -572,7 +555,7 @@ def sha512_hash(data: Span[UInt8, ...]) -> List[UInt8]: return sha512_final(ctx) -def sha512_final_to_buffer(mut ctx: SHA512Context, output: UnsafePointer[UInt8, MutAnyOrigin]): +def sha512_final_to_buffer(mut ctx: SHA512Context, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): var buf_ptr = ctx.buffer.unsafe_ptr() var final_low = ctx.count_low + UInt64(ctx.buffer_len) * 8 var final_high = ctx.count_high @@ -585,7 +568,7 @@ def sha512_final_to_buffer(mut ctx: SHA512Context, output: UnsafePointer[UInt8, if ctx.buffer_len > 112: memset_zero(buf_ptr + ctx.buffer_len, 128 - ctx.buffer_len) ctx.state = sha512_transform( - ctx.state, Span[UInt8, ...](ptr=buf_ptr, length=128) + ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=128) ) ctx.buffer_len = 0 @@ -598,7 +581,7 @@ def sha512_final_to_buffer(mut ctx: SHA512Context, output: UnsafePointer[UInt8, ctx.buffer[120 + i] = UInt8(UInt64(final_low >> UInt64(56 - i * 8)) & 0xFF) ctx.state = sha512_transform( - ctx.state, Span[UInt8, ...](ptr=buf_ptr, length=128) + ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=128) ) for i in range(8): @@ -614,7 +597,7 @@ def sha256_final_with_len(mut ctx: SHA256Context, output_len: Int) -> List[UInt8 if ctx.buffer_len > 56: memset_zero(buf_ptr + ctx.buffer_len, 64 - ctx.buffer_len) ctx.state = sha256_transform( - ctx.state, Span[UInt8, ...](ptr=buf_ptr, length=64) + ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=64) ) ctx.buffer_len = 0 @@ -625,7 +608,7 @@ def sha256_final_with_len(mut ctx: SHA256Context, output_len: Int) -> List[UInt8 ctx.buffer[56 + i] = UInt8(UInt64(bit_count >> UInt64(56 - i * 8)) & 0xFF) ctx.state = sha256_transform( - ctx.state, Span[UInt8, ...](ptr=buf_ptr, length=64) + ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=64) ) var output = List[UInt8](capacity=output_len) @@ -669,7 +652,7 @@ def sha256_final_partial( if ctx.buffer_len > 56: memset_zero(buf_ptr + ctx.buffer_len, 64 - ctx.buffer_len) ctx.state = sha256_transform( - ctx.state, Span[UInt8, ...](ptr=buf_ptr, length=64) + ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=64) ) ctx.buffer_len = 0 @@ -680,7 +663,7 @@ def sha256_final_partial( ctx.buffer[56 + i] = UInt8(UInt64(bit_count >> UInt64(56 - i * 8)) & 0xFF) ctx.state = sha256_transform( - ctx.state, Span[UInt8, ...](ptr=buf_ptr, length=64) + ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=64) ) var output = List[UInt8](capacity=output_len) @@ -730,7 +713,7 @@ def sha512_final_with_len(mut ctx: SHA512Context, output_len: Int) -> List[UInt8 if ctx.buffer_len > 112: memset_zero(buf_ptr + ctx.buffer_len, 128 - ctx.buffer_len) ctx.state = sha512_transform( - ctx.state, Span[UInt8, ...](ptr=buf_ptr, length=128) + ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=128) ) ctx.buffer_len = 0 @@ -743,7 +726,7 @@ def sha512_final_with_len(mut ctx: SHA512Context, output_len: Int) -> List[UInt8 ctx.buffer[120 + i] = UInt8(UInt64(final_low >> UInt64(56 - i * 8)) & 0xFF) ctx.state = sha512_transform( - ctx.state, Span[UInt8, ...](ptr=buf_ptr, length=128) + ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=128) ) var output = List[UInt8](capacity=output_len) @@ -778,7 +761,7 @@ def sha512_final_partial( if ctx.buffer_len > 112: memset_zero(buf_ptr + ctx.buffer_len, 128 - ctx.buffer_len) ctx.state = sha512_transform( - ctx.state, Span[UInt8, ...](ptr=buf_ptr, length=128) + ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=128) ) ctx.buffer_len = 0 @@ -791,7 +774,7 @@ def sha512_final_partial( ctx.buffer[120 + i] = UInt8(UInt64(final_low >> UInt64(56 - i * 8)) & 0xFF) ctx.state = sha512_transform( - ctx.state, Span[UInt8, ...](ptr=buf_ptr, length=128) + ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=128) ) var output = List[UInt8](capacity=output_len) diff --git a/src/thistle/sha3.mojo b/src/thistle/sha3.mojo index 84f54e8..c99fe4c 100644 --- a/src/thistle/sha3.mojo +++ b/src/thistle/sha3.mojo @@ -10,8 +10,9 @@ from std.bit import rotate_bits_left from std.builtin.simd import SIMD from std.builtin.dtype import DType from std.sys import llvm_intrinsic, CompilationTarget +from std.utils import StaticTuple -comptime KECCAK_RC = SIMD[DType.uint64, 24]( +comptime KECCAK_RC = StaticTuple[UInt64, 24]( 0x0000000000000001, 0x0000000000008082, 0x800000000000808A, 0x8000000080008000, 0x000000000000808B, 0x0000000080000001, @@ -60,7 +61,7 @@ def _bcax(a: _U64x2, b: _U64x2, c: _U64x2) -> _U64x2: return llvm_intrinsic["llvm.aarch64.crypto.bcaxu", _U64x2, has_side_effect=False](a, b, c) -def _keccak_f1600_hw(state: UnsafePointer[UInt64, MutAnyOrigin]): +def _keccak_f1600_hw(state: UnsafePointer[mut=True, UInt64, _, address_space=_]): var a0 = _U64x2(state[0], 0) var a1 = _U64x2(state[1], 0) var a2 = _U64x2(state[2], 0) @@ -180,14 +181,14 @@ def _keccak_f1600_hw(state: UnsafePointer[UInt64, MutAnyOrigin]): state[24] = a24[0] -def keccak_f1600(state: UnsafePointer[UInt64, MutAnyOrigin]): +def keccak_f1600(state: UnsafePointer[mut=True, UInt64, _, address_space=_]): comptime if _has_sha3_ext: _keccak_f1600_hw(state) return _keccak_f1600_scalar(state) -def _keccak_f1600_scalar(state: UnsafePointer[UInt64, MutAnyOrigin]): +def _keccak_f1600_scalar(state: UnsafePointer[mut=True, UInt64, _, address_space=_]): var a0 = state[0] var a1 = state[1] var a2 = state[2] @@ -341,6 +342,10 @@ struct SHA3Context(Movable): var buffer_len: Int def __init__(out self, rate_bits: Int): + debug_assert[assert_mode="safe"]( + 0 < rate_bits <= 1344 and rate_bits % 8 == 0, + "SHA-3 rate must be a positive multiple of 8 no larger than 1344 bits", + ) self.state = StackBuffer[UInt64, 25](fill=0) self.rate_bytes = rate_bits // 8 self.buffer = StackBuffer[UInt8, 168](fill=0) @@ -362,7 +367,7 @@ struct SHA3Context(Movable): @always_inline -def sha3_absorb_block(state: UnsafePointer[UInt64, MutAnyOrigin], block: UnsafePointer[UInt8, ImmutAnyOrigin], rate_bytes: Int): +def sha3_absorb_block(state: UnsafePointer[mut=True, UInt64, _, address_space=_], block: UnsafePointer[mut=False, UInt8, _, address_space=_], rate_bytes: Int): var full_lanes = rate_bytes // 8 for i in range(full_lanes): state[i] ^= (block + i * 8).bitcast[UInt64]().load[width=1, alignment=1]() @@ -376,20 +381,14 @@ def sha3_update(mut ctx: SHA3Context, data: Span[UInt8, ...]): if ctx.buffer_len > 0: var available = ctx.rate_bytes - ctx.buffer_len if total_len >= available: - memcpy( - dest=ctx.buffer.ptr() + ctx.buffer_len, - src=data.unsafe_ptr(), - count=available, - ) + for j in range(available): + ctx.buffer[ctx.buffer_len + j] = data[j] sha3_absorb_block(ctx.state.ptr(), ctx.buffer.ptr(), ctx.rate_bytes) ctx.buffer_len = 0 i += available else: - memcpy( - dest=ctx.buffer.ptr() + ctx.buffer_len, - src=data.unsafe_ptr(), - count=total_len, - ) + for j in range(total_len): + ctx.buffer[ctx.buffer_len + j] = data[j] ctx.buffer_len += total_len return @@ -399,15 +398,15 @@ def sha3_update(mut ctx: SHA3Context, data: Span[UInt8, ...]): if i < total_len: var remaining = total_len - i - memcpy( - dest=ctx.buffer.ptr(), - src=data.unsafe_ptr() + i, - count=remaining, - ) + for j in range(remaining): + ctx.buffer[j] = data[i + j] ctx.buffer_len = remaining def sha3_final(mut ctx: SHA3Context, output_len_bytes: Int) -> List[UInt8]: + debug_assert[assert_mode="safe"]( + output_len_bytes >= 0, "SHA-3 output length cannot be negative" + ) ctx.buffer[ctx.buffer_len] = 0x06 ctx.buffer_len += 1 @@ -444,6 +443,10 @@ def sha3_final(mut ctx: SHA3Context, output_len_bytes: Int) -> List[UInt8]: @always_inline def sha3_final_into(mut ctx: SHA3Context, mut output: StackBuffer[UInt8, ...], output_len_bytes: Int): + debug_assert[assert_mode="safe"]( + 0 <= output_len_bytes <= output.capacity(), + "SHA-3 output length exceeds destination capacity", + ) output.clear() ctx.buffer[ctx.buffer_len] = 0x06 ctx.buffer_len += 1 @@ -556,6 +559,10 @@ def shake_finalize(mut ctx: SHA3Context): @always_inline def shake_squeeze_prefix_into(mut ctx: SHA3Context, mut output: StackBuffer[UInt8, ...], output_len: Int): + debug_assert[assert_mode="safe"]( + 0 <= output_len <= output.capacity(), + "SHAKE output length exceeds destination capacity", + ) output.clear() output.set_len_unchecked(output_len) @@ -584,6 +591,9 @@ def shake_advance(mut ctx: SHA3Context): @always_inline def shake_final(mut ctx: SHA3Context, output_len: Int) -> List[UInt8]: + debug_assert[assert_mode="safe"]( + output_len >= 0, "SHAKE output length cannot be negative" + ) shake_finalize(ctx) var output = List[UInt8](capacity=output_len) @@ -611,6 +621,10 @@ def shake_final(mut ctx: SHA3Context, output_len: Int) -> List[UInt8]: @always_inline def shake_final_into(mut ctx: SHA3Context, mut output: StackBuffer[UInt8, ...], output_len: Int): + debug_assert[assert_mode="safe"]( + 0 <= output_len <= output.capacity(), + "SHAKE output length exceeds destination capacity", + ) output.clear() shake_finalize(ctx) diff --git a/src/thistle/sha_ni.mojo b/src/thistle/sha_ni.mojo index c0b3b3b..5ac813a 100644 --- a/src/thistle/sha_ni.mojo +++ b/src/thistle/sha_ni.mojo @@ -7,6 +7,7 @@ from std.memory import UnsafePointer, bitcast from .utils import StackBuffer from std.builtin.simd import SIMD from std.builtin.dtype import DType +from std.utils import StaticTuple from .sha2 import SHA256_IV, sha256_transform comptime SIMD128 = SIMD[DType.uint32, 4] @@ -111,12 +112,12 @@ def byte_swap32(v: SIMD128) -> SIMD128: @always_inline("nodebug") -def Load(ptr: UnsafePointer[UInt8, ImmutAnyOrigin]) -> SIMD128: +def Load(ptr: UnsafePointer[mut=False, UInt8, _, address_space=_]) -> SIMD128: return byte_swap32(ptr.bitcast[UInt32]().load[width=4, alignment=1]()) @always_inline("nodebug") -def prefetch_next_block(ptr: UnsafePointer[UInt8, ImmutAnyOrigin]): +def prefetch_next_block(ptr: UnsafePointer[mut=False, UInt8, _, address_space=_]): prefetch[PrefetchOptions().for_read().high_locality().to_data_cache()](ptr + 64) @@ -244,7 +245,7 @@ def sha256ni_hash(data: Span[UInt8, ...]) -> List[UInt8]: ctx.buffer_len += 1 if ctx.buffer_len > 56: - ctx.state = sha256ni_transform(ctx.state, Span[UInt8, ...](ptr=ctx.buffer.ptr(), length=64)) + ctx.state = sha256ni_transform(ctx.state, Span[UInt8, ...](unsafe_ptr=ctx.buffer.ptr(), length=64)) ctx.buffer_len = 0 while ctx.buffer_len < 56: @@ -254,7 +255,7 @@ def sha256ni_hash(data: Span[UInt8, ...]) -> List[UInt8]: for k in range(8): ctx.buffer[56 + k] = UInt8(UInt64(bit_count >> UInt64(56 - k * 8)) & 0xFF) - ctx.state = sha256ni_transform(ctx.state, Span[UInt8, ...](ptr=ctx.buffer.ptr(), length=64)) + ctx.state = sha256ni_transform(ctx.state, Span[UInt8, ...](unsafe_ptr=ctx.buffer.ptr(), length=64)) var output = List[UInt8](capacity=32) for k in range(8): @@ -271,7 +272,7 @@ def has_sha_ni() -> Bool: def sha256ni_transform_blocks( mut state: SIMD[DType.uint32, 8], - data: UnsafePointer[UInt8, ImmutAnyOrigin], + data: UnsafePointer[mut=False, UInt8, _, address_space=_], nblocks: Int, ): comptime if CompilationTarget.has_neon() and CompilationTarget._has_feature["sha2"]() and not CompilationTarget.is_x86(): @@ -312,13 +313,13 @@ def sha256ni_transform_blocks( return for blk in range(nblocks): - var span = Span[UInt8, ...](ptr=data + blk * 64, length=64) + var span = Span[UInt8, ...](unsafe_ptr=data + blk * 64, length=64) state = sha256ni_transform(state, span) comptime SIMD64x2 = SIMD[DType.uint64, 2] -comptime SHA512NI_K = SIMD[DType.uint64, 80]( +comptime SHA512NI_K = StaticTuple[UInt64, 80]( 0x428A2F98D728AE22, 0x7137449123EF65CD, 0xB5C0FBCFEC4D3B2F, 0xE9B5DBA58189DBBC, 0x3956C25BF348B538, 0x59F111F1B605D019, 0x923F82A4AF194F9B, 0xAB1C5ED5DA6D8118, 0xD807AA98A3030242, 0x12835B0145706FBE, 0x243185BE4EE4B28C, 0x550C7DC3D5FFB4E2, 0x72BE5D74F27B896F, 0x80DEB1FE3B1696B1, 0x9BDC06A725C71235, @@ -378,7 +379,7 @@ def byte_swap64(v: SIMD64x2) -> SIMD64x2: @always_inline("nodebug") -def Load512(ptr: UnsafePointer[UInt8, ImmutAnyOrigin]) -> SIMD64x2: +def Load512(ptr: UnsafePointer[mut=False, UInt8, _, address_space=_]) -> SIMD64x2: return byte_swap64(ptr.bitcast[UInt64]().load[width=2, alignment=1]()) @@ -388,13 +389,13 @@ def _ext1(a: SIMD64x2, b: SIMD64x2) -> SIMD64x2: @always_inline("nodebug") -def prefetch_next_block512(ptr: UnsafePointer[UInt8, ImmutAnyOrigin]): +def prefetch_next_block512(ptr: UnsafePointer[mut=False, UInt8, _, address_space=_]): prefetch[PrefetchOptions().for_read().high_locality().to_data_cache()](ptr + 128) def sha512ni_transform_blocks( mut state: SIMD[DType.uint64, 8], - data: UnsafePointer[UInt8, ImmutAnyOrigin], + data: UnsafePointer[mut=False, UInt8, _, address_space=_], nblocks: Int, ): var ab = SIMD64x2(state[0], state[1]) diff --git a/src/thistle/utils.mojo b/src/thistle/utils.mojo index 02c5655..6368a2e 100644 --- a/src/thistle/utils.mojo +++ b/src/thistle/utils.mojo @@ -1,7 +1,7 @@ from std.bit import byte_swap -struct StackInlineArray[ElementType: Copyable, size: Int](Copyable): +struct StackInlineArray[ElementType: Copyable & Deinitable, size: Int](Copyable): var _data: InlineArray[Self.ElementType, Self.size] @always_inline @@ -98,9 +98,7 @@ struct StackInlineArray[ElementType: Copyable, size: Int](Copyable): return self.unsafe_get(idx) @always_inline - def unsafe_set[ - _T: Copyable & ImplicitlyDestructible - ](mut self: StackInlineArray[_T, ...], idx: Int, var value: _T): + def unsafe_set(mut self, idx: Int, var value: Self.ElementType): debug_assert( 0 <= idx < Self.size, "The index provided must be within the range [0, len(List) -1] when using List.unsafe_set()", @@ -109,7 +107,7 @@ struct StackInlineArray[ElementType: Copyable, size: Int](Copyable): (self._data.unsafe_ptr() + idx).init_pointee_move(value^) -struct StackBuffer[T: Copyable & ImplicitlyDestructible, N: Int](Movable): +struct StackBuffer[T: Copyable & Deinitable, N: Int](Movable): var _data: InlineArray[Self.T, Self.N] var _len: Int @@ -143,7 +141,9 @@ struct StackBuffer[T: Copyable & ImplicitlyDestructible, N: Int](Movable): @always_inline def push(mut self, var val: Self.T): - debug_assert(self._len < Self.N, "StackBuffer overflow") + debug_assert[assert_mode="safe"]( + self._len < Self.N, "StackBuffer overflow" + ) self._data[self._len] = val^ self._len += 1 @@ -154,13 +154,17 @@ struct StackBuffer[T: Copyable & ImplicitlyDestructible, N: Int](Movable): @always_inline def pop(mut self) -> Self.T: - debug_assert(self._len > 0, "StackBuffer underflow") + debug_assert[assert_mode="safe"]( + self._len > 0, "StackBuffer underflow" + ) self._len -= 1 return self._data[self._len].copy() @always_inline def top(ref self) -> ref[self._data] Self.T: - debug_assert(self._len > 0, "StackBuffer empty") + debug_assert[assert_mode="safe"]( + self._len > 0, "StackBuffer empty" + ) return self._data[self._len - 1] @always_inline @@ -169,7 +173,10 @@ struct StackBuffer[T: Copyable & ImplicitlyDestructible, N: Int](Movable): @always_inline def set_len_unchecked(mut self, new_len: Int): - debug_assert(0 <= new_len <= Self.N, "StackBuffer set_len_unchecked out of bounds") + debug_assert[assert_mode="safe"]( + 0 <= new_len <= Self.N, + "StackBuffer set_len_unchecked out of bounds", + ) self._len = new_len @always_inline @@ -178,31 +185,44 @@ struct StackBuffer[T: Copyable & ImplicitlyDestructible, N: Int](Movable): @always_inline def __getitem__(ref self, i: Int) -> ref[self._data] Self.T: - debug_assert(0 <= i < Self.N, "StackBuffer index out of bounds") + debug_assert[assert_mode="safe"]( + 0 <= i < Self.N, "StackBuffer index out of bounds" + ) return self._data[i] @always_inline def __setitem__(mut self, i: Int, var val: Self.T): - debug_assert(0 <= i < Self.N, "StackBuffer index out of bounds") + debug_assert[assert_mode="safe"]( + 0 <= i < Self.N, "StackBuffer index out of bounds" + ) self._data[i] = val^ @always_inline - def ptr(mut self) -> UnsafePointer[Self.T, MutAnyOrigin]: - return self._data.unsafe_ptr().unsafe_origin_cast[MutAnyOrigin]() + def ptr[ + origin: Origin, address_space: AddressSpace, // + ](ref[origin, address_space] self) -> UnsafePointer[ + Self.T, origin, address_space=address_space + ]: + return ( + self._data.unsafe_ptr() + .unsafe_mut_cast[origin.mut]() + .unsafe_origin_cast[origin]() + .address_space_cast[address_space]() + ) @always_inline -def load_32be(ptr: UnsafePointer[UInt8, ImmutAnyOrigin], offset: Int) -> UInt32: +def load_32be(ptr: UnsafePointer[mut=False, UInt8, _, address_space=_], offset: Int) -> UInt32: return byte_swap((ptr + offset).bitcast[UInt32]().load[width=1, alignment=1]()) @always_inline -def load_64be(ptr: UnsafePointer[UInt8, ImmutAnyOrigin], offset: Int) -> UInt64: +def load_64be(ptr: UnsafePointer[mut=False, UInt8, _, address_space=_], offset: Int) -> UInt64: return byte_swap((ptr + offset).bitcast[UInt64]().load[width=1, alignment=1]()) @always_inline -def store_64be(p: UnsafePointer[UInt8, MutAnyOrigin], off: Int, v: UInt64): +def store_64be(p: UnsafePointer[mut=True, UInt8, _, address_space=_], off: Int, v: UInt64): for i in range(8): p[off + i] = UInt8((v >> UInt64(56 - 8 * i)) & 0xFF) @@ -248,7 +268,11 @@ def nibble_to_hex_char(nibble: UInt8) -> UInt8: @always_inline -def bytes_to_hex_simd(data: UnsafePointer[UInt8, ImmutAnyOrigin], len: Int) -> String: +def bytes_to_hex_simd(data: UnsafePointer[mut=False, UInt8, _, address_space=_], len: Int) -> String: + debug_assert[assert_mode="safe"]( + 0 <= len <= Int.MAX // 2, + "Hex input length cannot be negative or overflow the output size", + ) var result = String(capacity=len * 2) for i in range(len): var b = data[i] diff --git a/src/thistle/x25519.mojo b/src/thistle/x25519.mojo index 59f1864..7bea861 100644 --- a/src/thistle/x25519.mojo +++ b/src/thistle/x25519.mojo @@ -27,11 +27,17 @@ def _cswap_pair( _cswap_fe(swap, z_2, z_3) @no_inline -def x25519(scalar_in: Span[UInt8, ...], point: Span[UInt8, ...], output: UnsafePointer[UInt8, MutAnyOrigin]) raises: +def x25519( + scalar_in: Span[UInt8, ...], + point: Span[UInt8, ...], + output: Span[mut=True, UInt8, ...], +) raises: if len(scalar_in) < 32: raise Error("X25519 scalar must be 32 bytes") if len(point) < 32: raise Error("X25519 point must be 32 bytes") + if len(output) < 32: + raise Error("X25519 output needs at least 32 writable bytes") var scalar = StackInlineArray[UInt8, 32](uninitialized=True) for i in range(32): scalar[i] = scalar_in[i] @@ -74,14 +80,14 @@ def x25519(scalar_in: Span[UInt8, ...], point: Span[UInt8, ...], output: UnsafeP _cswap_pair(swap, x_2, x_3, z_2, z_3) var res = x_2 * z_2.invert() - res.to_bytes_into(output) + res.to_bytes_into(output.unsafe_ptr()) var scalar_ptr = scalar.unsafe_ptr() for i in range(32): scalar_ptr.store[volatile=True](i, UInt8(0)) def x25519_public_key( - private_key: Span[UInt8, ...], output: UnsafePointer[UInt8, MutAnyOrigin] + private_key: Span[UInt8, ...], output: Span[mut=True, UInt8, ...] ) raises: if len(private_key) != 32: raise Error("X25519 private key must be 32 bytes") @@ -89,11 +95,18 @@ def x25519_public_key( for i in range(32): base[i] = 0 base[0] = 9 - x25519(private_key, Span[UInt8, ...](ptr=base.unsafe_ptr(), length=32), output) + x25519( + private_key, + Span[UInt8, ...](unsafe_ptr=base.unsafe_ptr(), length=32), + output, + ) def x25519_keygen() raises -> Tuple[List[UInt8], List[UInt8]]: var private_key = random_bytes(32) var public_key = List[UInt8](unsafe_uninit_length=32) - x25519_public_key(Span[UInt8, ...](private_key), public_key.unsafe_ptr()) + x25519_public_key( + Span[UInt8, ...](private_key), + Span[mut=True, UInt8, ...](public_key), + ) return (private_key^, public_key^) diff --git a/tests/benchmark.mojo b/tests/benchmark.mojo index 875bcd8..74109dc 100644 --- a/tests/benchmark.mojo +++ b/tests/benchmark.mojo @@ -1,6 +1,6 @@ from std.time import perf_counter, perf_counter_ns from std.collections import List -from std.algorithm import parallelize +from max.algorithm import parallelize from std.random import random_ui64, seed from std.math import ceildiv from std.sys import has_accelerator @@ -47,18 +47,18 @@ def benchmark_x25519(duration_secs: Float64) raises -> String: for i in range(32): scalar[i] = UInt8(i + 1) point[i] = UInt8(9) if i == 0 else UInt8(0) - var scalar_span = Span[UInt8, ...](ptr=scalar.unsafe_ptr(), length=32) - var point_span = Span[UInt8, ...](ptr=point.unsafe_ptr(), length=32) - x25519(scalar_span, point_span, out.unsafe_ptr()) + var scalar_span = Span[UInt8, ...](unsafe_ptr=scalar.unsafe_ptr(), length=32) + var point_span = Span[UInt8, ...](unsafe_ptr=point.unsafe_ptr(), length=32) + x25519(scalar_span, point_span, Span[mut=True, UInt8, ...](out)) var count = 0 var start = perf_counter() while perf_counter() - start < duration_secs: - x25519(scalar_span, point_span, out.unsafe_ptr()) + x25519(scalar_span, point_span, Span[mut=True, UInt8, ...](out)) scalar.unsafe_ptr().store[volatile=True](0, out[0] | 8) count += 1 var duration = perf_counter() - start var ops = Float64(count) / duration - return "x25519 | throughput: " + String(ops)[byte=:8] + " ops/s, ops: " + String(count) + ", time: " + String(duration)[byte=:4] + "s" + return "x25519 | throughput: " + String(ops) + " ops/s, ops: " + String(count) + ", time: " + String(duration) + "s" def benchmark_p384(duration_secs: Float64) -> String: @@ -66,15 +66,15 @@ def benchmark_p384(duration_secs: Float64) -> String: var out = InlineArray[UInt8, 97](fill=0) scalar[47] = 7 var scalar_span = Span[UInt8, ...](scalar) - _ = p384_public_key(scalar_span, out.unsafe_ptr()) + _ = p384_public_key(scalar_span, Span[mut=True, UInt8, ...](out)) var count = 0 var start = perf_counter() while perf_counter() - start < duration_secs: - _ = p384_public_key(scalar_span, out.unsafe_ptr()) + _ = p384_public_key(scalar_span, Span[mut=True, UInt8, ...](out)) count += 1 var duration = perf_counter() - start var ops = Float64(count) / duration - return "p384-public-key | throughput: " + String(ops)[byte=:8] + " ops/s, ops: " + String(count) + ", time: " + String(duration)[byte=:4] + "s" + return "p384-public-key | throughput: " + String(ops) + " ops/s, ops: " + String(count) + ", time: " + String(duration) + "s" def benchmark_ecdsa(duration_secs: Float64) -> String: @@ -89,7 +89,9 @@ def benchmark_ecdsa(duration_secs: Float64) -> String: var start = perf_counter() while perf_counter() - start < duration_secs: _ = p256_ecdsa_sign( - Span[UInt8, ...](p256_key), msg, p256_sig.unsafe_ptr() + Span[UInt8, ...](p256_key), + msg, + Span[mut=True, UInt8, ...](unsafe_ptr=p256_sig.unsafe_ptr(), length=64), ) p256_count += 1 var p256_time = perf_counter() - start @@ -98,16 +100,18 @@ def benchmark_ecdsa(duration_secs: Float64) -> String: start = perf_counter() while perf_counter() - start < duration_secs: _ = p384_ecdsa_sign( - Span[UInt8, ...](p384_key), msg, p384_sig.unsafe_ptr() + Span[UInt8, ...](p384_key), + msg, + Span[mut=True, UInt8, ...](unsafe_ptr=p384_sig.unsafe_ptr(), length=96), ) p384_count += 1 var p384_time = perf_counter() - start return ( "p256-ecdsa-sign | throughput: " - + String(Float64(p256_count) / p256_time)[byte=:8] + " ops/s\n" + + String(Float64(p256_count) / p256_time) + " ops/s\n" + "p384-ecdsa-sign | throughput: " - + String(Float64(p384_count) / p384_time)[byte=:8] + " ops/s" + + String(Float64(p384_count) / p384_time) + " ops/s" ) @@ -120,10 +124,10 @@ def benchmark_ed25519(duration_secs: Float64) raises -> String: sk[i] = UInt8(i * 7 + 1) for i in range(64): msg[i] = UInt8(i) - var sk_span = Span[UInt8, ...](ptr=sk.unsafe_ptr(), length=32) - var msg_span = Span[UInt8, ...](ptr=msg.unsafe_ptr(), length=64) - var sig_span = Span[UInt8, ...](ptr=sig.unsafe_ptr(), length=64) - var pk_span = Span[UInt8, ...](ptr=pk.unsafe_ptr(), length=32) + var sk_span = Span[UInt8, ...](unsafe_ptr=sk.unsafe_ptr(), length=32) + var msg_span = Span[UInt8, ...](unsafe_ptr=msg.unsafe_ptr(), length=64) + var sig_span = Span[UInt8, ...](unsafe_ptr=sig.unsafe_ptr(), length=64) + var pk_span = Span[UInt8, ...](unsafe_ptr=pk.unsafe_ptr(), length=32) ed25519_generate_public_key(sk_span, pk.unsafe_ptr()) ed25519_sign(sk_span, msg_span, sig.unsafe_ptr()) @@ -146,8 +150,8 @@ def benchmark_ed25519(duration_secs: Float64) raises -> String: var verify_duration = perf_counter() - start var verify_ops = Float64(verify_count) / verify_duration - var result = "ed25519-sign | throughput: " + String(sign_ops)[byte=:8] + " ops/s, ops: " + String(sign_count) + ", time: " + String(sign_duration)[byte=:4] + "s\n" - result += "ed25519-verify | throughput: " + String(verify_ops)[byte=:8] + " ops/s, ops: " + String(verify_count) + ", time: " + String(verify_duration)[byte=:4] + "s" + var result = "ed25519-sign | throughput: " + String(sign_ops) + " ops/s, ops: " + String(sign_count) + ", time: " + String(sign_duration) + "s\n" + result += "ed25519-verify | throughput: " + String(verify_ops) + " ops/s, ops: " + String(verify_count) + ", time: " + String(verify_duration) + "s" if verify_failures > 0: result += " [" + String(verify_failures) + " FAILED VERIFICATIONS]" return result @@ -165,7 +169,7 @@ def benchmark_sha256(data: List[UInt8], duration_secs: Float64) -> String: var duration = end - start var mb = Float64(len(data) * count) / (1024 * 1024) var mbps = mb / duration - return "sha256 | throughput: " + String(mbps)[byte=:6] + " mb/s, hashes: " + String(count) + ", time: " + String(duration)[byte=:4] + "s" + return "sha256 | throughput: " + String(mbps) + " mb/s, hashes: " + String(count) + ", time: " + String(duration) + "s" def benchmark_sha256ni(data: List[UInt8], duration_secs: Float64) -> String: @@ -182,7 +186,7 @@ def benchmark_sha256ni(data: List[UInt8], duration_secs: Float64) -> String: var duration = end - start var mb = Float64(len(data) * count) / (1024 * 1024) var mbps = mb / duration - return "sha256-ni | throughput: " + String(mbps)[byte=:6] + " mb/s, hashes: " + String(count) + ", time: " + String(duration)[byte=:4] + "s" + return "sha256-ni | throughput: " + String(mbps) + " mb/s, hashes: " + String(count) + ", time: " + String(duration) + "s" def benchmark_sha512(data: List[UInt8], duration_secs: Float64) -> String: @@ -197,7 +201,7 @@ def benchmark_sha512(data: List[UInt8], duration_secs: Float64) -> String: var duration = end - start var mb = Float64(len(data) * count) / (1024 * 1024) var mbps = mb / duration - return "sha512 | throughput: " + String(mbps)[byte=:6] + " mb/s, hashes: " + String(count) + ", time: " + String(duration)[byte=:4] + "s" + return "sha512 | throughput: " + String(mbps) + " mb/s, hashes: " + String(count) + ", time: " + String(duration) + "s" def benchmark_sha3_256(data: List[UInt8], duration_secs: Float64) -> String: @@ -212,7 +216,7 @@ def benchmark_sha3_256(data: List[UInt8], duration_secs: Float64) -> String: var duration = end - start var mb = Float64(len(data) * count) / (1024 * 1024) var mbps = mb / duration - return "sha3-256 | throughput: " + String(mbps)[byte=:6] + " mb/s, hashes: " + String(count) + ", time: " + String(duration)[byte=:4] + "s" + return "sha3-256 | throughput: " + String(mbps) + " mb/s, hashes: " + String(count) + ", time: " + String(duration) + "s" def benchmark_blake2b(data: List[UInt8], duration_secs: Float64) raises -> String: @@ -228,7 +232,7 @@ def benchmark_blake2b(data: List[UInt8], duration_secs: Float64) raises -> Strin var duration = end - start var mb = Float64(len(data) * count) / (1024 * 1024) var mbps = mb / duration - return "blake2b | throughput: " + String(mbps)[byte=:6] + " mb/s, hashes: " + String(count) + ", time: " + String(duration)[byte=:4] + "s" + return "blake2b | throughput: " + String(mbps) + " mb/s, hashes: " + String(count) + ", time: " + String(duration) + "s" def benchmark_blake3(data: List[UInt8], duration_secs: Float64) raises -> String: @@ -243,7 +247,7 @@ def benchmark_blake3(data: List[UInt8], duration_secs: Float64) raises -> String var duration = end - start var mb = Float64(len(data) * count) / (1024 * 1024) var mbps = mb / duration - return "blake3 | throughput: " + String(mbps)[byte=:6] + " mb/s, hashes: " + String(count) + ", time: " + String(duration)[byte=:4] + "s" + return "blake3 | throughput: " + String(mbps) + " mb/s, hashes: " + String(count) + ", time: " + String(duration) + "s" def benchmark_camellia(data_size: Int, duration_secs: Float64) raises -> String: @@ -271,7 +275,7 @@ def benchmark_camellia(data_size: Int, duration_secs: Float64) raises -> String: blocks.free() var mbps = Float64(count * 16) / (1024 * 1024) / duration - return "camellia | throughput: " + String(mbps)[byte=:6] + " mb/s, blocks: " + String(count) + ", time: " + String(duration)[byte=:4] + "s" + return "camellia | throughput: " + String(mbps) + " mb/s, blocks: " + String(count) + ", time: " + String(duration) + "s" def benchmark_camellia_ctr(duration_secs: Float64) raises -> String: @@ -302,14 +306,14 @@ def benchmark_camellia_ctr(duration_secs: Float64) raises -> String: nonce.free() var mbps = Float64(count * size) / (1024 * 1024) / duration - return "camellia-ctr | throughput: " + String(mbps)[byte=:6] + " mb/s, chunks: " + String(count) + ", time: " + String(duration)[byte=:4] + "s" + return "camellia-ctr | throughput: " + String(mbps) + " mb/s, chunks: " + String(count) + ", time: " + String(duration) + "s" def benchmark_chacha20(data_size: Int, duration_secs: Float64) raises -> String: var key = SIMD[DType.uint8, 32](0) for i in range(32): key[i] = UInt8(i) - var nonce = SIMD[DType.uint8, 12](0) + var nonce = SIMD[DType.uint8, 16](0) var data = List[UInt8](capacity=data_size) for i in range(data_size): @@ -330,7 +334,7 @@ def benchmark_chacha20(data_size: Int, duration_secs: Float64) raises -> String: _ = checksum var mb = Float64(data_size * count) / (1024 * 1024) var mbps = mb / duration - return "chacha20 | throughput: " + String(mbps)[byte=:6] + " mb/s, encrypts: " + String(count) + ", time: " + String(duration)[byte=:4] + "s" + return "chacha20 | throughput: " + String(mbps) + " mb/s, encrypts: " + String(count) + ", time: " + String(duration) + "s" def benchmark_kcipher2(data_size: Int, duration_secs: Float64) -> String: @@ -353,7 +357,7 @@ def benchmark_kcipher2(data_size: Int, duration_secs: Float64) -> String: var duration = end - start var mb = Float64(data_size * count) / (1024 * 1024) var mbps = mb / duration - return "kcipher2 | throughput: " + String(mbps)[byte=:6] + " mb/s, encrypts: " + String(count) + ", time: " + String(duration)[byte=:4] + "s" + return "kcipher2 | throughput: " + String(mbps) + " mb/s, encrypts: " + String(count) + ", time: " + String(duration) + "s" def benchmark_argon2(duration_secs: Float64) raises -> String: @@ -371,7 +375,7 @@ def benchmark_argon2(duration_secs: Float64) raises -> String: var end = perf_counter() var duration = end - start var hps = Float64(count) / duration - return "argon2id | throughput: " + String(hps)[byte=:6] + " h/s, hashes: " + String(count) + ", time: " + String(duration)[byte=:4] + "s" + return "argon2id | throughput: " + String(hps) + " h/s, hashes: " + String(count) + ", time: " + String(duration) + "s" def benchmark_aes_cpu(duration_secs: Float64) raises -> String: @@ -396,7 +400,7 @@ def benchmark_aes_cpu(duration_secs: Float64) raises -> String: blocks.free() var mbps = Float64(count * 16) / (1024 * 1024) / duration - return "aes-128-cpu | throughput: " + String(mbps)[byte=:6] + " mb/s, blocks: " + String(count) + ", time: " + String(duration)[byte=:4] + "s" + return "aes-128-cpu | throughput: " + String(mbps) + " mb/s, blocks: " + String(count) + ", time: " + String(duration) + "s" def benchmark_aes_gpu_ecb() raises -> String: @@ -404,13 +408,15 @@ def benchmark_aes_gpu_ecb() raises -> String: if not has_accelerator(): return "aes-128-gpu-ecb | (GPU not available)" - from std.gpu.host import DeviceContext + from max.gpu.host import DeviceContext from thistle.aes_gpu import aes_gpu_kernel_ecb var key_ptr = alloc[UInt8](16) for i in range(16): key_ptr.store(i, TEST_KEY[i]) - var round_keys = expand_key_128(key_ptr) + var round_keys = expand_key_128( + Span[UInt8, ...](unsafe_ptr=key_ptr, length=16) + ) var num_blocks = 131072 var total_bytes = num_blocks * 16 @@ -423,7 +429,7 @@ def benchmark_aes_gpu_ecb() raises -> String: with DeviceContext() as ctx: var input_buffer = ctx.enqueue_create_buffer[DType.uint8](total_bytes) var output_buffer = ctx.enqueue_create_buffer[DType.uint8](total_bytes) - var skey_host = cpu_aes_ct_skey(round_keys, 10) + var skey_host = cpu_aes_ct_skey(round_keys.ptr(), 10) var skey_buffer = ctx.enqueue_create_buffer[DType.uint64](88) ctx.enqueue_copy(skey_buffer, skey_host.unsafe_ptr()) @@ -434,11 +440,11 @@ def benchmark_aes_gpu_ecb() raises -> String: var grid_dim = ceildiv(num_blocks, block_dim) ctx.enqueue_function[aes_gpu_kernel_ecb]( - input_buffer.unsafe_ptr(), - output_buffer.unsafe_ptr(), - skey_buffer.unsafe_ptr(), - num_blocks, - 10, + input_buffer, + output_buffer, + skey_buffer, + Int32(num_blocks), + Int32(10), grid_dim=grid_dim, block_dim=block_dim, ) @@ -448,11 +454,11 @@ def benchmark_aes_gpu_ecb() raises -> String: var start = perf_counter() for _ in range(iterations): ctx.enqueue_function[aes_gpu_kernel_ecb]( - input_buffer.unsafe_ptr(), - output_buffer.unsafe_ptr(), - skey_buffer.unsafe_ptr(), - num_blocks, - 10, + input_buffer, + output_buffer, + skey_buffer, + Int32(num_blocks), + Int32(10), grid_dim=grid_dim, block_dim=block_dim, ) @@ -465,10 +471,9 @@ def benchmark_aes_gpu_ecb() raises -> String: input_host.free() output_host.free() - round_keys.free() key_ptr.free() - return "aes-128-gpu-ecb | throughput: " + String(gbps)[byte=:6] + " gb/s, iterations: " + String(iterations) + return "aes-128-gpu-ecb | throughput: " + String(gbps) + " gb/s, iterations: " + String(iterations) def benchmark_aes_gpu_ctr() raises -> String: @@ -476,13 +481,15 @@ def benchmark_aes_gpu_ctr() raises -> String: if not has_accelerator(): return "aes-128-gpu-ctr | (GPU not available)" - from std.gpu.host import DeviceContext + from max.gpu.host import DeviceContext from thistle.aes_gpu import aes_gpu_kernel_ctr var key_ptr = alloc[UInt8](16) for i in range(16): key_ptr.store(i, TEST_KEY[i]) - var round_keys = expand_key_128(key_ptr) + var round_keys = expand_key_128( + Span[UInt8, ...](unsafe_ptr=key_ptr, length=16) + ) var num_blocks = 131072 var total_bytes = num_blocks * 16 @@ -498,7 +505,7 @@ def benchmark_aes_gpu_ctr() raises -> String: with DeviceContext() as ctx: var input_buffer = ctx.enqueue_create_buffer[DType.uint8](total_bytes) var output_buffer = ctx.enqueue_create_buffer[DType.uint8](total_bytes) - var skey_host = cpu_aes_ct_skey(round_keys, 10) + var skey_host = cpu_aes_ct_skey(round_keys.ptr(), 10) var skey_buffer = ctx.enqueue_create_buffer[DType.uint64](88) ctx.enqueue_copy(skey_buffer, skey_host.unsafe_ptr()) var nonce_buffer = ctx.enqueue_create_buffer[DType.uint8](16) @@ -511,12 +518,12 @@ def benchmark_aes_gpu_ctr() raises -> String: var grid_dim = ceildiv(num_blocks, block_dim) ctx.enqueue_function[aes_gpu_kernel_ctr]( - input_buffer.unsafe_ptr(), - output_buffer.unsafe_ptr(), - skey_buffer.unsafe_ptr(), - num_blocks, - nonce_buffer.unsafe_ptr(), - 10, + input_buffer, + output_buffer, + skey_buffer, + Int32(num_blocks), + nonce_buffer, + Int32(10), grid_dim=grid_dim, block_dim=block_dim, ) @@ -526,12 +533,12 @@ def benchmark_aes_gpu_ctr() raises -> String: var start = perf_counter() for _ in range(iterations): ctx.enqueue_function[aes_gpu_kernel_ctr]( - input_buffer.unsafe_ptr(), - output_buffer.unsafe_ptr(), - skey_buffer.unsafe_ptr(), - num_blocks, - nonce_buffer.unsafe_ptr(), - 10, + input_buffer, + output_buffer, + skey_buffer, + Int32(num_blocks), + nonce_buffer, + Int32(10), grid_dim=grid_dim, block_dim=block_dim, ) @@ -545,10 +552,9 @@ def benchmark_aes_gpu_ctr() raises -> String: input_host.free() output_host.free() nonce_host.free() - round_keys.free() key_ptr.free() - return "aes-128-gpu-ctr | throughput: " + String(gbps)[byte=:6] + " gb/s, iterations: " + String(iterations) + return "aes-128-gpu-ctr | throughput: " + String(gbps) + " gb/s, iterations: " + String(iterations) @@ -559,13 +565,15 @@ def benchmark_aes_gpu_gcm() raises -> String: if not has_accelerator(): return "aes-128-gpu-gcm | (GPU not available)" - from std.gpu.host import DeviceContext + from max.gpu.host import DeviceContext from thistle.aes_gpu import aes_gpu_kernel_gcm_ctr var key_ptr = alloc[UInt8](16) for i in range(16): key_ptr.store(i, TEST_KEY[i]) - var round_keys = expand_key_128(key_ptr) + var round_keys = expand_key_128( + Span[UInt8, ...](unsafe_ptr=key_ptr, length=16) + ) var num_blocks = 131072 var total_bytes = num_blocks * 16 @@ -582,7 +590,7 @@ def benchmark_aes_gpu_gcm() raises -> String: with DeviceContext() as ctx: var input_buffer = ctx.enqueue_create_buffer[DType.uint8](total_bytes) var output_buffer = ctx.enqueue_create_buffer[DType.uint8](total_bytes) - var skey_host = cpu_aes_ct_skey(round_keys, 10) + var skey_host = cpu_aes_ct_skey(round_keys.ptr(), 10) var skey_buffer = ctx.enqueue_create_buffer[DType.uint64](88) ctx.enqueue_copy(skey_buffer, skey_host.unsafe_ptr()) var nonce_buffer = ctx.enqueue_create_buffer[DType.uint8](16) @@ -595,12 +603,12 @@ def benchmark_aes_gpu_gcm() raises -> String: var grid_dim = ceildiv(num_blocks, block_dim) ctx.enqueue_function[aes_gpu_kernel_gcm_ctr]( - input_buffer.unsafe_ptr(), - output_buffer.unsafe_ptr(), - skey_buffer.unsafe_ptr(), - num_blocks, - nonce_buffer.unsafe_ptr(), - 10, + input_buffer, + output_buffer, + skey_buffer, + Int32(num_blocks), + nonce_buffer, + Int32(10), grid_dim=grid_dim, block_dim=block_dim, ) @@ -610,12 +618,12 @@ def benchmark_aes_gpu_gcm() raises -> String: var start = perf_counter() for _ in range(iterations): ctx.enqueue_function[aes_gpu_kernel_gcm_ctr]( - input_buffer.unsafe_ptr(), - output_buffer.unsafe_ptr(), - skey_buffer.unsafe_ptr(), - num_blocks, - nonce_buffer.unsafe_ptr(), - 10, + input_buffer, + output_buffer, + skey_buffer, + Int32(num_blocks), + nonce_buffer, + Int32(10), grid_dim=grid_dim, block_dim=block_dim, ) @@ -629,10 +637,9 @@ def benchmark_aes_gpu_gcm() raises -> String: input_host.free() output_host.free() nonce_host.free() - round_keys.free() key_ptr.free() - return "aes-128-gpu-gcm | throughput: " + String(gbps)[byte=:6] + " gb/s, iterations: " + String(iterations) + return "aes-128-gpu-gcm | throughput: " + String(gbps) + " gb/s, iterations: " + String(iterations) diff --git a/tests/dudect.mojo b/tests/dudect.mojo index 28aff2f..d0a0e1c 100644 --- a/tests/dudect.mojo +++ b/tests/dudect.mojo @@ -235,7 +235,7 @@ def run_aes_sw(mut rng: Rng) raises -> Bool: var kb = List[UInt8]() for i in range(16): kb.append(UInt8(i + 1)) - var rk = expand_key_128(kb.unsafe_ptr()) + var rk = expand_key_128(Span[UInt8, ...](kb)) var cls = _classes(N_FAST, rng) var inp = _fast_inputs(cls, 16, rng) var times = List[Float64](capacity=N_FAST) @@ -248,10 +248,9 @@ def run_aes_sw(mut rng: Rng) raises -> Bool: for _ in range(BATCH): for j in range(16): blk[j] = inp[i * 16 + j] - cpu_aes_encrypt(blk.unsafe_ptr(), rk, 10) + cpu_aes_encrypt(blk.unsafe_ptr(), rk.ptr(), 10) sink ^= blk[0] times.append(Float64(perf_counter_ns() - t0)) - rk.free() if sink == 42: print("") return _report("aes-128 software (fixed/random pt)", times, cls) @@ -260,7 +259,7 @@ def run_aes_sw(mut rng: Rng) raises -> Bool: def run_chacha20(mut rng: Rng) raises -> Bool: var cls = _classes(N_FAST, rng) var inp = _fast_inputs(cls, 32, rng) - var nonce = SIMD[DType.uint8, 12](0) + var nonce = SIMD[DType.uint8, 16](0) var times = List[Float64](capacity=N_FAST) var sink: UInt8 = 0 var data = List[UInt8]() @@ -402,7 +401,9 @@ def run_x25519(mut rng: Rng) raises -> Bool: sc[0] = 0x40 var t0 = perf_counter_ns() x25519( - Span[UInt8, ...](sc), Span[UInt8, ...](base), out.unsafe_ptr() + Span[UInt8, ...](sc), + Span[UInt8, ...](base), + Span[mut=True, UInt8, ...](out), ) times.append(Float64(perf_counter_ns() - t0)) sink ^= out[0] @@ -451,7 +452,9 @@ def run_p256(mut rng: Rng) raises -> Bool: for i in range(N_ASYM): var sc = _valid_scalar(32, rng, cls[i] == 0) var t0 = perf_counter_ns() - var ok = p256_public_key(Span[UInt8, ...](sc), out.unsafe_ptr()) + var ok = p256_public_key( + Span[UInt8, ...](sc), Span[mut=True, UInt8, ...](out) + ) times.append(Float64(perf_counter_ns() - t0)) sink ^= out[0] ^ (UInt8(1) if ok else UInt8(0)) if sink == 42: @@ -469,7 +472,9 @@ def run_p384(mut rng: Rng) raises -> Bool: for i in range(N_ASYM): var sc = _valid_scalar(48, rng, cls[i] == 0) var t0 = perf_counter_ns() - var ok = p384_public_key(Span[UInt8, ...](sc), out.unsafe_ptr()) + var ok = p384_public_key( + Span[UInt8, ...](sc), Span[mut=True, UInt8, ...](out) + ) times.append(Float64(perf_counter_ns() - t0)) sink ^= out[0] ^ (UInt8(1) if ok else UInt8(0)) if sink == 42: @@ -489,7 +494,7 @@ def run_p256_sign(mut rng: Rng) -> Bool: var ok = p256_ecdsa_sign( Span[UInt8, ...](private_key), Span[UInt8, ...](message), - signature.unsafe_ptr(), + Span[mut=True, UInt8, ...](signature), ) times.append(Float64(perf_counter_ns() - t0)) sink ^= signature[0] ^ (UInt8(1) if ok else UInt8(0)) @@ -510,7 +515,7 @@ def run_p384_sign(mut rng: Rng) -> Bool: var ok = p384_ecdsa_sign( Span[UInt8, ...](private_key), Span[UInt8, ...](message), - signature.unsafe_ptr(), + Span[mut=True, UInt8, ...](signature), ) times.append(Float64(perf_counter_ns() - t0)) sink ^= signature[0] ^ (UInt8(1) if ok else UInt8(0)) diff --git a/tests/gen_ed25519_table.mojo b/tests/gen_ed25519_table.mojo index 219b585..2fa4958 100644 --- a/tests/gen_ed25519_table.mojo +++ b/tests/gen_ed25519_table.mojo @@ -14,7 +14,7 @@ from thistle.curve25519 import FieldElement51 def _canonical(fe: FieldElement51) -> FieldElement51: var bytes = InlineArray[UInt8, 32](uninitialized=True) fe.to_bytes_into(bytes.unsafe_ptr()) - return fe_from_bytes(Span[UInt8, ...](ptr=bytes.unsafe_ptr(), length=32)) + return fe_from_bytes(Span[UInt8, ...](unsafe_ptr=bytes.unsafe_ptr(), length=32)) def _affine_niels_limbs(p: EdwardsPoint) -> InlineArray[UInt64, 15]: diff --git a/tests/test_aes_gpu.mojo b/tests/test_aes_gpu.mojo index 2efca96..6d9126c 100644 --- a/tests/test_aes_gpu.mojo +++ b/tests/test_aes_gpu.mojo @@ -3,10 +3,10 @@ from std.python import PythonObject from std.collections import List from std.sys import has_accelerator from thistle.sha2 import bytes_to_hex -from thistle.aes import cpu_aes_ct_skey, expand_key_128, expand_key_192, expand_key_256 +from thistle.aes import cpu_aes_ct_skey, AESExpandedKey from thistle.aes_gpu import aes_gpu_kernel_ecb, aes_gpu_kernel_ctr, aes_gpu_kernel_gcm_ctr from std.memory import alloc -from std.gpu.host import DeviceContext +from max.gpu.host import DeviceContext from std.memory.unsafe_pointer import UnsafePointer @@ -88,9 +88,6 @@ def test_aes_gpu_basic(json_data: PythonObject, py: PythonObject) raises -> Test var key_len = len(key_bytes) var rounds = 10 if key_len == 16 else (12 if key_len == 24 else 14) - var key_ptr = alloc[UInt8](key_len) - for j in range(key_len): - key_ptr.store(j, key_bytes[j]) var total_bytes = 64 var input_host = alloc[Scalar[DType.uint8]](total_bytes) @@ -103,14 +100,8 @@ def test_aes_gpu_basic(json_data: PythonObject, py: PythonObject) raises -> Test var input_buffer = ctx.enqueue_create_buffer[DType.uint8](total_bytes) var output_buffer = ctx.enqueue_create_buffer[DType.uint8](total_bytes) - var round_keys: UnsafePointer[UInt32, MutAnyOrigin] - if key_len == 16: - round_keys = expand_key_128(key_ptr) - elif key_len == 24: - round_keys = expand_key_192(key_ptr) - else: - round_keys = expand_key_256(key_ptr) - var skey_host = cpu_aes_ct_skey(round_keys, rounds) + var round_keys = AESExpandedKey(Span[UInt8, ...](key_bytes)) + var skey_host = cpu_aes_ct_skey(round_keys.ptr(), rounds) var skey_buffer = ctx.enqueue_create_buffer[DType.uint64]((rounds + 1) * 8) ctx.enqueue_copy(skey_buffer, skey_host.unsafe_ptr()) ctx.enqueue_copy(input_buffer, input_host) @@ -120,11 +111,11 @@ def test_aes_gpu_basic(json_data: PythonObject, py: PythonObject) raises -> Test var grid_dim = 1 ctx.enqueue_function[aes_gpu_kernel_ecb]( - input_buffer.unsafe_ptr(), - output_buffer.unsafe_ptr(), - skey_buffer.unsafe_ptr(), - 4, - rounds, + input_buffer, + output_buffer, + skey_buffer, + Int32(4), + Int32(rounds), grid_dim=grid_dim, block_dim=block_dim, ) @@ -157,8 +148,6 @@ def test_aes_gpu_basic(json_data: PythonObject, py: PythonObject) raises -> Test got_hex += byte_to_hex(output_host[j]) failures.append("AES-GPU " + name + ": expected " + expected + ", got " + got_hex) - key_ptr.free() - round_keys.free() input_host.free() output_host.free() @@ -183,28 +172,7 @@ def test_mode_gpu(json_data: PythonObject, mode: String) raises -> TestResult: var pt_bytes = hex_to_bytes(pt_hex) var key_len = len(key_bytes) - var key_ptr: UnsafePointer[UInt8, MutAnyOrigin] - var round_keys_size: Int - var round_keys: UnsafePointer[UInt32, MutAnyOrigin] - - if key_len == 16: - key_ptr = alloc[UInt8](16) - for j in range(16): - key_ptr.store(j, key_bytes[j]) - round_keys_size = 44 - round_keys = expand_key_128(key_ptr) - elif key_len == 24: - key_ptr = alloc[UInt8](24) - for j in range(24): - key_ptr.store(j, key_bytes[j]) - round_keys_size = 52 - round_keys = expand_key_192(key_ptr) - else: - key_ptr = alloc[UInt8](32) - for j in range(32): - key_ptr.store(j, key_bytes[j]) - round_keys_size = 60 - round_keys = expand_key_256(key_ptr) + var round_keys = AESExpandedKey(Span[UInt8, ...](key_bytes)) var total_bytes = len(pt_bytes) var n_blocks = total_bytes // 16 @@ -231,7 +199,7 @@ def test_mode_gpu(json_data: PythonObject, mode: String) raises -> TestResult: else: rounds = 14 - var skey_host = cpu_aes_ct_skey(round_keys, rounds) + var skey_host = cpu_aes_ct_skey(round_keys.ptr(), rounds) var skey_buffer = ctx.enqueue_create_buffer[DType.uint64]((rounds + 1) * 8) ctx.enqueue_copy(skey_buffer, skey_host.unsafe_ptr()) @@ -239,11 +207,11 @@ def test_mode_gpu(json_data: PythonObject, mode: String) raises -> TestResult: var nonce_buffer = ctx.enqueue_create_buffer[DType.uint8](16) if "ECB" in mode: ctx.enqueue_function[aes_gpu_kernel_ecb]( - input_buffer.unsafe_ptr(), - output_buffer.unsafe_ptr(), - skey_buffer.unsafe_ptr(), - n_blocks, - rounds, + input_buffer, + output_buffer, + skey_buffer, + Int32(n_blocks), + Int32(rounds), grid_dim=grid_dim, block_dim=block_dim, ) @@ -259,12 +227,12 @@ def test_mode_gpu(json_data: PythonObject, mode: String) raises -> TestResult: ctx.enqueue_copy(nonce_buffer, nonce_ptr) ctx.enqueue_function[aes_gpu_kernel_ctr]( - input_buffer.unsafe_ptr(), - output_buffer.unsafe_ptr(), - skey_buffer.unsafe_ptr(), - n_blocks, - nonce_buffer.unsafe_ptr(), - rounds, + input_buffer, + output_buffer, + skey_buffer, + Int32(n_blocks), + nonce_buffer, + Int32(rounds), grid_dim=grid_dim, block_dim=block_dim, ) @@ -284,22 +252,22 @@ def test_mode_gpu(json_data: PythonObject, mode: String) raises -> TestResult: ctx.enqueue_copy(nonce_buffer, nonce_ptr) ctx.enqueue_function[aes_gpu_kernel_gcm_ctr]( - input_buffer.unsafe_ptr(), - output_buffer.unsafe_ptr(), - skey_buffer.unsafe_ptr(), - n_blocks, - nonce_buffer.unsafe_ptr(), - rounds, + input_buffer, + output_buffer, + skey_buffer, + Int32(n_blocks), + nonce_buffer, + Int32(rounds), grid_dim=grid_dim, block_dim=block_dim, ) else: ctx.enqueue_function[aes_gpu_kernel_ecb]( - input_buffer.unsafe_ptr(), - output_buffer.unsafe_ptr(), - skey_buffer.unsafe_ptr(), - n_blocks, - rounds, + input_buffer, + output_buffer, + skey_buffer, + Int32(n_blocks), + Int32(rounds), grid_dim=grid_dim, block_dim=block_dim, ) @@ -332,8 +300,6 @@ def test_mode_gpu(json_data: PythonObject, mode: String) raises -> TestResult: got_hex += byte_to_hex(ct_ptr.load(j)) failures.append(mode + " " + String(i) + ": expected " + expected_ct + ", got " + got_hex) - key_ptr.free() - round_keys.free() pt_ptr.free() ct_ptr.free() nonce_ptr.free() diff --git a/tests/test_random.mojo b/tests/test_random.mojo index 1b35254..e53bb99 100644 --- a/tests/test_random.mojo +++ b/tests/test_random.mojo @@ -35,7 +35,7 @@ def main() raises: var large = List[UInt8](capacity=257) for _ in range(257): large.append(0) - random_fill(large.unsafe_ptr(), 257) + random_fill(Span[mut=True, UInt8, ...](large)) if _all_zero(large): raise Error("random_fill(257) returned all-zero output") @@ -43,7 +43,7 @@ def main() raises: var second = List[UInt8](capacity=257) for _ in range(257): second.append(0) - random_fill(second.unsafe_ptr(), 257) + random_fill(Span[mut=True, UInt8, ...](second)) if second[256] == 0: raise Error("random_fill(257) did not appear to write byte 256") diff --git a/tests/test_security_boundaries.mojo b/tests/test_security_boundaries.mojo new file mode 100644 index 0000000..1209826 --- /dev/null +++ b/tests/test_security_boundaries.mojo @@ -0,0 +1,119 @@ +from std.collections import List + +from thistle.aes import AESExpandedKey +from thistle.argon2 import variable_length_hash_into +from thistle.chacha20poly1305 import ( + chacha20_poly1305_encrypt, + hchacha20, +) +from thistle.p256 import p256_ecdsa_sign, p256_public_key +from thistle.p384 import p384_ecdsa_sign, p384_public_key +from thistle.x25519 import x25519 + + +def main() raises: + var empty = List[UInt8]() + var short = List[UInt8](length=15, fill=0) + var rejected = False + try: + var key = AESExpandedKey(Span[UInt8, ...](short)) + _ = key + except: + rejected = True + if not rejected: + raise Error("AES accepted an undersized key") + + var argon_output = List[UInt8](length=4, fill=0) + rejected = False + try: + variable_length_hash_into( + 5, + Span[UInt8, ...](empty), + Span[mut=True, UInt8, ...](argon_output), + ) + except: + rejected = True + if not rejected: + raise Error("Argon2 accepted an undersized destination") + + var key32 = List[UInt8](length=32, fill=1) + var point32 = List[UInt8](length=32, fill=0) + point32[0] = 9 + var output31 = List[UInt8](length=31, fill=0) + rejected = False + try: + x25519( + Span[UInt8, ...](key32), + Span[UInt8, ...](point32), + Span[mut=True, UInt8, ...](output31), + ) + except: + rejected = True + if not rejected: + raise Error("X25519 accepted an undersized destination") + + var input16 = List[UInt8](length=16, fill=0) + rejected = False + try: + hchacha20( + Span[UInt8, ...](key32), + Span[UInt8, ...](input16), + Span[mut=True, UInt8, ...](output31), + ) + except: + rejected = True + if not rejected: + raise Error("HChaCha20 accepted an undersized destination") + + var nonce = List[UInt8](length=12, fill=0) + var plaintext = List[UInt8](length=16, fill=0) + var ciphertext = List[UInt8](length=15, fill=0) + var tag = List[UInt8](length=16, fill=0) + rejected = False + try: + chacha20_poly1305_encrypt( + Span[UInt8, ...](key32), + Span[UInt8, ...](nonce), + Span[UInt8, ...](empty), + Span[UInt8, ...](plaintext), + Span[mut=True, UInt8, ...](ciphertext), + Span[mut=True, UInt8, ...](tag), + ) + except: + rejected = True + if not rejected: + raise Error("ChaCha20-Poly1305 accepted an undersized destination") + + var p256_private = List[UInt8](length=32, fill=0) + p256_private[31] = 1 + var p384_private = List[UInt8](length=48, fill=0) + p384_private[47] = 1 + var p256_output = List[UInt8](length=64, fill=0) + var p384_output = List[UInt8](length=96, fill=0) + if p256_public_key( + Span[UInt8, ...](p256_private), + Span[mut=True, UInt8, ...](p256_output), + ): + raise Error("P-256 public-key API accepted an undersized destination") + if p384_public_key( + Span[UInt8, ...](p384_private), + Span[mut=True, UInt8, ...](p384_output), + ): + raise Error("P-384 public-key API accepted an undersized destination") + + var short_signature = List[UInt8](length=63, fill=0) + var short_signature384 = List[UInt8](length=95, fill=0) + if p256_ecdsa_sign( + Span[UInt8, ...](p256_private), + Span[UInt8, ...](plaintext), + Span[mut=True, UInt8, ...](short_signature), + ): + raise Error("P-256 signing accepted an undersized destination") + if p384_ecdsa_sign( + Span[UInt8, ...](p384_private), + Span[UInt8, ...](plaintext), + Span[mut=True, UInt8, ...](short_signature384), + ): + raise Error("P-384 signing accepted an undersized destination") + + print("Security boundary tests passed") diff --git a/tests/test_sha3_capacity_guard.mojo b/tests/test_sha3_capacity_guard.mojo new file mode 100644 index 0000000..7b2ff05 --- /dev/null +++ b/tests/test_sha3_capacity_guard.mojo @@ -0,0 +1,10 @@ +from std.collections import List + +from thistle.sha3 import shake256_into +from thistle.utils import StackBuffer + + +def main(): + var input = List[UInt8]() + var output = StackBuffer[UInt8, 32]() + shake256_into(output, Span[UInt8, ...](input), 33) diff --git a/tests/test_sha3_capacity_guard.sh b/tests/test_sha3_capacity_guard.sh new file mode 100644 index 0000000..cc45839 --- /dev/null +++ b/tests/test_sha3_capacity_guard.sh @@ -0,0 +1,18 @@ +#!/bin/zsh + +set -uo pipefail + +guard_dir=$(mktemp -d) +guard_binary="$guard_dir/sha3-capacity-guard" +guard_log="$guard_dir/run.log" +trap 'rm -rf -- "$guard_dir"' EXIT + +mojo build -O3 -I src/ tests/test_sha3_capacity_guard.mojo \ + -o "$guard_binary" >/dev/null 2>&1 + +if "$guard_binary" >"$guard_log" 2>&1; then + echo "SHAKE capacity guard did not reject an oversized output" >&2 + exit 1 +fi + +grep -q "SHAKE output length exceeds destination capacity" "$guard_log" diff --git a/tests/test_signing.mojo b/tests/test_signing.mojo index adb1602..998fe1a 100644 --- a/tests/test_signing.mojo +++ b/tests/test_signing.mojo @@ -102,7 +102,7 @@ def test_p256() raises: if not p256_ecdsa_sign( Span[UInt8, ...](private_key), Span[UInt8, ...](message), - signature.unsafe_ptr(), + Span[mut=True, UInt8, ...](signature), ): raise Error("P-256 signing failed") if not equal(signature, expected): @@ -158,7 +158,7 @@ def test_p384() raises: if not p384_ecdsa_sign( Span[UInt8, ...](private_key), Span[UInt8, ...](message), - signature.unsafe_ptr(), + Span[mut=True, UInt8, ...](signature), ): raise Error("P-384 signing failed") if not equal(signature, expected): @@ -193,7 +193,8 @@ def test_keygen() raises: var p256_public = p256_pair[1].copy() var p256_check = List[UInt8](unsafe_uninit_length=65) if not p256_public_key( - Span[UInt8, ...](p256_private), p256_check.unsafe_ptr() + Span[UInt8, ...](p256_private), + Span[mut=True, UInt8, ...](p256_check), ) or not equal(p256_public, p256_check): raise Error("P-256 key generation failed") @@ -202,7 +203,8 @@ def test_keygen() raises: var p384_public = p384_pair[1].copy() var p384_check = List[UInt8](unsafe_uninit_length=97) if not p384_public_key( - Span[UInt8, ...](p384_private), p384_check.unsafe_ptr() + Span[UInt8, ...](p384_private), + Span[mut=True, UInt8, ...](p384_check), ) or not equal(p384_public, p384_check): raise Error("P-384 key generation failed") @@ -210,7 +212,9 @@ def test_keygen() raises: var x_private = x_pair[0].copy() var x_public = x_pair[1].copy() var x_check = List[UInt8](unsafe_uninit_length=32) - x25519_public_key(Span[UInt8, ...](x_private), x_check.unsafe_ptr()) + x25519_public_key( + Span[UInt8, ...](x_private), Span[mut=True, UInt8, ...](x_check) + ) if not equal(x_public, x_check): raise Error("X25519 key generation failed") diff --git a/tests/test_wycheproof_p256_ecdh.mojo b/tests/test_wycheproof_p256_ecdh.mojo index 2f44c19..795c061 100644 --- a/tests/test_wycheproof_p256_ecdh.mojo +++ b/tests/test_wycheproof_p256_ecdh.mojo @@ -132,7 +132,7 @@ def run_case( var got = p256_ecdh( Span[UInt8, ...](private_key), Span[UInt8, ...](public_key), - actual.unsafe_ptr(), + Span[mut=True, UInt8, ...](unsafe_ptr=actual.unsafe_ptr(), length=32), ) if is_valid and not got: print( @@ -179,4 +179,4 @@ def main() raises: "Wycheproof P-256 ECDH: ", ok_count, " passed, ", fail_count, " failed" ) if fail_count > 0: - raise Error("Wycheproof P-256 ECDH failed") \ No newline at end of file + raise Error("Wycheproof P-256 ECDH failed") diff --git a/tests/test_wycheproof_p384_ecdh.mojo b/tests/test_wycheproof_p384_ecdh.mojo index e8065fa..b076856 100644 --- a/tests/test_wycheproof_p384_ecdh.mojo +++ b/tests/test_wycheproof_p384_ecdh.mojo @@ -125,7 +125,7 @@ def run_case( var got = p384_ecdh( Span[UInt8, ...](private_key), Span[UInt8, ...](public_key), - actual.unsafe_ptr(), + Span[mut=True, UInt8, ...](unsafe_ptr=actual.unsafe_ptr(), length=48), ) if is_valid and not got: print( @@ -172,4 +172,4 @@ def main() raises: "Wycheproof P-384 ECDH: ", ok_count, " passed, ", fail_count, " failed" ) if fail_count > 0: - raise Error("Wycheproof P-384 ECDH failed") \ No newline at end of file + raise Error("Wycheproof P-384 ECDH failed") diff --git a/tests/test_wycheproof_x25519.mojo b/tests/test_wycheproof_x25519.mojo index 7a5b984..e05b21c 100644 --- a/tests/test_wycheproof_x25519.mojo +++ b/tests/test_wycheproof_x25519.mojo @@ -29,7 +29,13 @@ def run_case(tc_id: String, private_hex: String, public_hex: String, shared_hex: var public_key = hex_to_bytes(public_hex) var expected = hex_to_bytes(shared_hex) var actual = StackInlineArray[UInt8, 32](uninitialized=True) - x25519(Span[UInt8, ...](private_key), Span[UInt8, ...](public_key), actual.unsafe_ptr()) + x25519( + Span[UInt8, ...](private_key), + Span[UInt8, ...](public_key), + Span[mut=True, UInt8, ...]( + unsafe_ptr=actual.unsafe_ptr(), length=32 + ), + ) if not matches32(actual, expected): print("Test ", tc_id, " mismatch") return False diff --git a/tests/thistle_test_vectors.mojo b/tests/thistle_test_vectors.mojo index 335f77a..63cc9bc 100644 --- a/tests/thistle_test_vectors.mojo +++ b/tests/thistle_test_vectors.mojo @@ -2,7 +2,6 @@ from std.python import Python, PythonObject from std.collections import List from std.memory import alloc from std.memory.unsafe_pointer import UnsafePointer -from std.builtin.type_aliases import MutUntrackedOrigin from thistle.sha2 import ( bytes_to_hex, string_to_bytes, @@ -34,7 +33,7 @@ from thistle.kcipher2 import KCipher2 from thistle.pbkdf2 import pbkdf2_hmac_sha256, pbkdf2_hmac_sha512 from thistle.aes import ( cpu_aes_encrypt, cpu_aes_ecb_kernel, cpu_aes_cbc_kernel, cpu_aes_ctr_kernel, - cpu_aes_xts_kernel, expand_key_128, expand_key_192, expand_key_256, + cpu_aes_xts_kernel, AESExpandedKey, ) from thistle.aes_ni import aes_encrypt, has_aes_ni, has_x86_aes_ni, aes_gcm_encrypt, aes_gcm_decrypt from thistle.sha_ni import sha256ni_hash, has_sha_ni @@ -87,7 +86,7 @@ def byte_to_hex(b: UInt8) -> String: ) -def ptr_to_hex(ptr: UnsafePointer[UInt8, MutAnyOrigin], count: Int) -> String: +def ptr_to_hex(ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], count: Int) -> String: var s = String("") for j in range(count): s += byte_to_hex(ptr.load(j)) @@ -108,8 +107,8 @@ def list_to_simd32(lst: List[UInt8]) -> SIMD[DType.uint8, 32]: return result -def list_to_simd12(lst: List[UInt8]) -> SIMD[DType.uint8, 12]: - var result = SIMD[DType.uint8, 12](0) +def list_to_simd12(lst: List[UInt8]) -> SIMD[DType.uint8, 16]: + var result = SIMD[DType.uint8, 16](0) for i in range(min(len(lst), 12)): result[i] = lst[i] return result @@ -304,7 +303,11 @@ def test_camellia(data: PythonObject, py: PythonObject) raises -> TestResult: if ctr_out[idx] != (msg[idx] ^ ks_blk[j]): ctr_ok = False camellia_ctr_kernel( - ctr_out.unsafe_ptr(), ctr_out.unsafe_ptr(), cipher, ctr_nb, nonce.unsafe_ptr() + ctr_out.unsafe_ptr().unsafe_origin_cast[MutAnyOrigin](), + ctr_out.unsafe_ptr().unsafe_origin_cast[MutAnyOrigin](), + cipher, + ctr_nb, + nonce.unsafe_ptr(), ) for i2 in range(ctr_nb * 16): if ctr_out[i2] != msg[i2]: @@ -337,7 +340,11 @@ def test_camellia(data: PythonObject, py: PythonObject) raises -> TestResult: if cbc_out[bi * 16 + j] != prev[j]: cbc_ok = False camellia_cbc_decrypt_kernel( - cbc_out.unsafe_ptr(), cbc_out.unsafe_ptr(), cipher, cbc_nb, nonce.unsafe_ptr() + cbc_out.unsafe_ptr().unsafe_origin_cast[MutAnyOrigin](), + cbc_out.unsafe_ptr().unsafe_origin_cast[MutAnyOrigin](), + cipher, + cbc_nb, + nonce.unsafe_ptr(), ) for i2 in range(cbc_nb * 16): if cbc_out[i2] != cbc_msg[i2]: @@ -368,9 +375,9 @@ def test_chacha20(data: PythonObject, py: PythonObject) raises -> TestResult: UInt8, MutUntrackedOrigin ].unsafe_dangling() var out_span = Span[mut=True, UInt8, MutUntrackedOrigin]( - ptr=null_ptr, length=0 + unsafe_ptr=null_ptr, length=0 ) - cipher.encrypt_into(Span[UInt8, ...](ptr=null_ptr, length=0), out_span) + cipher.encrypt_into(Span[UInt8, ...](unsafe_ptr=null_ptr, length=0), out_span) if len(expected_ct) == 0: passed += 1 else: @@ -381,7 +388,7 @@ def test_chacha20(data: PythonObject, py: PythonObject) raises -> TestResult: else: var ct_ptr = alloc[UInt8](len(pt_bytes)) var ct_span = Span[mut=True, UInt8, MutUntrackedOrigin]( - ptr=ct_ptr, length=len(pt_bytes) + unsafe_ptr=ct_ptr, length=len(pt_bytes) ) cipher.encrypt_into(Span[UInt8, ...](pt_bytes), ct_span) var ok = len(pt_bytes) == len(expected_ct) @@ -609,7 +616,7 @@ def test_aes_cpu(data: PythonObject, py: PythonObject) raises -> TestResult: for j in range(16): pt_ptr.store(j, pt_bytes[j]) var round_keys = _expand_aes_key(key_bytes) - cpu_aes_encrypt(pt_ptr, round_keys, rounds) + cpu_aes_encrypt(pt_ptr, round_keys.ptr(), rounds) var got = ptr_to_hex(pt_ptr, 16) if got == expected_ct: passed += 1 @@ -623,7 +630,6 @@ def test_aes_cpu(data: PythonObject, py: PythonObject) raises -> TestResult: + ", got " + got ) - round_keys.free() pt_ptr.free() return TestResult(passed, failed, failures^) @@ -647,7 +653,7 @@ def test_aes_ni(data: PythonObject, py: PythonObject) raises -> TestResult: for j in range(16): pt_ptr.store(j, pt_bytes[j]) var round_keys = _expand_aes_key(key_bytes) - aes_encrypt(pt_ptr, round_keys, rounds) + aes_encrypt(pt_ptr, round_keys.ptr(), rounds) var got = ptr_to_hex(pt_ptr, 16) if got == expected_ct: passed += 1 @@ -661,7 +667,6 @@ def test_aes_ni(data: PythonObject, py: PythonObject) raises -> TestResult: + ", got " + got ) - round_keys.free() pt_ptr.free() return TestResult(passed, failed, failures^) @@ -784,12 +789,16 @@ def test_chacha20_poly1305(data: PythonObject, py: PythonObject, xchacha: Bool) if xchacha: xchacha20_poly1305_encrypt( Span[UInt8, ...](key), Span[UInt8, ...](iv), Span[UInt8, ...](aad), - Span[UInt8, ...](msg), out_ct.unsafe_ptr(), out_tag.unsafe_ptr(), + Span[UInt8, ...](msg), + Span[mut=True, UInt8, ...](out_ct), + Span[mut=True, UInt8, ...](out_tag), ) else: chacha20_poly1305_encrypt( Span[UInt8, ...](key), Span[UInt8, ...](iv), Span[UInt8, ...](aad), - Span[UInt8, ...](msg), out_ct.unsafe_ptr(), out_tag.unsafe_ptr(), + Span[UInt8, ...](msg), + Span[mut=True, UInt8, ...](out_ct), + Span[mut=True, UInt8, ...](out_tag), ) enc_matches = len(ct) == n and len(tag) == 16 for i in range(n): @@ -806,12 +815,14 @@ def test_chacha20_poly1305(data: PythonObject, py: PythonObject, xchacha: Bool) if xchacha: dec_ok = xchacha20_poly1305_decrypt( Span[UInt8, ...](key), Span[UInt8, ...](iv), Span[UInt8, ...](aad), - Span[UInt8, ...](ct), Span[UInt8, ...](tag), out_pt.unsafe_ptr(), + Span[UInt8, ...](ct), Span[UInt8, ...](tag), + Span[mut=True, UInt8, ...](out_pt), ) else: dec_ok = chacha20_poly1305_decrypt( Span[UInt8, ...](key), Span[UInt8, ...](iv), Span[UInt8, ...](aad), - Span[UInt8, ...](ct), Span[UInt8, ...](tag), out_pt.unsafe_ptr(), + Span[UInt8, ...](ct), Span[UInt8, ...](tag), + Span[mut=True, UInt8, ...](out_pt), ) if dec_ok: for i in range(len(ct)): @@ -833,19 +844,8 @@ def test_chacha20_poly1305(data: PythonObject, py: PythonObject, xchacha: Bool) return TestResult(passed, failed, failures^) -def _expand_aes_key(key: List[UInt8]) raises -> UnsafePointer[UInt32, MutAnyOrigin]: - var kp = alloc[UInt8](len(key)) - for i in range(len(key)): - kp[i] = key[i] - var rk: UnsafePointer[UInt32, MutAnyOrigin] - if len(key) == 16: - rk = expand_key_128(kp) - elif len(key) == 24: - rk = expand_key_192(kp) - else: - rk = expand_key_256(kp) - kp.free() - return rk +def _expand_aes_key(key: List[UInt8]) raises -> AESExpandedKey: + return AESExpandedKey(Span[UInt8, ...](key)) def test_aes_cpu_modes(data: PythonObject, py: PythonObject) raises -> TestResult: @@ -896,9 +896,9 @@ def test_aes_cpu_modes(data: PythonObject, py: PythonObject) raises -> TestResul var twp = alloc[UInt8](16) for j in range(16): twp[j] = tweak[j] - cpu_aes_xts_kernel(ip, op, rk1, rk2, nblocks, twp, rounds) - rk1.free() - rk2.free() + cpu_aes_xts_kernel( + ip, op, rk1.ptr(), rk2.ptr(), nblocks, twp, rounds + ) twp.free() else: var rk = _expand_aes_key(key) @@ -908,18 +908,23 @@ def test_aes_cpu_modes(data: PythonObject, py: PythonObject) raises -> TestResul var ivp = alloc[UInt8](16) for j in range(16): ivp[j] = iv[j] - cpu_aes_cbc_kernel(ip, op, rk, nblocks, ivp, rounds) + cpu_aes_cbc_kernel( + ip, op, rk.ptr(), nblocks, ivp, rounds + ) ivp.free() elif "CTR" in mode: var iv = hex_to_bytes(String(tv["iv"])) var ivp = alloc[UInt8](16) for j in range(16): ivp[j] = iv[j] - cpu_aes_ctr_kernel(ip, op, rk, nblocks, ivp, rounds) + cpu_aes_ctr_kernel( + ip, op, rk.ptr(), nblocks, ivp, rounds + ) ivp.free() else: - cpu_aes_ecb_kernel(ip, op, rk, nblocks, rounds) - rk.free() + cpu_aes_ecb_kernel( + ip, op, rk.ptr(), nblocks, rounds + ) ok = ptr_to_hex(op, n) == bytes_to_hex(ct_exp) ip.free() From f48f2eaf6166019f8fb3a880610febc2db010e2c Mon Sep 17 00:00:00 2001 From: libalpm64 <159192189+libalpm64@users.noreply.github.com> Date: Thu, 20 Aug 2026 07:14:47 -0400 Subject: [PATCH 2/4] ~ --- src/thistle/aes.mojo | 218 ++++++++--------- src/thistle/aes_gpu.mojo | 70 +++--- src/thistle/aes_ni.mojo | 366 ++++++++++++++-------------- src/thistle/argon2.mojo | 292 +++++++++++----------- src/thistle/blake2b.mojo | 54 ++-- src/thistle/blake3.mojo | 38 +-- src/thistle/camellia.mojo | 124 +++++----- src/thistle/chacha20.mojo | 42 ++-- src/thistle/chacha20poly1305.mojo | 58 ++--- src/thistle/curve25519.mojo | 22 +- src/thistle/ed25519.mojo | 98 ++++---- src/thistle/ed25519_table.mojo | 66 ++--- src/thistle/kcipher2.mojo | 22 +- src/thistle/ml_dsa.mojo | 68 +++--- src/thistle/ml_kem.mojo | 84 +++---- src/thistle/p256.mojo | 52 ++-- src/thistle/p256_table.mojo | 64 ++--- src/thistle/p384.mojo | 52 ++-- src/thistle/p384_table.mojo | 136 +++++------ src/thistle/pbkdf2.mojo | 100 +++++--- src/thistle/poly1305.mojo | 38 +-- src/thistle/random.mojo | 20 +- src/thistle/rsa.mojo | 104 ++++---- src/thistle/sha2.mojo | 100 ++++---- src/thistle/sha3.mojo | 270 ++++++++++---------- src/thistle/sha_ni.mojo | 50 ++-- src/thistle/utils.mojo | 42 ++-- src/thistle/x25519.mojo | 2 +- tests/benchmark.mojo | 94 +++---- tests/dudect.mojo | 6 +- tests/gen_ed25519_table.mojo | 6 +- tests/gen_p256_table.mojo | 2 +- tests/gen_p384_table.mojo | 2 +- tests/test_aes_gpu.mojo | 54 ++-- tests/test_security_boundaries.mojo | 69 ++++++ tests/test_wycheproof_ed25519.mojo | 2 +- tests/test_wycheproof_x25519.mojo | 2 +- tests/thistle_test_vectors.mojo | 82 +++---- 38 files changed, 1535 insertions(+), 1436 deletions(-) diff --git a/src/thistle/aes.mojo b/src/thistle/aes.mojo index e25d823..b108c7f 100644 --- a/src/thistle/aes.mojo +++ b/src/thistle/aes.mojo @@ -3,7 +3,7 @@ AES CPU implementation """ from std.bit import byte_swap -from std.memory import memset_zero +from std.memory import unsafe_memset_zero from std.utils import StaticTuple from .utils import StackBuffer @@ -11,29 +11,29 @@ comptime ROUNDS_128: Int = 10 @always_inline def _ct_encrypt1( - block: UnsafePointer[mut=True, UInt8, _, address_space=_], + block: Pointer[mut=True, UInt8, _, address_space=_], skey: List[UInt64], rounds: Int, ) -> None: var buf = InlineArray[UInt8, 64](fill=0) var bp = buf.unsafe_ptr() for i in range(16): - bp[i] = block[i] + bp[unsafe_offset=i] = block[unsafe_offset=i] cpu_aes_ct_encrypt4(bp, skey, rounds) for i in range(16): - block[i] = bp[i] + block[unsafe_offset=i] = bp[unsafe_offset=i] @always_inline def cpu_aes_encrypt( - pt_bytes: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], + pt_bytes: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_], ) -> None: cpu_aes_encrypt(pt_bytes, round_keys, 10) @always_inline def cpu_aes_encrypt( - pt_bytes: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], + pt_bytes: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_], rounds: Int, ) -> None: var skey = cpu_aes_ct_skey(round_keys, rounds) @@ -41,9 +41,9 @@ def cpu_aes_encrypt( @always_inline def cpu_aes_ecb_kernel( - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, rounds: Int ) -> None: @@ -56,68 +56,68 @@ def cpu_aes_ecb_kernel( if n > 16: n = 16 for j in range(n * 16): - sp[j] = input_ptr[i * 16 + j] + sp[unsafe_offset=j] = input_ptr[unsafe_offset=i * 16 + j] cpu_aes_ct_encrypt16(sp, skey, rounds) for j in range(n * 16): - output_ptr[i * 16 + j] = sp[j] + output_ptr[unsafe_offset=i * 16 + j] = sp[unsafe_offset=j] i += n @always_inline def cpu_aes_cbc_kernel( - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - iv_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + iv_ptr: Pointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: var skey = cpu_aes_ct_skey(round_keys, rounds) var prev_block = StaticTuple[UInt8, 16]( - iv_ptr[0], iv_ptr[1], iv_ptr[2], iv_ptr[3], - iv_ptr[4], iv_ptr[5], iv_ptr[6], iv_ptr[7], - iv_ptr[8], iv_ptr[9], iv_ptr[10], iv_ptr[11], - iv_ptr[12], iv_ptr[13], iv_ptr[14], iv_ptr[15] + iv_ptr[unsafe_offset=0], iv_ptr[unsafe_offset=1], iv_ptr[unsafe_offset=2], iv_ptr[unsafe_offset=3], + iv_ptr[unsafe_offset=4], iv_ptr[unsafe_offset=5], iv_ptr[unsafe_offset=6], iv_ptr[unsafe_offset=7], + iv_ptr[unsafe_offset=8], iv_ptr[unsafe_offset=9], iv_ptr[unsafe_offset=10], iv_ptr[unsafe_offset=11], + iv_ptr[unsafe_offset=12], iv_ptr[unsafe_offset=13], iv_ptr[unsafe_offset=14], iv_ptr[unsafe_offset=15] ) var i = 0 while i < num_blocks: - var block_ptr = input_ptr + i * 16 - var out_ptr = output_ptr + i * 16 + var block_ptr = input_ptr.unsafe_offset(i * 16) + var out_ptr = output_ptr.unsafe_offset(i * 16) for j in range(16): - out_ptr.store(j, block_ptr.load(j) ^ prev_block[j]) + out_ptr.unsafe_store(j, block_ptr.unsafe_load(j) ^ prev_block[j]) _ct_encrypt1(out_ptr, skey, rounds) for j in range(16): - prev_block[j] = out_ptr.load(j) + prev_block[j] = out_ptr.unsafe_load(j) i += 1 @always_inline def _ctr_write_block( - dst: UnsafePointer[mut=True, UInt8, _, address_space=_], - nonce_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + dst: Pointer[mut=True, UInt8, _, address_space=_], + nonce_ptr: Pointer[mut=True, UInt8, _, address_space=_], offset: Int, ) -> None: for j in range(16): - dst.store(j, nonce_ptr[j]) + dst.unsafe_store(j, nonce_ptr[unsafe_offset=j]) var carry = UInt64(offset) for j in range(15, -1, -1): if carry == 0: break - var total = UInt64(dst.load(j)) + (carry & 0xFF) - dst.store(j, UInt8(total & 0xFF)) + var total = UInt64(dst.unsafe_load(j)) + (carry & 0xFF) + dst.unsafe_store(j, UInt8(total & 0xFF)) carry = (carry >> 8) + (total >> 8) @always_inline def cpu_aes_ctr_kernel( - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - nonce_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + nonce_ptr: Pointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: var skey = cpu_aes_ct_skey(round_keys, rounds) @@ -129,33 +129,33 @@ def cpu_aes_ctr_kernel( if n > 16: n = 16 for k in range(16): - _ctr_write_block(kp + k * 16, nonce_ptr, i + (k if k < n else 0)) + _ctr_write_block(kp.unsafe_offset(k * 16), nonce_ptr, i + (k if k < n else 0)) cpu_aes_ct_encrypt16(kp, skey, rounds) for k in range(n): - var in_block = input_ptr + (i + k) * 16 - var out_block = output_ptr + (i + k) * 16 + var in_block = input_ptr.unsafe_offset((i + k) * 16) + var out_block = output_ptr.unsafe_offset((i + k) * 16) for j in range(16): - out_block.store(j, in_block.load(j) ^ kp.load(k * 16 + j)) + out_block.unsafe_store(j, in_block.unsafe_load(j) ^ kp.unsafe_load(k * 16 + j)) i += n @always_inline -def cpu_xts_mul_alpha_inplace(tweak_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_]) -> None: - var carry = (tweak_ptr.load(15) & 0x80) != 0 +def cpu_xts_mul_alpha_inplace(tweak_ptr: Pointer[mut=True, UInt8, _, address_space=_]) -> None: + var carry = (tweak_ptr.unsafe_load(15) & 0x80) != 0 for i in range(15, 0, -1): - tweak_ptr.store(i, (tweak_ptr.load(i) << UInt8(1)) | (tweak_ptr.load(i - 1) >> UInt8(7))) - var t0 = tweak_ptr.load(0) << UInt8(1) + tweak_ptr.unsafe_store(i, (tweak_ptr.unsafe_load(i) << UInt8(1)) | (tweak_ptr.unsafe_load(i - 1) >> UInt8(7))) + var t0 = tweak_ptr.unsafe_load(0) << UInt8(1) if carry: t0 = t0 ^ UInt8(0x87) - tweak_ptr.store(0, t0) + tweak_ptr.unsafe_store(0, t0) @always_inline def cpu_aes_xts_kernel( - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys1: UnsafePointer[mut=True, UInt32, _, address_space=_], - round_keys2: UnsafePointer[mut=True, UInt32, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], + round_keys1: Pointer[mut=True, UInt32, _, address_space=_], + round_keys2: Pointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - tweak_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + tweak_ptr: Pointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: var skey1 = cpu_aes_ct_skey(round_keys1, rounds) @@ -163,24 +163,24 @@ def cpu_aes_xts_kernel( var tweak = StackBuffer[UInt8, 16]() var wp = tweak.ptr() for j in range(16): - wp.store(j, tweak_ptr[j]) + wp.unsafe_store(j, tweak_ptr[unsafe_offset=j]) _ct_encrypt1(wp, skey2, rounds) var i = 0 while i < num_blocks: - var in_block = input_ptr + i * 16 - var out_block = output_ptr + i * 16 + var in_block = input_ptr.unsafe_offset(i * 16) + var out_block = output_ptr.unsafe_offset(i * 16) var xored = StackBuffer[UInt8, 16]() var xp = xored.ptr() for j in range(16): - xp.store(j, in_block.load(j) ^ wp.load(j)) + xp.unsafe_store(j, in_block.unsafe_load(j) ^ wp.unsafe_load(j)) _ct_encrypt1(xp, skey1, rounds) for j in range(16): - out_block.store(j, xp.load(j) ^ wp.load(j)) + out_block.unsafe_store(j, xp.unsafe_load(j) ^ wp.unsafe_load(j)) cpu_xts_mul_alpha_inplace(wp) i += 1 @@ -189,10 +189,10 @@ def cpu_aes_xts_kernel( def sub_word(w: UInt32) -> UInt32: var blk = InlineArray[UInt8, 16](fill=0) var bp = blk.unsafe_ptr() - bp[0] = UInt8((w >> 24) & 0xff) - bp[1] = UInt8((w >> 16) & 0xff) - bp[2] = UInt8((w >> 8) & 0xff) - bp[3] = UInt8(w & 0xff) + bp[unsafe_offset=0] = UInt8((w >> 24) & 0xff) + bp[unsafe_offset=1] = UInt8((w >> 16) & 0xff) + bp[unsafe_offset=2] = UInt8((w >> 8) & 0xff) + bp[unsafe_offset=3] = UInt8(w & 0xff) var q = InlineArray[SIMD[DType.uint64, 1], 8](fill=0) var pair = _ct_interleave_in[1](_ct_le32(bp, 0), 0, 0, 0) q[0] = pair[0] @@ -212,57 +212,57 @@ comptime RCON: StaticTuple[UInt8, 11] = StaticTuple[UInt8, 11]( ) def expand_key_128_into( - key_bytes: UnsafePointer[mut=False, UInt8, _, address_space=_], - w: UnsafePointer[mut=True, UInt32, _, address_space=_], + key_bytes: Pointer[mut=False, UInt8, _, address_space=_], + w: Pointer[mut=True, UInt32, _, address_space=_], ) raises -> None: for i in range(4): var key_val: UInt32 = 0 for j in range(4): - key_val |= UInt32(key_bytes.load(i * 4 + j)) << UInt32((3 - j) * 8) - w.store(i, key_val) + key_val |= UInt32(key_bytes.unsafe_load(i * 4 + j)) << UInt32((3 - j) * 8) + w.unsafe_store(i, key_val) for i in range(4, 44): - var temp = w.load(i - 1) + var temp = w.unsafe_load(i - 1) if i % 4 == 0: var rotated = (temp >> 24) | ((temp << 8) & 0xffffffff) temp = sub_word(rotated) temp ^= UInt32(RCON._unsafe_ref(i // 4 - 1)) << 24 - w.store(i, w.load(i - 4) ^ temp) + w.unsafe_store(i, w.unsafe_load(i - 4) ^ temp) def expand_key_192_into( - key_bytes: UnsafePointer[mut=False, UInt8, _, address_space=_], - w: UnsafePointer[mut=True, UInt32, _, address_space=_], + key_bytes: Pointer[mut=False, UInt8, _, address_space=_], + w: Pointer[mut=True, UInt32, _, address_space=_], ) raises -> None: for i in range(6): var key_val: UInt32 = 0 for j in range(4): - key_val |= UInt32(key_bytes.load(i * 4 + j)) << UInt32((3 - j) * 8) - w.store(i, key_val) + key_val |= UInt32(key_bytes.unsafe_load(i * 4 + j)) << UInt32((3 - j) * 8) + w.unsafe_store(i, key_val) for i in range(6, 52): - var temp = w.load(i - 1) + var temp = w.unsafe_load(i - 1) if i % 6 == 0: var rotated = (temp >> 24) | ((temp << 8) & 0xffffffff) temp = sub_word(rotated) temp ^= UInt32(RCON._unsafe_ref(i // 6 - 1)) << 24 - w.store(i, w.load(i - 6) ^ temp) + w.unsafe_store(i, w.unsafe_load(i - 6) ^ temp) def expand_key_256_into( - key_bytes: UnsafePointer[mut=False, UInt8, _, address_space=_], - w: UnsafePointer[mut=True, UInt32, _, address_space=_], + key_bytes: Pointer[mut=False, UInt8, _, address_space=_], + w: Pointer[mut=True, UInt32, _, address_space=_], ) raises -> None: for i in range(8): var key_val: UInt32 = 0 for j in range(4): - key_val |= UInt32(key_bytes.load(i * 4 + j)) << UInt32((3 - j) * 8) - w.store(i, key_val) + key_val |= UInt32(key_bytes.unsafe_load(i * 4 + j)) << UInt32((3 - j) * 8) + w.unsafe_store(i, key_val) for i in range(8, 60): - var temp = w.load(i - 1) + var temp = w.unsafe_load(i - 1) if i % 8 == 0: var rotated = (temp >> 24) | ((temp << 8) & 0xffffffff) temp = sub_word(rotated) temp ^= UInt32(RCON._unsafe_ref(i // 8 - 1)) << 24 elif i % 8 == 4: temp = sub_word(temp) - w.store(i, w.load(i - 8) ^ temp) + w.unsafe_store(i, w.unsafe_load(i - 8) ^ temp) struct AESExpandedKey(Movable): """Owned AES-128/192/256 round-key schedule, wiped on destruction.""" @@ -285,19 +285,19 @@ struct AESExpandedKey(Movable): raise Error("AES keys must contain exactly 16, 24, or 32 bytes") def __deinit__(deinit self): - memset_zero(self._round_keys.ptr(), 60) + unsafe_memset_zero(self._round_keys.ptr(), 60) @always_inline def ptr[ origin: Origin, address_space: AddressSpace, // - ](ref[origin, address_space] self) -> UnsafePointer[ + ](ref[origin, address_space] self) -> Pointer[ UInt32, origin, address_space=address_space ]: return ( self._round_keys.ptr() .unsafe_mut_cast[origin.mut]() .unsafe_origin_cast[origin]() - .address_space_cast[address_space]() + .unsafe_address_space_cast[address_space]() ) @@ -325,24 +325,24 @@ struct AESKey: def __init__(out self, key: StaticTuple[UInt8, 16]) raises: self._data = StackBuffer[UInt8, 16]() for i in range(16): - self._data.ptr().store(i, key[i]) + self._data.ptr().unsafe_store(i, key[i]) self._round_keys = StackBuffer[UInt32, 44]() expand_key_128_into(self._data.ptr(), self._round_keys.ptr()) def __deinit__(deinit self): - memset_zero(self._data.ptr(), 16) - memset_zero(self._round_keys.ptr(), 44) + unsafe_memset_zero(self._data.ptr(), 16) + unsafe_memset_zero(self._round_keys.ptr(), 44) def round_keys[ origin: Origin, address_space: AddressSpace, // - ](ref[origin, address_space] self) -> UnsafePointer[ + ](ref[origin, address_space] self) -> Pointer[ UInt32, origin, address_space=address_space ]: return ( self._round_keys.ptr() .unsafe_mut_cast[origin.mut]() .unsafe_origin_cast[origin]() - .address_space_cast[address_space]() + .unsafe_address_space_cast[address_space]() ) @@ -628,14 +628,14 @@ def _ct_mix_columns[W: Int](mut q: InlineArray[SIMD[DType.uint64, W], 8]): def cpu_aes_ct_skey( - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], rounds: Int + round_keys: Pointer[mut=True, UInt32, _, address_space=_], rounds: Int ) -> List[UInt64]: var skey = List[UInt64](capacity=(rounds + 1) * 8) for r in range(rounds + 1): - var w0 = SIMD[DType.uint64, 1](UInt64(byte_swap(round_keys.load(r * 4)))) - var w1 = SIMD[DType.uint64, 1](UInt64(byte_swap(round_keys.load(r * 4 + 1)))) - var w2 = SIMD[DType.uint64, 1](UInt64(byte_swap(round_keys.load(r * 4 + 2)))) - var w3 = SIMD[DType.uint64, 1](UInt64(byte_swap(round_keys.load(r * 4 + 3)))) + var w0 = SIMD[DType.uint64, 1](UInt64(byte_swap(round_keys.unsafe_load(r * 4)))) + var w1 = SIMD[DType.uint64, 1](UInt64(byte_swap(round_keys.unsafe_load(r * 4 + 1)))) + var w2 = SIMD[DType.uint64, 1](UInt64(byte_swap(round_keys.unsafe_load(r * 4 + 2)))) + var w3 = SIMD[DType.uint64, 1](UInt64(byte_swap(round_keys.unsafe_load(r * 4 + 3)))) var q = InlineArray[SIMD[DType.uint64, 1], 8](fill=0) for i in range(4): var pair = _ct_interleave_in[1](w0, w1, w2, w3) @@ -648,27 +648,27 @@ def cpu_aes_ct_skey( @always_inline -def _ct_le32(p: UnsafePointer[mut=True, UInt8, _, address_space=_], off: Int) -> UInt64: +def _ct_le32(p: Pointer[mut=True, UInt8, _, address_space=_], off: Int) -> UInt64: return ( - UInt64(p.load(off)) - | (UInt64(p.load(off + 1)) << 8) - | (UInt64(p.load(off + 2)) << 16) - | (UInt64(p.load(off + 3)) << 24) + UInt64(p.unsafe_load(off)) + | (UInt64(p.unsafe_load(off + 1)) << 8) + | (UInt64(p.unsafe_load(off + 2)) << 16) + | (UInt64(p.unsafe_load(off + 3)) << 24) ) @always_inline -def _ct_store_le32(p: UnsafePointer[mut=True, UInt8, _, address_space=_], off: Int, w: UInt64): - p.store(off, UInt8(w & 0xFF)) - p.store(off + 1, UInt8((w >> 8) & 0xFF)) - p.store(off + 2, UInt8((w >> 16) & 0xFF)) - p.store(off + 3, UInt8((w >> 24) & 0xFF)) +def _ct_store_le32(p: Pointer[mut=True, UInt8, _, address_space=_], off: Int, w: UInt64): + p.unsafe_store(off, UInt8(w & 0xFF)) + p.unsafe_store(off + 1, UInt8((w >> 8) & 0xFF)) + p.unsafe_store(off + 2, UInt8((w >> 16) & 0xFF)) + p.unsafe_store(off + 3, UInt8((w >> 24) & 0xFF)) @always_inline def _ct_encrypt_blocks[W: Int]( - blocks: UnsafePointer[mut=True, UInt8, _, address_space=_], - skp: UnsafePointer[mut=False, UInt64, _, address_space=_], + blocks: Pointer[mut=True, UInt8, _, address_space=_], + skp: Pointer[mut=False, UInt64, _, address_space=_], rounds: Int, ) -> None: var q = InlineArray[SIMD[DType.uint64, W], 8](fill=0) @@ -689,17 +689,17 @@ def _ct_encrypt_blocks[W: Int]( _ct_ortho(q) comptime for i in range(8): - q[i] ^= SIMD[DType.uint64, W](skp[i]) + q[i] ^= SIMD[DType.uint64, W](skp[unsafe_offset=i]) for r in range(1, rounds): _ct_sbox(q) _ct_shift_rows(q) _ct_mix_columns(q) comptime for i in range(8): - q[i] ^= SIMD[DType.uint64, W](skp[r * 8 + i]) + q[i] ^= SIMD[DType.uint64, W](skp[unsafe_offset=r * 8 + i]) _ct_sbox(q) _ct_shift_rows(q) comptime for i in range(8): - q[i] ^= SIMD[DType.uint64, W](skp[rounds * 8 + i]) + q[i] ^= SIMD[DType.uint64, W](skp[unsafe_offset=rounds * 8 + i]) _ct_ortho(q) for i in range(4): @@ -713,7 +713,7 @@ def _ct_encrypt_blocks[W: Int]( def cpu_aes_ct_encrypt4( - blocks: UnsafePointer[mut=True, UInt8, _, address_space=_], + blocks: Pointer[mut=True, UInt8, _, address_space=_], skey: List[UInt64], rounds: Int, ) -> None: @@ -721,7 +721,7 @@ def cpu_aes_ct_encrypt4( def cpu_aes_ct_encrypt16( - blocks: UnsafePointer[mut=True, UInt8, _, address_space=_], + blocks: Pointer[mut=True, UInt8, _, address_space=_], skey: List[UInt64], rounds: Int, ) -> None: @@ -729,15 +729,15 @@ def cpu_aes_ct_encrypt16( def cpu_aes_ct_encrypt( - pt_bytes: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], + pt_bytes: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_], rounds: Int = 10, ) -> None: var skey = cpu_aes_ct_skey(round_keys, rounds) var buf = InlineArray[UInt8, 64](fill=0) var bp = buf.unsafe_ptr() for i in range(16): - bp[i] = pt_bytes[i] + bp[unsafe_offset=i] = pt_bytes[unsafe_offset=i] cpu_aes_ct_encrypt4(bp, skey, rounds) for i in range(16): - pt_bytes[i] = bp[i] + pt_bytes[unsafe_offset=i] = bp[unsafe_offset=i] diff --git a/src/thistle/aes_gpu.mojo b/src/thistle/aes_gpu.mojo index b0b87c3..64c0d6d 100644 --- a/src/thistle/aes_gpu.mojo +++ b/src/thistle/aes_gpu.mojo @@ -4,46 +4,46 @@ AES-GPU implementation from std.gpu import global_idx from std.memory import stack_allocation -from std.memory.unsafe_pointer import UnsafePointer +from std.memory.unsafe_pointer import Pointer from .aes import _ct_encrypt_blocks @always_inline -def add_counter_offset(counter: UnsafePointer[mut=True, UInt8, _, address_space=_], offset: Int) -> None: +def add_counter_offset(counter: Pointer[mut=True, UInt8, _, address_space=_], offset: Int) -> None: var carry = offset for i in range(15, -1, -1): if carry == 0: break - var old = counter[i] + var old = counter[unsafe_offset=i] var addend = UInt8(carry & 0xff) var new_val = old + addend - counter[i] = new_val + counter[unsafe_offset=i] = new_val carry = carry >> 8 if new_val < old: carry += 1 @always_inline def _gcm_counter_from_j0( - j0: UnsafePointer[mut=True, UInt8, _, address_space=_], + j0: Pointer[mut=True, UInt8, _, address_space=_], block_index: Int, - counter: UnsafePointer[mut=True, UInt8, _, address_space=_], + counter: Pointer[mut=True, UInt8, _, address_space=_], ) -> None: for i in range(12): - counter[i] = j0[i] + counter[unsafe_offset=i] = j0[unsafe_offset=i] var base = ( - (UInt32(j0[12]) << 24) | (UInt32(j0[13]) << 16) - | (UInt32(j0[14]) << 8) | UInt32(j0[15]) + (UInt32(j0[unsafe_offset=12]) << 24) | (UInt32(j0[unsafe_offset=13]) << 16) + | (UInt32(j0[unsafe_offset=14]) << 8) | UInt32(j0[unsafe_offset=15]) ) var ctr = base + UInt32(1) + UInt32(block_index & 0xFFFFFFFF) - counter[12] = UInt8((ctr >> 24) & 0xFF) - counter[13] = UInt8((ctr >> 16) & 0xFF) - counter[14] = UInt8((ctr >> 8) & 0xFF) - counter[15] = UInt8(ctr & 0xFF) + counter[unsafe_offset=12] = UInt8((ctr >> 24) & 0xFF) + counter[unsafe_offset=13] = UInt8((ctr >> 16) & 0xFF) + counter[unsafe_offset=14] = UInt8((ctr >> 8) & 0xFF) + counter[unsafe_offset=15] = UInt8(ctr & 0xFF) @always_inline def aes_gpu_kernel_ecb( - input_data: UnsafePointer[mut=True, UInt8, MutUntrackedOrigin], - output_data: UnsafePointer[mut=True, UInt8, MutUntrackedOrigin], - skey: UnsafePointer[mut=True, UInt64, MutUntrackedOrigin], + input_data: Pointer[mut=True, UInt8, MutUntrackedOrigin], + output_data: Pointer[mut=True, UInt8, MutUntrackedOrigin], + skey: Pointer[mut=True, UInt64, MutUntrackedOrigin], n: Int32, rounds: Int32, ) -> None: @@ -58,21 +58,21 @@ def aes_gpu_kernel_ecb( if b >= num_blocks: b = base_block for j in range(16): - buf[k * 16 + j] = input_data[b * 16 + j] + buf[unsafe_offset=k * 16 + j] = input_data[unsafe_offset=b * 16 + j] _ct_encrypt_blocks[1](buf, skey, Int(rounds)) for k in range(4): var blk = base_block + k if blk < num_blocks: for j in range(16): - output_data[blk * 16 + j] = buf[k * 16 + j] + output_data[unsafe_offset=blk * 16 + j] = buf[unsafe_offset=k * 16 + j] @always_inline def aes_gpu_kernel_ctr( - input_data: UnsafePointer[mut=True, UInt8, MutUntrackedOrigin], - output_data: UnsafePointer[mut=True, UInt8, MutUntrackedOrigin], - skey: UnsafePointer[mut=True, UInt64, MutUntrackedOrigin], + input_data: Pointer[mut=True, UInt8, MutUntrackedOrigin], + output_data: Pointer[mut=True, UInt8, MutUntrackedOrigin], + skey: Pointer[mut=True, UInt64, MutUntrackedOrigin], n: Int32, - nonce: UnsafePointer[mut=True, UInt8, MutUntrackedOrigin], + nonce: Pointer[mut=True, UInt8, MutUntrackedOrigin], rounds: Int32, ) -> None: var tid = global_idx.x @@ -85,26 +85,26 @@ def aes_gpu_kernel_ctr( var b = base_block + k if b >= num_blocks: b = base_block - var dst = buf + k * 16 + var dst = buf.unsafe_offset(k * 16) for j in range(16): - dst[j] = nonce[j] + dst[unsafe_offset=j] = nonce[unsafe_offset=j] add_counter_offset(dst, b) _ct_encrypt_blocks[1](buf, skey, Int(rounds)) for k in range(4): var blk = base_block + k if blk < num_blocks: - var bp = input_data + blk * 16 - var op = output_data + blk * 16 + var bp = input_data.unsafe_offset(blk * 16) + var op = output_data.unsafe_offset(blk * 16) for j in range(16): - op[j] = bp[j] ^ buf[k * 16 + j] + op[unsafe_offset=j] = bp[unsafe_offset=j] ^ buf[unsafe_offset=k * 16 + j] @always_inline def aes_gpu_kernel_gcm_ctr( - input_data: UnsafePointer[mut=True, UInt8, MutUntrackedOrigin], - output_data: UnsafePointer[mut=True, UInt8, MutUntrackedOrigin], - skey: UnsafePointer[mut=True, UInt64, MutUntrackedOrigin], + input_data: Pointer[mut=True, UInt8, MutUntrackedOrigin], + output_data: Pointer[mut=True, UInt8, MutUntrackedOrigin], + skey: Pointer[mut=True, UInt64, MutUntrackedOrigin], n: Int32, - j0: UnsafePointer[mut=True, UInt8, MutUntrackedOrigin], + j0: Pointer[mut=True, UInt8, MutUntrackedOrigin], rounds: Int32, ) -> None: var tid = global_idx.x @@ -117,12 +117,12 @@ def aes_gpu_kernel_gcm_ctr( var b = base_block + k if b >= num_blocks: b = base_block - _gcm_counter_from_j0(j0, b, buf + k * 16) + _gcm_counter_from_j0(j0, b, buf.unsafe_offset(k * 16)) _ct_encrypt_blocks[1](buf, skey, Int(rounds)) for k in range(4): var blk = base_block + k if blk < num_blocks: - var bp = input_data + blk * 16 - var op = output_data + blk * 16 + var bp = input_data.unsafe_offset(blk * 16) + var op = output_data.unsafe_offset(blk * 16) for j in range(16): - op[j] = bp[j] ^ buf[k * 16 + j] + op[unsafe_offset=j] = bp[unsafe_offset=j] ^ buf[unsafe_offset=k * 16 + j] diff --git a/src/thistle/aes_ni.mojo b/src/thistle/aes_ni.mojo index 32c480b..871e296 100644 --- a/src/thistle/aes_ni.mojo +++ b/src/thistle/aes_ni.mojo @@ -4,7 +4,7 @@ AES-NI implementation from std.collections import List, InlineArray from std.sys import llvm_intrinsic, CompilationTarget -from std.memory import bitcast, memset_zero, memcpy, UnsafePointer +from std.memory import bitcast, unsafe_memset_zero, unsafe_memcpy, Pointer from std.utils import StaticTuple from .aes import cpu_aes_encrypt, cpu_aes_ct_encrypt, cpu_aes_ct_encrypt16, cpu_aes_ct_skey, expand_key_128_into, expand_key_192_into, expand_key_256_into from .utils import StackBuffer, load_64be, store_64be @@ -59,39 +59,39 @@ def _mm_aesenclast_si128(lhs: SIMD128, rhs: SIMD128) -> SIMD128: return SIMD128(0) @always_inline -def _mm_loadu_si128(ptr: UnsafePointer[mut=True, UInt8, _, address_space=_]) -> SIMD128: - return ptr.bitcast[UInt64]().load[width=2, alignment=1]() +def _mm_loadu_si128(ptr: Pointer[mut=True, UInt8, _, address_space=_]) -> SIMD128: + return ptr.unsafe_bitcast[UInt64]().unsafe_load[width=2, alignment=1]() @always_inline -def _mm_storeu_si128(ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], data: SIMD128) -> None: +def _mm_storeu_si128(ptr: Pointer[mut=True, UInt8, _, address_space=_], data: SIMD128) -> None: var bytes: SIMD[DType.uint8, 16] = bitcast[DType.uint8, 16](data) - ptr.store[width=16, alignment=1](0, bytes) + ptr.unsafe_store[width=16, alignment=1](0, bytes) @always_inline def _write_gcm_counter( - counter_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - j0_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + counter_ptr: Pointer[mut=True, UInt8, _, address_space=_], + j0_ptr: Pointer[mut=True, UInt8, _, address_space=_], block_index: Int, ) -> None: for j in range(12): - counter_ptr.store(j, j0_ptr.load(j)) + counter_ptr.unsafe_store(j, j0_ptr.unsafe_load(j)) var base = ( - (UInt32(j0_ptr.load(12)) << 24) | (UInt32(j0_ptr.load(13)) << 16) - | (UInt32(j0_ptr.load(14)) << 8) | UInt32(j0_ptr.load(15)) + (UInt32(j0_ptr.unsafe_load(12)) << 24) | (UInt32(j0_ptr.unsafe_load(13)) << 16) + | (UInt32(j0_ptr.unsafe_load(14)) << 8) | UInt32(j0_ptr.unsafe_load(15)) ) var ctr = base + UInt32(1) + UInt32(block_index & 0xFFFFFFFF) - counter_ptr.store(12, UInt8((ctr >> 24) & 0xFF)) - counter_ptr.store(13, UInt8((ctr >> 16) & 0xFF)) - counter_ptr.store(14, UInt8((ctr >> 8) & 0xFF)) - counter_ptr.store(15, UInt8(ctr & 0xFF)) + counter_ptr.unsafe_store(12, UInt8((ctr >> 24) & 0xFF)) + counter_ptr.unsafe_store(13, UInt8((ctr >> 16) & 0xFF)) + counter_ptr.unsafe_store(14, UInt8((ctr >> 8) & 0xFF)) + counter_ptr.unsafe_store(15, UInt8(ctr & 0xFF)) @always_inline -def _load_round_key(idx: Int, round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_]) -> SIMD128: - var w0 = round_keys.load(idx * 4) - var w1 = round_keys.load(idx * 4 + 1) - var w2 = round_keys.load(idx * 4 + 2) - var w3 = round_keys.load(idx * 4 + 3) +def _load_round_key(idx: Int, round_keys: Pointer[mut=True, UInt32, _, address_space=_]) -> SIMD128: + var w0 = round_keys.unsafe_load(idx * 4) + var w1 = round_keys.unsafe_load(idx * 4 + 1) + var w2 = round_keys.unsafe_load(idx * 4 + 2) + var w3 = round_keys.unsafe_load(idx * 4 + 3) var bytes = SIMD[DType.uint8, 16]( UInt8(w0 >> 24), UInt8(w0 >> 16), UInt8(w0 >> 8), UInt8(w0), UInt8(w1 >> 24), UInt8(w1 >> 16), UInt8(w1 >> 8), UInt8(w1), @@ -102,8 +102,8 @@ def _load_round_key(idx: Int, round_keys: UnsafePointer[mut=True, UInt32, _, add def x86_aes_encrypt_128( - pt: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_] + pt: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_] ) -> None: var state = _mm_loadu_si128(pt) x86_aes_encrypt_128_direct(state, round_keys) @@ -111,8 +111,8 @@ def x86_aes_encrypt_128( def x86_aes_encrypt_192( - pt: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_] + pt: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_] ) -> None: var state = _mm_loadu_si128(pt) x86_aes_encrypt_192_direct(state, round_keys) @@ -120,8 +120,8 @@ def x86_aes_encrypt_192( def x86_aes_encrypt_256( - pt: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_] + pt: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_] ) -> None: var state = _mm_loadu_si128(pt) x86_aes_encrypt_256_direct(state, round_keys) @@ -130,7 +130,7 @@ def x86_aes_encrypt_256( @always_inline def x86_aes_encrypt_128_direct( mut state: SIMD128, - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_] + round_keys: Pointer[mut=True, UInt32, _, address_space=_] ) -> None: var keys = StaticTuple[SIMD128, 11]() comptime for i in range(11): @@ -144,7 +144,7 @@ def x86_aes_encrypt_128_direct( @always_inline def x86_aes_encrypt_192_direct( mut state: SIMD128, - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_] + round_keys: Pointer[mut=True, UInt32, _, address_space=_] ) -> None: var keys = StaticTuple[SIMD128, 13]() comptime for i in range(13): @@ -158,7 +158,7 @@ def x86_aes_encrypt_192_direct( @always_inline def x86_aes_encrypt_256_direct( mut state: SIMD128, - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_] + round_keys: Pointer[mut=True, UInt32, _, address_space=_] ) -> None: var keys = StaticTuple[SIMD128, 15]() comptime for i in range(15): @@ -172,11 +172,11 @@ def x86_aes_encrypt_256_direct( @always_inline def _arm_load_keys[N: Int]( - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_] + round_keys: Pointer[mut=True, UInt32, _, address_space=_] ) -> InlineArray[SIMD16, N]: var keys = InlineArray[SIMD16, N](uninitialized=True) comptime for i in range(N): - var raw = (round_keys + i * 4).bitcast[UInt8]().load[ + var raw = (round_keys.unsafe_offset(i * 4)).unsafe_bitcast[UInt8]().unsafe_load[ width=16, alignment=1 ]() keys[i] = raw.shuffle[ @@ -197,46 +197,46 @@ def _arm_enc_block[NR: Int]( def arm_aes_encrypt_128( - pt: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_] + pt: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_] ) -> None: var keys = _arm_load_keys[11](round_keys) - var x = pt.load[width=16, alignment=1](0) - pt.store[alignment=1](0, _arm_enc_block[10](x, keys)) + var x = pt.unsafe_load[width=16, alignment=1](0) + pt.unsafe_store[alignment=1](0, _arm_enc_block[10](x, keys)) def arm_aes_encrypt_192( - pt: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_] + pt: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_] ) -> None: var keys = _arm_load_keys[13](round_keys) - var x = pt.load[width=16, alignment=1](0) - pt.store[alignment=1](0, _arm_enc_block[12](x, keys)) + var x = pt.unsafe_load[width=16, alignment=1](0) + pt.unsafe_store[alignment=1](0, _arm_enc_block[12](x, keys)) def arm_aes_encrypt_256( - pt: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_] + pt: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_] ) -> None: var keys = _arm_load_keys[15](round_keys) - var x = pt.load[width=16, alignment=1](0) - pt.store[alignment=1](0, _arm_enc_block[14](x, keys)) + var x = pt.unsafe_load[width=16, alignment=1](0) + pt.unsafe_store[alignment=1](0, _arm_enc_block[14](x, keys)) @always_inline def _arm_ecb_loop[NR: Int]( - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, ) -> None: var keys = _arm_load_keys[NR + 1](round_keys) var i = 0 while i + 4 <= num_blocks: var p = input_ptr + i * 16 - var b0 = p.load[width=16, alignment=1](0) - var b1 = p.load[width=16, alignment=1](16) - var b2 = p.load[width=16, alignment=1](32) - var b3 = p.load[width=16, alignment=1](48) + var b0 = p.unsafe_load[width=16, alignment=1](0) + var b1 = p.unsafe_load[width=16, alignment=1](16) + var b2 = p.unsafe_load[width=16, alignment=1](32) + var b3 = p.unsafe_load[width=16, alignment=1](48) comptime for r in range(NR - 1): b0 = _aesmc(_aese(b0, keys[r])) b1 = _aesmc(_aese(b1, keys[r])) @@ -247,22 +247,22 @@ def _arm_ecb_loop[NR: Int]( b2 = _aese(b2, keys[NR - 1]) ^ keys[NR] b3 = _aese(b3, keys[NR - 1]) ^ keys[NR] var q = output_ptr + i * 16 - q.store[alignment=1](0, b0) - q.store[alignment=1](16, b1) - q.store[alignment=1](32, b2) - q.store[alignment=1](48, b3) + q.unsafe_store[alignment=1](0, b0) + q.unsafe_store[alignment=1](16, b1) + q.unsafe_store[alignment=1](32, b2) + q.unsafe_store[alignment=1](48, b3) i += 4 while i < num_blocks: - var x = (input_ptr + i * 16).load[width=16, alignment=1](0) - (output_ptr + i * 16).store[alignment=1](0, _arm_enc_block[NR](x, keys)) + var x = (input_ptr + i * 16).unsafe_load[width=16, alignment=1](0) + (output_ptr + i * 16).unsafe_store[alignment=1](0, _arm_enc_block[NR](x, keys)) i += 1 @always_inline def arm_aes_ecb_kernel( - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, rounds: Int ) -> None: @@ -276,21 +276,21 @@ def arm_aes_ecb_kernel( @always_inline def _arm_cbc_loop[NR: Int]( - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - iv_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + iv_ptr: Pointer[mut=True, UInt8, _, address_space=_], ) -> None: var keys = _arm_load_keys[NR + 1](round_keys) - var prev = iv_ptr.load[width=16, alignment=1](0) + var prev = iv_ptr.unsafe_load[width=16, alignment=1](0) var src = input_ptr var dst = output_ptr var i = 0 while i < num_blocks: - var x = src.load[width=16, alignment=1](0) ^ prev + var x = src.unsafe_load[width=16, alignment=1](0) ^ prev prev = _arm_enc_block[NR](x, keys) - dst.store[alignment=1](0, prev) + dst.unsafe_store[alignment=1](0, prev) src += 16 dst += 16 i += 1 @@ -298,11 +298,11 @@ def _arm_cbc_loop[NR: Int]( @always_inline def arm_aes_cbc_kernel( - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - iv_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + iv_ptr: Pointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: if rounds == 10: @@ -314,15 +314,15 @@ def arm_aes_cbc_kernel( @always_inline def arm_aes_xts_kernel( - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys1: UnsafePointer[mut=True, UInt32, _, address_space=_], - round_keys2: UnsafePointer[mut=True, UInt32, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], + round_keys1: Pointer[mut=True, UInt32, _, address_space=_], + round_keys2: Pointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - tweak_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + tweak_ptr: Pointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: - var tweak = tweak_ptr.load[width=16, alignment=1](0) + var tweak = tweak_ptr.unsafe_load[width=16, alignment=1](0) if rounds == 10: tweak = _arm_enc_block[10](tweak, _arm_load_keys[11](round_keys2)) else: @@ -331,9 +331,9 @@ def arm_aes_xts_kernel( var keys = _arm_load_keys[11](round_keys1) var i = 0 while i < num_blocks: - var x = (input_ptr + i * 16).load[width=16, alignment=1](0) ^ tweak + var x = (input_ptr + i * 16).unsafe_load[width=16, alignment=1](0) ^ tweak x = _arm_enc_block[10](x, keys) ^ tweak - (output_ptr + i * 16).store[alignment=1](0, x) + (output_ptr + i * 16).unsafe_store[alignment=1](0, x) tweak = bitcast[DType.uint8, 16]( _gf_mul2_xts_simd(bitcast[DType.uint64, 2](tweak)) ) @@ -342,9 +342,9 @@ def arm_aes_xts_kernel( var keys = _arm_load_keys[15](round_keys1) var i = 0 while i < num_blocks: - var x = (input_ptr + i * 16).load[width=16, alignment=1](0) ^ tweak + var x = (input_ptr + i * 16).unsafe_load[width=16, alignment=1](0) ^ tweak x = _arm_enc_block[14](x, keys) ^ tweak - (output_ptr + i * 16).store[alignment=1](0, x) + (output_ptr + i * 16).unsafe_store[alignment=1](0, x) tweak = bitcast[DType.uint8, 16]( _gf_mul2_xts_simd(bitcast[DType.uint64, 2](tweak)) ) @@ -352,9 +352,9 @@ def arm_aes_xts_kernel( @always_inline def x86_aes_ecb_kernel( - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, rounds: Int ) -> None: @@ -374,11 +374,11 @@ def x86_aes_ecb_kernel( @always_inline def x86_aes_cbc_kernel( - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - iv_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + iv_ptr: Pointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: var prev_block = _mm_loadu_si128(iv_ptr) @@ -409,12 +409,12 @@ def _gf_mul2_xts_simd(val: SIMD128) -> SIMD128: @always_inline def x86_aes_xts_kernel( - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys1: UnsafePointer[mut=True, UInt32, _, address_space=_], - round_keys2: UnsafePointer[mut=True, UInt32, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], + round_keys1: Pointer[mut=True, UInt32, _, address_space=_], + round_keys2: Pointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - tweak_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + tweak_ptr: Pointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: var tweak = _mm_loadu_si128(tweak_ptr) @@ -442,8 +442,8 @@ def x86_aes_xts_kernel( @always_inline def aes_encrypt( - pt: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], + pt: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_], rounds: Int = 10 ) -> None: """Unified AES encryption function. @@ -471,11 +471,11 @@ def aes_encrypt( @always_inline def aes_gcm_ctr_kernel( - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - j0_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + j0_ptr: Pointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: comptime if CompilationTarget._has_feature["sse"]() and CompilationTarget._has_feature["aes"](): @@ -489,11 +489,11 @@ def aes_gcm_ctr_kernel( @always_inline def _hw_gcm_ctr_kernel( - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - j0_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + j0_ptr: Pointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: comptime if has_arm_crypto(): @@ -519,21 +519,21 @@ def _hw_gcm_ctr_kernel( _write_gcm_counter(cp, j0_ptr, i) aes_encrypt(cp, round_keys, rounds) - var in_block = input_ptr + i * 16 - var out_block = output_ptr + i * 16 + var in_block = input_ptr.unsafe_offset(i * 16) + var out_block = output_ptr.unsafe_offset(i * 16) for j in range(16): - out_block.store(j, in_block.load(j) ^ cp.load(j)) + out_block.unsafe_store(j, in_block.unsafe_load(j) ^ cp.unsafe_load(j)) i += 1 @always_inline def _arm_gcm_ctr_loop[NR: Int]( - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - j0_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + j0_ptr: Pointer[mut=True, UInt8, _, address_space=_], ) -> None: var keys = _arm_load_keys[NR + 1](round_keys) var ctr = StackBuffer[UInt8, 64]() @@ -541,13 +541,13 @@ def _arm_gcm_ctr_loop[NR: Int]( var i = 0 while i + 4 <= num_blocks: _write_gcm_counter(cp, j0_ptr, i) - _write_gcm_counter(cp + 16, j0_ptr, i + 1) - _write_gcm_counter(cp + 32, j0_ptr, i + 2) - _write_gcm_counter(cp + 48, j0_ptr, i + 3) - var b0 = cp.load[width=16, alignment=1](0) - var b1 = cp.load[width=16, alignment=1](16) - var b2 = cp.load[width=16, alignment=1](32) - var b3 = cp.load[width=16, alignment=1](48) + _write_gcm_counter(cp.unsafe_offset(16), j0_ptr, i + 1) + _write_gcm_counter(cp.unsafe_offset(32), j0_ptr, i + 2) + _write_gcm_counter(cp.unsafe_offset(48), j0_ptr, i + 3) + var b0 = cp.unsafe_load[width=16, alignment=1](0) + var b1 = cp.unsafe_load[width=16, alignment=1](16) + var b2 = cp.unsafe_load[width=16, alignment=1](32) + var b3 = cp.unsafe_load[width=16, alignment=1](48) comptime for r in range(NR - 1): b0 = _aesmc(_aese(b0, keys[r])) b1 = _aesmc(_aese(b1, keys[r])) @@ -557,36 +557,36 @@ def _arm_gcm_ctr_loop[NR: Int]( b1 = _aese(b1, keys[NR - 1]) ^ keys[NR] b2 = _aese(b2, keys[NR - 1]) ^ keys[NR] b3 = _aese(b3, keys[NR - 1]) ^ keys[NR] - var p = input_ptr + i * 16 - var q = output_ptr + i * 16 - q.store[alignment=1](0, p.load[width=16, alignment=1](0) ^ b0) - q.store[alignment=1](16, p.load[width=16, alignment=1](16) ^ b1) - q.store[alignment=1](32, p.load[width=16, alignment=1](32) ^ b2) - q.store[alignment=1](48, p.load[width=16, alignment=1](48) ^ b3) + var p = input_ptr.unsafe_offset(i * 16) + var q = output_ptr.unsafe_offset(i * 16) + q.unsafe_store[alignment=1](0, p.unsafe_load[width=16, alignment=1](0) ^ b0) + q.unsafe_store[alignment=1](16, p.unsafe_load[width=16, alignment=1](16) ^ b1) + q.unsafe_store[alignment=1](32, p.unsafe_load[width=16, alignment=1](32) ^ b2) + q.unsafe_store[alignment=1](48, p.unsafe_load[width=16, alignment=1](48) ^ b3) i += 4 while i < num_blocks: _write_gcm_counter(cp, j0_ptr, i) - var ks = _arm_enc_block[NR](cp.load[width=16, alignment=1](0), keys) - (output_ptr + i * 16).store[alignment=1]( - 0, (input_ptr + i * 16).load[width=16, alignment=1](0) ^ ks + var ks = _arm_enc_block[NR](cp.unsafe_load[width=16, alignment=1](0), keys) + (output_ptr.unsafe_offset(i * 16)).unsafe_store[alignment=1]( + 0, (input_ptr.unsafe_offset(i * 16)).unsafe_load[width=16, alignment=1](0) ^ ks ) i += 1 @always_inline def _arm_gcm_fused_loop[NR: Int]( - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - j0_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + j0_ptr: Pointer[mut=True, UInt8, _, address_space=_], mut gh: _GHash, ghash_ciphertext: Bool, ) -> None: var keys = _arm_load_keys[NR + 1](round_keys) var y = SIMD128(_bitrev64(gh.y_hi), _bitrev64(gh.y_lo)) - var j0u32 = bitcast[DType.uint32, 4](j0_ptr.load[width=16, alignment=1](0)) + var j0u32 = bitcast[DType.uint32, 4](j0_ptr.unsafe_load[width=16, alignment=1](0)) var ctr_base = llvm_intrinsic["llvm.bswap.i32", UInt32, has_side_effect=False](j0u32[3]) + 1 var i = 0 @@ -604,13 +604,13 @@ def _arm_gcm_fused_loop[NR: Int]( comptime for k in range(8): b[k] = _aese(b[k], keys[NR - 1]) ^ keys[NR] - var p = input_ptr + i * 16 - var q = output_ptr + i * 16 + var p = input_ptr.unsafe_offset(i * 16) + var q = output_ptr.unsafe_offset(i * 16) var g = InlineArray[SIMD16, 8](uninitialized=True) comptime for k in range(8): - var pt = p.load[width=16, alignment=1](k * 16) + var pt = p.unsafe_load[width=16, alignment=1](k * 16) var ct = pt ^ b[k] - q.store[alignment=1](k * 16, ct) + q.unsafe_store[alignment=1](k * 16, ct) g[k] = ct if ghash_ciphertext else pt var lo = SIMD128(0) @@ -640,13 +640,13 @@ def _arm_gcm_fused_loop[NR: Int]( comptime for k in range(4): b[k] = _aese(b[k], keys[NR - 1]) ^ keys[NR] - var p = input_ptr + i * 16 - var q = output_ptr + i * 16 + var p = input_ptr.unsafe_offset(i * 16) + var q = output_ptr.unsafe_offset(i * 16) var g = InlineArray[SIMD16, 4](uninitialized=True) comptime for k in range(4): - var pt = p.load[width=16, alignment=1](k * 16) + var pt = p.unsafe_load[width=16, alignment=1](k * 16) var ct = pt ^ b[k] - q.store[alignment=1](k * 16, ct) + q.unsafe_store[alignment=1](k * 16, ct) g[k] = ct if ghash_ciphertext else pt var lo = SIMD128(0) @@ -664,9 +664,9 @@ def _arm_gcm_fused_loop[NR: Int]( ctr_base + UInt32(i) ) var ks = _arm_enc_block[NR](bitcast[DType.uint8, 16](cv), keys) - var p0 = (input_ptr + i * 16).load[width=16, alignment=1](0) + var p0 = (input_ptr.unsafe_offset(i * 16)).unsafe_load[width=16, alignment=1](0) var c0 = p0 ^ ks - (output_ptr + i * 16).store[alignment=1](0, c0) + (output_ptr.unsafe_offset(i * 16)).unsafe_store[alignment=1](0, c0) var g0 = c0 if ghash_ciphertext else p0 y = _gf_mul_nat(_rev128(g0) ^ y, gh.hn) i += 1 @@ -677,11 +677,11 @@ def _arm_gcm_fused_loop[NR: Int]( @always_inline def _soft_gcm_ctr_kernel( - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - round_keys: UnsafePointer[mut=True, UInt32, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], + round_keys: Pointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - j0_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + j0_ptr: Pointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: var skey = cpu_aes_ct_skey(round_keys, rounds) @@ -693,13 +693,13 @@ def _soft_gcm_ctr_kernel( if n > 16: n = 16 for k in range(16): - _write_gcm_counter(kp + k * 16, j0_ptr, i + (k if k < n else 0)) + _write_gcm_counter(kp.unsafe_offset(k * 16), j0_ptr, i + (k if k < n else 0)) cpu_aes_ct_encrypt16(kp, skey, rounds) for k in range(n): - var in_block = input_ptr + (i + k) * 16 - var out_block = output_ptr + (i + k) * 16 + var in_block = input_ptr.unsafe_offset((i + k) * 16) + var out_block = output_ptr.unsafe_offset((i + k) * 16) for j in range(16): - out_block.store(j, in_block.load(j) ^ kp.load(k * 16 + j)) + out_block.unsafe_store(j, in_block.unsafe_load(j) ^ kp.unsafe_load(k * 16 + j)) i += n @always_inline @@ -734,8 +734,8 @@ def _rev128(v: SIMD16) -> SIMD128: @always_inline("nodebug") -def _load_nat128(p: UnsafePointer[mut=True, UInt8, _, address_space=_]) -> SIMD128: - return _rev128(p.load[width=16, alignment=1](0)) +def _load_nat128(p: Pointer[mut=True, UInt8, _, address_space=_]) -> SIMD128: + return _rev128(p.unsafe_load[width=16, alignment=1](0)) comptime _ZERO128 = SIMD128(0) @@ -843,7 +843,7 @@ struct _GHash(Copyable, Movable): self.y_lo = z_lo @always_inline - def update(mut self, data: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int): + def update(mut self, data: Pointer[mut=True, UInt8, _, address_space=_], length: Int): comptime if CompilationTarget.has_neon() and CompilationTarget._has_feature["aes"]() and not CompilationTarget.is_x86(): self._update_pmull(data, length) return @@ -851,35 +851,35 @@ struct _GHash(Copyable, Movable): self._update_soft(data, length) @always_inline - def _update_pmull(mut self, data: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int): + def _update_pmull(mut self, data: Pointer[mut=True, UInt8, _, address_space=_], length: Int): var y = SIMD128(_bitrev64(self.y_hi), _bitrev64(self.y_lo)) var off = 0 while off + 64 <= length: var lo = SIMD128(0) var hi = SIMD128(0) - _clmul_acc(_load_nat128(data + off) ^ y, self.hn4, lo, hi) - _clmul_acc(_load_nat128(data + off + 16), self.hn3, lo, hi) - _clmul_acc(_load_nat128(data + off + 32), self.hn2, lo, hi) - _clmul_acc(_load_nat128(data + off + 48), self.hn, lo, hi) + _clmul_acc(_load_nat128(data.unsafe_offset(off)) ^ y, self.hn4, lo, hi) + _clmul_acc(_load_nat128(data.unsafe_offset(off).unsafe_offset(16)), self.hn3, lo, hi) + _clmul_acc(_load_nat128(data.unsafe_offset(off).unsafe_offset(32)), self.hn2, lo, hi) + _clmul_acc(_load_nat128(data.unsafe_offset(off).unsafe_offset(48)), self.hn, lo, hi) y = _reduce_vec(lo, hi) off += 64 while off + 16 <= length: - y = _gf_mul_nat(_load_nat128(data + off) ^ y, self.hn) + y = _gf_mul_nat(_load_nat128(data.unsafe_offset(off)) ^ y, self.hn) off += 16 if off < length: var block = InlineArray[UInt8, 16](fill=0) for i in range(length - off): - block[i] = data[off + i] + block[i] = data[unsafe_offset=off + i] y = _gf_mul_nat(_load_nat128(block.unsafe_ptr()) ^ y, self.hn) self.y_hi = _bitrev64(y[0]) self.y_lo = _bitrev64(y[1]) @always_inline - def _update_soft(mut self, data: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int): + def _update_soft(mut self, data: Pointer[mut=True, UInt8, _, address_space=_], length: Int): var off = 0 while off < length: var block = InlineArray[UInt8, 16](fill=0) @@ -887,7 +887,7 @@ struct _GHash(Copyable, Movable): if n > 16: n = 16 for i in range(n): - block[i] = data[off + i] + block[i] = data[unsafe_offset=off + i] self.y_hi ^= load_64be(block.unsafe_ptr(), 0) self.y_lo ^= load_64be(block.unsafe_ptr(), 8) self._mul_y_by_h_soft() @@ -902,11 +902,11 @@ struct _GHash(Copyable, Movable): @always_inline def _encrypt_block( - rk: UnsafePointer[mut=True, UInt32, _, address_space=_], rounds: Int, - src: UnsafePointer[mut=True, UInt8, _, address_space=_], dst: UnsafePointer[mut=True, UInt8, _, address_space=_] + rk: Pointer[mut=True, UInt32, _, address_space=_], rounds: Int, + src: Pointer[mut=True, UInt8, _, address_space=_], dst: Pointer[mut=True, UInt8, _, address_space=_] ): for i in range(16): - dst[i] = src[i] + dst[unsafe_offset=i] = src[unsafe_offset=i] comptime if CompilationTarget._has_feature["sse"]() and CompilationTarget._has_feature["aes"](): aes_encrypt(dst, rk, rounds) else: @@ -939,10 +939,10 @@ def _derive_j0( def _gctr_and_ghash( - rk: UnsafePointer[mut=True, UInt32, _, address_space=_], rounds: Int, + rk: Pointer[mut=True, UInt32, _, address_space=_], rounds: Int, j0: InlineArray[UInt8, 16], - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], length: Int, mut gh: _GHash, ghash_ciphertext: Bool, @@ -981,15 +981,15 @@ def _gctr_and_ghash( _encrypt_block(rk, rounds, ctr.unsafe_ptr(), ks.unsafe_ptr()) var off = full_blocks * 16 for i in range(rem): - output_ptr[off + i] = input_ptr[off + i] ^ ks[i] + output_ptr[unsafe_offset=off + i] = input_ptr[unsafe_offset=off + i] ^ ks[i] if fused: if rem > 0: var off = full_blocks * 16 if ghash_ciphertext: - gh.update(output_ptr + off, rem) + gh.update(output_ptr.unsafe_offset(off), rem) else: - gh.update(input_ptr + off, rem) + gh.update(input_ptr.unsafe_offset(off), rem) elif ghash_ciphertext: gh.update(output_ptr, length) else: @@ -997,11 +997,11 @@ def _gctr_and_ghash( def _gcm_core_keyed( - rk: UnsafePointer[mut=True, UInt32, _, address_space=_], rounds: Int, + rk: Pointer[mut=True, UInt32, _, address_space=_], rounds: Int, mut gh: _GHash, iv: Span[UInt8, ...], aad: Span[UInt8, ...], - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], length: Int, mut tag: InlineArray[UInt8, 16], ghash_ciphertext: Bool, @@ -1044,7 +1044,7 @@ struct AESGCMContext(Copyable, Movable): expand_key_192_into(key_buf.unsafe_ptr(), self._rk.unsafe_ptr()) else: expand_key_256_into(key_buf.unsafe_ptr(), self._rk.unsafe_ptr()) - memset_zero(key_buf.unsafe_ptr(), 32) + unsafe_memset_zero(key_buf.unsafe_ptr(), 32) var zero_block = InlineArray[UInt8, 16](fill=0) var h_block = InlineArray[UInt8, 16](fill=0) @@ -1052,11 +1052,11 @@ struct AESGCMContext(Copyable, Movable): self._gh0 = _GHash( load_64be(h_block.unsafe_ptr(), 0), load_64be(h_block.unsafe_ptr(), 8) ) - memset_zero(h_block.unsafe_ptr(), 16) + unsafe_memset_zero(h_block.unsafe_ptr(), 16) - def __del__(deinit self): - memset_zero(self._rk.unsafe_ptr(), 60) - memset_zero(UnsafePointer(to=self._gh0), 1) + def __deinit__(deinit self): + unsafe_memset_zero(self._rk.unsafe_ptr(), 60) + unsafe_memset_zero(Pointer(to=self._gh0), 1) def encrypt( self, iv: Span[UInt8, ...], plaintext: Span[UInt8, ...], aad: Span[UInt8, ...] @@ -1081,8 +1081,8 @@ struct AESGCMContext(Copyable, Movable): var tag_out = List[UInt8](capacity=16) for i in range(16): tag_out.append(tag[i]) - memset_zero(rk.unsafe_ptr(), 60) - memset_zero(UnsafePointer(to=gh), 1) + unsafe_memset_zero(rk.unsafe_ptr(), 60) + unsafe_memset_zero(Pointer(to=gh), 1) return (ciphertext^, tag_out^) def decrypt( @@ -1115,13 +1115,13 @@ struct AESGCMContext(Copyable, Movable): if diff != 0: var pt_ptr = plaintext.unsafe_ptr() for i in range(n): - pt_ptr.store[volatile=True](i, UInt8(0)) - memset_zero(rk.unsafe_ptr(), 60) - memset_zero(UnsafePointer(to=gh), 1) + pt_ptr.unsafe_store[volatile=True](i, UInt8(0)) + unsafe_memset_zero(rk.unsafe_ptr(), 60) + unsafe_memset_zero(Pointer(to=gh), 1) return (List[UInt8](), False) - memset_zero(rk.unsafe_ptr(), 60) - memset_zero(UnsafePointer(to=gh), 1) + unsafe_memset_zero(rk.unsafe_ptr(), 60) + unsafe_memset_zero(Pointer(to=gh), 1) return (plaintext^, True) diff --git a/src/thistle/argon2.mojo b/src/thistle/argon2.mojo index cd6afbd..48b6361 100644 --- a/src/thistle/argon2.mojo +++ b/src/thistle/argon2.mojo @@ -4,7 +4,7 @@ RFC 9106 """ from std.collections import List -from std.memory import alloc, UnsafePointer, memcpy, memset_zero +from std.memory import Layout, Pointer, alloc, unsafe_memcpy, unsafe_memset_zero from max.algorithm import parallelize from std.bit import rotate_bits_left from .blake2b import Blake2b @@ -12,22 +12,22 @@ from .blake2b import Blake2b comptime MASK32 = 0xFFFFFFFF @always_inline -def zero_buffer(ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], len: Int): - memset_zero(ptr, len) +def zero_buffer(ptr: Pointer[mut=True, UInt8, _, address_space=_], len: Int): + unsafe_memset_zero(ptr, len) @always_inline -def zero_buffer_u64(ptr: UnsafePointer[mut=True, UInt64, _, address_space=_], len: Int): - memset_zero(ptr, len) +def zero_buffer_u64(ptr: Pointer[mut=True, UInt64, _, address_space=_], len: Int): + unsafe_memset_zero(ptr, len) @always_inline -def zero_and_free(ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], len: Int): +def zero_and_free(ptr: Pointer[mut=True, UInt8, _, address_space=_], len: Int): zero_buffer(ptr, len) - ptr.free() + ptr.unsafe_free() @always_inline -def zero_and_free_u64(ptr: UnsafePointer[mut=True, UInt64, _, address_space=_], len: Int): +def zero_and_free_u64(ptr: Pointer[mut=True, UInt64, _, address_space=_], len: Int): zero_buffer_u64(ptr, len) - ptr.free() + ptr.unsafe_free() @always_inline def f_bla_mka(x: UInt64, y: UInt64) -> UInt64: @@ -46,78 +46,78 @@ def gb(a: UInt64, b: UInt64, c: UInt64, d: UInt64) -> Tuple[UInt64, UInt64, UInt return (a_new, b_new, c_new, d_new) @always_inline -def _p_column(base: Int, v: UnsafePointer[mut=True, UInt64, _, address_space=_]): - var v0, v4, v8, v12 = gb(v[base + 0], v[base + 4], v[base + 8], v[base + 12]) - var v1, v5, v9, v13 = gb(v[base + 1], v[base + 5], v[base + 9], v[base + 13]) - var v2, v6, v10, v14 = gb(v[base + 2], v[base + 6], v[base + 10], v[base + 14]) - var v3, v7, v11, v15 = gb(v[base + 3], v[base + 7], v[base + 11], v[base + 15]) - v[base + 0] = v0 - v[base + 4] = v4 - v[base + 8] = v8 - v[base + 12] = v12 - v[base + 1] = v1 - v[base + 5] = v5 - v[base + 9] = v9 - v[base + 13] = v13 - v[base + 2] = v2 - v[base + 6] = v6 - v[base + 10] = v10 - v[base + 14] = v14 - v[base + 3] = v3 - v[base + 7] = v7 - v[base + 11] = v11 - v[base + 15] = v15 +def _p_column(base: Int, v: Pointer[mut=True, UInt64, _, address_space=_]): + var v0, v4, v8, v12 = gb(v[unsafe_offset=base + 0], v[unsafe_offset=base + 4], v[unsafe_offset=base + 8], v[unsafe_offset=base + 12]) + var v1, v5, v9, v13 = gb(v[unsafe_offset=base + 1], v[unsafe_offset=base + 5], v[unsafe_offset=base + 9], v[unsafe_offset=base + 13]) + var v2, v6, v10, v14 = gb(v[unsafe_offset=base + 2], v[unsafe_offset=base + 6], v[unsafe_offset=base + 10], v[unsafe_offset=base + 14]) + var v3, v7, v11, v15 = gb(v[unsafe_offset=base + 3], v[unsafe_offset=base + 7], v[unsafe_offset=base + 11], v[unsafe_offset=base + 15]) + v[unsafe_offset=base + 0] = v0 + v[unsafe_offset=base + 4] = v4 + v[unsafe_offset=base + 8] = v8 + v[unsafe_offset=base + 12] = v12 + v[unsafe_offset=base + 1] = v1 + v[unsafe_offset=base + 5] = v5 + v[unsafe_offset=base + 9] = v9 + v[unsafe_offset=base + 13] = v13 + v[unsafe_offset=base + 2] = v2 + v[unsafe_offset=base + 6] = v6 + v[unsafe_offset=base + 10] = v10 + v[unsafe_offset=base + 14] = v14 + v[unsafe_offset=base + 3] = v3 + v[unsafe_offset=base + 7] = v7 + v[unsafe_offset=base + 11] = v11 + v[unsafe_offset=base + 15] = v15 @always_inline -def _p_diagonal(base: Int, v: UnsafePointer[mut=True, UInt64, _, address_space=_]): - var v0, v5, v10, v15 = gb(v[base + 0], v[base + 5], v[base + 10], v[base + 15]) - var v1, v6, v11, v12 = gb(v[base + 1], v[base + 6], v[base + 11], v[base + 12]) - var v2, v7, v8, v13 = gb(v[base + 2], v[base + 7], v[base + 8], v[base + 13]) - var v3, v4, v9, v14 = gb(v[base + 3], v[base + 4], v[base + 9], v[base + 14]) - v[base + 0] = v0 - v[base + 5] = v5 - v[base + 10] = v10 - v[base + 15] = v15 - v[base + 1] = v1 - v[base + 6] = v6 - v[base + 11] = v11 - v[base + 12] = v12 - v[base + 2] = v2 - v[base + 7] = v7 - v[base + 8] = v8 - v[base + 13] = v13 - v[base + 3] = v3 - v[base + 4] = v4 - v[base + 9] = v9 - v[base + 14] = v14 +def _p_diagonal(base: Int, v: Pointer[mut=True, UInt64, _, address_space=_]): + var v0, v5, v10, v15 = gb(v[unsafe_offset=base + 0], v[unsafe_offset=base + 5], v[unsafe_offset=base + 10], v[unsafe_offset=base + 15]) + var v1, v6, v11, v12 = gb(v[unsafe_offset=base + 1], v[unsafe_offset=base + 6], v[unsafe_offset=base + 11], v[unsafe_offset=base + 12]) + var v2, v7, v8, v13 = gb(v[unsafe_offset=base + 2], v[unsafe_offset=base + 7], v[unsafe_offset=base + 8], v[unsafe_offset=base + 13]) + var v3, v4, v9, v14 = gb(v[unsafe_offset=base + 3], v[unsafe_offset=base + 4], v[unsafe_offset=base + 9], v[unsafe_offset=base + 14]) + v[unsafe_offset=base + 0] = v0 + v[unsafe_offset=base + 5] = v5 + v[unsafe_offset=base + 10] = v10 + v[unsafe_offset=base + 15] = v15 + v[unsafe_offset=base + 1] = v1 + v[unsafe_offset=base + 6] = v6 + v[unsafe_offset=base + 11] = v11 + v[unsafe_offset=base + 12] = v12 + v[unsafe_offset=base + 2] = v2 + v[unsafe_offset=base + 7] = v7 + v[unsafe_offset=base + 8] = v8 + v[unsafe_offset=base + 13] = v13 + v[unsafe_offset=base + 3] = v3 + v[unsafe_offset=base + 4] = v4 + v[unsafe_offset=base + 9] = v9 + v[unsafe_offset=base + 14] = v14 struct MemoryPool: - var block_buffer: UnsafePointer[UInt64, MutUntrackedOrigin] - var temp_buffer: UnsafePointer[UInt64, MutUntrackedOrigin] + var block_buffer: Pointer[UInt64, MutUntrackedOrigin] + var temp_buffer: Pointer[UInt64, MutUntrackedOrigin] var buffer_size: Int def __init__(out self, size: Int): self.buffer_size = size - self.block_buffer = alloc[UInt64](size) - self.temp_buffer = alloc[UInt64](size) + self.block_buffer = alloc(Layout[UInt64](count=size)).unsafe_leak() + self.temp_buffer = alloc(Layout[UInt64](count=size)).unsafe_leak() - def __del__(deinit self): + def __deinit__(deinit self): zero_and_free_u64(self.block_buffer, self.buffer_size) zero_and_free_u64(self.temp_buffer, self.buffer_size) @always_inline - def get_block(self) -> UnsafePointer[UInt64, MutUntrackedOrigin]: + def get_block(self) -> Pointer[UInt64, MutUntrackedOrigin]: return self.block_buffer @always_inline - def get_temp(self) -> UnsafePointer[UInt64, MutUntrackedOrigin]: + def get_temp(self) -> Pointer[UInt64, MutUntrackedOrigin]: return self.temp_buffer @always_inline def compression_g_with_pool( - out_ptr: UnsafePointer[mut=True, UInt64, _, address_space=_], - x_ptr: UnsafePointer[mut=False, UInt64, _, address_space=_], - y_ptr: UnsafePointer[mut=False, UInt64, _, address_space=_], + out_ptr: Pointer[mut=True, UInt64, _, address_space=_], + x_ptr: Pointer[mut=False, UInt64, _, address_space=_], + y_ptr: Pointer[mut=False, UInt64, _, address_space=_], with_xor: Bool, pool: MemoryPool, ): @@ -125,12 +125,12 @@ def compression_g_with_pool( var block_xy = pool.get_temp() for i in range(128): - var val = x_ptr[i] ^ y_ptr[i] - block[i] = val + var val = x_ptr[unsafe_offset=i] ^ y_ptr[unsafe_offset=i] + block[unsafe_offset=i] = val if with_xor: - block_xy[i] = val ^ out_ptr[i] + block_xy[unsafe_offset=i] = val ^ out_ptr[unsafe_offset=i] else: - block_xy[i] = val + block_xy[unsafe_offset=i] = val for i in range(8): var base = i * 16 @@ -138,22 +138,22 @@ def compression_g_with_pool( _p_diagonal(base, block) for col in range(8): - var v0 = block[col * 2 + 0] - var v1 = block[col * 2 + 1] - var v2 = block[col * 2 + 16] - var v3 = block[col * 2 + 17] - var v4 = block[col * 2 + 32] - var v5 = block[col * 2 + 33] - var v6 = block[col * 2 + 48] - var v7 = block[col * 2 + 49] - var v8 = block[col * 2 + 64] - var v9 = block[col * 2 + 65] - var v10 = block[col * 2 + 80] - var v11 = block[col * 2 + 81] - var v12 = block[col * 2 + 96] - var v13 = block[col * 2 + 97] - var v14 = block[col * 2 + 112] - var v15 = block[col * 2 + 113] + var v0 = block[unsafe_offset=col * 2 + 0] + var v1 = block[unsafe_offset=col * 2 + 1] + var v2 = block[unsafe_offset=col * 2 + 16] + var v3 = block[unsafe_offset=col * 2 + 17] + var v4 = block[unsafe_offset=col * 2 + 32] + var v5 = block[unsafe_offset=col * 2 + 33] + var v6 = block[unsafe_offset=col * 2 + 48] + var v7 = block[unsafe_offset=col * 2 + 49] + var v8 = block[unsafe_offset=col * 2 + 64] + var v9 = block[unsafe_offset=col * 2 + 65] + var v10 = block[unsafe_offset=col * 2 + 80] + var v11 = block[unsafe_offset=col * 2 + 81] + var v12 = block[unsafe_offset=col * 2 + 96] + var v13 = block[unsafe_offset=col * 2 + 97] + var v14 = block[unsafe_offset=col * 2 + 112] + var v15 = block[unsafe_offset=col * 2 + 113] v0, v4, v8, v12 = gb(v0, v4, v8, v12) v1, v5, v9, v13 = gb(v1, v5, v9, v13) @@ -165,32 +165,32 @@ def compression_g_with_pool( v2, v7, v8, v13 = gb(v2, v7, v8, v13) v3, v4, v9, v14 = gb(v3, v4, v9, v14) - block[col * 2 + 0] = v0 - block[col * 2 + 1] = v1 - block[col * 2 + 16] = v2 - block[col * 2 + 17] = v3 - block[col * 2 + 32] = v4 - block[col * 2 + 33] = v5 - block[col * 2 + 48] = v6 - block[col * 2 + 49] = v7 - block[col * 2 + 64] = v8 - block[col * 2 + 65] = v9 - block[col * 2 + 80] = v10 - block[col * 2 + 81] = v11 - block[col * 2 + 96] = v12 - block[col * 2 + 97] = v13 - block[col * 2 + 112] = v14 - block[col * 2 + 113] = v15 + block[unsafe_offset=col * 2 + 0] = v0 + block[unsafe_offset=col * 2 + 1] = v1 + block[unsafe_offset=col * 2 + 16] = v2 + block[unsafe_offset=col * 2 + 17] = v3 + block[unsafe_offset=col * 2 + 32] = v4 + block[unsafe_offset=col * 2 + 33] = v5 + block[unsafe_offset=col * 2 + 48] = v6 + block[unsafe_offset=col * 2 + 49] = v7 + block[unsafe_offset=col * 2 + 64] = v8 + block[unsafe_offset=col * 2 + 65] = v9 + block[unsafe_offset=col * 2 + 80] = v10 + block[unsafe_offset=col * 2 + 81] = v11 + block[unsafe_offset=col * 2 + 96] = v12 + block[unsafe_offset=col * 2 + 97] = v13 + block[unsafe_offset=col * 2 + 112] = v14 + block[unsafe_offset=col * 2 + 113] = v15 for i in range(128): - out_ptr[i] = block[i] ^ block_xy[i] + out_ptr[unsafe_offset=i] = block[unsafe_offset=i] ^ block_xy[unsafe_offset=i] @always_inline -def store_le32(ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], offset: Int, val: Int): - ptr[offset + 0] = UInt8(val & 0xFF) - ptr[offset + 1] = UInt8((val >> 8) & 0xFF) - ptr[offset + 2] = UInt8((val >> 16) & 0xFF) - ptr[offset + 3] = UInt8((val >> 24) & 0xFF) +def store_le32(ptr: Pointer[mut=True, UInt8, _, address_space=_], offset: Int, val: Int): + ptr[unsafe_offset=offset + 0] = UInt8(val & 0xFF) + ptr[unsafe_offset=offset + 1] = UInt8((val >> 8) & 0xFF) + ptr[unsafe_offset=offset + 2] = UInt8((val >> 16) & 0xFF) + ptr[unsafe_offset=offset + 3] = UInt8((val >> 24) & 0xFF) def variable_length_hash_into( @@ -206,7 +206,7 @@ def variable_length_hash_into( var out_ptr = output.unsafe_ptr() if t_len <= 64: - var le_buf = alloc[UInt8](4) + var le_buf = alloc(Layout[UInt8](count=4)).unsafe_leak() try: var ctx = Blake2b(t_len) store_le32(le_buf, 0, t_len) @@ -217,9 +217,9 @@ def variable_length_hash_into( zero_and_free(le_buf, 4) return - var le_buf = alloc[UInt8](4) + var le_buf = alloc(Layout[UInt8](count=4)).unsafe_leak() var r = (t_len + 31) // 32 - 2 - var v_buf = alloc[UInt8](64) + var v_buf = alloc(Layout[UInt8](count=64)).unsafe_leak() try: var ctx1 = Blake2b(64) store_le32(le_buf, 0, t_len) @@ -230,7 +230,7 @@ def variable_length_hash_into( var out_offset = 0 for _ in range(r - 1): for j in range(32): - out_ptr[out_offset + j] = v_buf[j] + out_ptr[unsafe_offset=out_offset + j] = v_buf[unsafe_offset=j] out_offset += 32 var ctx = Blake2b(64) @@ -238,13 +238,13 @@ def variable_length_hash_into( ctx.finalize_into(v_buf) for j in range(32): - out_ptr[out_offset + j] = v_buf[j] + out_ptr[unsafe_offset=out_offset + j] = v_buf[unsafe_offset=j] out_offset += 32 var last_len = t_len - 32 * r var ctx_last = Blake2b(last_len) ctx_last.update(Span[UInt8, ...](unsafe_ptr=v_buf, length=64)) - ctx_last.finalize_into(out_ptr + out_offset) + ctx_last.finalize_into(out_ptr.unsafe_offset(out_offset)) finally: zero_and_free(v_buf, 64) zero_and_free(le_buf, 4) @@ -252,7 +252,7 @@ def variable_length_hash_into( def variable_length_hash(t_len: Int, input: Span[UInt8, ...]) raises -> List[UInt8]: if t_len < 1: raise Error("Argon2 variable-length hash output must not be empty") - var out_buf = alloc[UInt8](t_len) + var out_buf = alloc(Layout[UInt8](count=t_len)).unsafe_leak() try: variable_length_hash_into( t_len, @@ -261,14 +261,14 @@ def variable_length_hash(t_len: Int, input: Span[UInt8, ...]) raises -> List[UIn ) var result = List[UInt8](capacity=t_len) for i in range(t_len): - result.append(out_buf[i]) + result.append(out_buf[unsafe_offset=i]) return result^ finally: zero_and_free(out_buf, t_len) @always_inline def _argon2_process_lane( - memory: UnsafePointer[mut=True, UInt64, _, address_space=_], + memory: Pointer[mut=True, UInt64, _, address_space=_], lane: Int, t: Int, slice_idx: Int, @@ -281,13 +281,13 @@ def _argon2_process_lane( type_code: Int, parallelism: Int, ): - var addressing_block = alloc[UInt64](128) - var z_u64 = alloc[UInt64](128) - var zero_u64 = alloc[UInt64](128) - var tmp_addr = alloc[UInt64](128) + var addressing_block = alloc(Layout[UInt64](count=128)).unsafe_leak() + var z_u64 = alloc(Layout[UInt64](count=128)).unsafe_leak() + var zero_u64 = alloc(Layout[UInt64](count=128)).unsafe_leak() + var tmp_addr = alloc(Layout[UInt64](count=128)).unsafe_leak() var pool = MemoryPool(128) var has_addressing_block = False - memset_zero(zero_u64, 128) + unsafe_memset_zero(zero_u64, 128) for index in range(seg_start, seg_end): if t == 0 and index < 2: @@ -303,24 +303,24 @@ def _argon2_process_lane( if not has_addressing_block or ( seg_offset % 128 == 0 ): - memset_zero(z_u64, 128) - z_u64[0] = UInt64(t) - z_u64[1] = UInt64(lane) - z_u64[2] = UInt64(slice_idx) - z_u64[3] = UInt64(m_prime_blocks) - z_u64[4] = UInt64(iterations) - z_u64[5] = UInt64(type_code) - z_u64[6] = UInt64((seg_offset // 128) + 1) + unsafe_memset_zero(z_u64, 128) + z_u64[unsafe_offset=0] = UInt64(t) + z_u64[unsafe_offset=1] = UInt64(lane) + z_u64[unsafe_offset=2] = UInt64(slice_idx) + z_u64[unsafe_offset=3] = UInt64(m_prime_blocks) + z_u64[unsafe_offset=4] = UInt64(iterations) + z_u64[unsafe_offset=5] = UInt64(type_code) + z_u64[unsafe_offset=6] = UInt64((seg_offset // 128) + 1) compression_g_with_pool(tmp_addr, zero_u64, z_u64, False, pool) compression_g_with_pool(addressing_block, zero_u64, tmp_addr, False, pool) has_addressing_block = True - var val = addressing_block[seg_offset % 128] + var val = addressing_block[unsafe_offset=seg_offset % 128] j1 = UInt32(val & 0xFFFFFFFF) j2 = UInt32(val >> 32) else: - var v0 = memory[lane * q * 128 + prev_index * 128] + var v0 = memory[unsafe_offset=lane * q * 128 + prev_index * 128] j1 = UInt32(v0 & 0xFFFFFFFF) j2 = UInt32(v0 >> 32) @@ -367,9 +367,9 @@ def _argon2_process_lane( ) * segment_length ref_index = (start_pos + Int(zz)) % q - var p_ptr = memory + (lane * q * 128 + prev_index * 128) - var r_ptr = memory + (ref_lane * q * 128 + ref_index * 128) - var c_ptr = memory + (lane * q * 128 + index * 128) + var p_ptr = memory.unsafe_offset((lane * q * 128 + prev_index * 128)) + var r_ptr = memory.unsafe_offset((ref_lane * q * 128 + ref_index * 128)) + var c_ptr = memory.unsafe_offset((lane * q * 128 + index * 128)) compression_g_with_pool( c_ptr, @@ -409,10 +409,10 @@ struct Argon2id: var secret: List[UInt8] var ad: List[UInt8] - def __del__(deinit self): + def __deinit__(deinit self): var secret_ptr = self.secret.unsafe_ptr() for i in range(len(self.secret)): - secret_ptr.store[volatile=True](i, UInt8(0)) + secret_ptr.unsafe_store[volatile=True](i, UInt8(0)) def __init__( out self, @@ -466,7 +466,7 @@ struct Argon2id: def hash(self, password: Span[UInt8, ...]) raises -> List[UInt8]: var h0_ctx = Blake2b(64) - var le_buf = alloc[UInt8](4) + var le_buf = alloc(Layout[UInt8](count=4)).unsafe_leak() store_le32(le_buf, 0, self.parallelism) h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf, length=4)) store_le32(le_buf, 0, self.tag_length) @@ -493,7 +493,7 @@ struct Argon2id: h0_ctx.update(Span[UInt8, ...](self.ad)) zero_and_free(le_buf, 4) - var h0_buf = alloc[UInt8](64) + var h0_buf = alloc(Layout[UInt8](count=64)).unsafe_leak() h0_ctx.finalize_into(h0_buf) var m_blocks = self.memory_size_kb @@ -505,10 +505,10 @@ struct Argon2id: var q = m_prime_blocks // self.parallelism var segment_length = q // 4 - var memory = alloc[UInt64](m_prime_blocks * 128) + var memory = alloc(Layout[UInt64](count=m_prime_blocks * 128)).unsafe_leak() - var h0_input = alloc[UInt8](72) - memcpy(dest=h0_input, src=h0_buf, count=64) + var h0_input = alloc(Layout[UInt8](count=72)).unsafe_leak() + unsafe_memcpy(dest=h0_input, src=h0_buf, count=64) zero_and_free(h0_buf, 64) for i in range(self.parallelism): @@ -516,7 +516,7 @@ struct Argon2id: store_le32(h0_input, 64, block_idx) store_le32(h0_input, 68, i) - var b_bytes = alloc[UInt8](1024) + var b_bytes = alloc(Layout[UInt8](count=1024)).unsafe_leak() variable_length_hash_into( 1024, Span[UInt8, ...](unsafe_ptr=h0_input, length=72), @@ -526,8 +526,8 @@ struct Argon2id: ) for k in range(128): - var word = (b_bytes + k * 8).bitcast[UInt64]().load[width=1, alignment=1]() - memory[i * q * 128 + block_idx * 128 + k] = word + var word = (b_bytes.unsafe_offset(k * 8)).unsafe_bitcast[UInt64]().unsafe_load[width=1, alignment=1]() + memory[unsafe_offset=i * q * 128 + block_idx * 128 + k] = word zero_and_free(b_bytes, 1024) zero_and_free(h0_input, 72) @@ -557,16 +557,16 @@ struct Argon2id: parallelize[process_lane](parallelism) - var c_block = alloc[UInt64](128) - memset_zero(c_block, 128) + var c_block = alloc(Layout[UInt64](count=128)).unsafe_leak() + unsafe_memset_zero(c_block, 128) for i in range(self.parallelism): - var last_ptr = memory + (i * q * 128 + (q - 1) * 128) + var last_ptr = memory.unsafe_offset((i * q * 128 + (q - 1) * 128)) for k in range(128): - c_block[k] ^= last_ptr[k] + c_block[unsafe_offset=k] ^= last_ptr[unsafe_offset=k] - var c_bytes = alloc[UInt8](1024) + var c_bytes = alloc(Layout[UInt8](count=1024)).unsafe_leak() for k in range(128): - (c_bytes + k * 8).bitcast[UInt64]().store[alignment=1](c_block[k]) + (c_bytes.unsafe_offset(k * 8)).unsafe_bitcast[UInt64]().unsafe_store[alignment=1](c_block[unsafe_offset=k]) zero_and_free_u64(c_block, 128) zero_and_free_u64(memory, m_prime_blocks * 128) diff --git a/src/thistle/blake2b.mojo b/src/thistle/blake2b.mojo index 2e63193..2704bd5 100644 --- a/src/thistle/blake2b.mojo +++ b/src/thistle/blake2b.mojo @@ -4,7 +4,7 @@ RFC 7693 """ from std.collections import List -from std.memory import UnsafePointer, memcpy, memset_zero +from std.memory import Pointer, unsafe_memcpy, unsafe_memset_zero from .utils import bytes_to_hex, string_to_bytes comptime BLAKE2B_IV = SIMD[DType.uint64, 8]( @@ -53,8 +53,8 @@ def g(a: UInt64, b: UInt64, c: UInt64, d: UInt64, x: UInt64, y: UInt64) -> Tuple @always_inline -def _mload(m: UnsafePointer[mut=False, UInt8, _, address_space=_], i: Int) -> UInt64: - return (m + i * 8).bitcast[UInt64]().load[width=1, alignment=1]() +def _mload(m: Pointer[mut=False, UInt8, _, address_space=_], i: Int) -> UInt64: + return (m.unsafe_offset(i * 8)).unsafe_bitcast[UInt64]().unsafe_load[width=1, alignment=1]() @always_inline @@ -63,7 +63,7 @@ def round_fn[r: Int]( mut v4: UInt64, mut v5: UInt64, mut v6: UInt64, mut v7: UInt64, mut v8: UInt64, mut v9: UInt64, mut v10: UInt64, mut v11: UInt64, mut v12: UInt64, mut v13: UInt64, mut v14: UInt64, mut v15: UInt64, - m: UnsafePointer[mut=False, UInt8, _, address_space=_], + m: Pointer[mut=False, UInt8, _, address_space=_], ) -> Tuple[UInt64, UInt64, UInt64, UInt64, UInt64, UInt64, UInt64, UInt64, UInt64, UInt64, UInt64, UInt64, UInt64, UInt64, UInt64, UInt64]: comptime s = SIGMA[r] @@ -90,8 +90,8 @@ struct Blake2b(Movable): var key_len: Int @always_inline - def _buf_ptr(mut self) -> UnsafePointer[UInt8, MutAnyOrigin]: - return self.buffer.unsafe_ptr().bitcast[UInt8]().unsafe_origin_cast[MutAnyOrigin]() + def _buf_ptr(mut self) -> Pointer[UInt8, MutAnyOrigin]: + return self.buffer.unsafe_ptr().unsafe_bitcast[UInt8]().unsafe_origin_cast[MutAnyOrigin]() def __init__(out self, out_len: Int = 64) raises: if out_len < 1 or out_len > 64: @@ -133,23 +133,23 @@ struct Blake2b(Movable): self.update(key) var buf = self._buf_ptr() while self.buffer_len < 128: - buf[self.buffer_len] = 0 + buf[unsafe_offset=self.buffer_len] = 0 self.buffer_len += 1 - def __init__(out self, *, deinit take: Self): - self.h = take.h - self.t_low = take.t_low - self.t_high = take.t_high - self.buffer = take.buffer^ - self.buffer_len = take.buffer_len - self.out_len = take.out_len - self.key_len = take.key_len - - def __del__(deinit self): - UnsafePointer(to=self.h).bitcast[UInt64]().store[volatile=True]( + def __init__(out self, *, deinit move: Self): + self.h = move.h + self.t_low = move.t_low + self.t_high = move.t_high + self.buffer = move.buffer^ + self.buffer_len = move.buffer_len + self.out_len = move.out_len + self.key_len = move.key_len + + def __deinit__(deinit self): + Pointer(to=self.h).unsafe_bitcast[UInt64]().unsafe_store[volatile=True]( 0, SIMD[DType.uint64, 8](0) ) - memset_zero(self.buffer.unsafe_ptr(), 16) + unsafe_memset_zero(self.buffer.unsafe_ptr(), 16) @always_inline def _inc_counter(mut self): @@ -157,7 +157,7 @@ struct Blake2b(Movable): if self.t_low < 128: self.t_high += 1 - def compress(mut self, m: UnsafePointer[mut=False, UInt8, _, address_space=_], is_last: Bool): + def compress(mut self, m: Pointer[mut=False, UInt8, _, address_space=_], is_last: Bool): var v0 = self.h[0] var v1 = self.h[1] var v2 = self.h[2] @@ -215,7 +215,7 @@ struct Blake2b(Movable): if total < to_copy: to_copy = total for j in range(to_copy): - self._buf_ptr()[self.buffer_len + j] = data[j] + self._buf_ptr()[unsafe_offset=self.buffer_len + j] = data[j] self.buffer_len += to_copy i += to_copy if i == total: @@ -226,29 +226,29 @@ struct Blake2b(Movable): while total - i > 128: self._inc_counter() - self.compress(data.unsafe_ptr() + i, False) + self.compress(data.unsafe_ptr().unsafe_offset(i), False) i += 128 for j in range(total - i): - self._buf_ptr()[j] = data[i + j] + self._buf_ptr()[unsafe_offset=j] = data[i + j] self.buffer_len = total - i - def finalize_into(mut self, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): + def finalize_into(mut self, output: Pointer[mut=True, UInt8, _, address_space=_]): var old_low = self.t_low self.t_low += UInt64(self.buffer_len) if self.t_low < old_low: self.t_high += 1 if self.buffer_len < 128: - memset_zero(self._buf_ptr() + self.buffer_len, 128 - self.buffer_len) + unsafe_memset_zero(self._buf_ptr().unsafe_offset(self.buffer_len), 128 - self.buffer_len) self.buffer_len = 128 self.compress(self._buf_ptr(), True) var h_copy = self.h - var h_bytes = UnsafePointer(to=h_copy).bitcast[UInt8]() + var h_bytes = Pointer(to=h_copy).unsafe_bitcast[UInt8]() for i in range(self.out_len): - output[i] = h_bytes[i] + output[unsafe_offset=i] = h_bytes[unsafe_offset=i] def finalize(mut self) -> List[UInt8]: var output = List[UInt8](capacity=self.out_len) diff --git a/src/thistle/blake3.mojo b/src/thistle/blake3.mojo index 2152422..b904e46 100644 --- a/src/thistle/blake3.mojo +++ b/src/thistle/blake3.mojo @@ -4,7 +4,7 @@ BLAKE3 cryptographic hash function from max.algorithm import parallelize from std.collections import List -from std.memory import UnsafePointer, bitcast +from std.memory import Pointer, bitcast from std.bit import count_trailing_zeros from std.utils import IndexList from thistle.utils import StackInlineArray @@ -112,7 +112,7 @@ def compress_internal[ counter: UInt64, blen: UInt8, flags: UInt8, - out_ptr: UnsafePointer[mut=True, SIMD[DType.uint32, w], _, address_space=_], + out_ptr: Pointer[mut=True, SIMD[DType.uint32, w], _, address_space=_], ): """BLAKE3 compression: 7 rounds of G with message permutation.""" # fmt: off @@ -153,8 +153,8 @@ def compress_internal[ transform() round() - out_ptr.bitcast[UInt32]().store( - v.unsafe_ptr().bitcast[UInt32]().load[width=w * 16]() + out_ptr.unsafe_bitcast[UInt32]().unsafe_store( + v.unsafe_ptr().unsafe_bitcast[UInt32]().unsafe_load[width=w * 16]() ) @always_inline @@ -191,7 +191,7 @@ def compress_internal_16way( base_counter: UInt64, blen: UInt8, flags: UInt8, - out_ptr: UnsafePointer[mut=True, SIMD[DType.uint32, 16], _, address_space=_], + out_ptr: Pointer[mut=True, SIMD[DType.uint32, 16], _, address_space=_], ): """16-way SIMD compression with per-lane sequential counters.""" var counters_low = SIMD[DType.uint32, 16]( @@ -236,14 +236,14 @@ def compress_internal_16way( transform() round() - out_ptr[0] = v[0] ^ v[8] - out_ptr[1] = v[1] ^ v[9] - out_ptr[2] = v[2] ^ v[10] - out_ptr[3] = v[3] ^ v[11] - out_ptr[4] = v[4] ^ v[12] - out_ptr[5] = v[5] ^ v[13] - out_ptr[6] = v[6] ^ v[14] - out_ptr[7] = v[7] ^ v[15] + out_ptr[unsafe_offset=0] = v[0] ^ v[8] + out_ptr[unsafe_offset=1] = v[1] ^ v[9] + out_ptr[unsafe_offset=2] = v[2] ^ v[10] + out_ptr[unsafe_offset=3] = v[3] ^ v[11] + out_ptr[unsafe_offset=4] = v[4] ^ v[12] + out_ptr[unsafe_offset=5] = v[5] ^ v[13] + out_ptr[unsafe_offset=6] = v[6] ^ v[14] + out_ptr[unsafe_offset=7] = v[7] ^ v[15] struct Hasher: var key: SIMD[DType.uint32, 8] @@ -273,7 +273,7 @@ struct Hasher: while len(d) > 0: if self.buf_len == 64: var blk = ( - self.buf.unsafe_ptr().bitcast[UInt32]().load[width=16, alignment=1]() + self.buf.unsafe_ptr().unsafe_bitcast[UInt32]().unsafe_load[width=16, alignment=1]() ) if self.blocks_compressed == 15: @@ -363,7 +363,7 @@ struct Hasher: for i in range(self.buf_len): temp_buf.unsafe_set(i, self.buf[i]) - var blk = temp_buf.unsafe_ptr().bitcast[UInt32]().load[width=16, alignment=1]() + var blk = temp_buf.unsafe_ptr().unsafe_bitcast[UInt32]().unsafe_load[width=16, alignment=1]() var flags = ( CHUNK_START if self.blocks_compressed == 0 else UInt8(0) @@ -408,7 +408,7 @@ struct Hasher: var to_copy = min(64, out_len - bytes_written) if to_copy == 64: - (out_buf.unsafe_ptr() + bytes_written).store(b) + (out_buf.unsafe_ptr().unsafe_offset(bytes_written)).unsafe_store(b) else: for i in range(to_copy): out_buf.unsafe_set(bytes_written + i, b[i]) @@ -439,7 +439,7 @@ def blake3_parallel_hash(input: Span[UInt8, ...], out_len: Int = 32) raises -> L @parameter def process_batch(tid: Int): var task_base = tid * BSIZE - var base_ptr = d.unsafe_ptr().bitcast[UInt32]() + var base_ptr = d.unsafe_ptr().unsafe_bitcast[UInt32]() var local_cvs = StackInlineArray[SIMD[DType.uint32, 8], 64]( uninitialized=True ) @@ -469,7 +469,7 @@ def blake3_parallel_hash(input: Span[UInt8, ...], out_len: Int = 32) raises -> L @parameter @always_inline def _load_idx(v: Int) -> SIMD[DType.uint32, 4]: - return base_ptr.load[width=4, alignment=1]( + return base_ptr.unsafe_load[width=4, alignment=1]( (base + v) * 256 + b * 16 + joff ) @@ -567,7 +567,7 @@ def blake3_parallel_hash(input: Span[UInt8, ...], out_len: Int = 32) raises -> L IV, left.join(right), 0, 64, PARENT ) local_cvs.unsafe_set(i, combined.slice[8]()) - batch_roots_ptr[tid] = local_cvs[0] + batch_roots_ptr[unsafe_offset=tid] = local_cvs[0] parallelize[process_batch](num_full_batches) diff --git a/src/thistle/camellia.mojo b/src/thistle/camellia.mojo index fe6f5a0..7dea04f 100644 --- a/src/thistle/camellia.mojo +++ b/src/thistle/camellia.mojo @@ -2,7 +2,7 @@ Camellia block cipher implementation per RFC 3713 """ -from std.memory import bitcast, UnsafePointer +from std.memory import bitcast, Pointer from std.bit import byte_swap, rotate_bits_left from std.collections import InlineArray from std.sys import llvm_intrinsic @@ -260,9 +260,9 @@ comptime _ONE_VALUE_LANES: UInt64 = 0xFF @always_inline -def _wipe_u64(ptr: UnsafePointer[mut=True, UInt64, _, address_space=_], count: Int): +def _wipe_u64(ptr: Pointer[mut=True, UInt64, _, address_space=_], count: Int): for i in range(count): - ptr.store[volatile=True](i, UInt64(0)) + ptr.unsafe_store[volatile=True](i, UInt64(0)) @always_inline @@ -570,7 +570,7 @@ struct CamelliaCipher: self.kwhw[r] = byte_swap(self.kw[r]) def wipe(mut self): - _wipe_u64(UnsafePointer(to=self.kw).bitcast[UInt64](), 4) + _wipe_u64(Pointer(to=self.kw).unsafe_bitcast[UInt64](), 4) _wipe_u64(self.k.unsafe_ptr(), 24) _wipe_u64(self.ke.unsafe_ptr(), 6) _wipe_u64(self.kp.unsafe_ptr(), 192) @@ -578,8 +578,8 @@ struct CamelliaCipher: _wipe_u64(self.khw.unsafe_ptr(), 24) _wipe_u64(self.kwhw.unsafe_ptr(), 4) - def __del__(deinit self): - _wipe_u64(UnsafePointer(to=self.kw).bitcast[UInt64](), 4) + def __deinit__(deinit self): + _wipe_u64(Pointer(to=self.kw).unsafe_bitcast[UInt64](), 4) _wipe_u64(self.k.unsafe_ptr(), 24) _wipe_u64(self.ke.unsafe_ptr(), 6) _wipe_u64(self.kp.unsafe_ptr(), 192) @@ -591,7 +591,7 @@ struct CamelliaCipher: def _bytes_to_u64_be(ref self, b: Span[UInt8, ...]) -> UInt64: return byte_swap( bitcast[DType.uint64, 1]( - b.unsafe_ptr().load[width=8, alignment=1](0) + b.unsafe_ptr().unsafe_load[width=8, alignment=1](0) )[0] ) @@ -703,7 +703,7 @@ struct CamelliaCipher: @always_inline def _load_half[W: Int]( - buf: UnsafePointer[mut=True, UInt8, _, address_space=_], off: Int, kw: UInt64 + buf: Pointer[mut=True, UInt8, _, address_space=_], off: Int, kw: UInt64 ) -> InlineArray[SIMD[DType.uint64, W], 8]: var q = InlineArray[SIMD[DType.uint64, W], 8](fill=0) var kwl = byte_swap(kw) @@ -712,7 +712,7 @@ def _load_half[W: Int]( var base = (e * 8 + j) * 16 + off q[j][e] = ( bitcast[DType.uint64, 1]( - buf.load[width=8, alignment=1](base) + buf.unsafe_load[width=8, alignment=1](base) )[0] ^ kwl ) @@ -722,7 +722,7 @@ def _load_half[W: Int]( @always_inline def _store_half[W: Int]( - buf: UnsafePointer[mut=True, UInt8, _, address_space=_], + buf: Pointer[mut=True, UInt8, _, address_space=_], off: Int, mut q: InlineArray[SIMD[DType.uint64, W], 8], kw: UInt64, @@ -732,7 +732,7 @@ def _store_half[W: Int]( comptime for e in range(W): comptime for j in range(8): var base = (e * 8 + j) * 16 + off - buf.store[alignment=1]( + buf.unsafe_store[alignment=1]( base, bitcast[DType.uint8, 8]( SIMD[DType.uint64, 1](q[j][e] ^ kwl) @@ -757,7 +757,7 @@ def _six_rounds[forward: Bool, W: Int]( @always_inline def _encrypt_batch[W: Int]( - cipher: CamelliaCipher, buf: UnsafePointer[mut=True, UInt8, _, address_space=_] + cipher: CamelliaCipher, buf: Pointer[mut=True, UInt8, _, address_space=_] ): var a = _load_half[W](buf, 0, cipher.kw[0]) var b = _load_half[W](buf, 8, cipher.kw[1]) @@ -780,7 +780,7 @@ def _encrypt_batch[W: Int]( @always_inline def _decrypt_batch[W: Int]( - cipher: CamelliaCipher, buf: UnsafePointer[mut=True, UInt8, _, address_space=_] + cipher: CamelliaCipher, buf: Pointer[mut=True, UInt8, _, address_space=_] ): var b = _load_half[W](buf, 0, cipher.kw[2]) var a = _load_half[W](buf, 8, cipher.kw[3]) @@ -803,7 +803,7 @@ def _decrypt_batch[W: Int]( @always_inline def _batch[encrypt: Bool, W: Int]( - cipher: CamelliaCipher, buf: UnsafePointer[mut=True, UInt8, _, address_space=_] + cipher: CamelliaCipher, buf: Pointer[mut=True, UInt8, _, address_space=_] ): comptime if encrypt: _encrypt_batch[W](cipher, buf) @@ -931,11 +931,11 @@ def _six_rounds_bs[forward: Bool]( @always_inline def _batch16_hw[encrypt: Bool]( - cipher: CamelliaCipher, buf: UnsafePointer[mut=True, UInt8, _, address_space=_] + cipher: CamelliaCipher, buf: Pointer[mut=True, UInt8, _, address_space=_] ): var m = InlineArray[_U8x16, 16](fill=_U8x16(0)) comptime for i in range(16): - m[i] = buf.load[width=16, alignment=1](i * 16) + m[i] = buf.unsafe_load[width=16, alignment=1](i * 16) _transpose16(m) var a = InlineArray[_U8x16, 8](fill=_U8x16(0)) @@ -980,7 +980,7 @@ def _batch16_hw[encrypt: Bool]( _transpose16(m) comptime for blk in range(16): - buf.store[alignment=1](blk * 16, m[_BITREV4[blk]]) + buf.unsafe_store[alignment=1](blk * 16, m[_BITREV4[blk]]) @always_inline @@ -1094,38 +1094,38 @@ def _camellia_block[encrypt: Bool]( def _camellia_blocks[encrypt: Bool]( cipher: CamelliaCipher, - data: UnsafePointer[mut=True, UInt8, _, address_space=_], + data: Pointer[mut=True, UInt8, _, address_space=_], num_blocks: Int, ): comptime if _has_hw_sbox(): var i = 0 while i + 16 <= num_blocks: - _batch16_hw[encrypt](cipher, data + i * 16) + _batch16_hw[encrypt](cipher, data.unsafe_offset(i * 16)) i += 16 if i < num_blocks: var scratch = InlineArray[UInt8, 256](fill=0) var sp = scratch.unsafe_ptr() for j in range((num_blocks - i) * 16): - sp[j] = data[i * 16 + j] + sp[unsafe_offset=j] = data[unsafe_offset=i * 16 + j] _batch16_hw[encrypt](cipher, sp) for j in range((num_blocks - i) * 16): - data[i * 16 + j] = sp[j] + data[unsafe_offset=i * 16 + j] = sp[unsafe_offset=j] else: var i = 0 while i + 32 <= num_blocks: - _batch[encrypt, 4](cipher, data + i * 16) + _batch[encrypt, 4](cipher, data.unsafe_offset(i * 16)) i += 32 while i + 8 <= num_blocks: - _batch[encrypt, 1](cipher, data + i * 16) + _batch[encrypt, 1](cipher, data.unsafe_offset(i * 16)) i += 8 if i < num_blocks: var scratch = InlineArray[UInt8, 128](fill=0) var sp = scratch.unsafe_ptr() for j in range((num_blocks - i) * 16): - sp[j] = data[i * 16 + j] + sp[unsafe_offset=j] = data[unsafe_offset=i * 16 + j] _batch[encrypt, 1](cipher, sp) for j in range((num_blocks - i) * 16): - data[i * 16 + j] = sp[j] + data[unsafe_offset=i * 16 + j] = sp[unsafe_offset=j] def camellia_encrypt_block( @@ -1142,7 +1142,7 @@ def camellia_decrypt_block( def camellia_encrypt_blocks( cipher: CamelliaCipher, - data: UnsafePointer[mut=True, UInt8, _, address_space=_], + data: Pointer[mut=True, UInt8, _, address_space=_], num_blocks: Int, ): _camellia_blocks[True](cipher, data, num_blocks) @@ -1150,61 +1150,61 @@ def camellia_encrypt_blocks( def camellia_decrypt_blocks( cipher: CamelliaCipher, - data: UnsafePointer[mut=True, UInt8, _, address_space=_], + data: Pointer[mut=True, UInt8, _, address_space=_], num_blocks: Int, ): _camellia_blocks[False](cipher, data, num_blocks) def camellia_cbc_encrypt_kernel( - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], cipher: CamelliaCipher, num_blocks: Int, - iv_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + iv_ptr: Pointer[mut=True, UInt8, _, address_space=_], ): - var prev = iv_ptr.load[width=16, alignment=1](0) + var prev = iv_ptr.unsafe_load[width=16, alignment=1](0) for i in range(num_blocks): - var x = input_ptr.load[width=16, alignment=1](i * 16) ^ prev + var x = input_ptr.unsafe_load[width=16, alignment=1](i * 16) ^ prev prev = camellia_encrypt_block(cipher, x) - output_ptr.store[alignment=1](i * 16, prev) + output_ptr.unsafe_store[alignment=1](i * 16, prev) def camellia_cbc_decrypt_kernel( - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], cipher: CamelliaCipher, num_blocks: Int, - iv_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + iv_ptr: Pointer[mut=True, UInt8, _, address_space=_], ): var ct = InlineArray[UInt8, 1024](fill=0) var pt = InlineArray[UInt8, 1024](fill=0) var ctp = ct.unsafe_ptr() var ptp = pt.unsafe_ptr() - var prev = iv_ptr.load[width=16, alignment=1](0) + var prev = iv_ptr.unsafe_load[width=16, alignment=1](0) var i = 0 while i < num_blocks: var n = num_blocks - i if n > 64: n = 64 for b in range(n): - var v = input_ptr.load[width=16, alignment=1]((i + b) * 16) - ctp.store[alignment=1](b * 16, v) - ptp.store[alignment=1](b * 16, v) + var v = input_ptr.unsafe_load[width=16, alignment=1]((i + b) * 16) + ctp.unsafe_store[alignment=1](b * 16, v) + ptp.unsafe_store[alignment=1](b * 16, v) camellia_decrypt_blocks(cipher, ptp, n) for b in range(n): - var out = ptp.load[width=16, alignment=1](b * 16) ^ prev - output_ptr.store[alignment=1]((i + b) * 16, out) - prev = ctp.load[width=16, alignment=1](b * 16) + var out = ptp.unsafe_load[width=16, alignment=1](b * 16) ^ prev + output_ptr.unsafe_store[alignment=1]((i + b) * 16, out) + prev = ctp.unsafe_load[width=16, alignment=1](b * 16) i += n def camellia_ctr_kernel( - input_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], - output_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + input_ptr: Pointer[mut=True, UInt8, _, address_space=_], + output_ptr: Pointer[mut=True, UInt8, _, address_space=_], cipher: CamelliaCipher, num_blocks: Int, - nonce_ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], + nonce_ptr: Pointer[mut=True, UInt8, _, address_space=_], ): var ks = InlineArray[UInt8, 512](fill=0) var kp = ks.unsafe_ptr() @@ -1212,42 +1212,42 @@ def camellia_ctr_kernel( comptime if _has_hw_sbox(): while i + 16 <= num_blocks: for k in range(16): - _ctr_write_block(kp + k * 16, nonce_ptr, i + k) + _ctr_write_block(kp.unsafe_offset(k * 16), nonce_ptr, i + k) _batch16_hw[True](cipher, kp) for b in range(16): var off = (i + b) * 16 - output_ptr.store[alignment=1]( + output_ptr.unsafe_store[alignment=1]( off, - input_ptr.load[width=16, alignment=1](off) - ^ kp.load[width=16, alignment=1](b * 16), + input_ptr.unsafe_load[width=16, alignment=1](off) + ^ kp.unsafe_load[width=16, alignment=1](b * 16), ) i += 16 while i < num_blocks: var n = num_blocks - i for k in range(16): _ctr_write_block( - kp + k * 16, nonce_ptr, i + (k if k < n else 0) + kp.unsafe_offset(k * 16), nonce_ptr, i + (k if k < n else 0) ) _batch16_hw[True](cipher, kp) for b in range(n): var off = (i + b) * 16 - output_ptr.store[alignment=1]( + output_ptr.unsafe_store[alignment=1]( off, - input_ptr.load[width=16, alignment=1](off) - ^ kp.load[width=16, alignment=1](b * 16), + input_ptr.unsafe_load[width=16, alignment=1](off) + ^ kp.unsafe_load[width=16, alignment=1](b * 16), ) i += n else: while i + 32 <= num_blocks: for k in range(32): - _ctr_write_block(kp + k * 16, nonce_ptr, i + k) + _ctr_write_block(kp.unsafe_offset(k * 16), nonce_ptr, i + k) _encrypt_batch[4](cipher, kp) for b in range(32): var off = (i + b) * 16 - output_ptr.store[alignment=1]( + output_ptr.unsafe_store[alignment=1]( off, - input_ptr.load[width=16, alignment=1](off) - ^ kp.load[width=16, alignment=1](b * 16), + input_ptr.unsafe_load[width=16, alignment=1](off) + ^ kp.unsafe_load[width=16, alignment=1](b * 16), ) i += 32 while i < num_blocks: @@ -1256,14 +1256,14 @@ def camellia_ctr_kernel( n = 8 for k in range(8): _ctr_write_block( - kp + k * 16, nonce_ptr, i + (k if k < n else 0) + kp.unsafe_offset(k * 16), nonce_ptr, i + (k if k < n else 0) ) _encrypt_batch[1](cipher, kp) for b in range(n): var off = (i + b) * 16 - output_ptr.store[alignment=1]( + output_ptr.unsafe_store[alignment=1]( off, - input_ptr.load[width=16, alignment=1](off) - ^ kp.load[width=16, alignment=1](b * 16), + input_ptr.unsafe_load[width=16, alignment=1](off) + ^ kp.unsafe_load[width=16, alignment=1](b * 16), ) i += n diff --git a/src/thistle/chacha20.mojo b/src/thistle/chacha20.mojo index 655a263..7113b42 100644 --- a/src/thistle/chacha20.mojo +++ b/src/thistle/chacha20.mojo @@ -1,7 +1,7 @@ # ChaCha20 stream cipher per RFC 7539. from std.memory import bitcast -from std.memory.unsafe_pointer import UnsafePointer +from std.memory.unsafe_pointer import Pointer from std.bit import rotate_bits_left comptime CHACHA_CONSTANTS = SIMD[DType.uint32, 4]( @@ -63,7 +63,7 @@ def simd_quarter_round( aa = aa + bb dd = dd ^ aa - dd = _rotl[16, dd.size](dd) + dd = _rotl[16, dd.length](dd) cc = cc + dd bb = bb ^ cc @@ -71,7 +71,7 @@ def simd_quarter_round( aa = aa + bb dd = dd ^ aa - dd = _rotl[8, dd.size](dd) + dd = _rotl[8, dd.length](dd) cc = cc + dd bb = bb ^ cc @@ -528,26 +528,36 @@ def _chacha20_block_scalar( return result +@always_inline +def _chacha20_nonce_words(nonce: Span[UInt8, ...]) raises -> SIMD[DType.uint32, 4]: + if len(nonce) != 12: + raise Error("ChaCha20 nonce must be exactly 12 bytes") + var padded = SIMD[DType.uint8, 16](0) + for i in range(12): + padded[i] = nonce[i] + return bitcast[DType.uint32, 4](padded) + + def chacha20_block( - key: SIMD[DType.uint8, 32], counter: UInt32, nonce: SIMD[DType.uint8, 16] -) -> SIMD[DType.uint8, 64]: + key: SIMD[DType.uint8, 32], counter: UInt32, nonce: Span[UInt8, ...] +) raises -> SIMD[DType.uint8, 64]: var key_words = bitcast[DType.uint32, 8](key) - var nonce_words = bitcast[DType.uint32, 4](nonce) + var nonce_words = _chacha20_nonce_words(nonce) var state = chacha20_block_core(key_words, counter, nonce_words) return bitcast[DType.uint8, 64](state) @always_inline def _xor_block64( - src: UnsafePointer[mut=True, UInt8, _, address_space=_], - dst: UnsafePointer[mut=True, UInt8, _, address_space=_], + src: Pointer[mut=True, UInt8, _, address_space=_], + dst: Pointer[mut=True, UInt8, _, address_space=_], keystream: SIMD[DType.uint32, 16], offset: Int, ): var ks = bitcast[DType.uint8, 64](keystream) - var v = (src + offset).load[width=64, alignment=1](0) - (dst + offset).store[alignment=1](0, v ^ ks) + var v = (src.unsafe_offset(offset)).unsafe_load[width=64, alignment=1](0) + (dst.unsafe_offset(offset)).unsafe_store[alignment=1](0, v ^ ks) struct ChaCha20: @@ -558,11 +568,11 @@ struct ChaCha20: def __init__( out self, key_bytes: SIMD[DType.uint8, 32], - nonce_bytes: SIMD[DType.uint8, 16], + nonce_bytes: Span[UInt8, ...], counter: UInt32 = 1, - ): + ) raises: self.key = bitcast[DType.uint32, 8](key_bytes) - self.nonce = bitcast[DType.uint32, 4](nonce_bytes) + self.nonce = _chacha20_nonce_words(nonce_bytes) self.counter = counter def _check_counter_space(self, data_len: Int) raises: @@ -574,8 +584,8 @@ struct ChaCha20: @always_inline def _stream_xor( mut self, - src: UnsafePointer[mut=True, UInt8, _, address_space=_], - dst: UnsafePointer[mut=True, UInt8, _, address_space=_], + src: Pointer[mut=True, UInt8, _, address_space=_], + dst: Pointer[mut=True, UInt8, _, address_space=_], length: Int, ) raises: self._check_counter_space(length) @@ -626,7 +636,7 @@ struct ChaCha20: ) var ks_u8 = bitcast[DType.uint8, 64](keystream) for i in range(length - offset): - (dst + offset + i).store(0, (src + offset + i).load(0) ^ ks_u8[i]) + (dst.unsafe_offset(offset).unsafe_offset(i)).unsafe_store(0, (src.unsafe_offset(offset).unsafe_offset(i)).unsafe_load(0) ^ ks_u8[i]) block_idx += 1 self.counter += UInt32(block_idx) diff --git a/src/thistle/chacha20poly1305.mojo b/src/thistle/chacha20poly1305.mojo index 3e4d268..dea0fb4 100644 --- a/src/thistle/chacha20poly1305.mojo +++ b/src/thistle/chacha20poly1305.mojo @@ -3,9 +3,9 @@ ChaCha20-Poly1305 and XChaCha20-Poly1305 AEAD RFC 8439 """ from std.memory import bitcast -from std.memory.unsafe_pointer import UnsafePointer +from std.memory.unsafe_pointer import Pointer from std.collections import InlineArray -from .chacha20 import ChaCha20, chacha20_block_core, simd_double_round, CHACHA_CONSTANTS +from .chacha20 import ChaCha20, chacha20_block_core, simd_double_round, CHACHA_CONSTANTS, _chacha20_nonce_words from .poly1305 import Poly1305 def hchacha20( @@ -18,8 +18,8 @@ def hchacha20( if len(output) < 32: raise Error("HChaCha20 output needs at least 32 writable bytes") var out_ptr = output.unsafe_ptr() - var kw = key.unsafe_ptr().bitcast[UInt32]().load[width=8, alignment=1]() - var iw = input16.unsafe_ptr().bitcast[UInt32]().load[width=4, alignment=1]() + var kw = key.unsafe_ptr().unsafe_bitcast[UInt32]().unsafe_load[width=8, alignment=1]() + var iw = input16.unsafe_ptr().unsafe_bitcast[UInt32]().unsafe_load[width=4, alignment=1]() var row0 = CHACHA_CONSTANTS var row1 = SIMD[DType.uint32, 4](kw[0], kw[1], kw[2], kw[3]) @@ -30,15 +30,15 @@ def hchacha20( var dr = simd_double_round(row0, row1, row2, row3) row0 = dr[0]; row1 = dr[1]; row2 = dr[2]; row3 = dr[3] - out_ptr.bitcast[UInt32]().store[alignment=1](0, row0) - (out_ptr + 16).bitcast[UInt32]().store[alignment=1](0, row3) + out_ptr.unsafe_bitcast[UInt32]().unsafe_store[alignment=1](0, row0) + (out_ptr.unsafe_offset(16)).unsafe_bitcast[UInt32]().unsafe_store[alignment=1](0, row3) def _aead_tag( poly_key: Span[UInt8, ...], aad: Span[UInt8, ...], ciphertext: Span[UInt8, ...], - output: UnsafePointer[mut=True, UInt8, _, address_space=_], + output: Pointer[mut=True, UInt8, _, address_space=_], ) raises: var p = Poly1305(poly_key) var zeros16 = InlineArray[UInt8, 16](fill=0) @@ -50,8 +50,8 @@ def _aead_tag( if len(ciphertext) % 16 != 0: p.update(Span[UInt8, ...](unsafe_ptr=zp, length=16 - len(ciphertext) % 16)) var lens = InlineArray[UInt8, 16](uninitialized=True) - lens.unsafe_ptr().bitcast[UInt64]().store[alignment=1](0, UInt64(len(aad))) - (lens.unsafe_ptr() + 8).bitcast[UInt64]().store[alignment=1](0, UInt64(len(ciphertext))) + lens.unsafe_ptr().unsafe_bitcast[UInt64]().unsafe_store[alignment=1](0, UInt64(len(aad))) + (lens.unsafe_ptr().unsafe_offset(8)).unsafe_bitcast[UInt64]().unsafe_store[alignment=1](0, UInt64(len(ciphertext))) p.update(Span[UInt8, ...](unsafe_ptr=lens.unsafe_ptr(), length=16)) p.finalize_into(output) @@ -61,24 +61,25 @@ def _aead_core[encrypt: Bool]( nonce: Span[UInt8, ...], aad: Span[UInt8, ...], input: Span[UInt8, ...], - output: UnsafePointer[mut=True, UInt8, _, address_space=_], - tag: UnsafePointer[mut=True, UInt8, _, address_space=_], + output: Pointer[mut=True, UInt8, _, address_space=_], + tag: Pointer[mut=True, UInt8, _, address_space=_], ) raises: - var key_bytes = key.unsafe_ptr().load[width=32, alignment=1](0) - var nonce_bytes = SIMD[DType.uint8, 16](0) + var key_bytes = key.unsafe_ptr().unsafe_load[width=32, alignment=1](0) + var nonce_bytes = InlineArray[UInt8, 12](fill=0) for i in range(12): nonce_bytes[i] = nonce[i] var kw = bitcast[DType.uint32, 8](key_bytes) - var nw = bitcast[DType.uint32, 4](nonce_bytes) + var nonce_span = Span[UInt8, ...](nonce_bytes) + var nw = _chacha20_nonce_words(nonce_span) var block0 = chacha20_block_core(kw, 0, nw) var poly_key = InlineArray[UInt8, 32](uninitialized=True) - poly_key.unsafe_ptr().store[alignment=1]( + poly_key.unsafe_ptr().unsafe_store[alignment=1]( 0, bitcast[DType.uint8, 64](block0).slice[32]() ) var poly_key_span = Span[UInt8, ...](unsafe_ptr=poly_key.unsafe_ptr(), length=32) - var cipher = ChaCha20(key_bytes, nonce_bytes, counter=1) + var cipher = ChaCha20(key_bytes, nonce_span, counter=1) var src = input.unsafe_ptr().unsafe_mut_cast[True]().unsafe_origin_cast[MutAnyOrigin]() cipher._stream_xor(src, output, len(input)) @@ -91,7 +92,7 @@ def _aead_core[encrypt: Bool]( _aead_tag(poly_key_span, aad, input, tag) var poly_key_ptr = poly_key.unsafe_ptr() for i in range(32): - poly_key_ptr.store[volatile=True](i, UInt8(0)) + poly_key_ptr.unsafe_store[volatile=True](i, UInt8(0)) def chacha20_poly1305_encrypt( @@ -132,15 +133,16 @@ def chacha20_poly1305_decrypt( if len(plaintext) < len(ciphertext): raise Error("ChaCha20-Poly1305 plaintext output is too small") - var key_bytes = key.unsafe_ptr().load[width=32, alignment=1](0) - var nonce_bytes = SIMD[DType.uint8, 16](0) + var key_bytes = key.unsafe_ptr().unsafe_load[width=32, alignment=1](0) + var nonce_bytes = InlineArray[UInt8, 12](fill=0) for i in range(12): nonce_bytes[i] = nonce[i] var kw = bitcast[DType.uint32, 8](key_bytes) - var nw = bitcast[DType.uint32, 4](nonce_bytes) + var nonce_span = Span[UInt8, ...](nonce_bytes) + var nw = _chacha20_nonce_words(nonce_span) var block0 = chacha20_block_core(kw, 0, nw) var poly_key = InlineArray[UInt8, 32](uninitialized=True) - poly_key.unsafe_ptr().store[alignment=1]( + poly_key.unsafe_ptr().unsafe_store[alignment=1]( 0, bitcast[DType.uint8, 64](block0).slice[32]() ) @@ -155,15 +157,15 @@ def chacha20_poly1305_decrypt( if diff != 0: var poly_key_ptr = poly_key.unsafe_ptr() for i in range(32): - poly_key_ptr.store[volatile=True](i, UInt8(0)) + poly_key_ptr.unsafe_store[volatile=True](i, UInt8(0)) return False - var cipher = ChaCha20(key_bytes, nonce_bytes, counter=1) + var cipher = ChaCha20(key_bytes, nonce_span, counter=1) var src = ciphertext.unsafe_ptr().unsafe_mut_cast[True]().unsafe_origin_cast[MutAnyOrigin]() cipher._stream_xor(src, plaintext.unsafe_ptr(), len(ciphertext)) var poly_key_ptr = poly_key.unsafe_ptr() for i in range(32): - poly_key_ptr.store[volatile=True](i, UInt8(0)) + poly_key_ptr.unsafe_store[volatile=True](i, UInt8(0)) return True @@ -183,11 +185,11 @@ def xchacha20_poly1305_encrypt( var sp = sub.unsafe_ptr().unsafe_origin_cast[MutAnyOrigin]() chacha20_poly1305_encrypt( Span[UInt8, ...](unsafe_ptr=sp, length=32), - Span[UInt8, ...](unsafe_ptr=sp + 32, length=12), + Span[UInt8, ...](unsafe_ptr=sp.unsafe_offset(32), length=12), aad, plaintext, ciphertext, tag, ) for i in range(44): - sp.store[volatile=True](i, UInt8(0)) + sp.unsafe_store[volatile=True](i, UInt8(0)) def xchacha20_poly1305_decrypt( @@ -206,11 +208,11 @@ def xchacha20_poly1305_decrypt( var sp = sub.unsafe_ptr().unsafe_origin_cast[MutAnyOrigin]() var ok = chacha20_poly1305_decrypt( Span[UInt8, ...](unsafe_ptr=sp, length=32), - Span[UInt8, ...](unsafe_ptr=sp + 32, length=12), + Span[UInt8, ...](unsafe_ptr=sp.unsafe_offset(32), length=12), aad, ciphertext, tag, plaintext, ) for i in range(44): - sp.store[volatile=True](i, UInt8(0)) + sp.unsafe_store[volatile=True](i, UInt8(0)) return ok diff --git a/src/thistle/curve25519.mojo b/src/thistle/curve25519.mojo index ef4bac2..671eef4 100644 --- a/src/thistle/curve25519.mojo +++ b/src/thistle/curve25519.mojo @@ -289,21 +289,21 @@ struct FieldElement51(Movable, Copyable, ImplicitlyCopyable): if len(bytes) < 32: raise Error("FieldElement51 input must be at least 32 bytes") @always_inline - def load8(ptr: UnsafePointer[mut=False, UInt8, _, address_space=_]) -> UInt64: - return ptr.bitcast[UInt64]().load[width=1, alignment=1]() + def load8(ptr: Pointer[mut=False, UInt8, _, address_space=_]) -> UInt64: + return ptr.unsafe_bitcast[UInt64]().unsafe_load[width=1, alignment=1]() var ptr = bytes.unsafe_ptr() var MASK = UInt64(0x7FFFFFFFFFFFF) var l0 = load8(ptr) & MASK - var l1 = (load8(ptr + 6) >> 3) & MASK - var l2 = (load8(ptr + 12) >> 6) & MASK - var l3 = (load8(ptr + 19) >> 1) & MASK - var l4 = (load8(ptr + 24) >> 12) & MASK + var l1 = (load8(ptr.unsafe_offset(6)) >> 3) & MASK + var l2 = (load8(ptr.unsafe_offset(12)) >> 6) & MASK + var l3 = (load8(ptr.unsafe_offset(19)) >> 1) & MASK + var l4 = (load8(ptr.unsafe_offset(24)) >> 12) & MASK return FieldElement51(l0, l1, l2, l3, l4) - def to_bytes_into(self, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): + def to_bytes_into(self, output: Pointer[mut=True, UInt8, _, address_space=_]): var res = self._reduce(self.limbs) var limbs = res.limbs @@ -327,7 +327,7 @@ struct FieldElement51(Movable, Copyable, ImplicitlyCopyable): var w2 = (limbs[2] >> 26) | (limbs[3] << 25) var w3 = (limbs[3] >> 39) | (limbs[4] << 12) - (output + 0).bitcast[UInt64]().store[alignment=1](w0) - (output + 8).bitcast[UInt64]().store[alignment=1](w1) - (output + 16).bitcast[UInt64]().store[alignment=1](w2) - (output + 24).bitcast[UInt64]().store[alignment=1](w3) + (output.unsafe_offset(0)).unsafe_bitcast[UInt64]().unsafe_store[alignment=1](w0) + (output.unsafe_offset(8)).unsafe_bitcast[UInt64]().unsafe_store[alignment=1](w1) + (output.unsafe_offset(16)).unsafe_bitcast[UInt64]().unsafe_store[alignment=1](w2) + (output.unsafe_offset(24)).unsafe_bitcast[UInt64]().unsafe_store[alignment=1](w3) diff --git a/src/thistle/ed25519.mojo b/src/thistle/ed25519.mojo index 278936e..febbada 100644 --- a/src/thistle/ed25519.mojo +++ b/src/thistle/ed25519.mojo @@ -87,7 +87,7 @@ def _encoded_y_lt_p(y: Span[UInt8, ...]) -> Bool: gt = 1 return lt == 1 -def _pack_limbs_into(limbs: SIMD[DType.uint64, 8], output: UnsafePointer[mut=True, UInt8, _, address_space=_]): +def _pack_limbs_into(limbs: SIMD[DType.uint64, 8], output: Pointer[mut=True, UInt8, _, address_space=_]): var words = SIMD[DType.uint64, 4](0, 0, 0, 0) words[0] = limbs[0] | (limbs[1] << 52) words[1] = (limbs[1] >> 12) | (limbs[2] << 40) @@ -95,11 +95,11 @@ def _pack_limbs_into(limbs: SIMD[DType.uint64, 8], output: UnsafePointer[mut=Tru words[3] = (limbs[3] >> 36) | (limbs[4] << 16) var bytes = bitcast[DType.uint8, 32](words) for i in range(32): - output[i] = bytes[i] + output[unsafe_offset=i] = bytes[i] def _unpack_limbs(bytes: Span[UInt8, ...]) -> SIMD[DType.uint64, 8]: # Input may be byte-aligned; use alignment=1 for the UInt64 wide load. - var words = bytes.unsafe_ptr().bitcast[UInt64]().load[width=4, alignment=1]() + var words = bytes.unsafe_ptr().unsafe_bitcast[UInt64]().unsafe_load[width=4, alignment=1]() comptime MASK = (UInt64(1) << 52) - 1 comptime TOP_MASK = (UInt64(1) << 48) - 1 var s = SIMD[DType.uint64, 8](0) @@ -114,7 +114,7 @@ def _from_512_raw(bytes: Span[UInt8, ...]) -> SIMD[DType.uint64, 8]: # RFC 8032 5.1.6: reduce 64-byte SHA-512 output modulo L. var ptr = bytes.unsafe_ptr() var lo_span = Span[UInt8, ...](unsafe_ptr=ptr, length=32) - var hi_span = Span[UInt8, ...](unsafe_ptr=ptr + 32, length=32) + var hi_span = Span[UInt8, ...](unsafe_ptr=ptr.unsafe_offset(32), length=32) var lo = Scalar.from_bytes(lo_span) var hi = Scalar.from_bytes(hi_span) var pow2_256 = Scalar(POW2_256_LIMBS) @@ -167,7 +167,7 @@ struct Scalar(Movable, Copyable, ImplicitlyCopyable): var raw = _unpack_limbs(bytes) return Scalar(raw)._montgomery_mul(Scalar(RR_LIMBS)) - def to_bytes_into(self, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): + def to_bytes_into(self, output: Pointer[mut=True, UInt8, _, address_space=_]): var raw = self._montgomery_mul( Scalar(SIMD[DType.uint64, 8](1, 0, 0, 0, 0, 0, 0, 0)) ) @@ -239,7 +239,7 @@ struct Scalar(Movable, Copyable, ImplicitlyCopyable): return Scalar(r)._sub(Scalar(L_LIMBS)) def wipe(mut self): - UnsafePointer(to=self.limbs).bitcast[UInt64]().store[volatile=True]( + Pointer(to=self.limbs).unsafe_bitcast[UInt64]().unsafe_store[volatile=True]( 0, SIMD[DType.uint64, 8](0) ) @@ -352,34 +352,34 @@ def _edwards_double_standalone(p: EdwardsPoint) -> EdwardsPoint: @no_inline def fe_from_bytes(bytes: Span[UInt8, ...]) -> FieldElement51: # Decode 255-bit little-endian field element; caller clears x-parity bit. - def load8(ptr: UnsafePointer[mut=False, UInt8, _, address_space=_]) -> UInt64: + def load8(ptr: Pointer[mut=False, UInt8, _, address_space=_]) -> UInt64: var v: UInt64 = 0 for j in range(8): - v |= UInt64(ptr[j]) << UInt64(j * 8) + v |= UInt64(ptr[unsafe_offset=j]) << UInt64(j * 8) return v var ptr = bytes.unsafe_ptr() var MASK = UInt64(0x7FFFFFFFFFFFF) var l0 = load8(ptr) & MASK - var l1 = (load8(ptr + 6) >> UInt64(3)) & MASK - var l2 = (load8(ptr + 12) >> UInt64(6)) & MASK - var l3 = (load8(ptr + 19) >> UInt64(1)) & MASK - var l4 = (load8(ptr + 24) >> UInt64(12)) & MASK + var l1 = (load8(ptr.unsafe_offset(6)) >> UInt64(3)) & MASK + var l2 = (load8(ptr.unsafe_offset(12)) >> UInt64(6)) & MASK + var l3 = (load8(ptr.unsafe_offset(19)) >> UInt64(1)) & MASK + var l4 = (load8(ptr.unsafe_offset(24)) >> UInt64(12)) & MASK return FieldElement51(l0, l1, l2, l3, l4) @no_inline -def edwards_encode_into(p: EdwardsPoint, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): +def edwards_encode_into(p: EdwardsPoint, output: Pointer[mut=True, UInt8, _, address_space=_]): # RFC 8032 5.1.2: encode y and store x parity in bit 255. _edwards_encode_with_zinv(p, p.Z.invert(), output) @no_inline -def _edwards_encode_with_zinv(p: EdwardsPoint, z_inv: FieldElement51, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): +def _edwards_encode_with_zinv(p: EdwardsPoint, z_inv: FieldElement51, output: Pointer[mut=True, UInt8, _, address_space=_]): var x = p.X * z_inv var y = p.Y * z_inv y.to_bytes_into(output) var x_bytes = InlineArray[UInt8, 32](uninitialized=True) x.to_bytes_into(x_bytes.unsafe_ptr()) var x_parity = x_bytes[0] & 1 - output[31] = output[31] | (x_parity << 7) + output[unsafe_offset=31] = output[unsafe_offset=31] | (x_parity << 7) @no_inline def edwards_decode(data: Span[UInt8, ...], strict: Bool = True) -> DecodeResult: @@ -605,16 +605,16 @@ def _radix16_digits(scalar: Span[UInt8, ...]) -> InlineArray[Int, 64]: return digits^ @always_inline -def _base_table_lookup(ptr: UnsafePointer[UInt64, _], j: Int, digit: Int) -> AffineNielsPoint: +def _base_table_lookup(ptr: Pointer[UInt64, _], j: Int, digit: Int) -> AffineNielsPoint: var d = Int64(digit) var sign = d >> 63 var absv = ((d ^ sign) - sign).cast[DType.uint64]() var acc = SIMD[DType.uint64, 16](0) acc[0] = 1 acc[5] = 1 - var base = ptr + j * 128 + var base = ptr.unsafe_offset(j * 128) for k in range(1, 9): - var cand = (base + (k - 1) * 16).load[width=16, alignment=8]() + var cand = (base.unsafe_offset((k - 1) * 16)).unsafe_load[width=16, alignment=8]() var diff = absv ^ UInt64(k) var m = UInt64(0) - ((diff - 1) >> 63) acc = acc ^ ((acc ^ cand) & SIMD[DType.uint64, 16](m)) @@ -643,9 +643,9 @@ def _mul_base_ct(scalar: Span[UInt8, ...]) -> EdwardsPoint: P = _edwards_double_standalone(P) for i in range(0, 64, 2): P = _add_affine_niels(P, _base_table_lookup(tptr, i >> 1, digits[i])) - var dptr = digits.unsafe_ptr().bitcast[UInt64]() + var dptr = digits.unsafe_ptr().unsafe_bitcast[UInt64]() for i in range(64): - dptr.store[volatile=True](i, UInt64(0)) + dptr.unsafe_store[volatile=True](i, UInt64(0)) return P def _naf5(scalar: Span[UInt8, ...]) -> InlineArray[Int, 256]: @@ -654,7 +654,7 @@ def _naf5(scalar: Span[UInt8, ...]) -> InlineArray[Int, 256]: words[4] = 0 var ptr = scalar.unsafe_ptr() for w in range(4): - words[w] = (ptr + 8 * w).bitcast[UInt64]().load[width=1, alignment=1]() + words[w] = (ptr.unsafe_offset(8 * w)).unsafe_bitcast[UInt64]().unsafe_load[width=1, alignment=1]() var pos = 0 var carry: UInt64 = 0 while pos < 256: @@ -677,12 +677,12 @@ def _naf5(scalar: Span[UInt8, ...]) -> InlineArray[Int, 256]: return naf^ @always_inline -def _b_odd_entry(ptr: UnsafePointer[UInt64, _], k: Int) -> AffineNielsPoint: - var base = ptr + k * 16 +def _b_odd_entry(ptr: Pointer[UInt64, _], k: Int) -> AffineNielsPoint: + var base = ptr.unsafe_offset(k * 16) return AffineNielsPoint( - FieldElement51(base[0], base[1], base[2], base[3], base[4]), - FieldElement51(base[5], base[6], base[7], base[8], base[9]), - FieldElement51(base[10], base[11], base[12], base[13], base[14]), + FieldElement51(base[unsafe_offset=0], base[unsafe_offset=1], base[unsafe_offset=2], base[unsafe_offset=3], base[unsafe_offset=4]), + FieldElement51(base[unsafe_offset=5], base[unsafe_offset=6], base[unsafe_offset=7], base[unsafe_offset=8], base[unsafe_offset=9]), + FieldElement51(base[unsafe_offset=10], base[unsafe_offset=11], base[unsafe_offset=12], base[unsafe_offset=13], base[unsafe_offset=14]), ) @no_inline @@ -718,7 +718,7 @@ def _double_scalar_mult_vartime(a: Span[UInt8, ...], A: EdwardsPoint, b: Span[UI return Q @no_inline -def ed25519_generate_public_key(private_key: Span[UInt8, ...], output: UnsafePointer[mut=True, UInt8, _, address_space=_]) raises: +def ed25519_generate_public_key(private_key: Span[UInt8, ...], output: Pointer[mut=True, UInt8, _, address_space=_]) raises: # RFC 8032 5.1.5: public key A = [pruned SHA512(secret)]B. if len(private_key) != 32: raise Error("Ed25519 private key must be 32 bytes") @@ -736,12 +736,12 @@ def ed25519_generate_public_key(private_key: Span[UInt8, ...], output: UnsafePoi var hash_ptr = hash.unsafe_ptr() var s_ptr = s_bytes.unsafe_ptr() for i in range(64): - hash_ptr.store[volatile=True](i, UInt8(0)) + hash_ptr.unsafe_store[volatile=True](i, UInt8(0)) for i in range(32): - s_ptr.store[volatile=True](i, UInt8(0)) + s_ptr.unsafe_store[volatile=True](i, UInt8(0)) @no_inline -def ed25519_sign(private_key: Span[UInt8, ...], message: Span[UInt8, ...], output: UnsafePointer[mut=True, UInt8, _, address_space=_]) raises: +def ed25519_sign(private_key: Span[UInt8, ...], message: Span[UInt8, ...], output: Pointer[mut=True, UInt8, _, address_space=_]) raises: # RFC 8032 5.1.6 pure Ed25519: # r = SHA512(prefix || M), R = [r]B, # k = SHA512(R || A || M), S = r + k*s mod L. @@ -759,7 +759,7 @@ def ed25519_sign(private_key: Span[UInt8, ...], message: Span[UInt8, ...], outpu var r_hash = InlineArray[UInt8, 64](uninitialized=True) var r_ctx = SHA512Context() - sha512_update(r_ctx, Span[UInt8, ...](unsafe_ptr=hash.unsafe_ptr() + 32, length=32)) + sha512_update(r_ctx, Span[UInt8, ...](unsafe_ptr=hash.unsafe_ptr().unsafe_offset(32), length=32)) sha512_update(r_ctx, message) sha512_final_to_buffer(r_ctx, r_hash.unsafe_ptr()) @@ -786,8 +786,8 @@ def ed25519_sign(private_key: Span[UInt8, ...], message: Span[UInt8, ...], outpu var S_bytes = InlineArray[UInt8, 32](uninitialized=True) S_scalar.to_bytes_into(S_bytes.unsafe_ptr()) - for i in range(32): output[i] = R_enc[i] - for i in range(32): output[32 + i] = S_bytes[i] + for i in range(32): output[unsafe_offset=i] = R_enc[i] + for i in range(32): output[unsafe_offset=32 + i] = S_bytes[i] ctx.wipe() r_ctx.wipe() s_scalar.wipe() @@ -799,13 +799,13 @@ def ed25519_sign(private_key: Span[UInt8, ...], message: Span[UInt8, ...], outpu var k_ptr = k_hash.unsafe_ptr() var S_ptr = S_bytes.unsafe_ptr() for i in range(64): - hash_ptr.store[volatile=True](i, UInt8(0)) - r_ptr.store[volatile=True](i, UInt8(0)) - k_ptr.store[volatile=True](i, UInt8(0)) + hash_ptr.unsafe_store[volatile=True](i, UInt8(0)) + r_ptr.unsafe_store[volatile=True](i, UInt8(0)) + k_ptr.unsafe_store[volatile=True](i, UInt8(0)) for i in range(32): - s_ptr.store[volatile=True](i, UInt8(0)) - r_bytes_ptr.store[volatile=True](i, UInt8(0)) - S_ptr.store[volatile=True](i, UInt8(0)) + s_ptr.unsafe_store[volatile=True](i, UInt8(0)) + r_bytes_ptr.unsafe_store[volatile=True](i, UInt8(0)) + S_ptr.unsafe_store[volatile=True](i, UInt8(0)) struct Ed25519SigningKey(Copyable, Movable): var _s: Scalar @@ -835,22 +835,22 @@ struct Ed25519SigningKey(Copyable, Movable): var hash_ptr = hash.unsafe_ptr() var s_ptr = s_bytes.unsafe_ptr() for i in range(64): - hash_ptr.store[volatile=True](i, UInt8(0)) + hash_ptr.unsafe_store[volatile=True](i, UInt8(0)) for i in range(32): - s_ptr.store[volatile=True](i, UInt8(0)) + s_ptr.unsafe_store[volatile=True](i, UInt8(0)) - def __del__(deinit self): + def __deinit__(deinit self): self._s.wipe() var p = self._prefix.unsafe_ptr() for i in range(32): - p.store[volatile=True](i, UInt8(0)) + p.unsafe_store[volatile=True](i, UInt8(0)) - def public_key_into(self, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): + def public_key_into(self, output: Pointer[mut=True, UInt8, _, address_space=_]): for i in range(32): output[i] = self._a_enc[i] @no_inline - def sign(self, message: Span[UInt8, ...], output: UnsafePointer[mut=True, UInt8, _, address_space=_]) raises: + def sign(self, message: Span[UInt8, ...], output: Pointer[mut=True, UInt8, _, address_space=_]) raises: var r_hash = InlineArray[UInt8, 64](uninitialized=True) var r_ctx = SHA512Context() sha512_update(r_ctx, Span[UInt8, ...](unsafe_ptr=self._prefix.unsafe_ptr(), length=32)) @@ -889,11 +889,11 @@ struct Ed25519SigningKey(Copyable, Movable): var r_bytes_ptr = r_bytes.unsafe_ptr() var S_ptr = S_bytes.unsafe_ptr() for i in range(64): - r_ptr.store[volatile=True](i, UInt8(0)) - k_ptr.store[volatile=True](i, UInt8(0)) + r_ptr.unsafe_store[volatile=True](i, UInt8(0)) + k_ptr.unsafe_store[volatile=True](i, UInt8(0)) for i in range(32): - r_bytes_ptr.store[volatile=True](i, UInt8(0)) - S_ptr.store[volatile=True](i, UInt8(0)) + r_bytes_ptr.unsafe_store[volatile=True](i, UInt8(0)) + S_ptr.unsafe_store[volatile=True](i, UInt8(0)) @no_inline diff --git a/src/thistle/ed25519_table.mojo b/src/thistle/ed25519_table.mojo index fad4e16..0ebaf69 100644 --- a/src/thistle/ed25519_table.mojo +++ b/src/thistle/ed25519_table.mojo @@ -369,42 +369,42 @@ comptime _ED25519_B_ODD = SIMD[DType.uint64, 128]( def ed25519_base_table() -> InlineArray[UInt64, 4096]: var t = InlineArray[UInt64, 4096](uninitialized=True) var p = t.unsafe_ptr() - p.store[alignment=8](0, _ED25519_BT0) - p.store[alignment=8](128, _ED25519_BT1) - p.store[alignment=8](256, _ED25519_BT2) - p.store[alignment=8](384, _ED25519_BT3) - p.store[alignment=8](512, _ED25519_BT4) - p.store[alignment=8](640, _ED25519_BT5) - p.store[alignment=8](768, _ED25519_BT6) - p.store[alignment=8](896, _ED25519_BT7) - p.store[alignment=8](1024, _ED25519_BT8) - p.store[alignment=8](1152, _ED25519_BT9) - p.store[alignment=8](1280, _ED25519_BT10) - p.store[alignment=8](1408, _ED25519_BT11) - p.store[alignment=8](1536, _ED25519_BT12) - p.store[alignment=8](1664, _ED25519_BT13) - p.store[alignment=8](1792, _ED25519_BT14) - p.store[alignment=8](1920, _ED25519_BT15) - p.store[alignment=8](2048, _ED25519_BT16) - p.store[alignment=8](2176, _ED25519_BT17) - p.store[alignment=8](2304, _ED25519_BT18) - p.store[alignment=8](2432, _ED25519_BT19) - p.store[alignment=8](2560, _ED25519_BT20) - p.store[alignment=8](2688, _ED25519_BT21) - p.store[alignment=8](2816, _ED25519_BT22) - p.store[alignment=8](2944, _ED25519_BT23) - p.store[alignment=8](3072, _ED25519_BT24) - p.store[alignment=8](3200, _ED25519_BT25) - p.store[alignment=8](3328, _ED25519_BT26) - p.store[alignment=8](3456, _ED25519_BT27) - p.store[alignment=8](3584, _ED25519_BT28) - p.store[alignment=8](3712, _ED25519_BT29) - p.store[alignment=8](3840, _ED25519_BT30) - p.store[alignment=8](3968, _ED25519_BT31) + p.unsafe_store[alignment=8](0, _ED25519_BT0) + p.unsafe_store[alignment=8](128, _ED25519_BT1) + p.unsafe_store[alignment=8](256, _ED25519_BT2) + p.unsafe_store[alignment=8](384, _ED25519_BT3) + p.unsafe_store[alignment=8](512, _ED25519_BT4) + p.unsafe_store[alignment=8](640, _ED25519_BT5) + p.unsafe_store[alignment=8](768, _ED25519_BT6) + p.unsafe_store[alignment=8](896, _ED25519_BT7) + p.unsafe_store[alignment=8](1024, _ED25519_BT8) + p.unsafe_store[alignment=8](1152, _ED25519_BT9) + p.unsafe_store[alignment=8](1280, _ED25519_BT10) + p.unsafe_store[alignment=8](1408, _ED25519_BT11) + p.unsafe_store[alignment=8](1536, _ED25519_BT12) + p.unsafe_store[alignment=8](1664, _ED25519_BT13) + p.unsafe_store[alignment=8](1792, _ED25519_BT14) + p.unsafe_store[alignment=8](1920, _ED25519_BT15) + p.unsafe_store[alignment=8](2048, _ED25519_BT16) + p.unsafe_store[alignment=8](2176, _ED25519_BT17) + p.unsafe_store[alignment=8](2304, _ED25519_BT18) + p.unsafe_store[alignment=8](2432, _ED25519_BT19) + p.unsafe_store[alignment=8](2560, _ED25519_BT20) + p.unsafe_store[alignment=8](2688, _ED25519_BT21) + p.unsafe_store[alignment=8](2816, _ED25519_BT22) + p.unsafe_store[alignment=8](2944, _ED25519_BT23) + p.unsafe_store[alignment=8](3072, _ED25519_BT24) + p.unsafe_store[alignment=8](3200, _ED25519_BT25) + p.unsafe_store[alignment=8](3328, _ED25519_BT26) + p.unsafe_store[alignment=8](3456, _ED25519_BT27) + p.unsafe_store[alignment=8](3584, _ED25519_BT28) + p.unsafe_store[alignment=8](3712, _ED25519_BT29) + p.unsafe_store[alignment=8](3840, _ED25519_BT30) + p.unsafe_store[alignment=8](3968, _ED25519_BT31) return t^ @no_inline def ed25519_b_odd_table() -> InlineArray[UInt64, 128]: var t = InlineArray[UInt64, 128](uninitialized=True) - t.unsafe_ptr().store[alignment=8](0, _ED25519_B_ODD) + t.unsafe_ptr().unsafe_store[alignment=8](0, _ED25519_B_ODD) return t^ diff --git a/src/thistle/kcipher2.mojo b/src/thistle/kcipher2.mojo index 8c9648f..c53e321 100644 --- a/src/thistle/kcipher2.mojo +++ b/src/thistle/kcipher2.mojo @@ -557,7 +557,7 @@ struct KCipher2: ](mut self, mut data: Span[mut=True, UInt8, origin]): var len_data = len(data) var data_ptr = data.unsafe_ptr() - var data_u64 = data_ptr.bitcast[UInt64]() + var data_u64 = data_ptr.unsafe_bitcast[UInt64]() var num_u64 = len_data // 8 var i = 0 @@ -570,25 +570,25 @@ struct KCipher2: self._next_normal() var ks3 = self.stream() self._next_normal() - (data_u64 + i).store[alignment=1]( - (data_u64 + i).load[width=1, alignment=1]() ^ ks0 + (data_u64.unsafe_offset(i)).unsafe_store[alignment=1]( + (data_u64.unsafe_offset(i)).unsafe_load[width=1, alignment=1]() ^ ks0 ) - (data_u64 + i + 1).store[alignment=1]( - (data_u64 + i + 1).load[width=1, alignment=1]() ^ ks1 + (data_u64.unsafe_offset(i).unsafe_offset(1)).unsafe_store[alignment=1]( + (data_u64.unsafe_offset(i).unsafe_offset(1)).unsafe_load[width=1, alignment=1]() ^ ks1 ) - (data_u64 + i + 2).store[alignment=1]( - (data_u64 + i + 2).load[width=1, alignment=1]() ^ ks2 + (data_u64.unsafe_offset(i).unsafe_offset(2)).unsafe_store[alignment=1]( + (data_u64.unsafe_offset(i).unsafe_offset(2)).unsafe_load[width=1, alignment=1]() ^ ks2 ) - (data_u64 + i + 3).store[alignment=1]( - (data_u64 + i + 3).load[width=1, alignment=1]() ^ ks3 + (data_u64.unsafe_offset(i).unsafe_offset(3)).unsafe_store[alignment=1]( + (data_u64.unsafe_offset(i).unsafe_offset(3)).unsafe_load[width=1, alignment=1]() ^ ks3 ) i += 4 while i < num_u64: var ks = self.stream() self._next_normal() - (data_u64 + i).store[alignment=1]( - (data_u64 + i).load[width=1, alignment=1]() ^ ks + (data_u64.unsafe_offset(i)).unsafe_store[alignment=1]( + (data_u64.unsafe_offset(i)).unsafe_load[width=1, alignment=1]() ^ ks ) i += 1 diff --git a/src/thistle/ml_dsa.mojo b/src/thistle/ml_dsa.mojo index 01729d5..70cda66 100644 --- a/src/thistle/ml_dsa.mojo +++ b/src/thistle/ml_dsa.mojo @@ -12,7 +12,7 @@ data, zeroize sensitive temporaries with volatile stores when practical. from std.collections import List from std.builtin.globals import global_constant -from std.memory import memset_zero +from std.memory import unsafe_memset_zero from thistle.sha3 import SHA3Context, sha3_update, shake_final, shake128, shake256 from thistle.random import random_bytes from thistle.utils import StackBuffer, zero_stack_u8 @@ -83,7 +83,7 @@ comptime ZETAS_TABLE: InlineArray[UInt32, 256] = [ def _zeta(i: Int) -> UInt32: debug_assert(i >= 0 and i < 256, "ML-DSA zeta index out of bounds") ref zetas = global_constant[ZETAS_TABLE]() - return zetas.unsafe_ptr()[i] + return zetas.unsafe_ptr()[unsafe_offset=i] @fieldwise_init @@ -133,25 +133,25 @@ struct MLDSAPrivateKey(Copyable, Movable): var t0: List[List[UInt32]] var k_seed: List[UInt8] - def __del__(deinit self): + def __deinit__(deinit self): var seed_ptr = self.seed.unsafe_ptr() for i in range(len(self.seed)): - seed_ptr.store[volatile=True](i, UInt8(0)) + seed_ptr.unsafe_store[volatile=True](i, UInt8(0)) var key_ptr = self.k_seed.unsafe_ptr() for i in range(len(self.k_seed)): - key_ptr.store[volatile=True](i, UInt8(0)) + key_ptr.unsafe_store[volatile=True](i, UInt8(0)) for row in range(len(self.s1)): var ptr = self.s1[row].unsafe_ptr() for i in range(len(self.s1[row])): - ptr.store[volatile=True](i, UInt32(0)) + ptr.unsafe_store[volatile=True](i, UInt32(0)) for row in range(len(self.s2)): var ptr = self.s2[row].unsafe_ptr() for i in range(len(self.s2[row])): - ptr.store[volatile=True](i, UInt32(0)) + ptr.unsafe_store[volatile=True](i, UInt32(0)) for row in range(len(self.t0)): var ptr = self.t0[row].unsafe_ptr() for i in range(len(self.t0[row])): - ptr.store[volatile=True](i, UInt32(0)) + ptr.unsafe_store[volatile=True](i, UInt32(0)) def params44() -> MLDSAParams: @@ -183,7 +183,7 @@ def private_key_size(p: MLDSAParams) -> Int: def _zero_poly() -> List[UInt32]: var r = List[UInt32](unsafe_uninit_length=N) - memset_zero(r.unsafe_ptr(), N) + unsafe_memset_zero(r.unsafe_ptr(), N) return r^ @@ -198,25 +198,25 @@ def _zero_dsa_poly_vec[ROWS: Int]() -> DSAPolyVec[ROWS]: def _zero_dsa_poly(mut p: DSAPoly): var ptr = p.unsafe_ptr() for i in range(N): - ptr.store[volatile=True](i, UInt32(0)) + ptr.unsafe_store[volatile=True](i, UInt32(0)) def _zero_dsa_poly_vec[ROWS: Int](mut v: DSAPolyVec[ROWS]): - var ptr = v.data.unsafe_ptr().bitcast[UInt32]() + var ptr = v.data.unsafe_ptr().unsafe_bitcast[UInt32]() for i in range(ROWS * N): - ptr.store[volatile=True](i, UInt32(0)) + ptr.unsafe_store[volatile=True](i, UInt32(0)) def _zero_dsa_hint_vec(mut h: DSAHintVec): - var ptr = h.unsafe_ptr().bitcast[UInt8]() + var ptr = h.unsafe_ptr().unsafe_bitcast[UInt8]() for i in range(MAX_K * N): - ptr.store[volatile=True](i, UInt8(0)) + ptr.unsafe_store[volatile=True](i, UInt8(0)) def _zero_dsa_ch(mut ch: InlineArray[UInt8, MLDSA_CRHBYTES]): var ptr = ch.unsafe_ptr() for i in range(MLDSA_CRHBYTES): - ptr.store[volatile=True](i, UInt8(0)) + ptr.unsafe_store[volatile=True](i, UInt8(0)) def _zero_poly_vec(count: Int) -> List[List[UInt32]]: @@ -255,13 +255,13 @@ def _ct_bool_to_u32(b: Bool) -> UInt32: def _zero_list_u8(mut data: List[UInt8]): var ptr = data.unsafe_ptr() for i in range(len(data)): - ptr.store[volatile=True](i, UInt8(0)) + ptr.unsafe_store[volatile=True](i, UInt8(0)) def _zero_list_u32(mut data: List[UInt32]): var ptr = data.unsafe_ptr() for i in range(len(data)): - ptr.store[volatile=True](i, UInt32(0)) + ptr.unsafe_store[volatile=True](i, UInt32(0)) def _zero_poly_vec_u32(mut data: List[List[UInt32]]): @@ -406,13 +406,13 @@ def _ntt_mul_into(mut r: List[UInt32], a: List[UInt32], b: List[UInt32]): @always_inline def _ntt_mul_ptrs( - r: UnsafePointer[mut=True, UInt32, _, address_space=_], - a: UnsafePointer[mut=False, UInt32, _, address_space=_], - b: UnsafePointer[mut=False, UInt32, _, address_space=_], + r: Pointer[mut=True, UInt32, _, address_space=_], + a: Pointer[mut=False, UInt32, _, address_space=_], + b: Pointer[mut=False, UInt32, _, address_space=_], ): var i = 0 while i < N: - r.store(i, _montgomery_mul_v(a.load[width=_VW](i), b.load[width=_VW](i))) + r.unsafe_store(i, _montgomery_mul_v(a.unsafe_load[width=_VW](i), b.unsafe_load[width=_VW](i))) i += _VW @@ -459,11 +459,11 @@ def _dsa_ntt_inplace(mut f: DSAPoly): var zv = _U32v(zeta).cast[DType.uint64]() var j = start while j < start + length: - var a = p.load[width=_VW](j) - var b = p.load[width=_VW](j + length) + var a = p.unsafe_load[width=_VW](j) + var b = p.unsafe_load[width=_VW](j + length) var t = _montgomery_reduce_v(zv * b.cast[DType.uint64]()) - p.store(j + length, _field_reduce_once_v(a - t + _U32v(Q))) - p.store(j, _field_reduce_once_v(a + t)) + p.unsafe_store(j + length, _field_reduce_once_v(a - t + _U32v(Q))) + p.unsafe_store(j, _field_reduce_once_v(a + t)) j += _VW else: var j = start @@ -514,10 +514,10 @@ def _dsa_inverse_ntt_inplace(mut f: DSAPoly): var zv = _U32v(zeta).cast[DType.uint64]() var j = start while j < start + length: - var a = p.load[width=_VW](j) - var b = p.load[width=_VW](j + length) - p.store(j, _field_reduce_once_v(a + b)) - p.store( + var a = p.unsafe_load[width=_VW](j) + var b = p.unsafe_load[width=_VW](j + length) + p.unsafe_store(j, _field_reduce_once_v(a + b)) + p.unsafe_store( j + length, _montgomery_reduce_v( zv * (b - a + _U32v(Q)).cast[DType.uint64]() @@ -535,7 +535,7 @@ def _dsa_inverse_ntt_inplace(mut f: DSAPoly): length *= 2 var i = 0 while i < N: - p.store(i, _montgomery_mul_v(p.load[width=_VW](i), _U32v(16382))) + p.unsafe_store(i, _montgomery_mul_v(p.unsafe_load[width=_VW](i), _U32v(16382))) i += _VW @@ -569,13 +569,13 @@ def _sample_ntt(rho: Span[UInt8, ...], s: UInt8, r: UInt8) raises -> List[UInt32 var limit = len(buf) - 2 while off < limit and n < N: var v = ( - UInt32(bp[off]) - | (UInt32(bp[off + 1]) << 8) - | (UInt32(bp[off + 2]) << 16) + UInt32(bp[unsafe_offset=off]) + | (UInt32(bp[unsafe_offset=off + 1]) << 8) + | (UInt32(bp[unsafe_offset=off + 2]) << 16) ) & 0x7FFFFF off += 3 if v < Q: - ap[n] = _field_to_montgomery_unchecked(v) + ap[unsafe_offset=n] = _field_to_montgomery_unchecked(v) n += 1 if n == N: return a^ diff --git a/src/thistle/ml_kem.mojo b/src/thistle/ml_kem.mojo index a15a3d3..27be353 100644 --- a/src/thistle/ml_kem.mojo +++ b/src/thistle/ml_kem.mojo @@ -110,7 +110,7 @@ comptime ZETAS_TABLE: InlineArray[Int16, 128] = [ def _zeta(i: Int) -> Int16: debug_assert(i >= 0 and i < 128, "ML-KEM zeta index out of bounds") ref zetas = global_constant[ZETAS_TABLE]() - return zetas.unsafe_ptr()[i] + return zetas.unsafe_ptr()[unsafe_offset=i] struct Poly(Copyable, Movable): @@ -188,7 +188,7 @@ def barrett_reduce_simd[w: Int](a: SIMD[DType.int16, w]) -> SIMD[DType.int16, w] @always_inline def _u24_le(buf: Span[UInt8, ...], offset: Int) -> UInt32: if offset + 4 <= len(buf): - return (buf.unsafe_ptr() + offset).bitcast[UInt32]().load[alignment=1]().cast[DType.uint32]() & 0x00FFFFFF + return buf.unsafe_ptr().unsafe_offset(offset).unsafe_bitcast[UInt32]().unsafe_load[alignment=1]().cast[DType.uint32]() & 0x00FFFFFF return UInt32(buf[offset]) | (UInt32(buf[offset + 1]) << 8) | (UInt32(buf[offset + 2]) << 16) @@ -199,7 +199,7 @@ def cbd3(mut r: Poly, buf: Span[UInt8, ...]) raises: var rp = r.coeffs.unsafe_ptr() for i in range(N // 4): var off = 3 * i - var t = UInt32(bp[off]) | (UInt32(bp[off + 1]) << 8) | (UInt32(bp[off + 2]) << 16) + var t = UInt32(bp[unsafe_offset=off]) | (UInt32(bp[unsafe_offset=off + 1]) << 8) | (UInt32(bp[unsafe_offset=off + 2]) << 16) var d = t & 0x00249249 d += (t >> 1) & 0x00249249 d += (t >> 2) & 0x00249249 @@ -212,10 +212,10 @@ def cbd3(mut r: Poly, buf: Span[UInt8, ...]) raises: var a3 = Int16((d >> 18) & 0x7) var b3 = Int16((d >> 21) & 0x7) var base = 4 * i - rp.store(base + 0, a0 - b0) - rp.store(base + 1, a1 - b1) - rp.store(base + 2, a2 - b2) - rp.store(base + 3, a3 - b3) + rp.unsafe_store(base + 0, a0 - b0) + rp.unsafe_store(base + 1, a1 - b1) + rp.unsafe_store(base + 2, a2 - b2) + rp.unsafe_store(base + 3, a3 - b3) def cbd2(mut r: Poly, buf: Span[UInt8, ...]) raises: @@ -225,7 +225,7 @@ def cbd2(mut r: Poly, buf: Span[UInt8, ...]) raises: var rp = r.coeffs.unsafe_ptr() for i in range(N // 8): var off = 4 * i - var t = UInt32(bp[off]) | (UInt32(bp[off + 1]) << 8) | (UInt32(bp[off + 2]) << 16) | (UInt32(bp[off + 3]) << 24) + var t = UInt32(bp[unsafe_offset=off]) | (UInt32(bp[unsafe_offset=off + 1]) << 8) | (UInt32(bp[unsafe_offset=off + 2]) << 16) | (UInt32(bp[unsafe_offset=off + 3]) << 24) var d = t & 0x55555555 d += (t >> 1) & 0x55555555 var base = 8 * i @@ -245,14 +245,14 @@ def cbd2(mut r: Poly, buf: Span[UInt8, ...]) raises: var b6 = Int16((d >> 26) & 0x3) var a7 = Int16((d >> 28) & 0x3) var b7 = Int16((d >> 30) & 0x3) - rp.store(base + 0, a0 - b0) - rp.store(base + 1, a1 - b1) - rp.store(base + 2, a2 - b2) - rp.store(base + 3, a3 - b3) - rp.store(base + 4, a4 - b4) - rp.store(base + 5, a5 - b5) - rp.store(base + 6, a6 - b6) - rp.store(base + 7, a7 - b7) + rp.unsafe_store(base + 0, a0 - b0) + rp.unsafe_store(base + 1, a1 - b1) + rp.unsafe_store(base + 2, a2 - b2) + rp.unsafe_store(base + 3, a3 - b3) + rp.unsafe_store(base + 4, a4 - b4) + rp.unsafe_store(base + 5, a5 - b5) + rp.unsafe_store(base + 6, a6 - b6) + rp.unsafe_store(base + 7, a7 - b7) def poly_cbd_eta1_512(mut r: Poly, buf: Span[UInt8, ...]) raises: @@ -280,11 +280,11 @@ def ntt(mut r: InlineArray[Int16, N]): var zeta = _zeta(zeta_idx) comptime if l >= W: def ntt_chunk[w: Int](off: Int) {ptr, zeta}: - var lo = ptr.load[width=w](start + off) - var hi = ptr.load[width=w](start + l + off) + var lo = ptr.unsafe_load[width=w](start + off) + var hi = ptr.unsafe_load[width=w](start + l + off) var t = fqmul_simd(SIMD[DType.int16, w](zeta), hi) - ptr.store[width=w](start + l + off, lo - t) - ptr.store[width=w](start + off, lo + t) + ptr.unsafe_store[width=w](start + l + off, lo - t) + ptr.unsafe_store[width=w](start + off, lo + t) vectorize[W, size=l](ntt_chunk) else: @@ -309,10 +309,10 @@ def invntt(mut r: InlineArray[Int16, N]): var zeta = _zeta(zeta_idx) comptime if l >= W: def invntt_chunk[w: Int](off: Int) {ptr, zeta}: - var lo = ptr.load[width=w](start + off) - var hi = ptr.load[width=w](start + l + off) - ptr.store[width=w](start + off, barrett_reduce_simd(lo + hi)) - ptr.store[width=w](start + l + off, fqmul_simd(SIMD[DType.int16, w](zeta), hi - lo)) + var lo = ptr.unsafe_load[width=w](start + off) + var hi = ptr.unsafe_load[width=w](start + l + off) + ptr.unsafe_store[width=w](start + off, barrett_reduce_simd(lo + hi)) + ptr.unsafe_store[width=w](start + l + off, fqmul_simd(SIMD[DType.int16, w](zeta), hi - lo)) vectorize[W, size=l](invntt_chunk) else: @@ -324,9 +324,9 @@ def invntt(mut r: InlineArray[Int16, N]): r[j + l] = fqmul(zeta, r[j + l]) def final_mul_chunk[w: Int](i: Int) {ptr}: - var a = ptr.load[width=w](i) + var a = ptr.unsafe_load[width=w](i) var f = SIMD[DType.int16, w](F) - ptr.store[width=w](i, fqmul_simd(a, f)) + ptr.unsafe_store[width=w](i, fqmul_simd(a, f)) vectorize[W, size=N](final_mul_chunk) @@ -346,11 +346,11 @@ def poly_reduce(mut r: Poly): var ptr = r.coeffs.unsafe_ptr().unsafe_origin_cast[MutAnyOrigin]() def reduce_chunk[w: Int](i: Int) {ptr}: - var a = ptr.load[width=w](i) + var a = ptr.unsafe_load[width=w](i) var v = SIMD[DType.int16, w](Int16(((1 << 26) + Q // 2) // Q)) var qv = SIMD[DType.int16, w](Int16(Q)) var t = ((a.cast[DType.int32]() * v.cast[DType.int32]() + (1 << 25)) >> 26).cast[DType.int16]() - ptr.store[width=w](i, a - t * qv) + ptr.unsafe_store[width=w](i, a - t * qv) vectorize[W, size=N](reduce_chunk) @@ -382,9 +382,9 @@ def poly_tomont(mut r: Poly): var ptr = r.coeffs.unsafe_ptr().unsafe_origin_cast[MutAnyOrigin]() def tomont_chunk[w: Int](i: Int) {ptr}: - var a = ptr.load[width=w](i) + var a = ptr.unsafe_load[width=w](i) var f = SIMD[DType.int16, w](F) - ptr.store[width=w](i, montgomery_reduce_simd(a.cast[DType.int32]() * f.cast[DType.int32]())) + ptr.unsafe_store[width=w](i, montgomery_reduce_simd(a.cast[DType.int32]() * f.cast[DType.int32]())) vectorize[W, size=N](tomont_chunk) @@ -399,7 +399,7 @@ def poly_add_inplace(mut r: Poly, ref b: Poly): ) def add_inplace_chunk[w: Int](i: Int) {rp, bp}: - rp.store[width=w](i, rp.load[width=w](i) + bp.load[width=w](i)) + rp.unsafe_store[width=w](i, rp.unsafe_load[width=w](i) + bp.unsafe_load[width=w](i)) vectorize[W, size=N](add_inplace_chunk) @@ -414,7 +414,7 @@ def poly_sub_from(mut r: Poly, ref a: Poly): ) def sub_from_chunk[w: Int](i: Int) {rp, ap}: - rp.store[width=w](i, ap.load[width=w](i) - rp.load[width=w](i)) + rp.unsafe_store[width=w](i, ap.unsafe_load[width=w](i) - rp.unsafe_load[width=w](i)) vectorize[W, size=N](sub_from_chunk) @@ -1005,14 +1005,14 @@ struct DecapsulationKey(Copyable, Movable): self.ek = EncapsulationKey() self.z = InlineArray[UInt8, SYMBYTES](fill=0) - def __del__(deinit self): + def __deinit__(deinit self): for row in range(K_MAX): var ptr = self.pke_dk.pv.vec[row].coeffs.unsafe_ptr() for i in range(N): - ptr.store[volatile=True](i, Int16(0)) + ptr.unsafe_store[volatile=True](i, Int16(0)) var z_ptr = self.z.unsafe_ptr() for i in range(SYMBYTES): - z_ptr.store[volatile=True](i, UInt8(0)) + z_ptr.unsafe_store[volatile=True](i, UInt8(0)) def pack_pk_stack(mut out: StackBuffer[UInt8, ...], ref pk: Polyvec, seed: InlineArray[UInt8, SYMBYTES], k: Int) -> Bool: @@ -1067,7 +1067,7 @@ def k_pke_keygen(mut ek: KPKEEncryptionKey, mut dk: KPKEDecapsulationKey, d: Spa var g = StackBuffer[UInt8, 2 * SYMBYTES]() hash_g_into(g, Span[UInt8, ...](unsafe_ptr=g_input.ptr(), length=g_input.len())) var rho = Span[UInt8, ...](unsafe_ptr=g.ptr(), length=SYMBYTES) - var sigma = Span[UInt8, ...](unsafe_ptr=g.ptr() + SYMBYTES, length=SYMBYTES) + var sigma = Span[UInt8, ...](unsafe_ptr=g.ptr().unsafe_offset(SYMBYTES), length=SYMBYTES) var a = InlineArray[Polyvec, K_MAX](fill=Polyvec()) gen_matrix(a, Span[UInt8, ...](rho), False, k) @@ -1118,7 +1118,7 @@ def k_pke_keygen_k[k: Int](mut ek: KPKEEncryptionKey, mut dk: KPKEDecapsulationK var g = StackBuffer[UInt8, 2 * SYMBYTES]() hash_g_into(g, Span[UInt8, ...](unsafe_ptr=g_input.ptr(), length=g_input.len())) var rho = Span[UInt8, ...](unsafe_ptr=g.ptr(), length=SYMBYTES) - var sigma = Span[UInt8, ...](unsafe_ptr=g.ptr() + SYMBYTES, length=SYMBYTES) + var sigma = Span[UInt8, ...](unsafe_ptr=g.ptr().unsafe_offset(SYMBYTES), length=SYMBYTES) var a = InlineArray[Polyvec, K_MAX](fill=Polyvec()) gen_matrix_k_static[k, False](a, Span[UInt8, ...](rho)) @@ -1426,7 +1426,7 @@ def _ct_select_u8(a: UInt8, b: UInt8, choice: UInt8) -> UInt8: def _zero_list(mut data: List[UInt8]): var ptr = data.unsafe_ptr() for i in range(len(data)): - ptr.store[volatile=True](i, UInt8(0)) + ptr.unsafe_store[volatile=True](i, UInt8(0)) @@ -1647,7 +1647,7 @@ def mlkem_encaps_seed_into(mut ciphertext_out: StackBuffer[UInt8, CIPHERTEXTBYTE zero_stack_u8(g_input) for i in range(SYMBYTES): shared_out.push_unchecked(g[i]) - if not k_pke_encrypt_into(ciphertext_out, ek, m, Span[UInt8, ...](unsafe_ptr=g.ptr() + SYMBYTES, length=SYMBYTES)): + if not k_pke_encrypt_into(ciphertext_out, ek, m, Span[UInt8, ...](unsafe_ptr=g.ptr().unsafe_offset(SYMBYTES), length=SYMBYTES)): zero_stack_u8(g) zero_stack_u8(shared_out) zero_stack_u8(ciphertext_out) @@ -1679,7 +1679,7 @@ def mlkem_encaps_seed_into_k[k: Int](mut ciphertext_out: StackBuffer[UInt8, CIPH zero_stack_u8(g_input) for i in range(SYMBYTES): shared_out.push_unchecked(g[i]) - if not k_pke_encrypt_into_k[k](ciphertext_out, ek, m, Span[UInt8, ...](unsafe_ptr=g.ptr() + SYMBYTES, length=SYMBYTES)): + if not k_pke_encrypt_into_k[k](ciphertext_out, ek, m, Span[UInt8, ...](unsafe_ptr=g.ptr().unsafe_offset(SYMBYTES), length=SYMBYTES)): zero_stack_u8(g) zero_stack_u8(shared_out) zero_stack_u8(ciphertext_out) @@ -1744,7 +1744,7 @@ def mlkem_decaps_into(mut shared_out: StackBuffer[UInt8, SYMBYTES], dk_bytes: Sp zero_stack_u8(g_input) var ct_check = StackBuffer[UInt8, CIPHERTEXTBYTES_MAX]() - if not k_pke_encrypt_into(ct_check, dk.ek.pke_ek, Span[UInt8, ...](unsafe_ptr=m.ptr(), length=m.len()), Span[UInt8, ...](unsafe_ptr=g.ptr() + SYMBYTES, length=SYMBYTES)): + if not k_pke_encrypt_into(ct_check, dk.ek.pke_ek, Span[UInt8, ...](unsafe_ptr=m.ptr(), length=m.len()), Span[UInt8, ...](unsafe_ptr=g.ptr().unsafe_offset(SYMBYTES), length=SYMBYTES)): zero_stack_u8(m) zero_stack_u8(g) zero_stack_u8(ct_check) @@ -1788,7 +1788,7 @@ def mlkem_decaps_into_k[k: Int](mut shared_out: StackBuffer[UInt8, SYMBYTES], dk zero_stack_u8(g_input) var ct_check = StackBuffer[UInt8, CIPHERTEXTBYTES_MAX]() - if not k_pke_encrypt_into_k[k](ct_check, dk.ek.pke_ek, Span[UInt8, ...](unsafe_ptr=m.ptr(), length=m.len()), Span[UInt8, ...](unsafe_ptr=g.ptr() + SYMBYTES, length=SYMBYTES)): + if not k_pke_encrypt_into_k[k](ct_check, dk.ek.pke_ek, Span[UInt8, ...](unsafe_ptr=m.ptr(), length=m.len()), Span[UInt8, ...](unsafe_ptr=g.ptr().unsafe_offset(SYMBYTES), length=SYMBYTES)): zero_stack_u8(m) zero_stack_u8(g) zero_stack_u8(ct_check) diff --git a/src/thistle/p256.mojo b/src/thistle/p256.mojo index 899f090..fc8a0f7 100644 --- a/src/thistle/p256.mojo +++ b/src/thistle/p256.mojo @@ -454,11 +454,11 @@ def _from_be(bytes: Span[UInt8, ...]) -> U256: return out -def _to_be(x: U256, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): +def _to_be(x: U256, output: Pointer[mut=True, UInt8, _, address_space=_]): for i in range(4): var limb = x.limbs[3 - i] for k in range(8): - output[i * 8 + k] = UInt8((limb >> UInt64(56 - 8 * k)) & 0xFF) + output[unsafe_offset=i * 8 + k] = UInt8((limb >> UInt64(56 - 8 * k)) & 0xFF) struct P256Point(Copyable, ImplicitlyCopyable, Movable): @@ -711,15 +711,15 @@ def _scalar_mult(k: U256, p: P256Point) -> P256Point: @always_inline def _base_table_entry( - tptr: UnsafePointer[UInt64, _], j: Int, d: UInt64 + tptr: Pointer[UInt64, _], j: Int, d: UInt64 ) -> P256Point: var qx = U256() var qy = U256() for t in range(1, 16): var hit = u64_zero_choice(UInt64(t) ^ d) var base = (j * 15 + (t - 1)) * 8 - var ex = U256(tptr[base], tptr[base + 1], tptr[base + 2], tptr[base + 3]) - var ey = U256(tptr[base + 4], tptr[base + 5], tptr[base + 6], tptr[base + 7]) + var ex = U256(tptr[unsafe_offset=base], tptr[unsafe_offset=base + 1], tptr[unsafe_offset=base + 2], tptr[unsafe_offset=base + 3]) + var ey = U256(tptr[unsafe_offset=base + 4], tptr[unsafe_offset=base + 5], tptr[unsafe_offset=base + 6], tptr[unsafe_offset=base + 7]) qx = _select_u256(qx, ex, hit) qy = _select_u256(qy, ey, hit) return P256Point(qx, qy, False) @@ -753,7 +753,7 @@ def _scalar_mult_base(k: U256) -> P256Point: def p256_decode_uncompressed(point: Span[UInt8, ...]) -> P256Point: if len(point) == 33 and (point[0] == 0x02 or point[0] == 0x03): var x = _from_be( - Span[UInt8, ...](unsafe_ptr=point.unsafe_ptr() + 1, length=32) + Span[UInt8, ...](unsafe_ptr=point.unsafe_ptr().unsafe_offset(1), length=32) ) if _cmp(x, _p()) >= 0: return P256Point() @@ -772,8 +772,8 @@ def p256_decode_uncompressed(point: Span[UInt8, ...]) -> P256Point: return p if len(point) != P256_POINT_SIZE or point[0] != 0x04: return P256Point() - var x = _from_be(Span[UInt8, ...](unsafe_ptr=point.unsafe_ptr() + 1, length=32)) - var y = _from_be(Span[UInt8, ...](unsafe_ptr=point.unsafe_ptr() + 33, length=32)) + var x = _from_be(Span[UInt8, ...](unsafe_ptr=point.unsafe_ptr().unsafe_offset(1), length=32)) + var y = _from_be(Span[UInt8, ...](unsafe_ptr=point.unsafe_ptr().unsafe_offset(33), length=32)) var p = P256Point(x, y, False) if not _is_on_curve(p): return P256Point() @@ -786,9 +786,9 @@ def p256_encode_uncompressed( if len(output) < P256_POINT_SIZE or point.infinity or not _is_on_curve(point): return False var out_ptr = output.unsafe_ptr() - out_ptr[0] = 0x04 - _to_be(point.x, out_ptr + 1) - _to_be(point.y, out_ptr + 33) + out_ptr[unsafe_offset=0] = 0x04 + _to_be(point.x, out_ptr.unsafe_offset(1)) + _to_be(point.y, out_ptr.unsafe_offset(33)) return True @@ -800,15 +800,15 @@ def p256_public_key( return False var d = _from_be(private_key) if d.is_zero() or _cmp(d, _n()) >= 0: - var dp = UnsafePointer(to=d.limbs[0]).unsafe_mut_cast[True]() + var dp = Pointer(to=d.limbs[0]).unsafe_mut_cast[True]() for i in range(4): - dp.store[volatile=True](i, UInt64(0)) + dp.unsafe_store[volatile=True](i, UInt64(0)) return False var q = _scalar_mult_base(d) var ok = p256_encode_uncompressed(q, output) - var dp = UnsafePointer(to=d.limbs[0]).unsafe_mut_cast[True]() + var dp = Pointer(to=d.limbs[0]).unsafe_mut_cast[True]() for i in range(4): - dp.store[volatile=True](i, UInt64(0)) + dp.unsafe_store[volatile=True](i, UInt64(0)) return ok @@ -822,20 +822,20 @@ def p256_ecdh( return False var d = _from_be(private_key) if d.is_zero() or _cmp(d, _n()) >= 0: - var dp = UnsafePointer(to=d.limbs[0]).unsafe_mut_cast[True]() + var dp = Pointer(to=d.limbs[0]).unsafe_mut_cast[True]() for i in range(4): - dp.store[volatile=True](i, UInt64(0)) + dp.unsafe_store[volatile=True](i, UInt64(0)) return False var q = p256_decode_uncompressed(public_key) if q.infinity: - var dp = UnsafePointer(to=d.limbs[0]).unsafe_mut_cast[True]() + var dp = Pointer(to=d.limbs[0]).unsafe_mut_cast[True]() for i in range(4): - dp.store[volatile=True](i, UInt64(0)) + dp.unsafe_store[volatile=True](i, UInt64(0)) return False var shared = _scalar_mult(d, q) - var dp = UnsafePointer(to=d.limbs[0]).unsafe_mut_cast[True]() + var dp = Pointer(to=d.limbs[0]).unsafe_mut_cast[True]() for i in range(4): - dp.store[volatile=True](i, UInt64(0)) + dp.unsafe_store[volatile=True](i, UInt64(0)) if shared.infinity: return False _to_be(shared.x, output.unsafe_ptr()) @@ -929,15 +929,15 @@ def _reduce_n(x: U256) -> U256: def _wipe_u256(mut x: U256): - var ptr = UnsafePointer(to=x.limbs[0]).unsafe_mut_cast[True]() + var ptr = Pointer(to=x.limbs[0]).unsafe_mut_cast[True]() for i in range(4): - ptr.store[volatile=True](i, UInt64(0)) + ptr.unsafe_store[volatile=True](i, UInt64(0)) def _wipe_list_u8(mut data: List[UInt8]): var ptr = data.unsafe_ptr() for i in range(len(data)): - ptr.store[volatile=True](i, UInt8(0)) + ptr.unsafe_store[volatile=True](i, UInt8(0)) def _rfc6979_p256(private_key: Span[UInt8, ...], digest: Span[UInt8, ...], skip: Int) -> U256: @@ -991,7 +991,7 @@ def _rfc6979_p256(private_key: Span[UInt8, ...], digest: Span[UInt8, ...], skip: _wipe_list_u8(seed) var hp = h1_bytes.unsafe_ptr() for i in range(32): - hp.store[volatile=True](i, UInt8(0)) + hp.unsafe_store[volatile=True](i, UInt8(0)) _wipe_u256(h1) return candidate accepted += 1 @@ -1046,7 +1046,7 @@ def p256_ecdsa_sign_digest( retry += 1 continue _to_be(r, signature_ptr) - _to_be(s, signature_ptr + 32) + _to_be(s, signature_ptr.unsafe_offset(32)) _wipe_u256(d) _wipe_u256(z) _wipe_u256(k) diff --git a/src/thistle/p256_table.mojo b/src/thistle/p256_table.mojo index d5e5d5f..e54f986 100644 --- a/src/thistle/p256_table.mojo +++ b/src/thistle/p256_table.mojo @@ -614,36 +614,36 @@ comptime _P256_BT31 = SIMD[DType.uint64, 128]( def p256_base_table() -> InlineArray[UInt64, 4096]: var t = InlineArray[UInt64, 4096](uninitialized=True) var p = t.unsafe_ptr() - p.store[alignment=8](0, _P256_BT0) - p.store[alignment=8](128, _P256_BT1) - p.store[alignment=8](256, _P256_BT2) - p.store[alignment=8](384, _P256_BT3) - p.store[alignment=8](512, _P256_BT4) - p.store[alignment=8](640, _P256_BT5) - p.store[alignment=8](768, _P256_BT6) - p.store[alignment=8](896, _P256_BT7) - p.store[alignment=8](1024, _P256_BT8) - p.store[alignment=8](1152, _P256_BT9) - p.store[alignment=8](1280, _P256_BT10) - p.store[alignment=8](1408, _P256_BT11) - p.store[alignment=8](1536, _P256_BT12) - p.store[alignment=8](1664, _P256_BT13) - p.store[alignment=8](1792, _P256_BT14) - p.store[alignment=8](1920, _P256_BT15) - p.store[alignment=8](2048, _P256_BT16) - p.store[alignment=8](2176, _P256_BT17) - p.store[alignment=8](2304, _P256_BT18) - p.store[alignment=8](2432, _P256_BT19) - p.store[alignment=8](2560, _P256_BT20) - p.store[alignment=8](2688, _P256_BT21) - p.store[alignment=8](2816, _P256_BT22) - p.store[alignment=8](2944, _P256_BT23) - p.store[alignment=8](3072, _P256_BT24) - p.store[alignment=8](3200, _P256_BT25) - p.store[alignment=8](3328, _P256_BT26) - p.store[alignment=8](3456, _P256_BT27) - p.store[alignment=8](3584, _P256_BT28) - p.store[alignment=8](3712, _P256_BT29) - p.store[alignment=8](3840, _P256_BT30) - p.store[alignment=8](3968, _P256_BT31) + p.unsafe_store[alignment=8](0, _P256_BT0) + p.unsafe_store[alignment=8](128, _P256_BT1) + p.unsafe_store[alignment=8](256, _P256_BT2) + p.unsafe_store[alignment=8](384, _P256_BT3) + p.unsafe_store[alignment=8](512, _P256_BT4) + p.unsafe_store[alignment=8](640, _P256_BT5) + p.unsafe_store[alignment=8](768, _P256_BT6) + p.unsafe_store[alignment=8](896, _P256_BT7) + p.unsafe_store[alignment=8](1024, _P256_BT8) + p.unsafe_store[alignment=8](1152, _P256_BT9) + p.unsafe_store[alignment=8](1280, _P256_BT10) + p.unsafe_store[alignment=8](1408, _P256_BT11) + p.unsafe_store[alignment=8](1536, _P256_BT12) + p.unsafe_store[alignment=8](1664, _P256_BT13) + p.unsafe_store[alignment=8](1792, _P256_BT14) + p.unsafe_store[alignment=8](1920, _P256_BT15) + p.unsafe_store[alignment=8](2048, _P256_BT16) + p.unsafe_store[alignment=8](2176, _P256_BT17) + p.unsafe_store[alignment=8](2304, _P256_BT18) + p.unsafe_store[alignment=8](2432, _P256_BT19) + p.unsafe_store[alignment=8](2560, _P256_BT20) + p.unsafe_store[alignment=8](2688, _P256_BT21) + p.unsafe_store[alignment=8](2816, _P256_BT22) + p.unsafe_store[alignment=8](2944, _P256_BT23) + p.unsafe_store[alignment=8](3072, _P256_BT24) + p.unsafe_store[alignment=8](3200, _P256_BT25) + p.unsafe_store[alignment=8](3328, _P256_BT26) + p.unsafe_store[alignment=8](3456, _P256_BT27) + p.unsafe_store[alignment=8](3584, _P256_BT28) + p.unsafe_store[alignment=8](3712, _P256_BT29) + p.unsafe_store[alignment=8](3840, _P256_BT30) + p.unsafe_store[alignment=8](3968, _P256_BT31) return t^ diff --git a/src/thistle/p384.mojo b/src/thistle/p384.mojo index 22fe48d..5d65764 100644 --- a/src/thistle/p384.mojo +++ b/src/thistle/p384.mojo @@ -399,11 +399,11 @@ def _from_be(bytes: Span[UInt8, ...]) -> U384: return out -def _to_be(x: U384, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): +def _to_be(x: U384, output: Pointer[mut=True, UInt8, _, address_space=_]): for i in range(6): var limb = x.limbs[5 - i] for k in range(8): - output[i * 8 + k] = UInt8((limb >> UInt64(56 - 8 * k)) & 0xFF) + output[unsafe_offset=i * 8 + k] = UInt8((limb >> UInt64(56 - 8 * k)) & 0xFF) struct P384Point(Copyable, ImplicitlyCopyable, Movable): @@ -653,7 +653,7 @@ def _scalar_mult(k: U384, p: P384Point) -> P384Point: @always_inline def _base_table_entry( - tptr: UnsafePointer[UInt64, _], j: Int, d: UInt64 + tptr: Pointer[UInt64, _], j: Int, d: UInt64 ) -> P384Point: # Scan the whole window. var qx = U384() @@ -661,8 +661,8 @@ def _base_table_entry( for t in range(1, 16): var hit = u64_zero_choice(UInt64(t) ^ d) var base = (j * 15 + (t - 1)) * 12 - var ex = U384(tptr[base], tptr[base + 1], tptr[base + 2], tptr[base + 3], tptr[base + 4], tptr[base + 5]) - var ey = U384(tptr[base + 6], tptr[base + 7], tptr[base + 8], tptr[base + 9], tptr[base + 10], tptr[base + 11]) + var ex = U384(tptr[unsafe_offset=base], tptr[unsafe_offset=base + 1], tptr[unsafe_offset=base + 2], tptr[unsafe_offset=base + 3], tptr[unsafe_offset=base + 4], tptr[unsafe_offset=base + 5]) + var ey = U384(tptr[unsafe_offset=base + 6], tptr[unsafe_offset=base + 7], tptr[unsafe_offset=base + 8], tptr[unsafe_offset=base + 9], tptr[unsafe_offset=base + 10], tptr[unsafe_offset=base + 11]) qx = _select_u384(qx, ex, hit) qy = _select_u384(qy, ey, hit) return P384Point(qx, qy, False) @@ -697,7 +697,7 @@ def _scalar_mult_base(k: U384) -> P384Point: def p384_decode_uncompressed(point: Span[UInt8, ...]) -> P384Point: if len(point) == 49 and (point[0] == 0x02 or point[0] == 0x03): var x = _from_be( - Span[UInt8, ...](unsafe_ptr=point.unsafe_ptr() + 1, length=48) + Span[UInt8, ...](unsafe_ptr=point.unsafe_ptr().unsafe_offset(1), length=48) ) if _cmp(x, _p()) >= 0: return P384Point() @@ -716,8 +716,8 @@ def p384_decode_uncompressed(point: Span[UInt8, ...]) -> P384Point: return p if len(point) != P384_POINT_SIZE or point[0] != 0x04: return P384Point() - var x = _from_be(Span[UInt8, ...](unsafe_ptr=point.unsafe_ptr() + 1, length=48)) - var y = _from_be(Span[UInt8, ...](unsafe_ptr=point.unsafe_ptr() + 49, length=48)) + var x = _from_be(Span[UInt8, ...](unsafe_ptr=point.unsafe_ptr().unsafe_offset(1), length=48)) + var y = _from_be(Span[UInt8, ...](unsafe_ptr=point.unsafe_ptr().unsafe_offset(49), length=48)) var p = P384Point(x, y, False) if not _is_on_curve(p): return P384Point() @@ -730,9 +730,9 @@ def p384_encode_uncompressed( if len(output) < P384_POINT_SIZE or point.infinity or not _is_on_curve(point): return False var out_ptr = output.unsafe_ptr() - out_ptr[0] = 0x04 - _to_be(point.x, out_ptr + 1) - _to_be(point.y, out_ptr + 49) + out_ptr[unsafe_offset=0] = 0x04 + _to_be(point.x, out_ptr.unsafe_offset(1)) + _to_be(point.y, out_ptr.unsafe_offset(49)) return True @@ -744,15 +744,15 @@ def p384_public_key( return False var d = _from_be(private_key) if d.is_zero() or _cmp(d, _n()) >= 0: - var dp = UnsafePointer(to=d.limbs[0]).unsafe_mut_cast[True]() + var dp = Pointer(to=d.limbs[0]).unsafe_mut_cast[True]() for i in range(6): - dp.store[volatile=True](i, UInt64(0)) + dp.unsafe_store[volatile=True](i, UInt64(0)) return False var q = _scalar_mult_base(d) var ok = p384_encode_uncompressed(q, output) - var dp = UnsafePointer(to=d.limbs[0]).unsafe_mut_cast[True]() + var dp = Pointer(to=d.limbs[0]).unsafe_mut_cast[True]() for i in range(6): - dp.store[volatile=True](i, UInt64(0)) + dp.unsafe_store[volatile=True](i, UInt64(0)) return ok @@ -766,20 +766,20 @@ def p384_ecdh( return False var d = _from_be(private_key) if d.is_zero() or _cmp(d, _n()) >= 0: - var dp = UnsafePointer(to=d.limbs[0]).unsafe_mut_cast[True]() + var dp = Pointer(to=d.limbs[0]).unsafe_mut_cast[True]() for i in range(6): - dp.store[volatile=True](i, UInt64(0)) + dp.unsafe_store[volatile=True](i, UInt64(0)) return False var q = p384_decode_uncompressed(public_key) if q.infinity: - var dp = UnsafePointer(to=d.limbs[0]).unsafe_mut_cast[True]() + var dp = Pointer(to=d.limbs[0]).unsafe_mut_cast[True]() for i in range(6): - dp.store[volatile=True](i, UInt64(0)) + dp.unsafe_store[volatile=True](i, UInt64(0)) return False var shared = _scalar_mult(d, q) - var dp = UnsafePointer(to=d.limbs[0]).unsafe_mut_cast[True]() + var dp = Pointer(to=d.limbs[0]).unsafe_mut_cast[True]() for i in range(6): - dp.store[volatile=True](i, UInt64(0)) + dp.unsafe_store[volatile=True](i, UInt64(0)) if shared.infinity: return False _to_be(shared.x, output.unsafe_ptr()) @@ -879,15 +879,15 @@ def _reduce_n(x: U384) -> U384: def _wipe_u384(mut x: U384): - var ptr = UnsafePointer(to=x.limbs[0]).unsafe_mut_cast[True]() + var ptr = Pointer(to=x.limbs[0]).unsafe_mut_cast[True]() for i in range(6): - ptr.store[volatile=True](i, UInt64(0)) + ptr.unsafe_store[volatile=True](i, UInt64(0)) def _wipe_list_u8(mut data: List[UInt8]): var ptr = data.unsafe_ptr() for i in range(len(data)): - ptr.store[volatile=True](i, UInt8(0)) + ptr.unsafe_store[volatile=True](i, UInt8(0)) def _rfc6979_p384(private_key: Span[UInt8, ...], digest: Span[UInt8, ...], skip: Int) -> U384: @@ -941,7 +941,7 @@ def _rfc6979_p384(private_key: Span[UInt8, ...], digest: Span[UInt8, ...], skip: _wipe_list_u8(seed) var hp = h1_bytes.unsafe_ptr() for i in range(48): - hp.store[volatile=True](i, UInt8(0)) + hp.unsafe_store[volatile=True](i, UInt8(0)) _wipe_u384(h1) return candidate accepted += 1 @@ -996,7 +996,7 @@ def p384_ecdsa_sign_digest( retry += 1 continue _to_be(r, signature_ptr) - _to_be(s, signature_ptr + 48) + _to_be(s, signature_ptr.unsafe_offset(48)) _wipe_u384(d) _wipe_u384(z) _wipe_u384(k) diff --git a/src/thistle/p384_table.mojo b/src/thistle/p384_table.mojo index 8c137ba..8fa6f67 100644 --- a/src/thistle/p384_table.mojo +++ b/src/thistle/p384_table.mojo @@ -1298,72 +1298,72 @@ comptime _P384_BT67 = SIMD[DType.uint64, 128]( def p384_base_table() -> InlineArray[UInt64, 8704]: var t = InlineArray[UInt64, 8704](uninitialized=True) var p = t.unsafe_ptr() - p.store[alignment=8](0, _P384_BT0) - p.store[alignment=8](128, _P384_BT1) - p.store[alignment=8](256, _P384_BT2) - p.store[alignment=8](384, _P384_BT3) - p.store[alignment=8](512, _P384_BT4) - p.store[alignment=8](640, _P384_BT5) - p.store[alignment=8](768, _P384_BT6) - p.store[alignment=8](896, _P384_BT7) - p.store[alignment=8](1024, _P384_BT8) - p.store[alignment=8](1152, _P384_BT9) - p.store[alignment=8](1280, _P384_BT10) - p.store[alignment=8](1408, _P384_BT11) - p.store[alignment=8](1536, _P384_BT12) - p.store[alignment=8](1664, _P384_BT13) - p.store[alignment=8](1792, _P384_BT14) - p.store[alignment=8](1920, _P384_BT15) - p.store[alignment=8](2048, _P384_BT16) - p.store[alignment=8](2176, _P384_BT17) - p.store[alignment=8](2304, _P384_BT18) - p.store[alignment=8](2432, _P384_BT19) - p.store[alignment=8](2560, _P384_BT20) - p.store[alignment=8](2688, _P384_BT21) - p.store[alignment=8](2816, _P384_BT22) - p.store[alignment=8](2944, _P384_BT23) - p.store[alignment=8](3072, _P384_BT24) - p.store[alignment=8](3200, _P384_BT25) - p.store[alignment=8](3328, _P384_BT26) - p.store[alignment=8](3456, _P384_BT27) - p.store[alignment=8](3584, _P384_BT28) - p.store[alignment=8](3712, _P384_BT29) - p.store[alignment=8](3840, _P384_BT30) - p.store[alignment=8](3968, _P384_BT31) - p.store[alignment=8](4096, _P384_BT32) - p.store[alignment=8](4224, _P384_BT33) - p.store[alignment=8](4352, _P384_BT34) - p.store[alignment=8](4480, _P384_BT35) - p.store[alignment=8](4608, _P384_BT36) - p.store[alignment=8](4736, _P384_BT37) - p.store[alignment=8](4864, _P384_BT38) - p.store[alignment=8](4992, _P384_BT39) - p.store[alignment=8](5120, _P384_BT40) - p.store[alignment=8](5248, _P384_BT41) - p.store[alignment=8](5376, _P384_BT42) - p.store[alignment=8](5504, _P384_BT43) - p.store[alignment=8](5632, _P384_BT44) - p.store[alignment=8](5760, _P384_BT45) - p.store[alignment=8](5888, _P384_BT46) - p.store[alignment=8](6016, _P384_BT47) - p.store[alignment=8](6144, _P384_BT48) - p.store[alignment=8](6272, _P384_BT49) - p.store[alignment=8](6400, _P384_BT50) - p.store[alignment=8](6528, _P384_BT51) - p.store[alignment=8](6656, _P384_BT52) - p.store[alignment=8](6784, _P384_BT53) - p.store[alignment=8](6912, _P384_BT54) - p.store[alignment=8](7040, _P384_BT55) - p.store[alignment=8](7168, _P384_BT56) - p.store[alignment=8](7296, _P384_BT57) - p.store[alignment=8](7424, _P384_BT58) - p.store[alignment=8](7552, _P384_BT59) - p.store[alignment=8](7680, _P384_BT60) - p.store[alignment=8](7808, _P384_BT61) - p.store[alignment=8](7936, _P384_BT62) - p.store[alignment=8](8064, _P384_BT63) - p.store[alignment=8](8192, _P384_BT64) - p.store[alignment=8](8320, _P384_BT65) - p.store[alignment=8](8448, _P384_BT66) - p.store[alignment=8](8576, _P384_BT67) + p.unsafe_store[alignment=8](0, _P384_BT0) + p.unsafe_store[alignment=8](128, _P384_BT1) + p.unsafe_store[alignment=8](256, _P384_BT2) + p.unsafe_store[alignment=8](384, _P384_BT3) + p.unsafe_store[alignment=8](512, _P384_BT4) + p.unsafe_store[alignment=8](640, _P384_BT5) + p.unsafe_store[alignment=8](768, _P384_BT6) + p.unsafe_store[alignment=8](896, _P384_BT7) + p.unsafe_store[alignment=8](1024, _P384_BT8) + p.unsafe_store[alignment=8](1152, _P384_BT9) + p.unsafe_store[alignment=8](1280, _P384_BT10) + p.unsafe_store[alignment=8](1408, _P384_BT11) + p.unsafe_store[alignment=8](1536, _P384_BT12) + p.unsafe_store[alignment=8](1664, _P384_BT13) + p.unsafe_store[alignment=8](1792, _P384_BT14) + p.unsafe_store[alignment=8](1920, _P384_BT15) + p.unsafe_store[alignment=8](2048, _P384_BT16) + p.unsafe_store[alignment=8](2176, _P384_BT17) + p.unsafe_store[alignment=8](2304, _P384_BT18) + p.unsafe_store[alignment=8](2432, _P384_BT19) + p.unsafe_store[alignment=8](2560, _P384_BT20) + p.unsafe_store[alignment=8](2688, _P384_BT21) + p.unsafe_store[alignment=8](2816, _P384_BT22) + p.unsafe_store[alignment=8](2944, _P384_BT23) + p.unsafe_store[alignment=8](3072, _P384_BT24) + p.unsafe_store[alignment=8](3200, _P384_BT25) + p.unsafe_store[alignment=8](3328, _P384_BT26) + p.unsafe_store[alignment=8](3456, _P384_BT27) + p.unsafe_store[alignment=8](3584, _P384_BT28) + p.unsafe_store[alignment=8](3712, _P384_BT29) + p.unsafe_store[alignment=8](3840, _P384_BT30) + p.unsafe_store[alignment=8](3968, _P384_BT31) + p.unsafe_store[alignment=8](4096, _P384_BT32) + p.unsafe_store[alignment=8](4224, _P384_BT33) + p.unsafe_store[alignment=8](4352, _P384_BT34) + p.unsafe_store[alignment=8](4480, _P384_BT35) + p.unsafe_store[alignment=8](4608, _P384_BT36) + p.unsafe_store[alignment=8](4736, _P384_BT37) + p.unsafe_store[alignment=8](4864, _P384_BT38) + p.unsafe_store[alignment=8](4992, _P384_BT39) + p.unsafe_store[alignment=8](5120, _P384_BT40) + p.unsafe_store[alignment=8](5248, _P384_BT41) + p.unsafe_store[alignment=8](5376, _P384_BT42) + p.unsafe_store[alignment=8](5504, _P384_BT43) + p.unsafe_store[alignment=8](5632, _P384_BT44) + p.unsafe_store[alignment=8](5760, _P384_BT45) + p.unsafe_store[alignment=8](5888, _P384_BT46) + p.unsafe_store[alignment=8](6016, _P384_BT47) + p.unsafe_store[alignment=8](6144, _P384_BT48) + p.unsafe_store[alignment=8](6272, _P384_BT49) + p.unsafe_store[alignment=8](6400, _P384_BT50) + p.unsafe_store[alignment=8](6528, _P384_BT51) + p.unsafe_store[alignment=8](6656, _P384_BT52) + p.unsafe_store[alignment=8](6784, _P384_BT53) + p.unsafe_store[alignment=8](6912, _P384_BT54) + p.unsafe_store[alignment=8](7040, _P384_BT55) + p.unsafe_store[alignment=8](7168, _P384_BT56) + p.unsafe_store[alignment=8](7296, _P384_BT57) + p.unsafe_store[alignment=8](7424, _P384_BT58) + p.unsafe_store[alignment=8](7552, _P384_BT59) + p.unsafe_store[alignment=8](7680, _P384_BT60) + p.unsafe_store[alignment=8](7808, _P384_BT61) + p.unsafe_store[alignment=8](7936, _P384_BT62) + p.unsafe_store[alignment=8](8064, _P384_BT63) + p.unsafe_store[alignment=8](8192, _P384_BT64) + p.unsafe_store[alignment=8](8320, _P384_BT65) + p.unsafe_store[alignment=8](8448, _P384_BT66) + p.unsafe_store[alignment=8](8576, _P384_BT67) return t^ diff --git a/src/thistle/pbkdf2.mojo b/src/thistle/pbkdf2.mojo index f76a3b9..5d5f411 100644 --- a/src/thistle/pbkdf2.mojo +++ b/src/thistle/pbkdf2.mojo @@ -3,7 +3,7 @@ PBKDF2 (Password-Based Key Derivation Function 2) Implementation in Mojo SP 800-132 / FIPS 140-2 / RFC 8018 """ from std.collections import List -from std.memory import memset_zero, memcpy, UnsafePointer +from std.memory import unsafe_memcpy, Pointer from .utils import StackBuffer from .sha2 import ( SHA384_IV, @@ -17,17 +17,27 @@ from .sha2 import ( sha512_final_with_len, ) +comptime PBKDF2_SHA256_MAX_DKLEN: Int = 0xFFFFFFFF * 32 +comptime PBKDF2_SHA512_MAX_DKLEN: Int = 0xFFFFFFFF * 64 + + +@always_inline +def _secure_zero(ptr: Pointer[mut=True, UInt8, _, address_space=_], count: Int): + """Wipe sensitive bytes with stores that cannot be optimized away.""" + for i in range(count): + ptr.unsafe_store[volatile=True](i, UInt8(0)) + @always_inline -def _xor_block_32(dst: UnsafePointer[mut=True, UInt8, _, address_space=_], src: UnsafePointer[mut=True, UInt8, _, address_space=_]): - var d = dst.bitcast[UInt64]().load[width=4, alignment=1]() - var s = src.bitcast[UInt64]().load[width=4, alignment=1]() - dst.bitcast[UInt64]().store[width=4, alignment=1](0, d ^ s) +def _xor_block_32(dst: Pointer[mut=True, UInt8, _, address_space=_], src: Pointer[mut=True, UInt8, _, address_space=_]): + var d = dst.unsafe_bitcast[UInt64]().unsafe_load[width=4, alignment=1]() + var s = src.unsafe_bitcast[UInt64]().unsafe_load[width=4, alignment=1]() + dst.unsafe_bitcast[UInt64]().unsafe_store[width=4, alignment=1](0, d ^ s) @always_inline -def _xor_block_64(dst: UnsafePointer[mut=True, UInt8, _, address_space=_], src: UnsafePointer[mut=True, UInt8, _, address_space=_]): - var d = dst.bitcast[UInt64]().load[width=8, alignment=1]() - var s = src.bitcast[UInt64]().load[width=8, alignment=1]() - dst.bitcast[UInt64]().store[width=8, alignment=1](0, d ^ s) +def _xor_block_64(dst: Pointer[mut=True, UInt8, _, address_space=_], src: Pointer[mut=True, UInt8, _, address_space=_]): + var d = dst.unsafe_bitcast[UInt64]().unsafe_load[width=8, alignment=1]() + var s = src.unsafe_bitcast[UInt64]().unsafe_load[width=8, alignment=1]() + dst.unsafe_bitcast[UInt64]().unsafe_store[width=8, alignment=1](0, d ^ s) struct PBKDF2SHA256(Movable): var ipad: StackBuffer[UInt8, 64] @@ -60,14 +70,14 @@ struct PBKDF2SHA256(Movable): for i in range(64): self.ipad[i] = k[i] ^ 0x36 self.opad[i] = k[i] ^ 0x5C - memset_zero(k.ptr(), 64) + _secure_zero(k.ptr(), 64) - def __del__(deinit self): - memset_zero(self.ipad.ptr(), 64) - memset_zero(self.opad.ptr(), 64) - memset_zero(self.inner_hash.ptr(), 32) - memset_zero(self.u_block.ptr(), 32) - memset_zero(self.counter_bytes.ptr(), 4) + def __deinit__(deinit self): + _secure_zero(self.ipad.ptr(), 64) + _secure_zero(self.opad.ptr(), 64) + _secure_zero(self.inner_hash.ptr(), 32) + _secure_zero(self.u_block.ptr(), 32) + _secure_zero(self.counter_bytes.ptr(), 4) @always_inline def hmac(mut self, data: Span[UInt8, ...]): @@ -100,9 +110,13 @@ struct PBKDF2SHA256(Movable): sha256_final_to_buffer(self.outer_ctx, self.u_block.ptr()) @always_inline - def derive(mut self, salt: Span[UInt8, ...], iterations: Int, dklen: Int) -> List[UInt8]: + def derive(mut self, salt: Span[UInt8, ...], iterations: Int, dklen: Int) raises -> List[UInt8]: + if dklen < 1 or dklen > PBKDF2_SHA256_MAX_DKLEN: + raise Error("PBKDF2-SHA256 dkLen exceeds the RFC 8018 limit") var hLen = 32 - var num_blocks = (dklen + hLen - 1) // hLen + var num_blocks = dklen // hLen + if dklen % hLen != 0: + num_blocks += 1 var derived_key = List[UInt8](capacity=dklen) var t_block = StackBuffer[UInt8, 32](fill=0) @@ -110,10 +124,10 @@ struct PBKDF2SHA256(Movable): for block_idx in range(1, num_blocks + 1): self.hmac_with_counter(salt, UInt32(block_idx)) - memcpy(dest=t_block.ptr(), src=self.u_block.ptr(), count=32) + unsafe_memcpy(dest=t_block.ptr(), src=self.u_block.ptr(), count=32) for _ in range(1, iterations): - memcpy( + unsafe_memcpy( dest=input_block.ptr(), src=self.u_block.ptr(), count=32, @@ -128,6 +142,8 @@ struct PBKDF2SHA256(Movable): for b in range(to_copy): derived_key.append(t_block[b]) + _secure_zero(t_block.ptr(), 32) + _secure_zero(input_block.ptr(), 32) return derived_key^ def pbkdf2_hmac_sha256( @@ -137,6 +153,8 @@ def pbkdf2_hmac_sha256( raise Error("PBKDF2 iterations must be at least 1") if dkLen < 1: raise Error("PBKDF2 dkLen must be at least 1") + if dkLen > PBKDF2_SHA256_MAX_DKLEN: + raise Error("PBKDF2-SHA256 dkLen exceeds the RFC 8018 limit") var ctx = PBKDF2SHA256(password) return ctx.derive(salt, iterations, dkLen) @@ -171,14 +189,14 @@ struct PBKDF2SHA512(Movable): for i in range(128): self.ipad[i] = k[i] ^ 0x36 self.opad[i] = k[i] ^ 0x5C - memset_zero(k.ptr(), 128) + _secure_zero(k.ptr(), 128) - def __del__(deinit self): - memset_zero(self.ipad.ptr(), 128) - memset_zero(self.opad.ptr(), 128) - memset_zero(self.inner_hash.ptr(), 64) - memset_zero(self.u_block.ptr(), 64) - memset_zero(self.counter_bytes.ptr(), 4) + def __deinit__(deinit self): + _secure_zero(self.ipad.ptr(), 128) + _secure_zero(self.opad.ptr(), 128) + _secure_zero(self.inner_hash.ptr(), 64) + _secure_zero(self.u_block.ptr(), 64) + _secure_zero(self.counter_bytes.ptr(), 4) @always_inline def hmac(mut self, data: Span[UInt8, ...]): @@ -211,9 +229,13 @@ struct PBKDF2SHA512(Movable): sha512_final_to_buffer(self.outer_ctx, self.u_block.ptr()) @always_inline - def derive(mut self, salt: Span[UInt8, ...], iterations: Int, dklen: Int) -> List[UInt8]: + def derive(mut self, salt: Span[UInt8, ...], iterations: Int, dklen: Int) raises -> List[UInt8]: + if dklen < 1 or dklen > PBKDF2_SHA512_MAX_DKLEN: + raise Error("PBKDF2-SHA512 dkLen exceeds the RFC 8018 limit") var hLen = 64 - var num_blocks = (dklen + hLen - 1) // hLen + var num_blocks = dklen // hLen + if dklen % hLen != 0: + num_blocks += 1 var derived_key = List[UInt8](capacity=dklen) var t_block = StackBuffer[UInt8, 64](fill=0) @@ -221,10 +243,10 @@ struct PBKDF2SHA512(Movable): for block_idx in range(1, num_blocks + 1): self.hmac_with_counter(salt, UInt32(block_idx)) - memcpy(dest=t_block.ptr(), src=self.u_block.ptr(), count=64) + unsafe_memcpy(dest=t_block.ptr(), src=self.u_block.ptr(), count=64) for _ in range(1, iterations): - memcpy( + unsafe_memcpy( dest=input_block.ptr(), src=self.u_block.ptr(), count=64, @@ -239,6 +261,8 @@ struct PBKDF2SHA512(Movable): for b in range(to_copy): derived_key.append(t_block[b]) + _secure_zero(t_block.ptr(), 64) + _secure_zero(input_block.ptr(), 64) return derived_key^ def pbkdf2_hmac_sha512( @@ -248,6 +272,8 @@ def pbkdf2_hmac_sha512( raise Error("PBKDF2 iterations must be at least 1") if dkLen < 1: raise Error("PBKDF2 dkLen must be at least 1") + if dkLen > PBKDF2_SHA512_MAX_DKLEN: + raise Error("PBKDF2-SHA512 dkLen exceeds the RFC 8018 limit") var ctx = PBKDF2SHA512(password) return ctx.derive(salt, iterations, dkLen) @@ -272,10 +298,10 @@ def hmac_sha384(key: Span[UInt8, ...], data: Span[UInt8, ...]) -> List[UInt8]: var k = StackBuffer[UInt8, 128](fill=0) if len(key) > 128: var kh = sha384_hash(key) - memcpy(dest=k.ptr(), src=kh.unsafe_ptr(), count=48) + unsafe_memcpy(dest=k.ptr(), src=kh.unsafe_ptr(), count=48) var khp = kh.unsafe_ptr() for i in range(48): - khp.store[volatile=True](i, UInt8(0)) + khp.unsafe_store[volatile=True](i, UInt8(0)) else: for i in range(len(key)): k[i] = key[i] @@ -300,12 +326,12 @@ def hmac_sha384(key: Span[UInt8, ...], data: Span[UInt8, ...]) -> List[UInt8]: var ip = ipad.ptr() var op = opad.ptr() for i in range(128): - kp.store[volatile=True](i, UInt8(0)) - ip.store[volatile=True](i, UInt8(0)) - op.store[volatile=True](i, UInt8(0)) + kp.unsafe_store[volatile=True](i, UInt8(0)) + ip.unsafe_store[volatile=True](i, UInt8(0)) + op.unsafe_store[volatile=True](i, UInt8(0)) var ihp = inner_hash.unsafe_ptr() for i in range(48): - ihp.store[volatile=True](i, UInt8(0)) + ihp.unsafe_store[volatile=True](i, UInt8(0)) inner.wipe() outer.wipe() return result^ diff --git a/src/thistle/poly1305.mojo b/src/thistle/poly1305.mojo index cd206cf..8a500b8 100644 --- a/src/thistle/poly1305.mojo +++ b/src/thistle/poly1305.mojo @@ -2,7 +2,7 @@ Poly1305 per RFC 8439 """ -from std.memory import UnsafePointer +from std.memory import Pointer from std.collections import InlineArray comptime _M44: UInt64 = 0xFFFFFFFFFFF @@ -10,8 +10,8 @@ comptime _M42: UInt64 = 0x3FFFFFFFFFF @always_inline -def _le64(ptr: UnsafePointer[mut=False, UInt8, _, address_space=_], offset: Int) -> UInt64: - return (ptr + offset).bitcast[UInt64]().load[width=1, alignment=1]() +def _le64(ptr: Pointer[mut=False, UInt8, _, address_space=_], offset: Int) -> UInt64: + return (ptr.unsafe_offset(offset)).unsafe_bitcast[UInt64]().unsafe_load[width=1, alignment=1]() struct _RPower(Movable, Copyable, ImplicitlyCopyable): @@ -67,7 +67,7 @@ def _reduce(mut h0: UInt64, mut h1: UInt64, mut h2: UInt64, d0: UInt128, d1: UIn @always_inline -def _limbs_at(ptr: UnsafePointer[mut=False, UInt8, _, address_space=_], offset: Int, hibit: UInt64) -> SIMD[DType.uint64, 4]: +def _limbs_at(ptr: Pointer[mut=False, UInt8, _, address_space=_], offset: Int, hibit: UInt64) -> SIMD[DType.uint64, 4]: var t0 = _le64(ptr, offset) var t1 = _le64(ptr, offset + 8) return SIMD[DType.uint64, 4]( @@ -175,7 +175,7 @@ struct Poly1305: _reduce(self.h0, self.h1, self.h2, d0, d1, d2) @no_inline - def _blocks8(mut self, ptr: UnsafePointer[mut=False, UInt8, _, address_space=_], count8: Int): + def _blocks8(mut self, ptr: Pointer[mut=False, UInt8, _, address_space=_], count8: Int): var h0 = self.h0 var h1 = self.h1 var h2 = self.h2 @@ -207,7 +207,7 @@ struct Poly1305: self.h2 = h2 @no_inline - def _blocks4(mut self, ptr: UnsafePointer[mut=False, UInt8, _, address_space=_], count4: Int): + def _blocks4(mut self, ptr: Pointer[mut=False, UInt8, _, address_space=_], count4: Int): var h0 = self.h0 var h1 = self.h1 var h2 = self.h2 @@ -239,7 +239,7 @@ struct Poly1305: if self.buf_len > 0: while self.buf_len < 16 and i < n: - self.buf[self.buf_len] = ptr[i] + self.buf[self.buf_len] = ptr[unsafe_offset=i] self.buf_len += 1 i += 1 if self.buf_len == 16: @@ -251,14 +251,14 @@ struct Poly1305: if octs > 0: if not self.powers8_ready: self._make_powers8() - self._blocks8(ptr + i, octs) + self._blocks8(ptr.unsafe_offset(i), octs) i += octs << 7 var quads = (n - i) >> 6 if quads > 0 and (self.powers4_ready or quads > 1): if not self.powers4_ready: self._make_powers4() - self._blocks4(ptr + i, quads) + self._blocks4(ptr.unsafe_offset(i), quads) i += quads << 6 while i + 16 <= n: @@ -266,11 +266,11 @@ struct Poly1305: i += 16 while i < n: - self.buf[self.buf_len] = ptr[i] + self.buf[self.buf_len] = ptr[unsafe_offset=i] self.buf_len += 1 i += 1 - def finalize_into(mut self, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): + def finalize_into(mut self, output: Pointer[mut=True, UInt8, _, address_space=_]): if self.buf_len > 0: self.buf[self.buf_len] = 1 for j in range(self.buf_len + 1, 16): @@ -324,17 +324,17 @@ struct Poly1305: var o0 = h0 | (h1 << 44) var o1 = (h1 >> 20) | (h2 << 24) - output.bitcast[UInt64]().store[alignment=1](0, o0) - (output + 8).bitcast[UInt64]().store[alignment=1](0, o1) + output.unsafe_bitcast[UInt64]().unsafe_store[alignment=1](0, o0) + (output.unsafe_offset(8)).unsafe_bitcast[UInt64]().unsafe_store[alignment=1](0, o1) self.wipe() def wipe(mut self): - var p = UnsafePointer(to=self.h0) - p.store[volatile=True](0, UInt64(0)) - UnsafePointer(to=self.h1).store[volatile=True](0, UInt64(0)) - UnsafePointer(to=self.h2).store[volatile=True](0, UInt64(0)) - UnsafePointer(to=self.pad0).store[volatile=True](0, UInt64(0)) - UnsafePointer(to=self.pad1).store[volatile=True](0, UInt64(0)) + var p = Pointer(to=self.h0) + p.unsafe_store[volatile=True](0, UInt64(0)) + Pointer(to=self.h1).unsafe_store[volatile=True](0, UInt64(0)) + Pointer(to=self.h2).unsafe_store[volatile=True](0, UInt64(0)) + Pointer(to=self.pad0).unsafe_store[volatile=True](0, UInt64(0)) + Pointer(to=self.pad1).unsafe_store[volatile=True](0, UInt64(0)) def poly1305_mac( diff --git a/src/thistle/random.mojo b/src/thistle/random.mojo index 4707e17..5115f47 100644 --- a/src/thistle/random.mojo +++ b/src/thistle/random.mojo @@ -1,10 +1,10 @@ from std.collections import List -from std.memory import UnsafePointer +from std.memory import Pointer from std.sys import CompilationTarget, inlined_assembly @always_inline -def _getrandom_linux_x86(buf: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int) -> Int: +def _getrandom_linux_x86(buf: Pointer[mut=True, UInt8, _, address_space=_], length: Int) -> Int: # Linux x86-64: getrandom(buf, len, flags=0), syscall 318. # rax is both syscall-number input and return-value output. return Int( @@ -18,7 +18,7 @@ def _getrandom_linux_x86(buf: UnsafePointer[mut=True, UInt8, _, address_space=_] @always_inline -def _getrandom_linux_arm(buf: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int) -> Int: +def _getrandom_linux_arm(buf: Pointer[mut=True, UInt8, _, address_space=_], length: Int) -> Int: # Linux aarch64: getrandom(buf, len, flags=0), syscall 278. # inputs go via scratch registers to avoid asm constraint conflicts return Int( @@ -38,7 +38,7 @@ def _getrandom_linux_arm(buf: UnsafePointer[mut=True, UInt8, _, address_space=_] @always_inline -def _getentropy_macos_arm(buf: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int) -> Int: +def _getentropy_macos_arm(buf: Pointer[mut=True, UInt8, _, address_space=_], length: Int) -> Int: # Darwin arm64: getentropy(buf, size), syscall 500. # x16 = syscall number, x0 = buffer, x1 = size. # Errors set carry and return errno in x0; convert that to -errno. @@ -61,10 +61,10 @@ def _getentropy_macos_arm(buf: UnsafePointer[mut=True, UInt8, _, address_space=_ ) -def _fill_linux_x86(buf: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int) raises: +def _fill_linux_x86(buf: Pointer[mut=True, UInt8, _, address_space=_], length: Int) raises: var offset = 0 while offset < length: - var ret = _getrandom_linux_x86(buf + offset, length - offset) + var ret = _getrandom_linux_x86(buf.unsafe_offset(offset), length - offset) if ret < 0: if ret == -4: # EINTR continue @@ -74,10 +74,10 @@ def _fill_linux_x86(buf: UnsafePointer[mut=True, UInt8, _, address_space=_], len offset += ret -def _fill_linux_arm(buf: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int) raises: +def _fill_linux_arm(buf: Pointer[mut=True, UInt8, _, address_space=_], length: Int) raises: var offset = 0 while offset < length: - var ret = _getrandom_linux_arm(buf + offset, length - offset) + var ret = _getrandom_linux_arm(buf.unsafe_offset(offset), length - offset) if ret < 0: if ret == -4: # EINTR continue @@ -87,11 +87,11 @@ def _fill_linux_arm(buf: UnsafePointer[mut=True, UInt8, _, address_space=_], len offset += ret -def _fill_macos_arm(buf: UnsafePointer[mut=True, UInt8, _, address_space=_], length: Int) raises: +def _fill_macos_arm(buf: Pointer[mut=True, UInt8, _, address_space=_], length: Int) raises: var offset = 0 while offset < length: var chunk = min(256, length - offset) - var ret = _getentropy_macos_arm(buf + offset, chunk) + var ret = _getentropy_macos_arm(buf.unsafe_offset(offset), chunk) if ret < 0: raise Error("getentropy syscall failed") if ret != 0: diff --git a/src/thistle/rsa.mojo b/src/thistle/rsa.mojo index 2b9d136..19c9923 100644 --- a/src/thistle/rsa.mojo +++ b/src/thistle/rsa.mojo @@ -3,7 +3,7 @@ RSASSA-PSS signature verification per RFC 8017 """ from std.collections import List, InlineArray -from std.memory import UnsafePointer +from std.memory import Pointer from std.bit import rotate_bits_left, count_leading_zeros from std.utils import StaticTuple from .random import random_bytes @@ -121,7 +121,7 @@ def _hash_len(alg: Int) raises -> Int: raise Error("unsupported hash for RSA-PSS") -def _hash_into(alg: Int, data: Span[UInt8, ...], output: UnsafePointer[mut=True, UInt8, _, address_space=_]) raises -> Int: +def _hash_into(alg: Int, data: Span[UInt8, ...], output: Pointer[mut=True, UInt8, _, address_space=_]) raises -> Int: if alg == SHA256: var ctx = SHA256Context() sha256_update(ctx, data) @@ -135,7 +135,7 @@ def _hash_into(alg: Int, data: Span[UInt8, ...], output: UnsafePointer[mut=True, if alg == SHA1: var d1 = _sha1(data) for i in range(20): - output[i] = d1[i] + output[unsafe_offset=i] = d1[i] return 20 var d: List[UInt8] if alg == SHA224: @@ -149,22 +149,22 @@ def _hash_into(alg: Int, data: Span[UInt8, ...], output: UnsafePointer[mut=True, else: raise Error("unsupported hash for RSA-PSS") for i in range(len(d)): - output[i] = d[i] + output[unsafe_offset=i] = d[i] return len(d) def _mgf1( alg: Int, - seed: UnsafePointer[UInt8, _], + seed: Pointer[UInt8, _], seed_len: Int, mask_len: Int, - output: UnsafePointer[mut=True, UInt8, _, address_space=_], + output: Pointer[mut=True, UInt8, _, address_space=_], ) raises: var h_len = _hash_len(alg) var block = InlineArray[UInt8, 128](fill=0) var digest = InlineArray[UInt8, 64](fill=0) for i in range(seed_len): - block[i] = seed[i] + block[i] = seed[unsafe_offset=i] var done = 0 var counter: UInt32 = 0 while done < mask_len: @@ -181,7 +181,7 @@ def _mgf1( if take > h_len: take = h_len for i in range(take): - output[done + i] = digest[i] + output[unsafe_offset=done + i] = digest[i] done += take counter += 1 @@ -192,7 +192,7 @@ def _emsa_pss_encode( sha: Int, mgf_sha: Int, em_bits: Int, - output: UnsafePointer[mut=True, UInt8, _, address_space=_], + output: Pointer[mut=True, UInt8, _, address_space=_], ) raises -> Bool: var h_len = _hash_len(sha) _ = _hash_len(mgf_sha) @@ -223,15 +223,15 @@ def _emsa_pss_encode( var mask = InlineArray[UInt8, 528](uninitialized=True) _mgf1(mgf_sha, h.unsafe_ptr(), h_len, db_len, mask.unsafe_ptr()) for i in range(db_len): - output[i] = db[i] ^ mask[i] - output[0] &= UInt8(0xFF) >> UInt8(8 * em_len - em_bits) + output[unsafe_offset=i] = db[i] ^ mask[i] + output[unsafe_offset=0] &= UInt8(0xFF) >> UInt8(8 * em_len - em_bits) for i in range(h_len): - output[db_len + i] = h[i] - output[em_len - 1] = 0xBC + output[unsafe_offset=db_len + i] = h[i] + output[unsafe_offset=em_len - 1] = 0xBC var mp = mprime.unsafe_ptr() for i in range(8 + h_len + len(salt)): - mp.store[volatile=True](i, UInt8(0)) + mp.unsafe_store[volatile=True](i, UInt8(0)) return True @@ -244,7 +244,7 @@ def _digest_info_prefix_len(alg: Int) raises -> Int: def _digest_info_prefix( - alg: Int, output: UnsafePointer[mut=True, UInt8, _, address_space=_] + alg: Int, output: Pointer[mut=True, UInt8, _, address_space=_] ) raises: var hex: String if alg == SHA1: @@ -265,7 +265,7 @@ def _digest_info_prefix( var lo = bytes[2 * i + 1] hi = hi - UInt8(48 if hi <= 57 else 87) lo = lo - UInt8(48 if lo <= 57 else 87) - output[i] = (hi << 4) | lo + output[unsafe_offset=i] = (hi << 4) | lo @always_inline @@ -315,7 +315,7 @@ def _bn_sub_copy( var d = (UInt128(1) << 64) + UInt128(a[i]) - UInt128(b[i]) - UInt128(borrow) out[i] = d.cast[DType.uint64]() borrow = 1 - (d >> 64).cast[DType.uint64]() - return out^, borrow + return out, borrow @always_inline @@ -329,7 +329,7 @@ def _bn_select( var mask = UInt64(0) - (choice & UInt64(1)) for i in range(k): out[i] = a[i] ^ (mask & (a[i] ^ b[i])) - return out^ + return out @always_inline @@ -621,7 +621,7 @@ struct RsaPublicKey: def _public_op( self, sig: Span[UInt8, ...], - output: UnsafePointer[mut=True, UInt8, _, address_space=_], + output: Pointer[mut=True, UInt8, _, address_space=_], ) raises -> Bool: var nb = self.nb var k = self.k @@ -656,7 +656,7 @@ struct RsaPublicKey: var m = _mont_mul(acc, one, self.n, self.n0, k) for i in range(nb): var limb = m[(nb - 1 - i) >> 3] - output[i] = UInt8((limb >> UInt64(8 * ((nb - 1 - i) & 7))) & 0xFF) + output[unsafe_offset=i] = UInt8((limb >> UInt64(8 * ((nb - 1 - i) & 7))) & 0xFF) return True def pss_verify( @@ -684,23 +684,23 @@ struct RsaPublicKey: for i in range(nb - em_len): if em[i] != 0: return False - var ep = em.unsafe_ptr() + (nb - em_len) + var ep = em.unsafe_ptr().unsafe_offset((nb - em_len)) - if ep[em_len - 1] != 0xBC: + if ep[unsafe_offset=em_len - 1] != 0xBC: return False var db_len = em_len - h_len - 1 - var h_ptr = ep + db_len + var h_ptr = ep.unsafe_offset(db_len) var top_mask = UInt8(0xFF) >> UInt8(8 * em_len - em_bits) - if (ep[0] & ~top_mask) != 0: + if (ep[unsafe_offset=0] & ~top_mask) != 0: return False var db_mask = InlineArray[UInt8, 528](uninitialized=True) _mgf1(mgf_sha, h_ptr, h_len, db_len, db_mask.unsafe_ptr()) var db = InlineArray[UInt8, 528](uninitialized=True) for i in range(db_len): - db[i] = ep[i] ^ db_mask[i] + db[i] = ep[unsafe_offset=i] ^ db_mask[i] db[0] &= top_mask var ps_len = db_len - salt_len - 1 @@ -729,14 +729,14 @@ struct RsaPublicKey: var diff: UInt8 = 0 for i in range(h_len): - diff |= h2[i] ^ h_ptr[i] + diff |= h2[i] ^ h_ptr[unsafe_offset=i] return diff == 0 def _wipe_bn(mut value: StaticTuple[UInt64, _NL], k: Int): - var ptr = UnsafePointer(to=value[0]).unsafe_mut_cast[True]() + var ptr = Pointer(to=value[0]).unsafe_mut_cast[True]() for i in range(k): - ptr.store[volatile=True](i, UInt64(0)) + ptr.unsafe_store[volatile=True](i, UInt64(0)) struct RsaPrivateKey: @@ -770,15 +770,15 @@ struct RsaPrivateKey: self.d[self.public.nb - d_len + i] = private_exponent[lead + i] _wipe_bn(dbn, self.public.k) - def __del__(deinit self): + def __deinit__(deinit self): var ptr = self.d.unsafe_ptr() for i in range(528): - ptr.store[volatile=True](i, UInt8(0)) + ptr.unsafe_store[volatile=True](i, UInt8(0)) def _private_op( self, encoded: Span[UInt8, ...], - signature: UnsafePointer[mut=True, UInt8, _, address_space=_], + signature: Pointer[mut=True, UInt8, _, address_space=_], ) raises -> Bool: var nb = self.public.nb var k = self.public.k @@ -819,7 +819,7 @@ struct RsaPrivateKey: var result = _mont_mul(acc, one, self.public.n, self.public.n0, k) for i in range(nb): var limb = result[(nb - 1 - i) >> 3] - signature[i] = UInt8((limb >> UInt64(8 * ((nb - 1 - i) & 7))) & 0xFF) + signature[unsafe_offset=i] = UInt8((limb >> UInt64(8 * ((nb - 1 - i) & 7))) & 0xFF) var recovered = InlineArray[UInt8, 528](uninitialized=True) var sig_span = Span[UInt8, ...](unsafe_ptr=signature, length=nb) @@ -830,7 +830,7 @@ struct RsaPrivateKey: valid = valid and diff == 0 if not valid: for i in range(nb): - signature.store[volatile=True](i, UInt8(0)) + signature.unsafe_store[volatile=True](i, UInt8(0)) _wipe_bn(input, k) _wipe_bn(base, k) @@ -846,14 +846,14 @@ struct RsaPrivateKey: salt: Span[UInt8, ...], sha: Int, mgf_sha: Int, - signature: UnsafePointer[mut=True, UInt8, _, address_space=_], + signature: Pointer[mut=True, UInt8, _, address_space=_], ) raises -> Bool: var em_bits = self.public.mod_bits - 1 var em_len = (em_bits + 7) // 8 var encoded = InlineArray[UInt8, 528](fill=0) if not _emsa_pss_encode( message, salt, sha, mgf_sha, em_bits, - encoded.unsafe_ptr() + (self.public.nb - em_len), + encoded.unsafe_ptr().unsafe_offset((self.public.nb - em_len)), ): return False var ok = self._private_op( @@ -862,7 +862,7 @@ struct RsaPrivateKey: ) var ep = encoded.unsafe_ptr() for i in range(self.public.nb): - ep.store[volatile=True](i, UInt8(0)) + ep.unsafe_store[volatile=True](i, UInt8(0)) return ok def pss_sign( @@ -886,7 +886,7 @@ struct RsaPrivateKey: ) var salt_ptr = salt.unsafe_ptr() for i in range(len(salt)): - salt_ptr.store[volatile=True](i, UInt8(0)) + salt_ptr.unsafe_store[volatile=True](i, UInt8(0)) if not ok: raise Error("RSA-PSS signing failed") return signature^ @@ -907,7 +907,7 @@ def _bn_reduce_bytes( var reduced, borrow = _bn_sub_copy(value, key.n, key.k) var take = carry | (borrow ^ UInt64(1)) value = _bn_select(value, reduced, take, key.k) - return value^ + return value def _private_pow( @@ -942,7 +942,7 @@ def _private_pow( _wipe_bn(acc, key.k) for i in range(16): _wipe_bn(table[i], key.k) - return result^ + return result def _bn_mul_parts( @@ -957,7 +957,7 @@ def _bn_mul_parts( out[i + j] = product.cast[DType.uint64]() carry = (product >> 64).cast[DType.uint64]() out[i + b_len] = carry - return out^ + return out def _bn_equal( @@ -1058,19 +1058,19 @@ struct RsaCrtPrivateKey: coefficient_diff |= check[i] var qbp = q_bytes.unsafe_ptr() for i in range(self.q.nb): - qbp.store[volatile=True](i, UInt8(0)) + qbp.unsafe_store[volatile=True](i, UInt8(0)) _wipe_bn(product, self.public.k) _wipe_bn(q_mod_p, self.p.k) _wipe_bn(check, self.p.k) if coefficient_diff != 0: raise Error("invalid RSA CRT coefficient") - def __del__(deinit self): + def __deinit__(deinit self): var dpp = self.dp.unsafe_ptr() var dqp = self.dq.unsafe_ptr() for i in range(528): - dpp.store[volatile=True](i, UInt8(0)) - dqp.store[volatile=True](i, UInt8(0)) + dpp.unsafe_store[volatile=True](i, UInt8(0)) + dqp.unsafe_store[volatile=True](i, UInt8(0)) _wipe_bn(self.qinv, self.p.k) _wipe_bn(self.p.n, self.p.k) _wipe_bn(self.p.rmod, self.p.k) @@ -1083,7 +1083,7 @@ struct RsaCrtPrivateKey: def _private_op( self, encoded: Span[UInt8, ...], - signature: UnsafePointer[mut=True, UInt8, _, address_space=_], + signature: Pointer[mut=True, UInt8, _, address_space=_], ) raises -> Bool: if len(encoded) != self.public.nb: return False @@ -1123,7 +1123,7 @@ struct RsaCrtPrivateKey: carry = (sum >> 64).cast[DType.uint64]() for i in range(self.public.nb): var limb = result[(self.public.nb - 1 - i) >> 3] - signature[i] = UInt8((limb >> UInt64(8 * ((self.public.nb - 1 - i) & 7))) & 0xFF) + signature[unsafe_offset=i] = UInt8((limb >> UInt64(8 * ((self.public.nb - 1 - i) & 7))) & 0xFF) var recovered = InlineArray[UInt8, 528](uninitialized=True) var valid = self.public._public_op( @@ -1135,7 +1135,7 @@ struct RsaCrtPrivateKey: valid = valid and diff == 0 if not valid: for i in range(self.public.nb): - signature.store[volatile=True](i, UInt8(0)) + signature.unsafe_store[volatile=True](i, UInt8(0)) _wipe_bn(mp, self.p.k) _wipe_bn(mq, self.q.k) @@ -1147,20 +1147,20 @@ struct RsaCrtPrivateKey: _wipe_bn(result, self.public.k) var m2p = m2_bytes.unsafe_ptr() for i in range(self.q.nb): - m2p.store[volatile=True](i, UInt8(0)) + m2p.unsafe_store[volatile=True](i, UInt8(0)) return valid def pss_sign_with_salt( self, message: Span[UInt8, ...], salt: Span[UInt8, ...], sha: Int, mgf_sha: Int, - signature: UnsafePointer[mut=True, UInt8, _, address_space=_], + signature: Pointer[mut=True, UInt8, _, address_space=_], ) raises -> Bool: var em_bits = self.public.mod_bits - 1 var em_len = (em_bits + 7) // 8 var encoded = InlineArray[UInt8, 528](fill=0) if not _emsa_pss_encode( message, salt, sha, mgf_sha, em_bits, - encoded.unsafe_ptr() + (self.public.nb - em_len), + encoded.unsafe_ptr().unsafe_offset((self.public.nb - em_len)), ): return False var ok = self._private_op( @@ -1168,7 +1168,7 @@ struct RsaCrtPrivateKey: ) var ep = encoded.unsafe_ptr() for i in range(self.public.nb): - ep.store[volatile=True](i, UInt8(0)) + ep.unsafe_store[volatile=True](i, UInt8(0)) return ok def pss_sign( @@ -1192,7 +1192,7 @@ struct RsaCrtPrivateKey: ) var salt_ptr = salt.unsafe_ptr() for i in range(len(salt)): - salt_ptr.store[volatile=True](i, UInt8(0)) + salt_ptr.unsafe_store[volatile=True](i, UInt8(0)) if not ok: raise Error("RSA-PSS signing failed") return signature^ diff --git a/src/thistle/sha2.mojo b/src/thistle/sha2.mojo index 1d8c06d..cd89850 100644 --- a/src/thistle/sha2.mojo +++ b/src/thistle/sha2.mojo @@ -4,7 +4,7 @@ RFC 6234 / FIPS 180-4 / CAVP validated """ from std.collections import List -from std.memory import UnsafePointer, memcpy, memset_zero +from std.memory import Pointer, unsafe_memcpy, unsafe_memset_zero from std.bit import rotate_bits_right, byte_swap from std.builtin.simd import SIMD from std.builtin.dtype import DType @@ -170,31 +170,31 @@ struct SHA256Context(Movable): self.buffer = InlineArray[UInt8, 64](fill=0) self.buffer_len = 0 - def __init__(out self, *, deinit take: Self): - self.state = take.state - self.count = take.count - self.buffer = take.buffer^ - self.buffer_len = take.buffer_len + def __init__(out self, *, deinit move: Self): + self.state = move.state + self.count = move.count + self.buffer = move.buffer^ + self.buffer_len = move.buffer_len - def __del__(deinit self): - memset_zero(self.buffer.unsafe_ptr(), 64) + def __deinit__(deinit self): + unsafe_memset_zero(self.buffer.unsafe_ptr(), 64) def reset(mut self): self.state = SHA256_IV self.count = 0 - memset_zero(self.buffer.unsafe_ptr(), 64) + unsafe_memset_zero(self.buffer.unsafe_ptr(), 64) self.buffer_len = 0 def reset(mut self, iv: SIMD[DType.uint32, 8]): self.state = iv self.count = 0 - memset_zero(self.buffer.unsafe_ptr(), 64) + unsafe_memset_zero(self.buffer.unsafe_ptr(), 64) self.buffer_len = 0 @always_inline def sha256_transform_blocks( mut state: SIMD[DType.uint32, 8], - data: UnsafePointer[mut=False, UInt8, _, address_space=_], + data: Pointer[mut=False, UInt8, _, address_space=_], nblocks: Int, ): comptime if (CompilationTarget.has_neon() and CompilationTarget._has_feature["sha2"]() and not CompilationTarget.is_x86()) or (CompilationTarget.is_x86() and CompilationTarget._has_feature["sse"]() and CompilationTarget._has_feature["sha"]()): @@ -211,7 +211,7 @@ def sha256_transform_blocks( var h0 = state[7] for blk in range(nblocks): - var block = data + blk * 64 + var block = data.unsafe_offset(blk * 64) var w = InlineArray[UInt32, 16](uninitialized=True) var a = a0 @@ -286,7 +286,7 @@ def sha256_update(mut ctx: SHA256Context, data: Span[UInt8, ...]): var available = 64 - ctx.buffer_len if total_len >= available: for j in range(available): - buf_ptr[ctx.buffer_len + j] = data[j] + buf_ptr[unsafe_offset=ctx.buffer_len + j] = data[j] ctx.state = sha256_transform( ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=64) ) @@ -295,20 +295,20 @@ def sha256_update(mut ctx: SHA256Context, data: Span[UInt8, ...]): ctx.buffer_len = 0 else: for j in range(total_len): - buf_ptr[ctx.buffer_len + j] = data[j] + buf_ptr[unsafe_offset=ctx.buffer_len + j] = data[j] ctx.buffer_len += total_len return var nblocks = (total_len - i) // 64 if nblocks > 0: - sha256_transform_blocks(ctx.state, data.unsafe_ptr() + i, nblocks) + sha256_transform_blocks(ctx.state, data.unsafe_ptr().unsafe_offset(i), nblocks) ctx.count += UInt64(nblocks) * 512 i += nblocks * 64 if i < total_len: var remaining = total_len - i for j in range(remaining): - buf_ptr[ctx.buffer_len + j] = data[i + j] + buf_ptr[unsafe_offset=ctx.buffer_len + j] = data[i + j] ctx.buffer_len += remaining @@ -325,7 +325,7 @@ def sha256_hash(data: Span[UInt8, ...]) -> List[UInt8]: return sha256_final(ctx) -def sha256_final_to_buffer(mut ctx: SHA256Context, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): +def sha256_final_to_buffer(mut ctx: SHA256Context, output: Pointer[mut=True, UInt8, _, address_space=_]): var buf_ptr = ctx.buffer.unsafe_ptr() var bit_count = ctx.count + UInt64(ctx.buffer_len) * 8 @@ -333,13 +333,13 @@ def sha256_final_to_buffer(mut ctx: SHA256Context, output: UnsafePointer[mut=Tru ctx.buffer_len += 1 if ctx.buffer_len > 56: - memset_zero(buf_ptr + ctx.buffer_len, 64 - ctx.buffer_len) + unsafe_memset_zero(buf_ptr.unsafe_offset(ctx.buffer_len), 64 - ctx.buffer_len) ctx.state = sha256_transform( ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=64) ) ctx.buffer_len = 0 - memset_zero(buf_ptr + ctx.buffer_len, 56 - ctx.buffer_len) + unsafe_memset_zero(buf_ptr.unsafe_offset(ctx.buffer_len), 56 - ctx.buffer_len) ctx.buffer_len = 56 for i in range(8): @@ -350,7 +350,7 @@ def sha256_final_to_buffer(mut ctx: SHA256Context, output: UnsafePointer[mut=Tru ) for i in range(8): - (output + i * 4).bitcast[UInt32]().store[alignment=1](byte_swap(ctx.state[i])) + (output.unsafe_offset(i * 4)).unsafe_bitcast[UInt32]().unsafe_store[alignment=1](byte_swap(ctx.state[i])) struct SHA512Context(Movable): var state: SIMD[DType.uint64, 8] @@ -374,23 +374,23 @@ struct SHA512Context(Movable): self.buffer = InlineArray[UInt8, 128](fill=0) self.buffer_len = 0 - def __init__(out self, *, deinit take: Self): - self.state = take.state - self.count_high = take.count_high - self.count_low = take.count_low - self.buffer = take.buffer^ - self.buffer_len = take.buffer_len + def __init__(out self, *, deinit move: Self): + self.state = move.state + self.count_high = move.count_high + self.count_low = move.count_low + self.buffer = move.buffer^ + self.buffer_len = move.buffer_len - def __del__(deinit self): - memset_zero(self.buffer.unsafe_ptr(), 128) + def __deinit__(deinit self): + unsafe_memset_zero(self.buffer.unsafe_ptr(), 128) def wipe(mut self): - UnsafePointer(to=self.state).bitcast[UInt64]().store[volatile=True]( + Pointer(to=self.state).unsafe_bitcast[UInt64]().unsafe_store[volatile=True]( 0, SIMD[DType.uint64, 8](0) ) var buf_ptr = self.buffer.unsafe_ptr() for i in range(128): - buf_ptr.store[volatile=True](i, UInt8(0)) + buf_ptr.unsafe_store[volatile=True](i, UInt8(0)) self.count_high = 0 self.count_low = 0 self.buffer_len = 0 @@ -400,21 +400,21 @@ struct SHA512Context(Movable): self.state = iv self.count_high = 0 self.count_low = 0 - memset_zero(self.buffer.unsafe_ptr(), 128) + unsafe_memset_zero(self.buffer.unsafe_ptr(), 128) self.buffer_len = 0 def reset(mut self, iv: SIMD[DType.uint64, 8]): self.state = iv self.count_high = 0 self.count_low = 0 - memset_zero(self.buffer.unsafe_ptr(), 128) + unsafe_memset_zero(self.buffer.unsafe_ptr(), 128) self.buffer_len = 0 @always_inline def sha512_transform_blocks( mut state: SIMD[DType.uint64, 8], - data: UnsafePointer[mut=False, UInt8, _, address_space=_], + data: Pointer[mut=False, UInt8, _, address_space=_], nblocks: Int, ): comptime if CompilationTarget.has_neon() and CompilationTarget._has_feature["sha3"]() and not CompilationTarget.is_x86(): @@ -431,7 +431,7 @@ def sha512_transform_blocks( var h0 = state[7] for blk in range(nblocks): - var block = data + blk * 128 + var block = data.unsafe_offset(blk * 128) var w = InlineArray[UInt64, 16](uninitialized=True) var a = a0 @@ -506,7 +506,7 @@ def sha512_update(mut ctx: SHA512Context, data: Span[UInt8, ...]): var available = 128 - ctx.buffer_len if total_len >= available: for j in range(available): - buf_ptr[ctx.buffer_len + j] = data[j] + buf_ptr[unsafe_offset=ctx.buffer_len + j] = data[j] ctx.state = sha512_transform( ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=128) ) @@ -520,13 +520,13 @@ def sha512_update(mut ctx: SHA512Context, data: Span[UInt8, ...]): ctx.buffer_len = 0 else: for j in range(total_len): - buf_ptr[ctx.buffer_len + j] = data[j] + buf_ptr[unsafe_offset=ctx.buffer_len + j] = data[j] ctx.buffer_len += total_len return var nblocks = (total_len - i) // 128 if nblocks > 0: - sha512_transform_blocks(ctx.state, data.unsafe_ptr() + i, nblocks) + sha512_transform_blocks(ctx.state, data.unsafe_ptr().unsafe_offset(i), nblocks) var old_low = ctx.count_low ctx.count_low += UInt64(nblocks) * 1024 @@ -538,7 +538,7 @@ def sha512_update(mut ctx: SHA512Context, data: Span[UInt8, ...]): if i < total_len: var remaining = total_len - i for j in range(remaining): - buf_ptr[ctx.buffer_len + j] = data[i + j] + buf_ptr[unsafe_offset=ctx.buffer_len + j] = data[i + j] ctx.buffer_len += remaining @@ -555,7 +555,7 @@ def sha512_hash(data: Span[UInt8, ...]) -> List[UInt8]: return sha512_final(ctx) -def sha512_final_to_buffer(mut ctx: SHA512Context, output: UnsafePointer[mut=True, UInt8, _, address_space=_]): +def sha512_final_to_buffer(mut ctx: SHA512Context, output: Pointer[mut=True, UInt8, _, address_space=_]): var buf_ptr = ctx.buffer.unsafe_ptr() var final_low = ctx.count_low + UInt64(ctx.buffer_len) * 8 var final_high = ctx.count_high @@ -566,13 +566,13 @@ def sha512_final_to_buffer(mut ctx: SHA512Context, output: UnsafePointer[mut=Tru ctx.buffer_len += 1 if ctx.buffer_len > 112: - memset_zero(buf_ptr + ctx.buffer_len, 128 - ctx.buffer_len) + unsafe_memset_zero(buf_ptr.unsafe_offset(ctx.buffer_len), 128 - ctx.buffer_len) ctx.state = sha512_transform( ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=128) ) ctx.buffer_len = 0 - memset_zero(buf_ptr + ctx.buffer_len, 112 - ctx.buffer_len) + unsafe_memset_zero(buf_ptr.unsafe_offset(ctx.buffer_len), 112 - ctx.buffer_len) ctx.buffer_len = 112 for i in range(8): @@ -585,7 +585,7 @@ def sha512_final_to_buffer(mut ctx: SHA512Context, output: UnsafePointer[mut=Tru ) for i in range(8): - (output + i * 8).bitcast[UInt64]().store[alignment=1](byte_swap(ctx.state[i])) + (output.unsafe_offset(i * 8)).unsafe_bitcast[UInt64]().unsafe_store[alignment=1](byte_swap(ctx.state[i])) def sha256_final_with_len(mut ctx: SHA256Context, output_len: Int) -> List[UInt8]: var buf_ptr = ctx.buffer.unsafe_ptr() @@ -595,13 +595,13 @@ def sha256_final_with_len(mut ctx: SHA256Context, output_len: Int) -> List[UInt8 ctx.buffer_len += 1 if ctx.buffer_len > 56: - memset_zero(buf_ptr + ctx.buffer_len, 64 - ctx.buffer_len) + unsafe_memset_zero(buf_ptr.unsafe_offset(ctx.buffer_len), 64 - ctx.buffer_len) ctx.state = sha256_transform( ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=64) ) ctx.buffer_len = 0 - memset_zero(buf_ptr + ctx.buffer_len, 56 - ctx.buffer_len) + unsafe_memset_zero(buf_ptr.unsafe_offset(ctx.buffer_len), 56 - ctx.buffer_len) ctx.buffer_len = 56 for i in range(8): @@ -650,13 +650,13 @@ def sha256_final_partial( ctx.buffer_len += 1 if ctx.buffer_len > 56: - memset_zero(buf_ptr + ctx.buffer_len, 64 - ctx.buffer_len) + unsafe_memset_zero(buf_ptr.unsafe_offset(ctx.buffer_len), 64 - ctx.buffer_len) ctx.state = sha256_transform( ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=64) ) ctx.buffer_len = 0 - memset_zero(buf_ptr + ctx.buffer_len, 56 - ctx.buffer_len) + unsafe_memset_zero(buf_ptr.unsafe_offset(ctx.buffer_len), 56 - ctx.buffer_len) ctx.buffer_len = 56 for i in range(8): @@ -711,13 +711,13 @@ def sha512_final_with_len(mut ctx: SHA512Context, output_len: Int) -> List[UInt8 ctx.buffer_len += 1 if ctx.buffer_len > 112: - memset_zero(buf_ptr + ctx.buffer_len, 128 - ctx.buffer_len) + unsafe_memset_zero(buf_ptr.unsafe_offset(ctx.buffer_len), 128 - ctx.buffer_len) ctx.state = sha512_transform( ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=128) ) ctx.buffer_len = 0 - memset_zero(buf_ptr + ctx.buffer_len, 112 - ctx.buffer_len) + unsafe_memset_zero(buf_ptr.unsafe_offset(ctx.buffer_len), 112 - ctx.buffer_len) ctx.buffer_len = 112 for i in range(8): @@ -759,13 +759,13 @@ def sha512_final_partial( ctx.buffer_len += 1 if ctx.buffer_len > 112: - memset_zero(buf_ptr + ctx.buffer_len, 128 - ctx.buffer_len) + unsafe_memset_zero(buf_ptr.unsafe_offset(ctx.buffer_len), 128 - ctx.buffer_len) ctx.state = sha512_transform( ctx.state, Span[UInt8, ...](unsafe_ptr=buf_ptr, length=128) ) ctx.buffer_len = 0 - memset_zero(buf_ptr + ctx.buffer_len, 112 - ctx.buffer_len) + unsafe_memset_zero(buf_ptr.unsafe_offset(ctx.buffer_len), 112 - ctx.buffer_len) ctx.buffer_len = 112 for i in range(8): diff --git a/src/thistle/sha3.mojo b/src/thistle/sha3.mojo index c99fe4c..ac93ef5 100644 --- a/src/thistle/sha3.mojo +++ b/src/thistle/sha3.mojo @@ -4,7 +4,7 @@ FIPS 202 """ from std.collections import List -from std.memory import UnsafePointer, memcpy, memset_zero +from std.memory import Pointer, unsafe_memcpy, unsafe_memset_zero from .utils import StackBuffer, bytes_to_hex, string_to_bytes from std.bit import rotate_bits_left from std.builtin.simd import SIMD @@ -61,32 +61,32 @@ def _bcax(a: _U64x2, b: _U64x2, c: _U64x2) -> _U64x2: return llvm_intrinsic["llvm.aarch64.crypto.bcaxu", _U64x2, has_side_effect=False](a, b, c) -def _keccak_f1600_hw(state: UnsafePointer[mut=True, UInt64, _, address_space=_]): - var a0 = _U64x2(state[0], 0) - var a1 = _U64x2(state[1], 0) - var a2 = _U64x2(state[2], 0) - var a3 = _U64x2(state[3], 0) - var a4 = _U64x2(state[4], 0) - var a5 = _U64x2(state[5], 0) - var a6 = _U64x2(state[6], 0) - var a7 = _U64x2(state[7], 0) - var a8 = _U64x2(state[8], 0) - var a9 = _U64x2(state[9], 0) - var a10 = _U64x2(state[10], 0) - var a11 = _U64x2(state[11], 0) - var a12 = _U64x2(state[12], 0) - var a13 = _U64x2(state[13], 0) - var a14 = _U64x2(state[14], 0) - var a15 = _U64x2(state[15], 0) - var a16 = _U64x2(state[16], 0) - var a17 = _U64x2(state[17], 0) - var a18 = _U64x2(state[18], 0) - var a19 = _U64x2(state[19], 0) - var a20 = _U64x2(state[20], 0) - var a21 = _U64x2(state[21], 0) - var a22 = _U64x2(state[22], 0) - var a23 = _U64x2(state[23], 0) - var a24 = _U64x2(state[24], 0) +def _keccak_f1600_hw(state: Pointer[mut=True, UInt64, _, address_space=_]): + var a0 = _U64x2(state[unsafe_offset=0], 0) + var a1 = _U64x2(state[unsafe_offset=1], 0) + var a2 = _U64x2(state[unsafe_offset=2], 0) + var a3 = _U64x2(state[unsafe_offset=3], 0) + var a4 = _U64x2(state[unsafe_offset=4], 0) + var a5 = _U64x2(state[unsafe_offset=5], 0) + var a6 = _U64x2(state[unsafe_offset=6], 0) + var a7 = _U64x2(state[unsafe_offset=7], 0) + var a8 = _U64x2(state[unsafe_offset=8], 0) + var a9 = _U64x2(state[unsafe_offset=9], 0) + var a10 = _U64x2(state[unsafe_offset=10], 0) + var a11 = _U64x2(state[unsafe_offset=11], 0) + var a12 = _U64x2(state[unsafe_offset=12], 0) + var a13 = _U64x2(state[unsafe_offset=13], 0) + var a14 = _U64x2(state[unsafe_offset=14], 0) + var a15 = _U64x2(state[unsafe_offset=15], 0) + var a16 = _U64x2(state[unsafe_offset=16], 0) + var a17 = _U64x2(state[unsafe_offset=17], 0) + var a18 = _U64x2(state[unsafe_offset=18], 0) + var a19 = _U64x2(state[unsafe_offset=19], 0) + var a20 = _U64x2(state[unsafe_offset=20], 0) + var a21 = _U64x2(state[unsafe_offset=21], 0) + var a22 = _U64x2(state[unsafe_offset=22], 0) + var a23 = _U64x2(state[unsafe_offset=23], 0) + var a24 = _U64x2(state[unsafe_offset=24], 0) comptime for round in range(24): var c0 = _eor3(a0, a5, _eor3(a10, a15, a20)) @@ -154,66 +154,66 @@ def _keccak_f1600_hw(state: UnsafePointer[mut=True, UInt64, _, address_space=_]) a23 = _bcax(b23, b20, b24) a24 = _bcax(b24, b21, b20) - state[0] = a0[0] - state[1] = a1[0] - state[2] = a2[0] - state[3] = a3[0] - state[4] = a4[0] - state[5] = a5[0] - state[6] = a6[0] - state[7] = a7[0] - state[8] = a8[0] - state[9] = a9[0] - state[10] = a10[0] - state[11] = a11[0] - state[12] = a12[0] - state[13] = a13[0] - state[14] = a14[0] - state[15] = a15[0] - state[16] = a16[0] - state[17] = a17[0] - state[18] = a18[0] - state[19] = a19[0] - state[20] = a20[0] - state[21] = a21[0] - state[22] = a22[0] - state[23] = a23[0] - state[24] = a24[0] - - -def keccak_f1600(state: UnsafePointer[mut=True, UInt64, _, address_space=_]): + state[unsafe_offset=0] = a0[0] + state[unsafe_offset=1] = a1[0] + state[unsafe_offset=2] = a2[0] + state[unsafe_offset=3] = a3[0] + state[unsafe_offset=4] = a4[0] + state[unsafe_offset=5] = a5[0] + state[unsafe_offset=6] = a6[0] + state[unsafe_offset=7] = a7[0] + state[unsafe_offset=8] = a8[0] + state[unsafe_offset=9] = a9[0] + state[unsafe_offset=10] = a10[0] + state[unsafe_offset=11] = a11[0] + state[unsafe_offset=12] = a12[0] + state[unsafe_offset=13] = a13[0] + state[unsafe_offset=14] = a14[0] + state[unsafe_offset=15] = a15[0] + state[unsafe_offset=16] = a16[0] + state[unsafe_offset=17] = a17[0] + state[unsafe_offset=18] = a18[0] + state[unsafe_offset=19] = a19[0] + state[unsafe_offset=20] = a20[0] + state[unsafe_offset=21] = a21[0] + state[unsafe_offset=22] = a22[0] + state[unsafe_offset=23] = a23[0] + state[unsafe_offset=24] = a24[0] + + +def keccak_f1600(state: Pointer[mut=True, UInt64, _, address_space=_]): comptime if _has_sha3_ext: _keccak_f1600_hw(state) return _keccak_f1600_scalar(state) -def _keccak_f1600_scalar(state: UnsafePointer[mut=True, UInt64, _, address_space=_]): - var a0 = state[0] - var a1 = state[1] - var a2 = state[2] - var a3 = state[3] - var a4 = state[4] - var a5 = state[5] - var a6 = state[6] - var a7 = state[7] - var a8 = state[8] - var a9 = state[9] - var a10 = state[10] - var a11 = state[11] - var a12 = state[12] - var a13 = state[13] - var a14 = state[14] - var a15 = state[15] - var a16 = state[16] - var a17 = state[17] - var a18 = state[18] - var a19 = state[19] - var a20 = state[20] - var a21 = state[21] - var a22 = state[22] - var a23 = state[23] - var a24 = state[24] +def _keccak_f1600_scalar(state: Pointer[mut=True, UInt64, _, address_space=_]): + var a0 = state[unsafe_offset=0] + var a1 = state[unsafe_offset=1] + var a2 = state[unsafe_offset=2] + var a3 = state[unsafe_offset=3] + var a4 = state[unsafe_offset=4] + var a5 = state[unsafe_offset=5] + var a6 = state[unsafe_offset=6] + var a7 = state[unsafe_offset=7] + var a8 = state[unsafe_offset=8] + var a9 = state[unsafe_offset=9] + var a10 = state[unsafe_offset=10] + var a11 = state[unsafe_offset=11] + var a12 = state[unsafe_offset=12] + var a13 = state[unsafe_offset=13] + var a14 = state[unsafe_offset=14] + var a15 = state[unsafe_offset=15] + var a16 = state[unsafe_offset=16] + var a17 = state[unsafe_offset=17] + var a18 = state[unsafe_offset=18] + var a19 = state[unsafe_offset=19] + var a20 = state[unsafe_offset=20] + var a21 = state[unsafe_offset=21] + var a22 = state[unsafe_offset=22] + var a23 = state[unsafe_offset=23] + var a24 = state[unsafe_offset=24] comptime for round in range(24): var c0 = a0 ^ a5 ^ a10 ^ a15 ^ a20 @@ -308,31 +308,31 @@ def _keccak_f1600_scalar(state: UnsafePointer[mut=True, UInt64, _, address_space a0 ^= KECCAK_RC[round] - state[0] = a0 - state[1] = a1 - state[2] = a2 - state[3] = a3 - state[4] = a4 - state[5] = a5 - state[6] = a6 - state[7] = a7 - state[8] = a8 - state[9] = a9 - state[10] = a10 - state[11] = a11 - state[12] = a12 - state[13] = a13 - state[14] = a14 - state[15] = a15 - state[16] = a16 - state[17] = a17 - state[18] = a18 - state[19] = a19 - state[20] = a20 - state[21] = a21 - state[22] = a22 - state[23] = a23 - state[24] = a24 + state[unsafe_offset=0] = a0 + state[unsafe_offset=1] = a1 + state[unsafe_offset=2] = a2 + state[unsafe_offset=3] = a3 + state[unsafe_offset=4] = a4 + state[unsafe_offset=5] = a5 + state[unsafe_offset=6] = a6 + state[unsafe_offset=7] = a7 + state[unsafe_offset=8] = a8 + state[unsafe_offset=9] = a9 + state[unsafe_offset=10] = a10 + state[unsafe_offset=11] = a11 + state[unsafe_offset=12] = a12 + state[unsafe_offset=13] = a13 + state[unsafe_offset=14] = a14 + state[unsafe_offset=15] = a15 + state[unsafe_offset=16] = a16 + state[unsafe_offset=17] = a17 + state[unsafe_offset=18] = a18 + state[unsafe_offset=19] = a19 + state[unsafe_offset=20] = a20 + state[unsafe_offset=21] = a21 + state[unsafe_offset=22] = a22 + state[unsafe_offset=23] = a23 + state[unsafe_offset=24] = a24 struct SHA3Context(Movable): @@ -351,26 +351,26 @@ struct SHA3Context(Movable): self.buffer = StackBuffer[UInt8, 168](fill=0) self.buffer_len = 0 - def __init__(out self, *, deinit take: Self): - self.state = take.state^ - self.rate_bytes = take.rate_bytes - self.buffer = take.buffer^ - self.buffer_len = take.buffer_len + def __init__(out self, *, deinit move: Self): + self.state = move.state^ + self.rate_bytes = move.rate_bytes + self.buffer = move.buffer^ + self.buffer_len = move.buffer_len - def __del__(deinit self): + def __deinit__(deinit self): var state_ptr = self.state.ptr() for i in range(25): - state_ptr.store[volatile=True](i, UInt64(0)) + state_ptr.unsafe_store[volatile=True](i, UInt64(0)) var buffer_ptr = self.buffer.ptr() for i in range(168): - buffer_ptr.store[volatile=True](i, UInt8(0)) + buffer_ptr.unsafe_store[volatile=True](i, UInt8(0)) @always_inline -def sha3_absorb_block(state: UnsafePointer[mut=True, UInt64, _, address_space=_], block: UnsafePointer[mut=False, UInt8, _, address_space=_], rate_bytes: Int): +def sha3_absorb_block(state: Pointer[mut=True, UInt64, _, address_space=_], block: Pointer[mut=False, UInt8, _, address_space=_], rate_bytes: Int): var full_lanes = rate_bytes // 8 for i in range(full_lanes): - state[i] ^= (block + i * 8).bitcast[UInt64]().load[width=1, alignment=1]() + state[unsafe_offset=i] ^= block.unsafe_offset(i * 8).unsafe_bitcast[UInt64]().unsafe_load[width=1, alignment=1]() keccak_f1600(state) @@ -393,7 +393,7 @@ def sha3_update(mut ctx: SHA3Context, data: Span[UInt8, ...]): return while i + ctx.rate_bytes <= total_len: - sha3_absorb_block(ctx.state.ptr(), data.unsafe_ptr() + i, ctx.rate_bytes) + sha3_absorb_block(ctx.state.ptr(), data.unsafe_ptr().unsafe_offset(i), ctx.rate_bytes) i += ctx.rate_bytes if i < total_len: @@ -412,7 +412,7 @@ def sha3_final(mut ctx: SHA3Context, output_len_bytes: Int) -> List[UInt8]: var pad_len = ctx.rate_bytes - ctx.buffer_len if pad_len > 0: - memset_zero(ctx.buffer.ptr() + ctx.buffer_len, pad_len) + unsafe_memset_zero(ctx.buffer.ptr().unsafe_offset(ctx.buffer_len), pad_len) ctx.buffer_len = ctx.rate_bytes ctx.buffer[ctx.rate_bytes - 1] |= 0x80 @@ -428,9 +428,9 @@ def sha3_final(mut ctx: SHA3Context, output_len_bytes: Int) -> List[UInt8]: if output_len_bytes - offset < limit: limit = output_len_bytes - offset - memcpy( - dest=output.unsafe_ptr() + offset, - src=ctx.state.ptr().bitcast[UInt8](), + unsafe_memcpy( + dest=output.unsafe_ptr().unsafe_offset(offset), + src=ctx.state.ptr().unsafe_bitcast[UInt8](), count=limit, ) @@ -453,7 +453,7 @@ def sha3_final_into(mut ctx: SHA3Context, mut output: StackBuffer[UInt8, ...], o var pad_len = ctx.rate_bytes - ctx.buffer_len if pad_len > 0: - memset_zero(ctx.buffer.ptr() + ctx.buffer_len, pad_len) + unsafe_memset_zero(ctx.buffer.ptr().unsafe_offset(ctx.buffer_len), pad_len) ctx.buffer_len = ctx.rate_bytes ctx.buffer[ctx.rate_bytes - 1] |= 0x80 @@ -467,9 +467,9 @@ def sha3_final_into(mut ctx: SHA3Context, mut output: StackBuffer[UInt8, ...], o if output_len_bytes - offset < limit: limit = output_len_bytes - offset - memcpy( - dest=output.ptr() + offset, - src=ctx.state.ptr().bitcast[UInt8](), + unsafe_memcpy( + dest=output.ptr().unsafe_offset(offset), + src=ctx.state.ptr().unsafe_bitcast[UInt8](), count=limit, ) @@ -549,7 +549,7 @@ def shake_finalize(mut ctx: SHA3Context): var pad_len = ctx.rate_bytes - ctx.buffer_len if pad_len > 0: - memset_zero(ctx.buffer.ptr() + ctx.buffer_len, pad_len) + unsafe_memset_zero(ctx.buffer.ptr().unsafe_offset(ctx.buffer_len), pad_len) ctx.buffer_len = ctx.rate_bytes ctx.buffer[ctx.rate_bytes - 1] |= 0x80 @@ -572,9 +572,9 @@ def shake_squeeze_prefix_into(mut ctx: SHA3Context, mut output: StackBuffer[UInt if output_len - offset < limit: limit = output_len - offset - memcpy( - dest=output.ptr() + offset, - src=ctx.state.ptr().bitcast[UInt8](), + unsafe_memcpy( + dest=output.ptr().unsafe_offset(offset), + src=ctx.state.ptr().unsafe_bitcast[UInt8](), count=limit, ) @@ -606,9 +606,9 @@ def shake_final(mut ctx: SHA3Context, output_len: Int) -> List[UInt8]: if output_len - offset < limit: limit = output_len - offset - memcpy( - dest=output.unsafe_ptr() + offset, - src=ctx.state.ptr().bitcast[UInt8](), + unsafe_memcpy( + dest=output.unsafe_ptr().unsafe_offset(offset), + src=ctx.state.ptr().unsafe_bitcast[UInt8](), count=limit, ) @@ -636,9 +636,9 @@ def shake_final_into(mut ctx: SHA3Context, mut output: StackBuffer[UInt8, ...], if output_len - offset < limit: limit = output_len - offset - memcpy( - dest=output.ptr() + offset, - src=ctx.state.ptr().bitcast[UInt8](), + unsafe_memcpy( + dest=output.ptr().unsafe_offset(offset), + src=ctx.state.ptr().unsafe_bitcast[UInt8](), count=limit, ) diff --git a/src/thistle/sha_ni.mojo b/src/thistle/sha_ni.mojo index 5ac813a..aacda35 100644 --- a/src/thistle/sha_ni.mojo +++ b/src/thistle/sha_ni.mojo @@ -3,7 +3,7 @@ SHA-NI implementation In Mojo. """ from std.sys import llvm_intrinsic, CompilationTarget, prefetch, PrefetchOptions -from std.memory import UnsafePointer, bitcast +from std.memory import Pointer, bitcast from .utils import StackBuffer from std.builtin.simd import SIMD from std.builtin.dtype import DType @@ -112,13 +112,13 @@ def byte_swap32(v: SIMD128) -> SIMD128: @always_inline("nodebug") -def Load(ptr: UnsafePointer[mut=False, UInt8, _, address_space=_]) -> SIMD128: - return byte_swap32(ptr.bitcast[UInt32]().load[width=4, alignment=1]()) +def Load(ptr: Pointer[mut=False, UInt8, _, address_space=_]) -> SIMD128: + return byte_swap32(ptr.unsafe_bitcast[UInt32]().unsafe_load[width=4, alignment=1]()) @always_inline("nodebug") -def prefetch_next_block(ptr: UnsafePointer[mut=False, UInt8, _, address_space=_]): - prefetch[PrefetchOptions().for_read().high_locality().to_data_cache()](ptr + 64) +def prefetch_next_block(ptr: Pointer[mut=False, UInt8, _, address_space=_]): + prefetch[PrefetchOptions().for_read().high_locality().to_data_cache()](ptr.unsafe_offset(64)) def sha256ni_transform(state: SIMD[DType.uint32, 8], block: Span[UInt8, ...]) -> SIMD[DType.uint32, 8]: @@ -142,9 +142,9 @@ def _sha256ni_transform_arm(state: SIMD[DType.uint32, 8], block: Span[UInt8, ... var w = InlineArray[SIMD128, 4](uninitialized=True) w[0] = Load(ptr) - w[1] = Load(ptr + 16) - w[2] = Load(ptr + 32) - w[3] = Load(ptr + 48) + w[1] = Load(ptr.unsafe_offset(16)) + w[2] = Load(ptr.unsafe_offset(32)) + w[3] = Load(ptr.unsafe_offset(48)) comptime for i in range(16): var wk = w[i & 3] + SIMD128(SHA256_K[4 * i], SHA256_K[4 * i + 1], SHA256_K[4 * i + 2], SHA256_K[4 * i + 3]) @@ -178,9 +178,9 @@ def _sha256ni_transform_x86(state: SIMD[DType.uint32, 8], block: Span[UInt8, ... # expand all 64 words to 16 simd registers var w0 = Load(ptr) - var w1 = Load(ptr + 16) - var w2 = Load(ptr + 32) - var w3 = Load(ptr + 48) + var w1 = Load(ptr.unsafe_offset(16)) + var w2 = Load(ptr.unsafe_offset(32)) + var w3 = Load(ptr.unsafe_offset(48)) var w = InlineArray[SIMD128, 16](uninitialized=True) w[0] = w0 @@ -228,7 +228,7 @@ def sha256ni_hash(data: Span[UInt8, ...]) -> List[UInt8]: while i + 64 <= total_len: if i + 64 < total_len: - prefetch_next_block(data.unsafe_ptr() + i) + prefetch_next_block(data.unsafe_ptr().unsafe_offset(i)) ctx.state = sha256ni_transform(ctx.state, data[i:i+64]) ctx.count += 512 i += 64 @@ -272,7 +272,7 @@ def has_sha_ni() -> Bool: def sha256ni_transform_blocks( mut state: SIMD[DType.uint32, 8], - data: UnsafePointer[mut=False, UInt8, _, address_space=_], + data: Pointer[mut=False, UInt8, _, address_space=_], nblocks: Int, ): comptime if CompilationTarget.has_neon() and CompilationTarget._has_feature["sha2"]() and not CompilationTarget.is_x86(): @@ -280,7 +280,7 @@ def sha256ni_transform_blocks( var st1 = SIMD128(state[4], state[5], state[6], state[7]) for blk in range(nblocks): - var ptr = data + blk * 64 + var ptr = data.unsafe_offset(blk * 64) if blk + 1 < nblocks: prefetch_next_block(ptr) var old_st0 = st0 @@ -288,9 +288,9 @@ def sha256ni_transform_blocks( var w = InlineArray[SIMD128, 4](uninitialized=True) w[0] = Load(ptr) - w[1] = Load(ptr + 16) - w[2] = Load(ptr + 32) - w[3] = Load(ptr + 48) + w[1] = Load(ptr.unsafe_offset(16)) + w[2] = Load(ptr.unsafe_offset(32)) + w[3] = Load(ptr.unsafe_offset(48)) comptime for i in range(16): var wk = w[i & 3] + SIMD128(SHA256_K[4 * i], SHA256_K[4 * i + 1], SHA256_K[4 * i + 2], SHA256_K[4 * i + 3]) @@ -313,7 +313,7 @@ def sha256ni_transform_blocks( return for blk in range(nblocks): - var span = Span[UInt8, ...](unsafe_ptr=data + blk * 64, length=64) + var span = Span[UInt8, ...](unsafe_ptr=data.unsafe_offset(blk * 64), length=64) state = sha256ni_transform(state, span) @@ -379,8 +379,8 @@ def byte_swap64(v: SIMD64x2) -> SIMD64x2: @always_inline("nodebug") -def Load512(ptr: UnsafePointer[mut=False, UInt8, _, address_space=_]) -> SIMD64x2: - return byte_swap64(ptr.bitcast[UInt64]().load[width=2, alignment=1]()) +def Load512(ptr: Pointer[mut=False, UInt8, _, address_space=_]) -> SIMD64x2: + return byte_swap64(ptr.unsafe_bitcast[UInt64]().unsafe_load[width=2, alignment=1]()) @always_inline("nodebug") @@ -389,13 +389,13 @@ def _ext1(a: SIMD64x2, b: SIMD64x2) -> SIMD64x2: @always_inline("nodebug") -def prefetch_next_block512(ptr: UnsafePointer[mut=False, UInt8, _, address_space=_]): - prefetch[PrefetchOptions().for_read().high_locality().to_data_cache()](ptr + 128) +def prefetch_next_block512(ptr: Pointer[mut=False, UInt8, _, address_space=_]): + prefetch[PrefetchOptions().for_read().high_locality().to_data_cache()](ptr.unsafe_offset(128)) def sha512ni_transform_blocks( mut state: SIMD[DType.uint64, 8], - data: UnsafePointer[mut=False, UInt8, _, address_space=_], + data: Pointer[mut=False, UInt8, _, address_space=_], nblocks: Int, ): var ab = SIMD64x2(state[0], state[1]) @@ -404,13 +404,13 @@ def sha512ni_transform_blocks( var gh = SIMD64x2(state[6], state[7]) for blk in range(nblocks): - var block = data + blk * 128 + var block = data.unsafe_offset(blk * 128) if blk + 1 < nblocks: prefetch_next_block512(block) var w = InlineArray[SIMD64x2, 8](uninitialized=True) comptime for i in range(8): - w[i] = Load512(block + i * 16) + w[i] = Load512(block.unsafe_offset(i * 16)) var v = InlineArray[SIMD64x2, 4](uninitialized=True) v[0] = ab diff --git a/src/thistle/utils.mojo b/src/thistle/utils.mojo index 6368a2e..9672176 100644 --- a/src/thistle/utils.mojo +++ b/src/thistle/utils.mojo @@ -39,17 +39,17 @@ struct StackInlineArray[ElementType: Copyable & Deinitable, size: Int](Copyable) var ptr = self.unsafe_ptr() comptime for i in range(Self.size): - ptr.init_pointee_move_from( - UnsafePointer(to=storage[i]).unsafe_mut_cast[True]() + ptr.unsafe_write_move_from( + Pointer(to=storage[i]).unsafe_mut_cast[True]() ) - ptr += 1 + ptr = ptr.unsafe_offset(1) storage^._annihilate() @always_inline def unsafe_ptr[ origin: Origin, address_space: AddressSpace, // - ](ref[origin, address_space] self) -> UnsafePointer[ + ](ref[origin, address_space] self) -> Pointer[ Self.ElementType, origin, address_space=address_space @@ -58,7 +58,7 @@ struct StackInlineArray[ElementType: Copyable & Deinitable, size: Int](Copyable) self._data.unsafe_ptr() .unsafe_mut_cast[origin.mut]() .unsafe_origin_cast[origin]() - .address_space_cast[address_space]() + .unsafe_address_space_cast[address_space]() ) @always_inline @@ -103,8 +103,8 @@ struct StackInlineArray[ElementType: Copyable & Deinitable, size: Int](Copyable) 0 <= idx < Self.size, "The index provided must be within the range [0, len(List) -1] when using List.unsafe_set()", ) - (self._data.unsafe_ptr() + idx).destroy_pointee() - (self._data.unsafe_ptr() + idx).init_pointee_move(value^) + (self._data.unsafe_ptr().unsafe_offset(idx)).unsafe_deinit_pointee() + (self._data.unsafe_ptr().unsafe_offset(idx)).unsafe_write(value^) struct StackBuffer[T: Copyable & Deinitable, N: Int](Movable): @@ -123,9 +123,9 @@ struct StackBuffer[T: Copyable & Deinitable, N: Int](Movable): self._len = 0 @always_inline - def __init__(out self, *, deinit take: Self): - self._data = take._data^ - self._len = take._len + def __init__(out self, *, deinit move: Self): + self._data = move._data^ + self._len = move._len @always_inline def len(self) -> Int: @@ -200,31 +200,31 @@ struct StackBuffer[T: Copyable & Deinitable, N: Int](Movable): @always_inline def ptr[ origin: Origin, address_space: AddressSpace, // - ](ref[origin, address_space] self) -> UnsafePointer[ + ](ref[origin, address_space] self) -> Pointer[ Self.T, origin, address_space=address_space ]: return ( self._data.unsafe_ptr() .unsafe_mut_cast[origin.mut]() .unsafe_origin_cast[origin]() - .address_space_cast[address_space]() + .unsafe_address_space_cast[address_space]() ) @always_inline -def load_32be(ptr: UnsafePointer[mut=False, UInt8, _, address_space=_], offset: Int) -> UInt32: - return byte_swap((ptr + offset).bitcast[UInt32]().load[width=1, alignment=1]()) +def load_32be(ptr: Pointer[mut=False, UInt8, _, address_space=_], offset: Int) -> UInt32: + return byte_swap((ptr.unsafe_offset(offset)).unsafe_bitcast[UInt32]().unsafe_load[width=1, alignment=1]()) @always_inline -def load_64be(ptr: UnsafePointer[mut=False, UInt8, _, address_space=_], offset: Int) -> UInt64: - return byte_swap((ptr + offset).bitcast[UInt64]().load[width=1, alignment=1]()) +def load_64be(ptr: Pointer[mut=False, UInt8, _, address_space=_], offset: Int) -> UInt64: + return byte_swap((ptr.unsafe_offset(offset)).unsafe_bitcast[UInt64]().unsafe_load[width=1, alignment=1]()) @always_inline -def store_64be(p: UnsafePointer[mut=True, UInt8, _, address_space=_], off: Int, v: UInt64): +def store_64be(p: Pointer[mut=True, UInt8, _, address_space=_], off: Int, v: UInt64): for i in range(8): - p[off + i] = UInt8((v >> UInt64(56 - 8 * i)) & 0xFF) + p[unsafe_offset=off + i] = UInt8((v >> UInt64(56 - 8 * i)) & 0xFF) @always_inline @@ -254,7 +254,7 @@ def u64_zero_choice(x: UInt64) -> UInt64: def zero_stack_u8(mut data: StackBuffer[UInt8, ...]): var ptr = data.ptr() for i in range(data.len()): - ptr.store[volatile=True](i, UInt8(0)) + ptr.unsafe_store[volatile=True](i, UInt8(0)) data.clear() @@ -268,14 +268,14 @@ def nibble_to_hex_char(nibble: UInt8) -> UInt8: @always_inline -def bytes_to_hex_simd(data: UnsafePointer[mut=False, UInt8, _, address_space=_], len: Int) -> String: +def bytes_to_hex_simd(data: Pointer[mut=False, UInt8, _, address_space=_], len: Int) -> String: debug_assert[assert_mode="safe"]( 0 <= len <= Int.MAX // 2, "Hex input length cannot be negative or overflow the output size", ) var result = String(capacity=len * 2) for i in range(len): - var b = data[i] + var b = data[unsafe_offset=i] result += chr(Int(nibble_to_hex_char((b >> 4) & 0x0F))) result += chr(Int(nibble_to_hex_char(b & 0x0F))) return result diff --git a/src/thistle/x25519.mojo b/src/thistle/x25519.mojo index 7bea861..1ba6611 100644 --- a/src/thistle/x25519.mojo +++ b/src/thistle/x25519.mojo @@ -83,7 +83,7 @@ def x25519( res.to_bytes_into(output.unsafe_ptr()) var scalar_ptr = scalar.unsafe_ptr() for i in range(32): - scalar_ptr.store[volatile=True](i, UInt8(0)) + scalar_ptr.unsafe_store[volatile=True](i, UInt8(0)) def x25519_public_key( diff --git a/tests/benchmark.mojo b/tests/benchmark.mojo index 74109dc..5298bb5 100644 --- a/tests/benchmark.mojo +++ b/tests/benchmark.mojo @@ -4,6 +4,7 @@ from max.algorithm import parallelize from std.random import random_ui64, seed from std.math import ceildiv from std.sys import has_accelerator +from std.memory import Layout, alloc from thistle.argon2 import Argon2id from thistle.blake2b import Blake2b @@ -20,7 +21,6 @@ from thistle.ed25519 import ed25519_sign, ed25519_verify, ed25519_generate_publi from thistle.p256 import p256_ecdsa_sign from thistle.p384 import p384_public_key, p384_ecdsa_sign from thistle.utils import StackInlineArray -from std.memory import alloc from std.utils import StaticTuple comptime TEST_KEY: StaticTuple[UInt8, 16] = StaticTuple[UInt8, 16]( @@ -54,7 +54,7 @@ def benchmark_x25519(duration_secs: Float64) raises -> String: var start = perf_counter() while perf_counter() - start < duration_secs: x25519(scalar_span, point_span, Span[mut=True, UInt8, ...](out)) - scalar.unsafe_ptr().store[volatile=True](0, out[0] | 8) + scalar.unsafe_ptr().unsafe_store[volatile=True](0, out[0] | 8) count += 1 var duration = perf_counter() - start var ops = Float64(count) / duration @@ -257,9 +257,9 @@ def benchmark_camellia(data_size: Int, duration_secs: Float64) raises -> String: var cipher = CamelliaCipher(Span[UInt8, ...](key)) var nb = 32 - var blocks = alloc[UInt8](nb * 16) + var blocks = alloc(Layout[UInt8](count=nb * 16)).unsafe_leak() for i in range(nb * 16): - blocks.store(i, UInt8(i % 256)) + blocks.unsafe_store(i, UInt8(i % 256)) for _ in range(100): camellia_encrypt_blocks(cipher, blocks, nb) @@ -272,7 +272,7 @@ def benchmark_camellia(data_size: Int, duration_secs: Float64) raises -> String: var end = perf_counter() var duration = end - start - blocks.free() + blocks.unsafe_free() var mbps = Float64(count * 16) / (1024 * 1024) / duration return "camellia | throughput: " + String(mbps) + " mb/s, blocks: " + String(count) + ", time: " + String(duration) + "s" @@ -285,12 +285,12 @@ def benchmark_camellia_ctr(duration_secs: Float64) raises -> String: var cipher = CamelliaCipher(Span[UInt8, ...](key)) var size = 64 * 1024 - var buf = alloc[UInt8](size) + var buf = alloc(Layout[UInt8](count=size)).unsafe_leak() for i in range(size): - buf.store(i, UInt8(i % 256)) - var nonce = alloc[UInt8](16) + buf.unsafe_store(i, UInt8(i % 256)) + var nonce = alloc(Layout[UInt8](count=16)).unsafe_leak() for i in range(16): - nonce.store(i, UInt8(i * 3)) + nonce.unsafe_store(i, UInt8(i * 3)) camellia_ctr_kernel(buf, buf, cipher, size // 16, nonce) @@ -302,8 +302,8 @@ def benchmark_camellia_ctr(duration_secs: Float64) raises -> String: var end = perf_counter() var duration = end - start - buf.free() - nonce.free() + buf.unsafe_free() + nonce.unsafe_free() var mbps = Float64(count * size) / (1024 * 1024) / duration return "camellia-ctr | throughput: " + String(mbps) + " mb/s, chunks: " + String(count) + ", time: " + String(duration) + "s" @@ -313,14 +313,14 @@ def benchmark_chacha20(data_size: Int, duration_secs: Float64) raises -> String: var key = SIMD[DType.uint8, 32](0) for i in range(32): key[i] = UInt8(i) - var nonce = SIMD[DType.uint8, 16](0) + var nonce = InlineArray[UInt8, 12](fill=0) var data = List[UInt8](capacity=data_size) for i in range(data_size): data.append(UInt8(i % 256)) var span = Span[mut=True, UInt8](data) - var cipher = ChaCha20(key, nonce) + var cipher = ChaCha20(key, Span[UInt8, ...](nonce)) var checksum: UInt64 = 0 var count = 0 @@ -382,9 +382,9 @@ def benchmark_aes_cpu(duration_secs: Float64) raises -> String: var key = AESKey(TEST_KEY) var round_keys = key.round_keys() var skey = cpu_aes_ct_skey(round_keys, ROUNDS_128) - var blocks = alloc[UInt8](256) + var blocks = alloc(Layout[UInt8](count=256)).unsafe_leak() for i in range(256): - blocks.store(i, TEST_PT[i % 16]) + blocks.unsafe_store(i, TEST_PT[i % 16]) for _ in range(100): cpu_aes_ct_encrypt16(blocks, skey, ROUNDS_128) @@ -397,7 +397,7 @@ def benchmark_aes_cpu(duration_secs: Float64) raises -> String: var end = perf_counter() var duration = end - start - blocks.free() + blocks.unsafe_free() var mbps = Float64(count * 16) / (1024 * 1024) / duration return "aes-128-cpu | throughput: " + String(mbps) + " mb/s, blocks: " + String(count) + ", time: " + String(duration) + "s" @@ -411,20 +411,20 @@ def benchmark_aes_gpu_ecb() raises -> String: from max.gpu.host import DeviceContext from thistle.aes_gpu import aes_gpu_kernel_ecb - var key_ptr = alloc[UInt8](16) + var key_ptr = alloc(Layout[UInt8](count=16)).unsafe_leak() for i in range(16): - key_ptr.store(i, TEST_KEY[i]) + key_ptr.unsafe_store(i, TEST_KEY[i]) var round_keys = expand_key_128( Span[UInt8, ...](unsafe_ptr=key_ptr, length=16) ) var num_blocks = 131072 var total_bytes = num_blocks * 16 - var input_host = alloc[Scalar[DType.uint8]](total_bytes) - var output_host = alloc[Scalar[DType.uint8]](total_bytes) + var input_host = alloc(Layout[Scalar[DType.uint8]](count=total_bytes)).unsafe_leak() + var output_host = alloc(Layout[Scalar[DType.uint8]](count=total_bytes)).unsafe_leak() for i in range(total_bytes): - input_host[i] = TEST_PT[i % 16] + input_host[unsafe_offset=i] = TEST_PT[i % 16] with DeviceContext() as ctx: var input_buffer = ctx.enqueue_create_buffer[DType.uint8](total_bytes) @@ -469,9 +469,9 @@ def benchmark_aes_gpu_ecb() raises -> String: var total_gb = Float64(iterations * total_bytes) / 1024.0 / 1024.0 / 1024.0 var gbps = total_gb / duration - input_host.free() - output_host.free() - key_ptr.free() + input_host.unsafe_free() + output_host.unsafe_free() + key_ptr.unsafe_free() return "aes-128-gpu-ecb | throughput: " + String(gbps) + " gb/s, iterations: " + String(iterations) @@ -484,23 +484,23 @@ def benchmark_aes_gpu_ctr() raises -> String: from max.gpu.host import DeviceContext from thistle.aes_gpu import aes_gpu_kernel_ctr - var key_ptr = alloc[UInt8](16) + var key_ptr = alloc(Layout[UInt8](count=16)).unsafe_leak() for i in range(16): - key_ptr.store(i, TEST_KEY[i]) + key_ptr.unsafe_store(i, TEST_KEY[i]) var round_keys = expand_key_128( Span[UInt8, ...](unsafe_ptr=key_ptr, length=16) ) var num_blocks = 131072 var total_bytes = num_blocks * 16 - var input_host = alloc[Scalar[DType.uint8]](total_bytes) - var output_host = alloc[Scalar[DType.uint8]](total_bytes) - var nonce_host = alloc[Scalar[DType.uint8]](16) + var input_host = alloc(Layout[Scalar[DType.uint8]](count=total_bytes)).unsafe_leak() + var output_host = alloc(Layout[Scalar[DType.uint8]](count=total_bytes)).unsafe_leak() + var nonce_host = alloc(Layout[Scalar[DType.uint8]](count=16)).unsafe_leak() for i in range(total_bytes): - input_host[i] = TEST_PT[i % 16] + input_host[unsafe_offset=i] = TEST_PT[i % 16] for i in range(16): - nonce_host[i] = 0 + nonce_host[unsafe_offset=i] = 0 with DeviceContext() as ctx: var input_buffer = ctx.enqueue_create_buffer[DType.uint8](total_bytes) @@ -549,10 +549,10 @@ def benchmark_aes_gpu_ctr() raises -> String: var total_gb = Float64(iterations * total_bytes) / 1024.0 / 1024.0 / 1024.0 var gbps = total_gb / duration - input_host.free() - output_host.free() - nonce_host.free() - key_ptr.free() + input_host.unsafe_free() + output_host.unsafe_free() + nonce_host.unsafe_free() + key_ptr.unsafe_free() return "aes-128-gpu-ctr | throughput: " + String(gbps) + " gb/s, iterations: " + String(iterations) @@ -568,24 +568,24 @@ def benchmark_aes_gpu_gcm() raises -> String: from max.gpu.host import DeviceContext from thistle.aes_gpu import aes_gpu_kernel_gcm_ctr - var key_ptr = alloc[UInt8](16) + var key_ptr = alloc(Layout[UInt8](count=16)).unsafe_leak() for i in range(16): - key_ptr.store(i, TEST_KEY[i]) + key_ptr.unsafe_store(i, TEST_KEY[i]) var round_keys = expand_key_128( Span[UInt8, ...](unsafe_ptr=key_ptr, length=16) ) var num_blocks = 131072 var total_bytes = num_blocks * 16 - var input_host = alloc[Scalar[DType.uint8]](total_bytes) - var output_host = alloc[Scalar[DType.uint8]](total_bytes) - var nonce_host = alloc[Scalar[DType.uint8]](16) + var input_host = alloc(Layout[Scalar[DType.uint8]](count=total_bytes)).unsafe_leak() + var output_host = alloc(Layout[Scalar[DType.uint8]](count=total_bytes)).unsafe_leak() + var nonce_host = alloc(Layout[Scalar[DType.uint8]](count=16)).unsafe_leak() for i in range(total_bytes): - input_host[i] = TEST_PT[i % 16] + input_host[unsafe_offset=i] = TEST_PT[i % 16] for i in range(16): - nonce_host[i] = 0 - nonce_host[15] = 1 + nonce_host[unsafe_offset=i] = 0 + nonce_host[unsafe_offset=15] = 1 with DeviceContext() as ctx: var input_buffer = ctx.enqueue_create_buffer[DType.uint8](total_bytes) @@ -634,10 +634,10 @@ def benchmark_aes_gpu_gcm() raises -> String: var total_gb = Float64(iterations * total_bytes) / 1024.0 / 1024.0 / 1024.0 var gbps = total_gb / duration - input_host.free() - output_host.free() - nonce_host.free() - key_ptr.free() + input_host.unsafe_free() + output_host.unsafe_free() + nonce_host.unsafe_free() + key_ptr.unsafe_free() return "aes-128-gpu-gcm | throughput: " + String(gbps) + " gb/s, iterations: " + String(iterations) diff --git a/tests/dudect.mojo b/tests/dudect.mojo index d0a0e1c..b54c821 100644 --- a/tests/dudect.mojo +++ b/tests/dudect.mojo @@ -3,7 +3,7 @@ Dudect test (Reparaz/Balasch/Verbauwhede). """ from std.time import perf_counter_ns -from std.memory import UnsafePointer +from std.memory import Pointer from thistle.camellia import CamelliaCipher, camellia_encrypt_block from thistle.kcipher2 import KCipher2 @@ -259,7 +259,7 @@ def run_aes_sw(mut rng: Rng) raises -> Bool: def run_chacha20(mut rng: Rng) raises -> Bool: var cls = _classes(N_FAST, rng) var inp = _fast_inputs(cls, 32, rng) - var nonce = SIMD[DType.uint8, 16](0) + var nonce = InlineArray[UInt8, 12](fill=0) var times = List[Float64](capacity=N_FAST) var sink: UInt8 = 0 var data = List[UInt8]() @@ -271,7 +271,7 @@ def run_chacha20(mut rng: Rng) raises -> Bool: key[j] = inp[i * 32 + j] var t0 = perf_counter_ns() for _ in range(BATCH): - var c = ChaCha20(key, nonce) + var c = ChaCha20(key, Span[UInt8, ...](nonce)) var span = Span[mut=True, UInt8](data) c.encrypt_inplace(span) sink ^= data[0] diff --git a/tests/gen_ed25519_table.mojo b/tests/gen_ed25519_table.mojo index 2fa4958..c4098dc 100644 --- a/tests/gen_ed25519_table.mojo +++ b/tests/gen_ed25519_table.mojo @@ -29,7 +29,7 @@ def _affine_niels_limbs(p: EdwardsPoint) -> InlineArray[UInt64, 15]: out[i] = y_plus_x.limbs[i] out[5 + i] = y_minus_x.limbs[i] out[10 + i] = xy2d.limbs[i] - return out + return out^ def _print_chunk(name: String, points: InlineArray[EdwardsPoint, 8]): @@ -81,11 +81,11 @@ def main() raises: print(" var t = InlineArray[UInt64, 4096](uninitialized=True)") print(" var p = t.unsafe_ptr()") for j in range(32): - print(" p.store[alignment=8](" + String(j * 128) + ", _ED25519_BT" + String(j) + ")") + print(" p.unsafe_store[alignment=8](" + String(j * 128) + ", _ED25519_BT" + String(j) + ")") print(" return t") print() print("@no_inline") print("def ed25519_b_odd_table() -> InlineArray[UInt64, 128]:") print(" var t = InlineArray[UInt64, 128](uninitialized=True)") - print(" t.unsafe_ptr().store[alignment=8](0, _ED25519_B_ODD)") + print(" t.unsafe_ptr().unsafe_store[alignment=8](0, _ED25519_B_ODD)") print(" return t") diff --git a/tests/gen_p256_table.mojo b/tests/gen_p256_table.mojo index dc3548e..13bdfb0 100644 --- a/tests/gen_p256_table.mojo +++ b/tests/gen_p256_table.mojo @@ -39,5 +39,5 @@ def main() raises: print(" var t = InlineArray[UInt64, 4096](uninitialized=True)") print(" var p = t.unsafe_ptr()") for c in range(32): - print(" p.store[alignment=8](" + String(c * 128) + ", _P256_BT" + String(c) + ")") + print(" p.unsafe_store[alignment=8](" + String(c * 128) + ", _P256_BT" + String(c) + ")") print(" return t^") diff --git a/tests/gen_p384_table.mojo b/tests/gen_p384_table.mojo index 4262171..c499cb5 100644 --- a/tests/gen_p384_table.mojo +++ b/tests/gen_p384_table.mojo @@ -39,5 +39,5 @@ def main() raises: print(" var t = InlineArray[UInt64, 8704](uninitialized=True)") print(" var p = t.unsafe_ptr()") for c in range(68): - print(" p.store[alignment=8](" + String(c * 128) + ", _P384_BT" + String(c) + ")") + print(" p.unsafe_store[alignment=8](" + String(c * 128) + ", _P384_BT" + String(c) + ")") print(" return t^") diff --git a/tests/test_aes_gpu.mojo b/tests/test_aes_gpu.mojo index 6d9126c..e5a8171 100644 --- a/tests/test_aes_gpu.mojo +++ b/tests/test_aes_gpu.mojo @@ -5,9 +5,9 @@ from std.sys import has_accelerator from thistle.sha2 import bytes_to_hex from thistle.aes import cpu_aes_ct_skey, AESExpandedKey from thistle.aes_gpu import aes_gpu_kernel_ecb, aes_gpu_kernel_ctr, aes_gpu_kernel_gcm_ctr -from std.memory import alloc from max.gpu.host import DeviceContext -from std.memory.unsafe_pointer import UnsafePointer +from std.memory.unsafe_pointer import Pointer +from std.memory import Layout, alloc def byte_to_hex(b: UInt8) -> String: @@ -81,7 +81,6 @@ def test_aes_gpu_basic(json_data: PythonObject, py: PythonObject) raises -> Test var name = String(v["name"]) var key_hex = String(v["key"]) var pt_hex = String(v["plaintext"]) - var expected_ct_hex = String(v["ciphertext"]) var key_bytes = hex_to_bytes(key_hex) var pt_bytes_data = hex_to_bytes(pt_hex) @@ -90,12 +89,12 @@ def test_aes_gpu_basic(json_data: PythonObject, py: PythonObject) raises -> Test var rounds = 10 if key_len == 16 else (12 if key_len == 24 else 14) var total_bytes = 64 - var input_host = alloc[Scalar[DType.uint8]](total_bytes) - var output_host = alloc[Scalar[DType.uint8]](total_bytes) + var input_host = alloc(Layout[Scalar[DType.uint8]](count=total_bytes)).unsafe_leak() + var output_host = alloc(Layout[Scalar[DType.uint8]](count=total_bytes)).unsafe_leak() for block in range(4): for j in range(16): - input_host[block * 16 + j] = pt_bytes_data[j] + input_host[unsafe_offset=block * 16 + j] = pt_bytes_data[j] var input_buffer = ctx.enqueue_create_buffer[DType.uint8](total_bytes) var output_buffer = ctx.enqueue_create_buffer[DType.uint8](total_bytes) @@ -125,14 +124,13 @@ def test_aes_gpu_basic(json_data: PythonObject, py: PythonObject) raises -> Test ctx.enqueue_copy(output_host, output_buffer) ctx.synchronize() - var correct = True var expected = String(v["ciphertext"]) var all_passed = True for block in range(4): for j in range(16): var expected_byte = hex_char_to_val(Int(expected.as_bytes()[j * 2])) << 4 expected_byte |= hex_char_to_val(Int(expected.as_bytes()[j * 2 + 1])) - var actual_byte = output_host[block * 16 + j] + var actual_byte = output_host[unsafe_offset=block * 16 + j] if actual_byte != expected_byte: all_passed = False break @@ -145,11 +143,11 @@ def test_aes_gpu_basic(json_data: PythonObject, py: PythonObject) raises -> Test failed += 1 var got_hex = String("") for j in range(16): - got_hex += byte_to_hex(output_host[j]) + got_hex += byte_to_hex(output_host[unsafe_offset=j]) failures.append("AES-GPU " + name + ": expected " + expected + ", got " + got_hex) - input_host.free() - output_host.free() + input_host.unsafe_free() + output_host.unsafe_free() return TestResult(passed, failed, failures^) @@ -176,11 +174,11 @@ def test_mode_gpu(json_data: PythonObject, mode: String) raises -> TestResult: var total_bytes = len(pt_bytes) var n_blocks = total_bytes // 16 - var pt_ptr = alloc[UInt8](total_bytes) + var pt_ptr = alloc(Layout[UInt8](count=total_bytes)).unsafe_leak() for j in range(total_bytes): - pt_ptr.store(j, pt_bytes[j]) + pt_ptr.unsafe_store(j, pt_bytes[j]) - var ct_ptr = alloc[UInt8](total_bytes) + var ct_ptr = alloc(Layout[UInt8](count=total_bytes)).unsafe_leak() var input_buffer = ctx.enqueue_create_buffer[DType.uint8](total_bytes) var output_buffer = ctx.enqueue_create_buffer[DType.uint8](total_bytes) @@ -203,7 +201,7 @@ def test_mode_gpu(json_data: PythonObject, mode: String) raises -> TestResult: var skey_buffer = ctx.enqueue_create_buffer[DType.uint64]((rounds + 1) * 8) ctx.enqueue_copy(skey_buffer, skey_host.unsafe_ptr()) - var nonce_ptr = alloc[UInt8](16) + var nonce_ptr = alloc(Layout[UInt8](count=16)).unsafe_leak() var nonce_buffer = ctx.enqueue_create_buffer[DType.uint8](16) if "ECB" in mode: ctx.enqueue_function[aes_gpu_kernel_ecb]( @@ -219,11 +217,11 @@ def test_mode_gpu(json_data: PythonObject, mode: String) raises -> TestResult: var iv_hex = String(tv.get("iv", PythonObject())) if iv_hex.byte_length() == 0: for j in range(16): - nonce_ptr[j] = 0 + nonce_ptr[unsafe_offset=j] = 0 else: var iv_bytes = hex_to_bytes(iv_hex) for j in range(16): - nonce_ptr.store(j, iv_bytes[j]) + nonce_ptr.unsafe_store(j, iv_bytes[j]) ctx.enqueue_copy(nonce_buffer, nonce_ptr) ctx.enqueue_function[aes_gpu_kernel_ctr]( @@ -240,15 +238,15 @@ def test_mode_gpu(json_data: PythonObject, mode: String) raises -> TestResult: var nonce_hex = String(tv.get("nonce", PythonObject())) if nonce_hex.byte_length() == 0: for j in range(12): - nonce_ptr[j] = 0 + nonce_ptr[unsafe_offset=j] = 0 else: var nonce_bytes = hex_to_bytes(nonce_hex) for j in range(12): - nonce_ptr.store(j, nonce_bytes[j]) - nonce_ptr[12] = 0 - nonce_ptr[13] = 0 - nonce_ptr[14] = 0 - nonce_ptr[15] = 1 + nonce_ptr.unsafe_store(j, nonce_bytes[j]) + nonce_ptr[unsafe_offset=12] = 0 + nonce_ptr[unsafe_offset=13] = 0 + nonce_ptr[unsafe_offset=14] = 0 + nonce_ptr[unsafe_offset=15] = 1 ctx.enqueue_copy(nonce_buffer, nonce_ptr) ctx.enqueue_function[aes_gpu_kernel_gcm_ctr]( @@ -287,7 +285,7 @@ def test_mode_gpu(json_data: PythonObject, mode: String) raises -> TestResult: for j in range(total_bytes): var expected_byte = hex_char_to_val(Int(expected_ct.as_bytes()[j * 2])) << 4 expected_byte = expected_byte | hex_char_to_val(Int(expected_ct.as_bytes()[j * 2 + 1])) - if ct_ptr.load(j) != expected_byte: + if ct_ptr.unsafe_load(j) != expected_byte: correct = False break @@ -297,12 +295,12 @@ def test_mode_gpu(json_data: PythonObject, mode: String) raises -> TestResult: failed += 1 var got_hex = String("") for j in range(total_bytes): - got_hex += byte_to_hex(ct_ptr.load(j)) + got_hex += byte_to_hex(ct_ptr.unsafe_load(j)) failures.append(mode + " " + String(i) + ": expected " + expected_ct + ", got " + got_hex) - pt_ptr.free() - ct_ptr.free() - nonce_ptr.free() + pt_ptr.unsafe_free() + ct_ptr.unsafe_free() + nonce_ptr.unsafe_free() return TestResult(passed, failed, failures^) diff --git a/tests/test_security_boundaries.mojo b/tests/test_security_boundaries.mojo index 1209826..c22a6b6 100644 --- a/tests/test_security_boundaries.mojo +++ b/tests/test_security_boundaries.mojo @@ -6,8 +6,15 @@ from thistle.chacha20poly1305 import ( chacha20_poly1305_encrypt, hchacha20, ) +from thistle.chacha20 import ChaCha20 from thistle.p256 import p256_ecdsa_sign, p256_public_key from thistle.p384 import p384_ecdsa_sign, p384_public_key +from thistle.pbkdf2 import ( + PBKDF2_SHA256_MAX_DKLEN, + PBKDF2_SHA512_MAX_DKLEN, + pbkdf2_hmac_sha256, + pbkdf2_hmac_sha512, +) from thistle.x25519 import x25519 @@ -65,6 +72,68 @@ def main() raises: if not rejected: raise Error("HChaCha20 accepted an undersized destination") + var chacha_key = SIMD[DType.uint8, 32](0) + var oversized_nonce = List[UInt8](length=16, fill=0) + rejected = False + try: + var invalid_cipher = ChaCha20( + chacha_key, Span[UInt8, ...](oversized_nonce) + ) + _ = invalid_cipher + except: + rejected = True + if not rejected: + raise Error("ChaCha20 accepted a nonce longer than 12 bytes") + + var nonce_a = List[UInt8](length=12, fill=0) + var nonce_b = List[UInt8](length=12, fill=0) + nonce_b[11] = 1 + var zeros = List[UInt8](length=64, fill=0) + var stream_a = List[UInt8](length=64, fill=0) + var stream_b = List[UInt8](length=64, fill=0) + var cipher_a = ChaCha20(chacha_key, Span[UInt8, ...](nonce_a)) + var cipher_b = ChaCha20(chacha_key, Span[UInt8, ...](nonce_b)) + var stream_a_span = Span[mut=True, UInt8, ...](stream_a) + var stream_b_span = Span[mut=True, UInt8, ...](stream_b) + cipher_a.encrypt_into( + Span[UInt8, ...](zeros), stream_a_span + ) + cipher_b.encrypt_into( + Span[UInt8, ...](zeros), stream_b_span + ) + var nonce_affects_stream = False + for i in range(64): + nonce_affects_stream |= stream_a[i] != stream_b[i] + if not nonce_affects_stream: + raise Error("ChaCha20 ignored the final nonce byte") + + var empty_salt = List[UInt8]() + rejected = False + try: + _ = pbkdf2_hmac_sha256( + Span[UInt8, ...](empty), + Span[UInt8, ...](empty_salt), + 1, + PBKDF2_SHA256_MAX_DKLEN + 1, + ) + except: + rejected = True + if not rejected: + raise Error("PBKDF2-SHA256 accepted an oversized derived key") + + rejected = False + try: + _ = pbkdf2_hmac_sha512( + Span[UInt8, ...](empty), + Span[UInt8, ...](empty_salt), + 1, + PBKDF2_SHA512_MAX_DKLEN + 1, + ) + except: + rejected = True + if not rejected: + raise Error("PBKDF2-SHA512 accepted an oversized derived key") + var nonce = List[UInt8](length=12, fill=0) var plaintext = List[UInt8](length=16, fill=0) var ciphertext = List[UInt8](length=15, fill=0) diff --git a/tests/test_wycheproof_ed25519.mojo b/tests/test_wycheproof_ed25519.mojo index 4d33a0a..34074e5 100644 --- a/tests/test_wycheproof_ed25519.mojo +++ b/tests/test_wycheproof_ed25519.mojo @@ -1,5 +1,5 @@ """ -Wycheproof Ed25519 test suite +Wycheproof Ed25519 test suite. """ from std.collections import List diff --git a/tests/test_wycheproof_x25519.mojo b/tests/test_wycheproof_x25519.mojo index e05b21c..944ddc2 100644 --- a/tests/test_wycheproof_x25519.mojo +++ b/tests/test_wycheproof_x25519.mojo @@ -1,5 +1,5 @@ """ -Wycheproof X25519 test suite +Wycheproof X25519 test suite. """ from std.collections import List from std.python import Python diff --git a/tests/thistle_test_vectors.mojo b/tests/thistle_test_vectors.mojo index 63cc9bc..0877c5c 100644 --- a/tests/thistle_test_vectors.mojo +++ b/tests/thistle_test_vectors.mojo @@ -1,7 +1,7 @@ from std.python import Python, PythonObject from std.collections import List -from std.memory import alloc -from std.memory.unsafe_pointer import UnsafePointer +from std.memory.unsafe_pointer import Pointer +from std.memory import Layout, alloc from thistle.sha2 import ( bytes_to_hex, string_to_bytes, @@ -86,10 +86,10 @@ def byte_to_hex(b: UInt8) -> String: ) -def ptr_to_hex(ptr: UnsafePointer[mut=True, UInt8, _, address_space=_], count: Int) -> String: +def ptr_to_hex(ptr: Pointer[mut=True, UInt8, _, address_space=_], count: Int) -> String: var s = String("") for j in range(count): - s += byte_to_hex(ptr.load(j)) + s += byte_to_hex(ptr.unsafe_load(j)) return s @@ -107,13 +107,6 @@ def list_to_simd32(lst: List[UInt8]) -> SIMD[DType.uint8, 32]: return result -def list_to_simd12(lst: List[UInt8]) -> SIMD[DType.uint8, 16]: - var result = SIMD[DType.uint8, 16](0) - for i in range(min(len(lst), 12)): - result[i] = lst[i] - return result - - @fieldwise_init struct TestResult(Copyable, Movable): var passed: Int @@ -221,8 +214,8 @@ def test_camellia(data: PythonObject, py: PythonObject) raises -> TestResult: var pt_bytes = hex_to_bytes(pt_hex) var ct_bytes = hex_to_bytes(ct_hex) - var pt_blk = pt_bytes.unsafe_ptr().load[width=16, alignment=1](0) - var ct_blk = ct_bytes.unsafe_ptr().load[width=16, alignment=1](0) + var pt_blk = pt_bytes.unsafe_ptr().unsafe_load[width=16, alignment=1](0) + var ct_blk = ct_bytes.unsafe_ptr().unsafe_load[width=16, alignment=1](0) var got_ct = bytes_to_hex(camellia_encrypt_block(cipher, pt_blk)) if got_ct == ct_hex: @@ -250,7 +243,7 @@ def test_camellia(data: PythonObject, py: PythonObject) raises -> TestResult: var expected = List[UInt8](capacity=nb * 16) for bi in range(nb): var one = camellia_encrypt_block( - cipher, buf.unsafe_ptr().load[width=16, alignment=1](bi * 16) + cipher, buf.unsafe_ptr().unsafe_load[width=16, alignment=1](bi * 16) ) for j in range(16): expected.append(one[j]) @@ -296,7 +289,7 @@ def test_camellia(data: PythonObject, py: PythonObject) raises -> TestResult: ctr_blk[j] = UInt8(total & 0xFF) carry = (carry >> 8) + (total >> 8) var ks_blk = camellia_encrypt_block( - cipher, ctr_blk.unsafe_ptr().load[width=16, alignment=1](0) + cipher, ctr_blk.unsafe_ptr().unsafe_load[width=16, alignment=1](0) ) for j in range(16): var idx = bi * 16 + j @@ -334,7 +327,7 @@ def test_camellia(data: PythonObject, py: PythonObject) raises -> TestResult: for j in range(16): prev[j] = nonce[j] for bi in range(cbc_nb): - var x = cbc_msg.unsafe_ptr().load[width=16, alignment=1](bi * 16) ^ prev + var x = cbc_msg.unsafe_ptr().unsafe_load[width=16, alignment=1](bi * 16) ^ prev prev = camellia_encrypt_block(cipher, x) for j in range(16): if cbc_out[bi * 16 + j] != prev[j]: @@ -363,15 +356,16 @@ def test_chacha20(data: PythonObject, py: PythonObject) raises -> TestResult: for i in range(Int(py=data.__len__())): var v = data[i] var name = String(v["name"]) + var nonce = hex_to_bytes(String(v["nonce"])) var cipher = ChaCha20( list_to_simd32(hex_to_bytes(String(v["key"]))), - list_to_simd12(hex_to_bytes(String(v["nonce"]))), + Span[UInt8, ...](nonce), UInt32(Int(py=v["counter"])), ) var pt_bytes = hex_to_bytes(String(v["plaintext"])) var expected_ct = hex_to_bytes(String(v["ciphertext"])) if len(pt_bytes) == 0: - var null_ptr = UnsafePointer[ + var null_ptr = Pointer[ UInt8, MutUntrackedOrigin ].unsafe_dangling() var out_span = Span[mut=True, UInt8, MutUntrackedOrigin]( @@ -386,7 +380,7 @@ def test_chacha20(data: PythonObject, py: PythonObject) raises -> TestResult: "ChaCha20 " + name + ": empty plaintext test failed" ) else: - var ct_ptr = alloc[UInt8](len(pt_bytes)) + var ct_ptr = alloc(Layout[UInt8](count=len(pt_bytes))).unsafe_leak() var ct_span = Span[mut=True, UInt8, MutUntrackedOrigin]( unsafe_ptr=ct_ptr, length=len(pt_bytes) ) @@ -394,7 +388,7 @@ def test_chacha20(data: PythonObject, py: PythonObject) raises -> TestResult: var ok = len(pt_bytes) == len(expected_ct) if ok: for j in range(len(pt_bytes)): - if ct_ptr[j] != expected_ct[j]: + if ct_ptr[unsafe_offset=j] != expected_ct[j]: ok = False break if ok: @@ -410,7 +404,7 @@ def test_chacha20(data: PythonObject, py: PythonObject) raises -> TestResult: + ptr_to_hex(ct_ptr, min(len(pt_bytes), 16)) + "..." ) - ct_ptr.free() + ct_ptr.unsafe_free() return TestResult(passed, failed, failures^) @@ -612,9 +606,9 @@ def test_aes_cpu(data: PythonObject, py: PythonObject) raises -> TestResult: var pt_bytes = hex_to_bytes(String(v["plaintext"])) var expected_ct = String(v["ciphertext"]) var rounds = 10 if len(key_bytes) == 16 else (12 if len(key_bytes) == 24 else 14) - var pt_ptr = alloc[UInt8](16) + var pt_ptr = alloc(Layout[UInt8](count=16)).unsafe_leak() for j in range(16): - pt_ptr.store(j, pt_bytes[j]) + pt_ptr.unsafe_store(j, pt_bytes[j]) var round_keys = _expand_aes_key(key_bytes) cpu_aes_encrypt(pt_ptr, round_keys.ptr(), rounds) var got = ptr_to_hex(pt_ptr, 16) @@ -630,7 +624,7 @@ def test_aes_cpu(data: PythonObject, py: PythonObject) raises -> TestResult: + ", got " + got ) - pt_ptr.free() + pt_ptr.unsafe_free() return TestResult(passed, failed, failures^) @@ -649,9 +643,9 @@ def test_aes_ni(data: PythonObject, py: PythonObject) raises -> TestResult: var pt_bytes = hex_to_bytes(String(v["plaintext"])) var expected_ct = String(v["ciphertext"]) var rounds = 10 if len(key_bytes) == 16 else (12 if len(key_bytes) == 24 else 14) - var pt_ptr = alloc[UInt8](16) + var pt_ptr = alloc(Layout[UInt8](count=16)).unsafe_leak() for j in range(16): - pt_ptr.store(j, pt_bytes[j]) + pt_ptr.unsafe_store(j, pt_bytes[j]) var round_keys = _expand_aes_key(key_bytes) aes_encrypt(pt_ptr, round_keys.ptr(), rounds) var got = ptr_to_hex(pt_ptr, 16) @@ -667,7 +661,7 @@ def test_aes_ni(data: PythonObject, py: PythonObject) raises -> TestResult: + ", got " + got ) - pt_ptr.free() + pt_ptr.unsafe_free() return TestResult(passed, failed, failures^) @@ -783,8 +777,8 @@ def test_chacha20_poly1305(data: PythonObject, py: PythonObject, xchacha: Bool) for _ in range(16): out_tag.append(0) - var enc_matches = False - var dec_ok = False + var enc_matches: Bool + var dec_ok: Bool try: if xchacha: xchacha20_poly1305_encrypt( @@ -876,11 +870,11 @@ def test_aes_cpu_modes(data: PythonObject, py: PythonObject) raises -> TestResul ok = ok and bytes_to_hex(enc[1]) == bytes_to_hex(tag_exp) else: var nblocks = n // 16 - var ip = alloc[UInt8](n) - var op = alloc[UInt8](n) + var ip = alloc(Layout[UInt8](count=n)).unsafe_leak() + var op = alloc(Layout[UInt8](count=n)).unsafe_leak() for j in range(n): - ip[j] = pt[j] - op[j] = 0 + ip[unsafe_offset=j] = pt[j] + op[unsafe_offset=j] = 0 if "XTS" in mode: var half = len(key) // 2 @@ -893,42 +887,42 @@ def test_aes_cpu_modes(data: PythonObject, py: PythonObject) raises -> TestResul var rk2 = _expand_aes_key(k2) var rounds = 10 if half == 16 else 14 var tweak = hex_to_bytes(String(tv["tweak"])) - var twp = alloc[UInt8](16) + var twp = alloc(Layout[UInt8](count=16)).unsafe_leak() for j in range(16): - twp[j] = tweak[j] + twp[unsafe_offset=j] = tweak[j] cpu_aes_xts_kernel( ip, op, rk1.ptr(), rk2.ptr(), nblocks, twp, rounds ) - twp.free() + twp.unsafe_free() else: var rk = _expand_aes_key(key) var rounds = 10 if len(key) == 16 else (12 if len(key) == 24 else 14) if "CBC" in mode: var iv = hex_to_bytes(String(tv["iv"])) - var ivp = alloc[UInt8](16) + var ivp = alloc(Layout[UInt8](count=16)).unsafe_leak() for j in range(16): - ivp[j] = iv[j] + ivp[unsafe_offset=j] = iv[j] cpu_aes_cbc_kernel( ip, op, rk.ptr(), nblocks, ivp, rounds ) - ivp.free() + ivp.unsafe_free() elif "CTR" in mode: var iv = hex_to_bytes(String(tv["iv"])) - var ivp = alloc[UInt8](16) + var ivp = alloc(Layout[UInt8](count=16)).unsafe_leak() for j in range(16): - ivp[j] = iv[j] + ivp[unsafe_offset=j] = iv[j] cpu_aes_ctr_kernel( ip, op, rk.ptr(), nblocks, ivp, rounds ) - ivp.free() + ivp.unsafe_free() else: cpu_aes_ecb_kernel( ip, op, rk.ptr(), nblocks, rounds ) ok = ptr_to_hex(op, n) == bytes_to_hex(ct_exp) - ip.free() - op.free() + ip.unsafe_free() + op.unsafe_free() if ok: passed += 1 From a590f14e0285a7141fc2bd7f849f3d306e54fc8a Mon Sep 17 00:00:00 2001 From: libalpm64 <159192189+libalpm64@users.noreply.github.com> Date: Thu, 20 Aug 2026 17:50:48 -0400 Subject: [PATCH 3/4] ~ --- src/thistle/__init__.mojo | 3 +- src/thistle/aes_ni.mojo | 17 ++-- src/thistle/argon2.mojo | 20 +++-- src/thistle/blake2b.mojo | 13 ++- src/thistle/blake3.mojo | 18 ++-- src/thistle/chacha20poly1305.mojo | 12 +-- src/thistle/ed25519.mojo | 118 ++++++++++++++++----------- src/thistle/ed25519_table.mojo | 4 +- src/thistle/p256.mojo | 10 +-- src/thistle/p256_table.mojo | 2 +- src/thistle/p384.mojo | 10 +-- src/thistle/p384_table.mojo | 2 +- src/thistle/poly1305.mojo | 13 ++- src/thistle/rsa.mojo | 57 +++++++------ src/thistle/sha2.mojo | 4 +- src/thistle/sha3.mojo | 37 ++++----- src/thistle/sha_ni.mojo | 10 +-- src/thistle/utils.mojo | 84 +++++++------------ src/thistle/x25519.mojo | 4 +- tests/benchmark.mojo | 32 +++++--- tests/dudect.mojo | 3 +- tests/gen_ed25519_table.mojo | 12 +-- tests/gen_p256_table.mojo | 2 +- tests/gen_p384_table.mojo | 2 +- tests/test_security_boundaries.mojo | 71 ++++++++++++++++ tests/test_sha3_capacity_guard.sh | 22 +++-- tests/test_signing.mojo | 22 ++++- tests/test_wycheproof_p256_ecdh.mojo | 2 +- tests/test_wycheproof_p384_ecdh.mojo | 2 +- tests/test_wycheproof_x25519.mojo | 2 +- 30 files changed, 377 insertions(+), 233 deletions(-) diff --git a/src/thistle/__init__.mojo b/src/thistle/__init__.mojo index 0c25d74..e6c4527 100644 --- a/src/thistle/__init__.mojo +++ b/src/thistle/__init__.mojo @@ -14,8 +14,7 @@ from .pbkdf2 import pbkdf2_hmac_sha256, pbkdf2_hmac_sha512 from .argon2 import Argon2id, argon2id_hash_string from .aes import AESKey, AESExpandedKey, expand_key_128, expand_key_192, expand_key_256 -from .aes_ni import has_aes_ni, aes_gcm_ctr_kernel, aes_gcm_encrypt, aes_gcm_decrypt, AESGCMContext -from .aes_gpu import aes_gpu_kernel_ecb, aes_gpu_kernel_ctr, aes_gpu_kernel_gcm_ctr +from .aes_ni import has_aes_ni, aes_gcm_encrypt, aes_gcm_decrypt, AESGCMContext from .camellia import CamelliaCipher from .chacha20 import ChaCha20, chacha20_block from .poly1305 import Poly1305, poly1305_mac diff --git a/src/thistle/aes_ni.mojo b/src/thistle/aes_ni.mojo index 871e296..9a4171b 100644 --- a/src/thistle/aes_ni.mojo +++ b/src/thistle/aes_ni.mojo @@ -5,6 +5,7 @@ AES-NI implementation from std.collections import List, InlineArray from std.sys import llvm_intrinsic, CompilationTarget from std.memory import bitcast, unsafe_memset_zero, unsafe_memcpy, Pointer +from std.os import abort from std.utils import StaticTuple from .aes import cpu_aes_encrypt, cpu_aes_ct_encrypt, cpu_aes_ct_encrypt16, cpu_aes_ct_skey, expand_key_128_into, expand_key_192_into, expand_key_256_into from .utils import StackBuffer, load_64be, store_64be @@ -174,7 +175,7 @@ def x86_aes_encrypt_256_direct( def _arm_load_keys[N: Int]( round_keys: Pointer[mut=True, UInt32, _, address_space=_] ) -> InlineArray[SIMD16, N]: - var keys = InlineArray[SIMD16, N](uninitialized=True) + var keys = InlineArray[SIMD16, N](fill=SIMD16(0)) comptime for i in range(N): var raw = (round_keys.unsafe_offset(i * 4)).unsafe_bitcast[UInt8]().unsafe_load[ width=16, alignment=1 @@ -451,6 +452,8 @@ def aes_encrypt( Uses X86 AES-NI or ARM Crypto Extensions on supported hardware, falls back to software implementation otherwise. """ + if rounds != 10 and rounds != 12 and rounds != 14: + abort("AES round count must be 10, 12, or 14") comptime if CompilationTarget._has_feature["sse"]() and CompilationTarget._has_feature["aes"](): if rounds == 10: x86_aes_encrypt_128(pt, round_keys) @@ -478,6 +481,10 @@ def aes_gcm_ctr_kernel( j0_ptr: Pointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: + if rounds != 10 and rounds != 12 and rounds != 14: + abort("AES round count must be 10, 12, or 14") + if num_blocks < 0: + abort("AES-GCM block count cannot be negative") comptime if CompilationTarget._has_feature["sse"]() and CompilationTarget._has_feature["aes"](): _hw_gcm_ctr_kernel(input_ptr, output_ptr, round_keys, num_blocks, j0_ptr, rounds) else: @@ -591,7 +598,7 @@ def _arm_gcm_fused_loop[NR: Int]( var i = 0 while i + 8 <= num_blocks: - var b = InlineArray[SIMD16, 8](uninitialized=True) + var b = InlineArray[SIMD16, 8](fill=SIMD16(0)) comptime for k in range(8): var cv = j0u32 cv[3] = llvm_intrinsic["llvm.bswap.i32", UInt32, has_side_effect=False]( @@ -606,7 +613,7 @@ def _arm_gcm_fused_loop[NR: Int]( var p = input_ptr.unsafe_offset(i * 16) var q = output_ptr.unsafe_offset(i * 16) - var g = InlineArray[SIMD16, 8](uninitialized=True) + var g = InlineArray[SIMD16, 8](fill=SIMD16(0)) comptime for k in range(8): var pt = p.unsafe_load[width=16, alignment=1](k * 16) var ct = pt ^ b[k] @@ -627,7 +634,7 @@ def _arm_gcm_fused_loop[NR: Int]( i += 8 if i + 4 <= num_blocks: - var b = InlineArray[SIMD16, 4](uninitialized=True) + var b = InlineArray[SIMD16, 4](fill=SIMD16(0)) comptime for k in range(4): var cv = j0u32 cv[3] = llvm_intrinsic["llvm.bswap.i32", UInt32, has_side_effect=False]( @@ -642,7 +649,7 @@ def _arm_gcm_fused_loop[NR: Int]( var p = input_ptr.unsafe_offset(i * 16) var q = output_ptr.unsafe_offset(i * 16) - var g = InlineArray[SIMD16, 4](uninitialized=True) + var g = InlineArray[SIMD16, 4](fill=SIMD16(0)) comptime for k in range(4): var pt = p.unsafe_load[width=16, alignment=1](k * 16) var ct = pt ^ b[k] diff --git a/src/thistle/argon2.mojo b/src/thistle/argon2.mojo index 48b6361..38f8104 100644 --- a/src/thistle/argon2.mojo +++ b/src/thistle/argon2.mojo @@ -212,7 +212,7 @@ def variable_length_hash_into( store_le32(le_buf, 0, t_len) ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf, length=4)) ctx.update(input) - ctx.finalize_into(out_ptr) + ctx.finalize_into(output) finally: zero_and_free(le_buf, 4) return @@ -225,7 +225,9 @@ def variable_length_hash_into( store_le32(le_buf, 0, t_len) ctx1.update(Span[UInt8, ...](unsafe_ptr=le_buf, length=4)) ctx1.update(input) - ctx1.finalize_into(v_buf) + ctx1.finalize_into( + Span[mut=True, UInt8, ...](unsafe_ptr=v_buf, length=64) + ) var out_offset = 0 for _ in range(r - 1): @@ -235,7 +237,9 @@ def variable_length_hash_into( var ctx = Blake2b(64) ctx.update(Span[UInt8, ...](unsafe_ptr=v_buf, length=64)) - ctx.finalize_into(v_buf) + ctx.finalize_into( + Span[mut=True, UInt8, ...](unsafe_ptr=v_buf, length=64) + ) for j in range(32): out_ptr[unsafe_offset=out_offset + j] = v_buf[unsafe_offset=j] @@ -244,7 +248,11 @@ def variable_length_hash_into( var last_len = t_len - 32 * r var ctx_last = Blake2b(last_len) ctx_last.update(Span[UInt8, ...](unsafe_ptr=v_buf, length=64)) - ctx_last.finalize_into(out_ptr.unsafe_offset(out_offset)) + ctx_last.finalize_into( + Span[mut=True, UInt8, ...]( + unsafe_ptr=out_ptr.unsafe_offset(out_offset), length=last_len + ) + ) finally: zero_and_free(v_buf, 64) zero_and_free(le_buf, 4) @@ -494,7 +502,9 @@ struct Argon2id: zero_and_free(le_buf, 4) var h0_buf = alloc(Layout[UInt8](count=64)).unsafe_leak() - h0_ctx.finalize_into(h0_buf) + h0_ctx.finalize_into( + Span[mut=True, UInt8, ...](unsafe_ptr=h0_buf, length=64) + ) var m_blocks = self.memory_size_kb var m_prime_blocks = ( diff --git a/src/thistle/blake2b.mojo b/src/thistle/blake2b.mojo index 2704bd5..a7c3e86 100644 --- a/src/thistle/blake2b.mojo +++ b/src/thistle/blake2b.mojo @@ -233,7 +233,9 @@ struct Blake2b(Movable): self._buf_ptr()[unsafe_offset=j] = data[i + j] self.buffer_len = total - i - def finalize_into(mut self, output: Pointer[mut=True, UInt8, _, address_space=_]): + def _finalize_into_unchecked( + mut self, output: Pointer[mut=True, UInt8, _, address_space=_] + ): var old_low = self.t_low self.t_low += UInt64(self.buffer_len) if self.t_low < old_low: @@ -250,11 +252,18 @@ struct Blake2b(Movable): for i in range(self.out_len): output[unsafe_offset=i] = h_bytes[unsafe_offset=i] + def finalize_into( + mut self, output: Span[mut=True, UInt8, ...] + ) raises: + if len(output) < self.out_len: + raise Error("BLAKE2b output buffer is too small") + self._finalize_into_unchecked(output.unsafe_ptr()) + def finalize(mut self) -> List[UInt8]: var output = List[UInt8](capacity=self.out_len) for _ in range(self.out_len): output.append(0) - self.finalize_into(output.unsafe_ptr()) + self._finalize_into_unchecked(output.unsafe_ptr()) return output^ diff --git a/src/thistle/blake3.mojo b/src/thistle/blake3.mojo index b904e46..c3c770f 100644 --- a/src/thistle/blake3.mojo +++ b/src/thistle/blake3.mojo @@ -175,7 +175,9 @@ def compress_core( ] # fmt: on - var res = StackInlineArray[SIMD[DType.uint32, 1], 16](uninitialized=True) + var res = StackInlineArray[SIMD[DType.uint32, 1], 16]( + fill=SIMD[DType.uint32, 1](0) + ) compress_internal[1](cv, m^, counter, blen, flags, res.unsafe_ptr()) var final = SIMD[DType.uint32, 16]() @@ -258,11 +260,11 @@ struct Hasher: def __init__(out self): self.key = IV self.original_key = IV - self.cv_stack = {uninitialized=True} + self.cv_stack = {fill=SIMD[DType.uint32, 8](0)} self.stack_len = 0 - self.buf = {uninitialized=True} + self.buf = {fill=0} self.buf_len = 0 self.chunk_counter = 0 @@ -356,7 +358,7 @@ struct Hasher: if out_len < 0: raise Error("BLAKE3 output length must be non-negative") out_buf = {unsafe_uninit_length=out_len} - var temp_buf = StackInlineArray[UInt8, 64](uninitialized=True) + var temp_buf = StackInlineArray[UInt8, 64](fill=0) for i in range(64): temp_buf.unsafe_set(i, 0) @@ -441,7 +443,7 @@ def blake3_parallel_hash(input: Span[UInt8, ...], out_len: Int = 32) raises -> L var task_base = tid * BSIZE var base_ptr = d.unsafe_ptr().unsafe_bitcast[UInt32]() var local_cvs = StackInlineArray[SIMD[DType.uint32, 8], 64]( - uninitialized=True + fill=SIMD[DType.uint32, 8](0) ) for i in range(0, BSIZE, 16): var base = task_base + i @@ -457,10 +459,10 @@ def blake3_parallel_hash(input: Span[UInt8, ...], out_len: Int = 32) raises -> L CHUNK_END if b == 15 else UInt8(0) ) var ma = StackInlineArray[SIMD[DType.uint32, 8], 16]( - uninitialized=True + fill=SIMD[DType.uint32, 8](0) ) var mb = StackInlineArray[SIMD[DType.uint32, 8], 16]( - uninitialized=True + fill=SIMD[DType.uint32, 8](0) ) for j in range(4): @@ -531,7 +533,7 @@ def blake3_parallel_hash(input: Span[UInt8, ...], out_len: Int = 32) raises -> L ] var res = StackInlineArray[SIMD[DType.uint32, 16], 8]( - uninitialized=True + fill=SIMD[DType.uint32, 16](0) ) compress_internal_16way( c, diff --git a/src/thistle/chacha20poly1305.mojo b/src/thistle/chacha20poly1305.mojo index dea0fb4..c687d20 100644 --- a/src/thistle/chacha20poly1305.mojo +++ b/src/thistle/chacha20poly1305.mojo @@ -49,11 +49,13 @@ def _aead_tag( p.update(ciphertext) if len(ciphertext) % 16 != 0: p.update(Span[UInt8, ...](unsafe_ptr=zp, length=16 - len(ciphertext) % 16)) - var lens = InlineArray[UInt8, 16](uninitialized=True) + var lens = InlineArray[UInt8, 16](fill=0) lens.unsafe_ptr().unsafe_bitcast[UInt64]().unsafe_store[alignment=1](0, UInt64(len(aad))) (lens.unsafe_ptr().unsafe_offset(8)).unsafe_bitcast[UInt64]().unsafe_store[alignment=1](0, UInt64(len(ciphertext))) p.update(Span[UInt8, ...](unsafe_ptr=lens.unsafe_ptr(), length=16)) - p.finalize_into(output) + p.finalize_into( + Span[mut=True, UInt8, ...](unsafe_ptr=output, length=16) + ) def _aead_core[encrypt: Bool]( @@ -73,7 +75,7 @@ def _aead_core[encrypt: Bool]( var nonce_span = Span[UInt8, ...](nonce_bytes) var nw = _chacha20_nonce_words(nonce_span) var block0 = chacha20_block_core(kw, 0, nw) - var poly_key = InlineArray[UInt8, 32](uninitialized=True) + var poly_key = InlineArray[UInt8, 32](fill=0) poly_key.unsafe_ptr().unsafe_store[alignment=1]( 0, bitcast[DType.uint8, 64](block0).slice[32]() ) @@ -141,12 +143,12 @@ def chacha20_poly1305_decrypt( var nonce_span = Span[UInt8, ...](nonce_bytes) var nw = _chacha20_nonce_words(nonce_span) var block0 = chacha20_block_core(kw, 0, nw) - var poly_key = InlineArray[UInt8, 32](uninitialized=True) + var poly_key = InlineArray[UInt8, 32](fill=0) poly_key.unsafe_ptr().unsafe_store[alignment=1]( 0, bitcast[DType.uint8, 64](block0).slice[32]() ) - var expected = InlineArray[UInt8, 16](uninitialized=True) + var expected = InlineArray[UInt8, 16](fill=0) _aead_tag( Span[UInt8, ...](unsafe_ptr=poly_key.unsafe_ptr(), length=32), aad, ciphertext, expected.unsafe_ptr(), diff --git a/src/thistle/ed25519.mojo b/src/thistle/ed25519.mojo index febbada..321c57a 100644 --- a/src/thistle/ed25519.mojo +++ b/src/thistle/ed25519.mojo @@ -181,7 +181,7 @@ struct Scalar(Movable, Copyable, ImplicitlyCopyable): @staticmethod def from_bytes_clamped(bytes: Span[UInt8, ...]) -> Scalar: # RFC 8032 5.1.5: prune SHA512(secret)[0..31] into the secret scalar. - var s = InlineArray[UInt8, 32](uninitialized=True) + var s = InlineArray[UInt8, 32](fill=0) for i in range(32): s[i] = bytes[i] s[0] &= 0xF8 @@ -209,13 +209,13 @@ struct Scalar(Movable, Copyable, ImplicitlyCopyable): @staticmethod def _montgomery_mul_raw(a: SIMD[DType.uint64, 8], b: SIMD[DType.uint64, 8]) -> SIMD[DType.uint64, 8]: - var z = InlineArray[UInt128, 9](uninitialized=True) + var z = InlineArray[UInt128, 9](fill=0) for i in range(9): z[i] = 0 for i in range(5): for j in range(5): z[i + j] += UInt128(a[i]) * UInt128(b[j]) var carry: UInt128 = 0 - var n = InlineArray[UInt64, 5](uninitialized=True) + var n = InlineArray[UInt64, 5](fill=0) for i in range(5): var sum = carry + z[i] for j in range(i): @@ -376,7 +376,7 @@ def _edwards_encode_with_zinv(p: EdwardsPoint, z_inv: FieldElement51, output: Po var x = p.X * z_inv var y = p.Y * z_inv y.to_bytes_into(output) - var x_bytes = InlineArray[UInt8, 32](uninitialized=True) + var x_bytes = InlineArray[UInt8, 32](fill=0) x.to_bytes_into(x_bytes.unsafe_ptr()) var x_parity = x_bytes[0] & 1 output[unsafe_offset=31] = output[unsafe_offset=31] | (x_parity << 7) @@ -385,7 +385,7 @@ def _edwards_encode_with_zinv(p: EdwardsPoint, z_inv: FieldElement51, output: Po def edwards_decode(data: Span[UInt8, ...], strict: Bool = True) -> DecodeResult: # RFC 8032 5.1.3: strict point decoding. # Reject y >= p, invalid square roots, and x == 0 with sign bit set. - var y_bytes = InlineArray[UInt8, 32](uninitialized=True) + var y_bytes = InlineArray[UInt8, 32](fill=0) for i in range(32): y_bytes[i] = data[i] var sign = (y_bytes[31] >> 7) & 1 @@ -404,7 +404,7 @@ def edwards_decode(data: Span[UInt8, ...], strict: Bool = True) -> DecodeResult: var x = x_opt.unsafe_value() # x = 0 has no odd/negative alternate root. - var x_zero_bytes = InlineArray[UInt8, 32](uninitialized=True) + var x_zero_bytes = InlineArray[UInt8, 32](fill=0) x.to_bytes_into(x_zero_bytes.unsafe_ptr()) var x_is_zero = True for i in range(32): @@ -414,14 +414,14 @@ def edwards_decode(data: Span[UInt8, ...], strict: Bool = True) -> DecodeResult: return DecodeResult(False, EdwardsPoint()) var x_try = x - var x_try_bytes = InlineArray[UInt8, 32](uninitialized=True) + var x_try_bytes = InlineArray[UInt8, 32](fill=0) x_try.to_bytes_into(x_try_bytes.unsafe_ptr()) if (x_try_bytes[0] & 1) != sign: # Choose the root matching the encoded x parity. x_try = FieldElement51.ZERO() - x_try var chk = x_try.square() * v - u - var chk_bytes = InlineArray[UInt8, 32](uninitialized=True) + var chk_bytes = InlineArray[UInt8, 32](fill=0) chk.to_bytes_into(chk_bytes.unsafe_ptr()) var ok = True for i in range(32): @@ -442,8 +442,8 @@ def _is_small_order(p: EdwardsPoint) -> Bool: var q = _edwards_double_standalone(p) q = _edwards_double_standalone(q) q = _edwards_double_standalone(q) - var x = InlineArray[UInt8, 32](uninitialized=True) - var yz = InlineArray[UInt8, 32](uninitialized=True) + var x = InlineArray[UInt8, 32](fill=0) + var yz = InlineArray[UInt8, 32](fill=0) q.X.to_bytes_into(x.unsafe_ptr()) (q.Y - q.Z).to_bytes_into(yz.unsafe_ptr()) var diff = UInt8(0) @@ -466,8 +466,8 @@ def sqrt_ratio_checked(u: FieldElement51, v: FieldElement51) -> Optional[FieldEl var vx2 = x.square() * v var diff = vx2 - u var diff2 = vx2 + u - var diff_bytes = InlineArray[UInt8, 32](uninitialized=True) - var diff2_bytes = InlineArray[UInt8, 32](uninitialized=True) + var diff_bytes = InlineArray[UInt8, 32](fill=0) + var diff2_bytes = InlineArray[UInt8, 32](fill=0) diff.to_bytes_into(diff_bytes.unsafe_ptr()) diff2.to_bytes_into(diff2_bytes.unsafe_ptr()) var is_zero = True @@ -516,6 +516,13 @@ struct ProjectiveNielsPoint(Movable, Copyable, ImplicitlyCopyable): var Z: FieldElement51 var T2d: FieldElement51 + @always_inline + def __init__(out self): + self.Y_plus_X = FieldElement51() + self.Y_minus_X = FieldElement51() + self.Z = FieldElement51() + self.T2d = FieldElement51() + @always_inline def __init__(out self, Y_plus_X: FieldElement51, Y_minus_X: FieldElement51, Z: FieldElement51, T2d: FieldElement51): self.Y_plus_X = Y_plus_X @@ -594,7 +601,7 @@ def _sub_projective_niels(p: EdwardsPoint, n: ProjectiveNielsPoint) -> EdwardsPo @always_inline def _radix16_digits(scalar: Span[UInt8, ...]) -> InlineArray[Int, 64]: - var digits = InlineArray[Int, 64](uninitialized=True) + var digits = InlineArray[Int, 64](fill=0) for i in range(32): digits[2 * i] = Int(scalar[i] & 15) digits[2 * i + 1] = Int((scalar[i] >> 4) & 15) @@ -650,7 +657,7 @@ def _mul_base_ct(scalar: Span[UInt8, ...]) -> EdwardsPoint: def _naf5(scalar: Span[UInt8, ...]) -> InlineArray[Int, 256]: var naf = InlineArray[Int, 256](fill=0) - var words = InlineArray[UInt64, 5](uninitialized=True) + var words = InlineArray[UInt64, 5](fill=0) words[4] = 0 var ptr = scalar.unsafe_ptr() for w in range(4): @@ -691,7 +698,7 @@ def _double_scalar_mult_vartime(a: Span[UInt8, ...], A: EdwardsPoint, b: Span[UI var naf_b = _naf5(b) var d2 = ed25519_d2() var A2n = _to_projective_niels(_edwards_double_standalone(A), d2) - var Ai = InlineArray[ProjectiveNielsPoint, 8](uninitialized=True) + var Ai = InlineArray[ProjectiveNielsPoint, 8](fill=ProjectiveNielsPoint()) var cur = A Ai[0] = _to_projective_niels(A, d2) for k in range(1, 8): @@ -718,19 +725,24 @@ def _double_scalar_mult_vartime(a: Span[UInt8, ...], A: EdwardsPoint, b: Span[UI return Q @no_inline -def ed25519_generate_public_key(private_key: Span[UInt8, ...], output: Pointer[mut=True, UInt8, _, address_space=_]) raises: +def ed25519_generate_public_key( + private_key: Span[UInt8, ...], output: Span[mut=True, UInt8, ...] +) raises: # RFC 8032 5.1.5: public key A = [pruned SHA512(secret)]B. if len(private_key) != 32: raise Error("Ed25519 private key must be 32 bytes") - var hash = InlineArray[UInt8, 64](uninitialized=True) + if len(output) < 32: + raise Error("Ed25519 public-key output must be at least 32 bytes") + var output_ptr = output.unsafe_ptr() + var hash = InlineArray[UInt8, 64](fill=0) var ctx = SHA512Context() sha512_update(ctx, private_key) sha512_final_to_buffer(ctx, hash.unsafe_ptr()) var s = Scalar.from_bytes_clamped(Span[UInt8, ...](unsafe_ptr=hash.unsafe_ptr(), length=32)) - var s_bytes = InlineArray[UInt8, 32](uninitialized=True) + var s_bytes = InlineArray[UInt8, 32](fill=0) s.to_bytes_into(s_bytes.unsafe_ptr()) var pub_point = _mul_base_ct(Span[UInt8, ...](unsafe_ptr=s_bytes.unsafe_ptr(), length=32)) - edwards_encode_into(pub_point, output) + edwards_encode_into(pub_point, output_ptr) ctx.wipe() s.wipe() var hash_ptr = hash.unsafe_ptr() @@ -741,40 +753,47 @@ def ed25519_generate_public_key(private_key: Span[UInt8, ...], output: Pointer[m s_ptr.unsafe_store[volatile=True](i, UInt8(0)) @no_inline -def ed25519_sign(private_key: Span[UInt8, ...], message: Span[UInt8, ...], output: Pointer[mut=True, UInt8, _, address_space=_]) raises: +def ed25519_sign( + private_key: Span[UInt8, ...], + message: Span[UInt8, ...], + output: Span[mut=True, UInt8, ...], +) raises: # RFC 8032 5.1.6 pure Ed25519: # r = SHA512(prefix || M), R = [r]B, # k = SHA512(R || A || M), S = r + k*s mod L. if len(private_key) != 32: raise Error("Ed25519 private key must be 32 bytes") - var hash = InlineArray[UInt8, 64](uninitialized=True) + if len(output) < 64: + raise Error("Ed25519 signature output must be at least 64 bytes") + var output_ptr = output.unsafe_ptr() + var hash = InlineArray[UInt8, 64](fill=0) var ctx = SHA512Context() sha512_update(ctx, private_key) sha512_final_to_buffer(ctx, hash.unsafe_ptr()) var s_scalar = Scalar.from_bytes_clamped(Span[UInt8, ...](unsafe_ptr=hash.unsafe_ptr(), length=32)) - var s_bytes = InlineArray[UInt8, 32](uninitialized=True) + var s_bytes = InlineArray[UInt8, 32](fill=0) s_scalar.to_bytes_into(s_bytes.unsafe_ptr()) var A_point = _mul_base_ct(Span[UInt8, ...](unsafe_ptr=s_bytes.unsafe_ptr(), length=32)) - var r_hash = InlineArray[UInt8, 64](uninitialized=True) + var r_hash = InlineArray[UInt8, 64](fill=0) var r_ctx = SHA512Context() sha512_update(r_ctx, Span[UInt8, ...](unsafe_ptr=hash.unsafe_ptr().unsafe_offset(32), length=32)) sha512_update(r_ctx, message) sha512_final_to_buffer(r_ctx, r_hash.unsafe_ptr()) var r_scalar = Scalar.from_bytes_wide(Span[UInt8, ...](unsafe_ptr=r_hash.unsafe_ptr(), length=64)) - var r_bytes = InlineArray[UInt8, 32](uninitialized=True) + var r_bytes = InlineArray[UInt8, 32](fill=0) r_scalar.to_bytes_into(r_bytes.unsafe_ptr()) var R_point = _mul_base_ct(Span[UInt8, ...](unsafe_ptr=r_bytes.unsafe_ptr(), length=32)) var zz_inv = (A_point.Z * R_point.Z).invert() - var A_enc = InlineArray[UInt8, 32](uninitialized=True) - var R_enc = InlineArray[UInt8, 32](uninitialized=True) + var A_enc = InlineArray[UInt8, 32](fill=0) + var R_enc = InlineArray[UInt8, 32](fill=0) _edwards_encode_with_zinv(A_point, zz_inv * R_point.Z, A_enc.unsafe_ptr()) _edwards_encode_with_zinv(R_point, zz_inv * A_point.Z, R_enc.unsafe_ptr()) - var k_hash = InlineArray[UInt8, 64](uninitialized=True) + var k_hash = InlineArray[UInt8, 64](fill=0) var k_ctx = SHA512Context() sha512_update(k_ctx, Span[UInt8, ...](unsafe_ptr=R_enc.unsafe_ptr(), length=32)) sha512_update(k_ctx, Span[UInt8, ...](unsafe_ptr=A_enc.unsafe_ptr(), length=32)) @@ -783,11 +802,11 @@ def ed25519_sign(private_key: Span[UInt8, ...], message: Span[UInt8, ...], outpu var k_scalar = Scalar.from_bytes_wide(Span[UInt8, ...](unsafe_ptr=k_hash.unsafe_ptr(), length=64)) var S_scalar = r_scalar + k_scalar * s_scalar - var S_bytes = InlineArray[UInt8, 32](uninitialized=True) + var S_bytes = InlineArray[UInt8, 32](fill=0) S_scalar.to_bytes_into(S_bytes.unsafe_ptr()) - for i in range(32): output[unsafe_offset=i] = R_enc[i] - for i in range(32): output[unsafe_offset=32 + i] = S_bytes[i] + for i in range(32): output_ptr[unsafe_offset=i] = R_enc[i] + for i in range(32): output_ptr[unsafe_offset=32 + i] = S_bytes[i] ctx.wipe() r_ctx.wipe() s_scalar.wipe() @@ -815,20 +834,20 @@ struct Ed25519SigningKey(Copyable, Movable): def __init__(out self, private_key: Span[UInt8, ...]) raises: if len(private_key) != 32: raise Error("Ed25519 private key must be 32 bytes") - var hash = InlineArray[UInt8, 64](uninitialized=True) + var hash = InlineArray[UInt8, 64](fill=0) var ctx = SHA512Context() sha512_update(ctx, private_key) sha512_final_to_buffer(ctx, hash.unsafe_ptr()) self._s = Scalar.from_bytes_clamped(Span[UInt8, ...](unsafe_ptr=hash.unsafe_ptr(), length=32)) - self._prefix = InlineArray[UInt8, 32](uninitialized=True) + self._prefix = InlineArray[UInt8, 32](fill=0) for i in range(32): self._prefix[i] = hash[32 + i] - var s_bytes = InlineArray[UInt8, 32](uninitialized=True) + var s_bytes = InlineArray[UInt8, 32](fill=0) self._s.to_bytes_into(s_bytes.unsafe_ptr()) var A_point = _mul_base_ct(Span[UInt8, ...](unsafe_ptr=s_bytes.unsafe_ptr(), length=32)) - self._a_enc = InlineArray[UInt8, 32](uninitialized=True) + self._a_enc = InlineArray[UInt8, 32](fill=0) edwards_encode_into(A_point, self._a_enc.unsafe_ptr()) ctx.wipe() @@ -845,27 +864,34 @@ struct Ed25519SigningKey(Copyable, Movable): for i in range(32): p.unsafe_store[volatile=True](i, UInt8(0)) - def public_key_into(self, output: Pointer[mut=True, UInt8, _, address_space=_]): + def public_key_into(self, output: Span[mut=True, UInt8, ...]) -> Bool: + if len(output) < 32: + return False for i in range(32): output[i] = self._a_enc[i] + return True @no_inline - def sign(self, message: Span[UInt8, ...], output: Pointer[mut=True, UInt8, _, address_space=_]) raises: - var r_hash = InlineArray[UInt8, 64](uninitialized=True) + def sign( + self, message: Span[UInt8, ...], output: Span[mut=True, UInt8, ...] + ) raises: + if len(output) < 64: + raise Error("Ed25519 signature output must be at least 64 bytes") + var r_hash = InlineArray[UInt8, 64](fill=0) var r_ctx = SHA512Context() sha512_update(r_ctx, Span[UInt8, ...](unsafe_ptr=self._prefix.unsafe_ptr(), length=32)) sha512_update(r_ctx, message) sha512_final_to_buffer(r_ctx, r_hash.unsafe_ptr()) var r_scalar = Scalar.from_bytes_wide(Span[UInt8, ...](unsafe_ptr=r_hash.unsafe_ptr(), length=64)) - var r_bytes = InlineArray[UInt8, 32](uninitialized=True) + var r_bytes = InlineArray[UInt8, 32](fill=0) r_scalar.to_bytes_into(r_bytes.unsafe_ptr()) var R_point = _mul_base_ct(Span[UInt8, ...](unsafe_ptr=r_bytes.unsafe_ptr(), length=32)) - var R_enc = InlineArray[UInt8, 32](uninitialized=True) + var R_enc = InlineArray[UInt8, 32](fill=0) edwards_encode_into(R_point, R_enc.unsafe_ptr()) - var k_hash = InlineArray[UInt8, 64](uninitialized=True) + var k_hash = InlineArray[UInt8, 64](fill=0) var k_ctx = SHA512Context() sha512_update(k_ctx, Span[UInt8, ...](unsafe_ptr=R_enc.unsafe_ptr(), length=32)) sha512_update(k_ctx, Span[UInt8, ...](unsafe_ptr=self._a_enc.unsafe_ptr(), length=32)) @@ -874,7 +900,7 @@ struct Ed25519SigningKey(Copyable, Movable): var k_scalar = Scalar.from_bytes_wide(Span[UInt8, ...](unsafe_ptr=k_hash.unsafe_ptr(), length=64)) var S_scalar = r_scalar + k_scalar * self._s - var S_bytes = InlineArray[UInt8, 32](uninitialized=True) + var S_bytes = InlineArray[UInt8, 32](fill=0) S_scalar.to_bytes_into(S_bytes.unsafe_ptr()) for i in range(32): @@ -909,16 +935,16 @@ def ed25519_verify(public_key: Span[UInt8, ...], message: Span[UInt8, ...], sign if _is_small_order(A): return False - var R_enc = InlineArray[UInt8, 32](uninitialized=True) + var R_enc = InlineArray[UInt8, 32](fill=0) for i in range(32): R_enc[i] = signature[i] - var S_bytes = InlineArray[UInt8, 32](uninitialized=True) + var S_bytes = InlineArray[UInt8, 32](fill=0) for i in range(32): S_bytes[i] = signature[32 + i] var S_bytes_span = Span[UInt8, ...](unsafe_ptr=S_bytes.unsafe_ptr(), length=32) if not _s_lt_l(S_bytes_span): return False - var k_hash = InlineArray[UInt8, 64](uninitialized=True) + var k_hash = InlineArray[UInt8, 64](fill=0) var k_ctx = SHA512Context() sha512_update(k_ctx, Span[UInt8, ...](unsafe_ptr=R_enc.unsafe_ptr(), length=32)) sha512_update(k_ctx, public_key) @@ -928,12 +954,12 @@ def ed25519_verify(public_key: Span[UInt8, ...], message: Span[UInt8, ...], sign var k_hash_span = Span[UInt8, ...](unsafe_ptr=k_hash.unsafe_ptr(), length=64) var k_scalar = Scalar.from_bytes_wide(k_hash_span) - var k_bytes = InlineArray[UInt8, 32](uninitialized=True) + var k_bytes = InlineArray[UInt8, 32](fill=0) k_scalar.to_bytes_into(k_bytes.unsafe_ptr()) var k_bytes_span = Span[UInt8, ...](unsafe_ptr=k_bytes.unsafe_ptr(), length=32) var Q = _double_scalar_mult_vartime(k_bytes_span, edwards_negate(A), S_bytes_span) - var Q_enc = InlineArray[UInt8, 32](uninitialized=True) + var Q_enc = InlineArray[UInt8, 32](fill=0) edwards_encode_into(Q, Q_enc.unsafe_ptr()) for i in range(32): if Q_enc[i] != R_enc[i]: diff --git a/src/thistle/ed25519_table.mojo b/src/thistle/ed25519_table.mojo index 0ebaf69..a2222ba 100644 --- a/src/thistle/ed25519_table.mojo +++ b/src/thistle/ed25519_table.mojo @@ -367,7 +367,7 @@ comptime _ED25519_B_ODD = SIMD[DType.uint64, 128]( @no_inline def ed25519_base_table() -> InlineArray[UInt64, 4096]: - var t = InlineArray[UInt64, 4096](uninitialized=True) + var t = InlineArray[UInt64, 4096](fill=0) var p = t.unsafe_ptr() p.unsafe_store[alignment=8](0, _ED25519_BT0) p.unsafe_store[alignment=8](128, _ED25519_BT1) @@ -405,6 +405,6 @@ def ed25519_base_table() -> InlineArray[UInt64, 4096]: @no_inline def ed25519_b_odd_table() -> InlineArray[UInt64, 128]: - var t = InlineArray[UInt64, 128](uninitialized=True) + var t = InlineArray[UInt64, 128](fill=0) t.unsafe_ptr().unsafe_store[alignment=8](0, _ED25519_B_ODD) return t^ diff --git a/src/thistle/p256.mojo b/src/thistle/p256.mojo index fc8a0f7..d197c61 100644 --- a/src/thistle/p256.mojo +++ b/src/thistle/p256.mojo @@ -661,21 +661,21 @@ def _jacobian_to_affine(p: P256JacobianPoint) -> P256Point: def _scalar_mult(k: U256, p: P256Point) -> P256Point: # Fixed-window scan keeps scalar access constant-time. var pm = P256Point(_to_mont(p.x), _to_mont(p.y), False) - var jac = InlineArray[P256JacobianPoint, 15](uninitialized=True) + var jac = InlineArray[P256JacobianPoint, 15](fill=P256JacobianPoint()) jac[0] = P256JacobianPoint(pm.x, pm.y, _one_mont(), False) jac[1] = _jacobian_double_ct(jac[0]) for i in range(2, 15): jac[i] = _jacobian_add_affine_non_equal_ct(jac[i - 1], pm) # Normalize the table with one inversion. - var prefix = InlineArray[U256, 15](uninitialized=True) + var prefix = InlineArray[U256, 15](fill=U256()) prefix[0] = jac[0].z for i in range(1, 15): prefix[i] = _mont_mul(prefix[i - 1], jac[i].z) var inv_acc = _inv_p(prefix[14]) - var tx = InlineArray[U256, 15](uninitialized=True) - var ty = InlineArray[U256, 15](uninitialized=True) + var tx = InlineArray[U256, 15](fill=U256()) + var ty = InlineArray[U256, 15](fill=U256()) for jj in range(15): var j = 14 - jj var zinv = inv_acc @@ -942,7 +942,7 @@ def _wipe_list_u8(mut data: List[UInt8]): def _rfc6979_p256(private_key: Span[UInt8, ...], digest: Span[UInt8, ...], skip: Int) -> U256: var h1 = _reduce_n(_from_be(digest)) - var h1_bytes = InlineArray[UInt8, 32](uninitialized=True) + var h1_bytes = InlineArray[UInt8, 32](fill=0) _to_be(h1, h1_bytes.unsafe_ptr()) var k = List[UInt8](length=32, fill=0) var v = List[UInt8](length=32, fill=1) diff --git a/src/thistle/p256_table.mojo b/src/thistle/p256_table.mojo index e54f986..5da3149 100644 --- a/src/thistle/p256_table.mojo +++ b/src/thistle/p256_table.mojo @@ -612,7 +612,7 @@ comptime _P256_BT31 = SIMD[DType.uint64, 128]( ) def p256_base_table() -> InlineArray[UInt64, 4096]: - var t = InlineArray[UInt64, 4096](uninitialized=True) + var t = InlineArray[UInt64, 4096](fill=0) var p = t.unsafe_ptr() p.unsafe_store[alignment=8](0, _P256_BT0) p.unsafe_store[alignment=8](128, _P256_BT1) diff --git a/src/thistle/p384.mojo b/src/thistle/p384.mojo index 5d65764..8a5aa63 100644 --- a/src/thistle/p384.mojo +++ b/src/thistle/p384.mojo @@ -604,20 +604,20 @@ def _jacobian_to_affine(p: P384JacobianPoint) -> P384Point: def _scalar_mult(k: U384, p: P384Point) -> P384Point: var pm = P384Point(_to_mont(p.x), _to_mont(p.y), False) - var jac = InlineArray[P384JacobianPoint, 15](uninitialized=True) + var jac = InlineArray[P384JacobianPoint, 15](fill=P384JacobianPoint()) jac[0] = P384JacobianPoint(pm.x, pm.y, _one_mont(), False) jac[1] = _jacobian_double_ct(jac[0]) for i in range(2, 15): jac[i] = _jacobian_add_affine_non_equal_ct(jac[i - 1], pm) - var prefix = InlineArray[U384, 15](uninitialized=True) + var prefix = InlineArray[U384, 15](fill=U384()) prefix[0] = jac[0].z for i in range(1, 15): prefix[i] = _mont_mul(prefix[i - 1], jac[i].z) var inv_acc = _inv_p(prefix[14]) - var tx = InlineArray[U384, 15](uninitialized=True) - var ty = InlineArray[U384, 15](uninitialized=True) + var tx = InlineArray[U384, 15](fill=U384()) + var ty = InlineArray[U384, 15](fill=U384()) for jj in range(15): var j = 14 - jj var zinv = inv_acc @@ -892,7 +892,7 @@ def _wipe_list_u8(mut data: List[UInt8]): def _rfc6979_p384(private_key: Span[UInt8, ...], digest: Span[UInt8, ...], skip: Int) -> U384: var h1 = _reduce_n(_from_be(digest)) - var h1_bytes = InlineArray[UInt8, 48](uninitialized=True) + var h1_bytes = InlineArray[UInt8, 48](fill=0) _to_be(h1, h1_bytes.unsafe_ptr()) var k = List[UInt8](length=48, fill=0) var v = List[UInt8](length=48, fill=1) diff --git a/src/thistle/p384_table.mojo b/src/thistle/p384_table.mojo index 8fa6f67..961dc7c 100644 --- a/src/thistle/p384_table.mojo +++ b/src/thistle/p384_table.mojo @@ -1296,7 +1296,7 @@ comptime _P384_BT67 = SIMD[DType.uint64, 128]( ) def p384_base_table() -> InlineArray[UInt64, 8704]: - var t = InlineArray[UInt64, 8704](uninitialized=True) + var t = InlineArray[UInt64, 8704](fill=0) var p = t.unsafe_ptr() p.unsafe_store[alignment=8](0, _P384_BT0) p.unsafe_store[alignment=8](128, _P384_BT1) diff --git a/src/thistle/poly1305.mojo b/src/thistle/poly1305.mojo index 8a500b8..d8daad4 100644 --- a/src/thistle/poly1305.mojo +++ b/src/thistle/poly1305.mojo @@ -270,7 +270,9 @@ struct Poly1305: self.buf_len += 1 i += 1 - def finalize_into(mut self, output: Pointer[mut=True, UInt8, _, address_space=_]): + def _finalize_into_unchecked( + mut self, output: Pointer[mut=True, UInt8, _, address_space=_] + ): if self.buf_len > 0: self.buf[self.buf_len] = 1 for j in range(self.buf_len + 1, 16): @@ -328,6 +330,13 @@ struct Poly1305: (output.unsafe_offset(8)).unsafe_bitcast[UInt64]().unsafe_store[alignment=1](0, o1) self.wipe() + def finalize_into( + mut self, output: Span[mut=True, UInt8, ...] + ) raises: + if len(output) < 16: + raise Error("Poly1305 output needs at least 16 writable bytes") + self._finalize_into_unchecked(output.unsafe_ptr()) + def wipe(mut self): var p = Pointer(to=self.h0) p.unsafe_store[volatile=True](0, UInt64(0)) @@ -346,4 +355,4 @@ def poly1305_mac( raise Error("Poly1305 output needs at least 16 writable bytes") var p = Poly1305(key) p.update(message) - p.finalize_into(output.unsafe_ptr()) + p.finalize_into(output) diff --git a/src/thistle/rsa.mojo b/src/thistle/rsa.mojo index 19c9923..0d84ca1 100644 --- a/src/thistle/rsa.mojo +++ b/src/thistle/rsa.mojo @@ -48,7 +48,7 @@ def _sha1(data: Span[UInt8, ...]) -> InlineArray[UInt8, 20]: for i in range(7, -1, -1): padded.append(UInt8((bits >> UInt64(8 * i)) & 0xFF)) - var w = InlineArray[UInt32, 80](uninitialized=True) + var w = InlineArray[UInt32, 80](fill=0) var off = 0 while off < len(padded): for t in range(16): @@ -92,8 +92,8 @@ def _sha1(data: Span[UInt8, ...]) -> InlineArray[UInt8, 20]: h4 += e off += 64 - var out = InlineArray[UInt8, 20](uninitialized=True) - var hs = InlineArray[UInt32, 5](uninitialized=True) + var out = InlineArray[UInt8, 20](fill=0) + var hs = InlineArray[UInt32, 5](fill=0) hs[0] = h0 hs[1] = h1 hs[2] = h2 @@ -200,14 +200,14 @@ def _emsa_pss_encode( if em_len < h_len + 2 or len(salt) > em_len - h_len - 2: return False - var m_hash = InlineArray[UInt8, 64](uninitialized=True) + var m_hash = InlineArray[UInt8, 64](fill=0) _ = _hash_into(sha, message, m_hash.unsafe_ptr()) var mprime = InlineArray[UInt8, 534](fill=0) for i in range(h_len): mprime[8 + i] = m_hash[i] for i in range(len(salt)): mprime[8 + h_len + i] = salt[i] - var h = InlineArray[UInt8, 64](uninitialized=True) + var h = InlineArray[UInt8, 64](fill=0) _ = _hash_into( sha, Span[UInt8, ...](unsafe_ptr=mprime.unsafe_ptr(), length=8 + h_len + len(salt)), @@ -220,7 +220,7 @@ def _emsa_pss_encode( db[ps_len] = 1 for i in range(len(salt)): db[ps_len + 1 + i] = salt[i] - var mask = InlineArray[UInt8, 528](uninitialized=True) + var mask = InlineArray[UInt8, 528](fill=0) _mgf1(mgf_sha, h.unsafe_ptr(), h_len, db_len, mask.unsafe_ptr()) for i in range(db_len): output[unsafe_offset=i] = db[i] ^ mask[i] @@ -447,7 +447,7 @@ def _mont_sqr_k[K: Int]( n: StaticTuple[UInt64, _NL], n0: UInt64, ) -> StaticTuple[UInt64, _NL]: - var t = InlineArray[UInt64, 2 * _NL + 2](uninitialized=True) + var t = InlineArray[UInt64, 2 * _NL + 2](fill=0) comptime for z in range(2 * K + 1): t[z] = 0 @@ -678,7 +678,7 @@ struct RsaPublicKey: if em_len < h_len + 2 or salt_len > em_len - h_len - 2: return False - var em = InlineArray[UInt8, 528](uninitialized=True) + var em = InlineArray[UInt8, 528](fill=0) if not self._public_op(signature, em.unsafe_ptr()): return False for i in range(nb - em_len): @@ -696,9 +696,9 @@ struct RsaPublicKey: if (ep[unsafe_offset=0] & ~top_mask) != 0: return False - var db_mask = InlineArray[UInt8, 528](uninitialized=True) + var db_mask = InlineArray[UInt8, 528](fill=0) _mgf1(mgf_sha, h_ptr, h_len, db_len, db_mask.unsafe_ptr()) - var db = InlineArray[UInt8, 528](uninitialized=True) + var db = InlineArray[UInt8, 528](fill=0) for i in range(db_len): db[i] = ep[unsafe_offset=i] ^ db_mask[i] db[0] &= top_mask @@ -710,17 +710,17 @@ struct RsaPublicKey: if db[ps_len] != 0x01: return False - var m_hash = InlineArray[UInt8, 64](uninitialized=True) + var m_hash = InlineArray[UInt8, 64](fill=0) _ = _hash_into(sha, message, m_hash.unsafe_ptr()) - var mprime = InlineArray[UInt8, 534](uninitialized=True) + var mprime = InlineArray[UInt8, 534](fill=0) for i in range(8): mprime[i] = 0 for i in range(h_len): mprime[8 + i] = m_hash[i] for i in range(salt_len): mprime[8 + h_len + i] = db[ps_len + 1 + i] - var h2 = InlineArray[UInt8, 64](uninitialized=True) + var h2 = InlineArray[UInt8, 64](fill=0) _ = _hash_into( sha, Span[UInt8, ...](unsafe_ptr=mprime.unsafe_ptr(), length=8 + h_len + salt_len), @@ -821,7 +821,7 @@ struct RsaPrivateKey: var limb = result[(nb - 1 - i) >> 3] signature[unsafe_offset=i] = UInt8((limb >> UInt64(8 * ((nb - 1 - i) & 7))) & 0xFF) - var recovered = InlineArray[UInt8, 528](uninitialized=True) + var recovered = InlineArray[UInt8, 528](fill=0) var sig_span = Span[UInt8, ...](unsafe_ptr=signature, length=nb) var valid = self.public._public_op(sig_span, recovered.unsafe_ptr()) var diff = UInt8(0) @@ -846,8 +846,10 @@ struct RsaPrivateKey: salt: Span[UInt8, ...], sha: Int, mgf_sha: Int, - signature: Pointer[mut=True, UInt8, _, address_space=_], + signature: Span[mut=True, UInt8, ...], ) raises -> Bool: + if len(signature) < self.public.nb: + return False var em_bits = self.public.mod_bits - 1 var em_len = (em_bits + 7) // 8 var encoded = InlineArray[UInt8, 528](fill=0) @@ -858,7 +860,7 @@ struct RsaPrivateKey: return False var ok = self._private_op( Span[UInt8, ...](unsafe_ptr=encoded.unsafe_ptr(), length=self.public.nb), - signature, + signature.unsafe_ptr(), ) var ep = encoded.unsafe_ptr() for i in range(self.public.nb): @@ -882,7 +884,8 @@ struct RsaPrivateKey: var salt = random_bytes(salt_len) var signature = List[UInt8](unsafe_uninit_length=self.public.nb) var ok = self.pss_sign_with_salt( - message, Span[UInt8, ...](salt), sha, mgf_sha, signature.unsafe_ptr() + message, Span[UInt8, ...](salt), sha, mgf_sha, + Span[mut=True, UInt8, ...](signature), ) var salt_ptr = salt.unsafe_ptr() for i in range(len(salt)): @@ -1125,7 +1128,7 @@ struct RsaCrtPrivateKey: var limb = result[(self.public.nb - 1 - i) >> 3] signature[unsafe_offset=i] = UInt8((limb >> UInt64(8 * ((self.public.nb - 1 - i) & 7))) & 0xFF) - var recovered = InlineArray[UInt8, 528](uninitialized=True) + var recovered = InlineArray[UInt8, 528](fill=0) var valid = self.public._public_op( Span[UInt8, ...](unsafe_ptr=signature, length=self.public.nb), recovered.unsafe_ptr() ) @@ -1153,8 +1156,10 @@ struct RsaCrtPrivateKey: def pss_sign_with_salt( self, message: Span[UInt8, ...], salt: Span[UInt8, ...], sha: Int, mgf_sha: Int, - signature: Pointer[mut=True, UInt8, _, address_space=_], + signature: Span[mut=True, UInt8, ...], ) raises -> Bool: + if len(signature) < self.public.nb: + return False var em_bits = self.public.mod_bits - 1 var em_len = (em_bits + 7) // 8 var encoded = InlineArray[UInt8, 528](fill=0) @@ -1164,7 +1169,8 @@ struct RsaCrtPrivateKey: ): return False var ok = self._private_op( - Span[UInt8, ...](unsafe_ptr=encoded.unsafe_ptr(), length=self.public.nb), signature + Span[UInt8, ...](unsafe_ptr=encoded.unsafe_ptr(), length=self.public.nb), + signature.unsafe_ptr(), ) var ep = encoded.unsafe_ptr() for i in range(self.public.nb): @@ -1188,7 +1194,8 @@ struct RsaCrtPrivateKey: var salt = random_bytes(salt_len) var signature = List[UInt8](unsafe_uninit_length=self.public.nb) var ok = self.pss_sign_with_salt( - message, Span[UInt8, ...](salt), sha, mgf_sha, signature.unsafe_ptr() + message, Span[UInt8, ...](salt), sha, mgf_sha, + Span[mut=True, UInt8, ...](signature), ) var salt_ptr = salt.unsafe_ptr() for i in range(len(salt)): @@ -1209,7 +1216,7 @@ def _pkcs1_v15_verify( if key.nb < t_len + 11: return False - var em = InlineArray[UInt8, 528](uninitialized=True) + var em = InlineArray[UInt8, 528](fill=0) if not key._public_op(signature, em.unsafe_ptr()): return False @@ -1224,7 +1231,7 @@ def _pkcs1_v15_verify( for i in range(prefix_len): diff |= em[3 + ps_len + i] ^ prefix[i] - var digest = InlineArray[UInt8, 64](uninitialized=True) + var digest = InlineArray[UInt8, 64](fill=0) _ = _hash_into(sha, message, digest.unsafe_ptr()) for i in range(h_len): diff |= em[3 + ps_len + prefix_len + i] ^ digest[i] @@ -1260,7 +1267,7 @@ def rsa_pss_sign_with_salt( var key = RsaPrivateKey(modulus, exponent, private_exponent) var signature = List[UInt8](unsafe_uninit_length=key.public.nb) if not key.pss_sign_with_salt( - message, salt, sha, mgf_sha, signature.unsafe_ptr() + message, salt, sha, mgf_sha, Span[mut=True, UInt8, ...](signature) ): raise Error("RSA-PSS signing failed") return signature^ @@ -1291,7 +1298,7 @@ def rsa_pss_crt_sign_with_salt( ) var signature = List[UInt8](unsafe_uninit_length=key.public.nb) if not key.pss_sign_with_salt( - message, salt, sha, mgf_sha, signature.unsafe_ptr() + message, salt, sha, mgf_sha, Span[mut=True, UInt8, ...](signature) ): raise Error("RSA-PSS signing failed") return signature^ diff --git a/src/thistle/sha2.mojo b/src/thistle/sha2.mojo index cd89850..37d819f 100644 --- a/src/thistle/sha2.mojo +++ b/src/thistle/sha2.mojo @@ -212,7 +212,7 @@ def sha256_transform_blocks( for blk in range(nblocks): var block = data.unsafe_offset(blk * 64) - var w = InlineArray[UInt32, 16](uninitialized=True) + var w = InlineArray[UInt32, 16](fill=0) var a = a0 var b = b0 @@ -432,7 +432,7 @@ def sha512_transform_blocks( for blk in range(nblocks): var block = data.unsafe_offset(blk * 128) - var w = InlineArray[UInt64, 16](uninitialized=True) + var w = InlineArray[UInt64, 16](fill=0) var a = a0 var b = b0 diff --git a/src/thistle/sha3.mojo b/src/thistle/sha3.mojo index ac93ef5..650d0fd 100644 --- a/src/thistle/sha3.mojo +++ b/src/thistle/sha3.mojo @@ -5,6 +5,7 @@ FIPS 202 from std.collections import List from std.memory import Pointer, unsafe_memcpy, unsafe_memset_zero +from std.os import abort from .utils import StackBuffer, bytes_to_hex, string_to_bytes from std.bit import rotate_bits_left from std.builtin.simd import SIMD @@ -342,10 +343,10 @@ struct SHA3Context(Movable): var buffer_len: Int def __init__(out self, rate_bits: Int): - debug_assert[assert_mode="safe"]( - 0 < rate_bits <= 1344 and rate_bits % 8 == 0, - "SHA-3 rate must be a positive multiple of 8 no larger than 1344 bits", - ) + if not (0 < rate_bits <= 1344 and rate_bits % 8 == 0): + abort( + "SHA-3 rate must be a positive multiple of 8 no larger than 1344 bits" + ) self.state = StackBuffer[UInt64, 25](fill=0) self.rate_bytes = rate_bits // 8 self.buffer = StackBuffer[UInt8, 168](fill=0) @@ -404,9 +405,8 @@ def sha3_update(mut ctx: SHA3Context, data: Span[UInt8, ...]): def sha3_final(mut ctx: SHA3Context, output_len_bytes: Int) -> List[UInt8]: - debug_assert[assert_mode="safe"]( - output_len_bytes >= 0, "SHA-3 output length cannot be negative" - ) + if output_len_bytes < 0: + abort("SHA-3 output length cannot be negative") ctx.buffer[ctx.buffer_len] = 0x06 ctx.buffer_len += 1 @@ -443,10 +443,8 @@ def sha3_final(mut ctx: SHA3Context, output_len_bytes: Int) -> List[UInt8]: @always_inline def sha3_final_into(mut ctx: SHA3Context, mut output: StackBuffer[UInt8, ...], output_len_bytes: Int): - debug_assert[assert_mode="safe"]( - 0 <= output_len_bytes <= output.capacity(), - "SHA-3 output length exceeds destination capacity", - ) + if output_len_bytes < 0 or output_len_bytes > output.capacity(): + abort("SHA-3 output length exceeds destination capacity") output.clear() ctx.buffer[ctx.buffer_len] = 0x06 ctx.buffer_len += 1 @@ -559,10 +557,8 @@ def shake_finalize(mut ctx: SHA3Context): @always_inline def shake_squeeze_prefix_into(mut ctx: SHA3Context, mut output: StackBuffer[UInt8, ...], output_len: Int): - debug_assert[assert_mode="safe"]( - 0 <= output_len <= output.capacity(), - "SHAKE output length exceeds destination capacity", - ) + if output_len < 0 or output_len > output.capacity(): + abort("SHAKE output length exceeds destination capacity") output.clear() output.set_len_unchecked(output_len) @@ -591,9 +587,8 @@ def shake_advance(mut ctx: SHA3Context): @always_inline def shake_final(mut ctx: SHA3Context, output_len: Int) -> List[UInt8]: - debug_assert[assert_mode="safe"]( - output_len >= 0, "SHAKE output length cannot be negative" - ) + if output_len < 0: + abort("SHAKE output length cannot be negative") shake_finalize(ctx) var output = List[UInt8](capacity=output_len) @@ -621,10 +616,8 @@ def shake_final(mut ctx: SHA3Context, output_len: Int) -> List[UInt8]: @always_inline def shake_final_into(mut ctx: SHA3Context, mut output: StackBuffer[UInt8, ...], output_len: Int): - debug_assert[assert_mode="safe"]( - 0 <= output_len <= output.capacity(), - "SHAKE output length exceeds destination capacity", - ) + if output_len < 0 or output_len > output.capacity(): + abort("SHAKE output length exceeds destination capacity") output.clear() shake_finalize(ctx) diff --git a/src/thistle/sha_ni.mojo b/src/thistle/sha_ni.mojo index aacda35..503edac 100644 --- a/src/thistle/sha_ni.mojo +++ b/src/thistle/sha_ni.mojo @@ -140,7 +140,7 @@ def _sha256ni_transform_arm(state: SIMD[DType.uint32, 8], block: Span[UInt8, ... var old_st0 = st0 var old_st1 = st1 - var w = InlineArray[SIMD128, 4](uninitialized=True) + var w = InlineArray[SIMD128, 4](fill=SIMD128(0)) w[0] = Load(ptr) w[1] = Load(ptr.unsafe_offset(16)) w[2] = Load(ptr.unsafe_offset(32)) @@ -182,7 +182,7 @@ def _sha256ni_transform_x86(state: SIMD[DType.uint32, 8], block: Span[UInt8, ... var w2 = Load(ptr.unsafe_offset(32)) var w3 = Load(ptr.unsafe_offset(48)) - var w = InlineArray[SIMD128, 16](uninitialized=True) + var w = InlineArray[SIMD128, 16](fill=SIMD128(0)) w[0] = w0 w[1] = w1 w[2] = w2 @@ -286,7 +286,7 @@ def sha256ni_transform_blocks( var old_st0 = st0 var old_st1 = st1 - var w = InlineArray[SIMD128, 4](uninitialized=True) + var w = InlineArray[SIMD128, 4](fill=SIMD128(0)) w[0] = Load(ptr) w[1] = Load(ptr.unsafe_offset(16)) w[2] = Load(ptr.unsafe_offset(32)) @@ -408,11 +408,11 @@ def sha512ni_transform_blocks( if blk + 1 < nblocks: prefetch_next_block512(block) - var w = InlineArray[SIMD64x2, 8](uninitialized=True) + var w = InlineArray[SIMD64x2, 8](fill=SIMD64x2(0)) comptime for i in range(8): w[i] = Load512(block.unsafe_offset(i * 16)) - var v = InlineArray[SIMD64x2, 4](uninitialized=True) + var v = InlineArray[SIMD64x2, 4](fill=SIMD64x2(0)) v[0] = ab v[1] = cd v[2] = ef diff --git a/src/thistle/utils.mojo b/src/thistle/utils.mojo index 9672176..c0cc9e8 100644 --- a/src/thistle/utils.mojo +++ b/src/thistle/utils.mojo @@ -1,20 +1,18 @@ from std.bit import byte_swap +from std.os import abort struct StackInlineArray[ElementType: Copyable & Deinitable, size: Int](Copyable): var _data: InlineArray[Self.ElementType, Self.size] @always_inline - def __init__(out self, *, uninitialized: Bool): - self._data = InlineArray[Self.ElementType, Self.size]( - uninitialized=True - ) + def __init__(out self, *, var fill: Self.ElementType): + self._data = InlineArray[Self.ElementType, Self.size](fill=fill^) @always_inline def __init__(out self, var *elems: Self.ElementType, __list_literal__: NoneType): - debug_assert( - len(elems) == Self.size, "No. of elems must match array size" - ) + if len(elems) != Self.size: + abort("StackInlineArray literal length must match its size") self = Self(storage=elems^) @always_inline @@ -27,14 +25,13 @@ struct StackInlineArray[ElementType: Copyable & Deinitable, size: Int](Copyable) elt_is_mutable=True, origin=origin, Self.ElementType, is_owned=True ], ): - debug_assert( - len(storage) == Self.size, - "Expected variadic list of length ", - Self.size, - ", received ", - len(storage), + if len(storage) != Self.size: + abort("StackInlineArray storage length must match its size") + # Each owned variadic element below move-initializes exactly one slot. + # No slot is read, assigned, or deinitialized before that initialization. + self._data = InlineArray[Self.ElementType, Self.size]( + uninitialized=True ) - self = {uninitialized=True} var ptr = self.unsafe_ptr() @@ -64,13 +61,8 @@ struct StackInlineArray[ElementType: Copyable & Deinitable, size: Int](Copyable) @always_inline def unsafe_get[I: Indexer](ref self, idx: I) -> ref[self._data] Self.ElementType: var i = index(idx) - debug_assert( - 0 <= i < Self.size, - " InlineArray.unsafe_get() index out of bounds: ", - i, - " should be greater than or equal to 0 and less than ", - Self.size, - ) + if i < 0 or i >= Self.size: + abort("StackInlineArray index out of bounds") return self._data.unsafe_get(i) @always_inline @@ -91,36 +83,29 @@ struct StackInlineArray[ElementType: Copyable & Deinitable, size: Int](Copyable) @always_inline def __getitem__(ref self, idx: Int) -> ref[self._data] Self.ElementType: - debug_assert( - 0 <= idx < Self.size, - "Index out of bounds: ", idx, " should be in [0, ", Self.size, ")", - ) return self.unsafe_get(idx) @always_inline def unsafe_set(mut self, idx: Int, var value: Self.ElementType): - debug_assert( - 0 <= idx < Self.size, - "The index provided must be within the range [0, len(List) -1] when using List.unsafe_set()", - ) - (self._data.unsafe_ptr().unsafe_offset(idx)).unsafe_deinit_pointee() - (self._data.unsafe_ptr().unsafe_offset(idx)).unsafe_write(value^) + if idx < 0 or idx >= Self.size: + abort("StackInlineArray index out of bounds") + self._data[idx] = value^ -struct StackBuffer[T: Copyable & Deinitable, N: Int](Movable): +struct StackBuffer[T: Copyable & Deinitable & Defaultable, N: Int](Movable): var _data: InlineArray[Self.T, Self.N] var _len: Int @always_inline def __init__(out self): comptime assert Self.T.__del__is_trivial, "StackBuffer requires trivially destructible types (UInt8, UInt32, UInt64, etc)" - self._data = InlineArray[Self.T, Self.N](uninitialized=True) + self._data = InlineArray[Self.T, Self.N](fill=Self.T()) self._len = 0 @always_inline def __init__(out self, *, var fill: Self.T): self._data = InlineArray[Self.T, Self.N](fill=fill^) - self._len = 0 + self._len = Self.N @always_inline def __init__(out self, *, deinit move: Self): @@ -141,9 +126,8 @@ struct StackBuffer[T: Copyable & Deinitable, N: Int](Movable): @always_inline def push(mut self, var val: Self.T): - debug_assert[assert_mode="safe"]( - self._len < Self.N, "StackBuffer overflow" - ) + if self._len >= Self.N: + abort("StackBuffer overflow") self._data[self._len] = val^ self._len += 1 @@ -154,17 +138,15 @@ struct StackBuffer[T: Copyable & Deinitable, N: Int](Movable): @always_inline def pop(mut self) -> Self.T: - debug_assert[assert_mode="safe"]( - self._len > 0, "StackBuffer underflow" - ) + if self._len <= 0: + abort("StackBuffer underflow") self._len -= 1 return self._data[self._len].copy() @always_inline def top(ref self) -> ref[self._data] Self.T: - debug_assert[assert_mode="safe"]( - self._len > 0, "StackBuffer empty" - ) + if self._len <= 0: + abort("StackBuffer is empty") return self._data[self._len - 1] @always_inline @@ -173,10 +155,8 @@ struct StackBuffer[T: Copyable & Deinitable, N: Int](Movable): @always_inline def set_len_unchecked(mut self, new_len: Int): - debug_assert[assert_mode="safe"]( - 0 <= new_len <= Self.N, - "StackBuffer set_len_unchecked out of bounds", - ) + if new_len < 0 or new_len > Self.N: + abort("StackBuffer length out of bounds") self._len = new_len @always_inline @@ -185,16 +165,14 @@ struct StackBuffer[T: Copyable & Deinitable, N: Int](Movable): @always_inline def __getitem__(ref self, i: Int) -> ref[self._data] Self.T: - debug_assert[assert_mode="safe"]( - 0 <= i < Self.N, "StackBuffer index out of bounds" - ) + if i < 0 or i >= self._len: + abort("StackBuffer index out of bounds") return self._data[i] @always_inline def __setitem__(mut self, i: Int, var val: Self.T): - debug_assert[assert_mode="safe"]( - 0 <= i < Self.N, "StackBuffer index out of bounds" - ) + if i < 0 or i >= self._len: + abort("StackBuffer index out of bounds") self._data[i] = val^ @always_inline diff --git a/src/thistle/x25519.mojo b/src/thistle/x25519.mojo index 1ba6611..72743f0 100644 --- a/src/thistle/x25519.mojo +++ b/src/thistle/x25519.mojo @@ -38,7 +38,7 @@ def x25519( raise Error("X25519 point must be 32 bytes") if len(output) < 32: raise Error("X25519 output needs at least 32 writable bytes") - var scalar = StackInlineArray[UInt8, 32](uninitialized=True) + var scalar = StackInlineArray[UInt8, 32](fill=0) for i in range(32): scalar[i] = scalar_in[i] scalar[0] &= 248 @@ -91,7 +91,7 @@ def x25519_public_key( ) raises: if len(private_key) != 32: raise Error("X25519 private key must be 32 bytes") - var base = StackInlineArray[UInt8, 32](uninitialized=True) + var base = StackInlineArray[UInt8, 32](fill=0) for i in range(32): base[i] = 0 base[0] = 9 diff --git a/tests/benchmark.mojo b/tests/benchmark.mojo index 5298bb5..f072693 100644 --- a/tests/benchmark.mojo +++ b/tests/benchmark.mojo @@ -41,9 +41,9 @@ def generate_data(length: Int) -> List[UInt8]: def benchmark_x25519(duration_secs: Float64) raises -> String: - var scalar = InlineArray[UInt8, 32](uninitialized=True) - var point = InlineArray[UInt8, 32](uninitialized=True) - var out = InlineArray[UInt8, 32](uninitialized=True) + var scalar = InlineArray[UInt8, 32](fill=0) + var point = InlineArray[UInt8, 32](fill=0) + var out = InlineArray[UInt8, 32](fill=0) for i in range(32): scalar[i] = UInt8(i + 1) point[i] = UInt8(9) if i == 0 else UInt8(0) @@ -81,8 +81,8 @@ def benchmark_ecdsa(duration_secs: Float64) -> String: var p256_key = InlineArray[UInt8, 32](fill=1) var p384_key = InlineArray[UInt8, 48](fill=1) var message = InlineArray[UInt8, 64](fill=7) - var p256_sig = InlineArray[UInt8, 64](uninitialized=True) - var p384_sig = InlineArray[UInt8, 96](uninitialized=True) + var p256_sig = InlineArray[UInt8, 64](fill=0) + var p384_sig = InlineArray[UInt8, 96](fill=0) var msg = Span[UInt8, ...](message) var p256_count = 0 @@ -116,10 +116,10 @@ def benchmark_ecdsa(duration_secs: Float64) -> String: def benchmark_ed25519(duration_secs: Float64) raises -> String: - var sk = InlineArray[UInt8, 32](uninitialized=True) - var pk = InlineArray[UInt8, 32](uninitialized=True) - var msg = InlineArray[UInt8, 64](uninitialized=True) - var sig = InlineArray[UInt8, 64](uninitialized=True) + var sk = InlineArray[UInt8, 32](fill=0) + var pk = InlineArray[UInt8, 32](fill=0) + var msg = InlineArray[UInt8, 64](fill=0) + var sig = InlineArray[UInt8, 64](fill=0) for i in range(32): sk[i] = UInt8(i * 7 + 1) for i in range(64): @@ -127,19 +127,25 @@ def benchmark_ed25519(duration_secs: Float64) raises -> String: var sk_span = Span[UInt8, ...](unsafe_ptr=sk.unsafe_ptr(), length=32) var msg_span = Span[UInt8, ...](unsafe_ptr=msg.unsafe_ptr(), length=64) var sig_span = Span[UInt8, ...](unsafe_ptr=sig.unsafe_ptr(), length=64) + var sig_out = Span[mut=True, UInt8, ...]( + unsafe_ptr=sig.unsafe_ptr(), length=64 + ) var pk_span = Span[UInt8, ...](unsafe_ptr=pk.unsafe_ptr(), length=32) - ed25519_generate_public_key(sk_span, pk.unsafe_ptr()) - ed25519_sign(sk_span, msg_span, sig.unsafe_ptr()) + var pk_out = Span[mut=True, UInt8, ...]( + unsafe_ptr=pk.unsafe_ptr(), length=32 + ) + ed25519_generate_public_key(sk_span, pk_out) + ed25519_sign(sk_span, msg_span, sig_out) var sign_count = 0 var start = perf_counter() while perf_counter() - start < duration_secs: - ed25519_sign(sk_span, msg_span, sig.unsafe_ptr()) + ed25519_sign(sk_span, msg_span, sig_out) sign_count += 1 var sign_duration = perf_counter() - start var sign_ops = Float64(sign_count) / sign_duration - ed25519_sign(sk_span, msg_span, sig.unsafe_ptr()) + ed25519_sign(sk_span, msg_span, sig_out) var verify_count = 0 var verify_failures = 0 start = perf_counter() diff --git a/tests/dudect.mojo b/tests/dudect.mojo index b54c821..a22c819 100644 --- a/tests/dudect.mojo +++ b/tests/dudect.mojo @@ -433,7 +433,8 @@ def run_ed25519(mut rng: Rng) raises -> Bool: sk[0] = 0x33 var t0 = perf_counter_ns() ed25519_sign( - Span[UInt8, ...](sk), Span[UInt8, ...](msg), sig.unsafe_ptr() + Span[UInt8, ...](sk), Span[UInt8, ...](msg), + Span[mut=True, UInt8, ...](sig), ) times.append(Float64(perf_counter_ns() - t0)) sink ^= sig[0] diff --git a/tests/gen_ed25519_table.mojo b/tests/gen_ed25519_table.mojo index c4098dc..2da77fb 100644 --- a/tests/gen_ed25519_table.mojo +++ b/tests/gen_ed25519_table.mojo @@ -12,7 +12,7 @@ from thistle.curve25519 import FieldElement51 def _canonical(fe: FieldElement51) -> FieldElement51: - var bytes = InlineArray[UInt8, 32](uninitialized=True) + var bytes = InlineArray[UInt8, 32](fill=0) fe.to_bytes_into(bytes.unsafe_ptr()) return fe_from_bytes(Span[UInt8, ...](unsafe_ptr=bytes.unsafe_ptr(), length=32)) @@ -24,7 +24,7 @@ def _affine_niels_limbs(p: EdwardsPoint) -> InlineArray[UInt64, 15]: var y_plus_x = _canonical(y + x) var y_minus_x = _canonical(y - x) var xy2d = _canonical(x * y * ed25519_d2()) - var out = InlineArray[UInt64, 15](uninitialized=True) + var out = InlineArray[UInt64, 15](fill=0) for i in range(5): out[i] = y_plus_x.limbs[i] out[5 + i] = y_minus_x.limbs[i] @@ -59,7 +59,7 @@ def main() raises: var P = B for j in range(32): - var row = InlineArray[EdwardsPoint, 8](uninitialized=True) + var row = InlineArray[EdwardsPoint, 8](fill=EdwardsPoint()) row[0] = P for k in range(1, 8): row[k] = edwards_add(row[k - 1], P) @@ -69,7 +69,7 @@ def main() raises: P = edwards_double(P) var B2 = edwards_double(B) - var odd = InlineArray[EdwardsPoint, 8](uninitialized=True) + var odd = InlineArray[EdwardsPoint, 8](fill=EdwardsPoint()) odd[0] = B for k in range(1, 8): odd[k] = edwards_add(odd[k - 1], B2) @@ -78,7 +78,7 @@ def main() raises: print("@no_inline") print("def ed25519_base_table() -> InlineArray[UInt64, 4096]:") - print(" var t = InlineArray[UInt64, 4096](uninitialized=True)") + print(" var t = InlineArray[UInt64, 4096](fill=0)") print(" var p = t.unsafe_ptr()") for j in range(32): print(" p.unsafe_store[alignment=8](" + String(j * 128) + ", _ED25519_BT" + String(j) + ")") @@ -86,6 +86,6 @@ def main() raises: print() print("@no_inline") print("def ed25519_b_odd_table() -> InlineArray[UInt64, 128]:") - print(" var t = InlineArray[UInt64, 128](uninitialized=True)") + print(" var t = InlineArray[UInt64, 128](fill=0)") print(" t.unsafe_ptr().unsafe_store[alignment=8](0, _ED25519_B_ODD)") print(" return t") diff --git a/tests/gen_p256_table.mojo b/tests/gen_p256_table.mojo index 13bdfb0..3015322 100644 --- a/tests/gen_p256_table.mojo +++ b/tests/gen_p256_table.mojo @@ -36,7 +36,7 @@ def main() raises: print(")") print("") print("def p256_base_table() -> InlineArray[UInt64, 4096]:") - print(" var t = InlineArray[UInt64, 4096](uninitialized=True)") + print(" var t = InlineArray[UInt64, 4096](fill=0)") print(" var p = t.unsafe_ptr()") for c in range(32): print(" p.unsafe_store[alignment=8](" + String(c * 128) + ", _P256_BT" + String(c) + ")") diff --git a/tests/gen_p384_table.mojo b/tests/gen_p384_table.mojo index c499cb5..34033ca 100644 --- a/tests/gen_p384_table.mojo +++ b/tests/gen_p384_table.mojo @@ -36,7 +36,7 @@ def main() raises: print(")") print("") print("def p384_base_table() -> InlineArray[UInt64, 8704]:") - print(" var t = InlineArray[UInt64, 8704](uninitialized=True)") + print(" var t = InlineArray[UInt64, 8704](fill=0)") print(" var p = t.unsafe_ptr()") for c in range(68): print(" p.unsafe_store[alignment=8](" + String(c * 128) + ", _P384_BT" + String(c) + ")") diff --git a/tests/test_security_boundaries.mojo b/tests/test_security_boundaries.mojo index c22a6b6..9488163 100644 --- a/tests/test_security_boundaries.mojo +++ b/tests/test_security_boundaries.mojo @@ -2,11 +2,17 @@ from std.collections import List from thistle.aes import AESExpandedKey from thistle.argon2 import variable_length_hash_into +from thistle.blake2b import Blake2b from thistle.chacha20poly1305 import ( chacha20_poly1305_encrypt, hchacha20, ) from thistle.chacha20 import ChaCha20 +from thistle.ed25519 import ( + Ed25519SigningKey, + ed25519_generate_public_key, + ed25519_sign, +) from thistle.p256 import p256_ecdsa_sign, p256_public_key from thistle.p384 import p384_ecdsa_sign, p384_public_key from thistle.pbkdf2 import ( @@ -15,6 +21,7 @@ from thistle.pbkdf2 import ( pbkdf2_hmac_sha256, pbkdf2_hmac_sha512, ) +from thistle.poly1305 import Poly1305 from thistle.x25519 import x25519 @@ -43,6 +50,28 @@ def main() raises: if not rejected: raise Error("Argon2 accepted an undersized destination") + var blake_output = List[UInt8](length=31, fill=0) + var blake = Blake2b(32) + blake.update(Span[UInt8, ...](empty)) + rejected = False + try: + blake.finalize_into(Span[mut=True, UInt8, ...](blake_output)) + except: + rejected = True + if not rejected: + raise Error("BLAKE2b accepted an undersized destination") + + var poly_key = List[UInt8](length=32, fill=0) + var poly_output = List[UInt8](length=15, fill=0) + var poly = Poly1305(Span[UInt8, ...](poly_key)) + rejected = False + try: + poly.finalize_into(Span[mut=True, UInt8, ...](poly_output)) + except: + rejected = True + if not rejected: + raise Error("Poly1305 accepted an undersized destination") + var key32 = List[UInt8](length=32, fill=1) var point32 = List[UInt8](length=32, fill=0) point32[0] = 9 @@ -107,6 +136,48 @@ def main() raises: if not nonce_affects_stream: raise Error("ChaCha20 ignored the final nonce byte") + var ed_private = List[UInt8](length=32, fill=1) + var ed_public_short = List[UInt8](length=31, fill=0) + rejected = False + try: + ed25519_generate_public_key( + Span[UInt8, ...](ed_private), + Span[mut=True, UInt8, ...](ed_public_short), + ) + except: + rejected = True + if not rejected: + raise Error("Ed25519 accepted an undersized public-key destination") + + var ed_signature_short = List[UInt8](length=63, fill=0) + rejected = False + try: + ed25519_sign( + Span[UInt8, ...](ed_private), + Span[UInt8, ...](empty), + Span[mut=True, UInt8, ...](ed_signature_short), + ) + except: + rejected = True + if not rejected: + raise Error("Ed25519 accepted an undersized signature destination") + + var ed_key = Ed25519SigningKey(Span[UInt8, ...](ed_private)) + if ed_key.public_key_into( + Span[mut=True, UInt8, ...](ed_public_short) + ): + raise Error("Ed25519 key object accepted an undersized public-key destination") + rejected = False + try: + ed_key.sign( + Span[UInt8, ...](empty), + Span[mut=True, UInt8, ...](ed_signature_short), + ) + except: + rejected = True + if not rejected: + raise Error("Ed25519 key object accepted an undersized signature destination") + var empty_salt = List[UInt8]() rejected = False try: diff --git a/tests/test_sha3_capacity_guard.sh b/tests/test_sha3_capacity_guard.sh index cc45839..189ed9d 100644 --- a/tests/test_sha3_capacity_guard.sh +++ b/tests/test_sha3_capacity_guard.sh @@ -3,16 +3,20 @@ set -uo pipefail guard_dir=$(mktemp -d) -guard_binary="$guard_dir/sha3-capacity-guard" -guard_log="$guard_dir/run.log" trap 'rm -rf -- "$guard_dir"' EXIT -mojo build -O3 -I src/ tests/test_sha3_capacity_guard.mojo \ - -o "$guard_binary" >/dev/null 2>&1 +for assert_mode in safe none; do + guard_binary="$guard_dir/sha3-capacity-guard-$assert_mode" + guard_log="$guard_dir/run-$assert_mode.log" -if "$guard_binary" >"$guard_log" 2>&1; then - echo "SHAKE capacity guard did not reject an oversized output" >&2 - exit 1 -fi + mojo build -O3 -D ASSERT="$assert_mode" -I src/ \ + tests/test_sha3_capacity_guard.mojo -o "$guard_binary" \ + >/dev/null 2>&1 -grep -q "SHAKE output length exceeds destination capacity" "$guard_log" + if "$guard_binary" >"$guard_log" 2>&1; then + echo "SHAKE capacity guard failed with ASSERT=$assert_mode" >&2 + exit 1 + fi + + grep -q "SHAKE output length exceeds destination capacity" "$guard_log" +done diff --git a/tests/test_signing.mojo b/tests/test_signing.mojo index 998fe1a..6902520 100644 --- a/tests/test_signing.mojo +++ b/tests/test_signing.mojo @@ -229,6 +229,18 @@ def test_rsa_pss_signing() raises: ) var msg = hex_bytes("74686973746c65") var salt = List[UInt8](length=32, fill=0xA5) + var raw_key = RsaPrivateKey( + Span[UInt8, ...](n), Span[UInt8, ...](e), Span[UInt8, ...](d) + ) + var short_rsa_signature = List[UInt8](length=255, fill=0) + if raw_key.pss_sign_with_salt( + Span[UInt8, ...](msg), + Span[UInt8, ...](salt), + SHA256, + SHA256, + Span[mut=True, UInt8, ...](short_rsa_signature), + ): + raise Error("RSA-PSS accepted an undersized signature destination") var sig = rsa_pss_sign_with_salt( Span[UInt8, ...](n), Span[UInt8, ...](e), @@ -272,13 +284,21 @@ def test_rsa_pss_signing() raises: Span[UInt8, ...](dq), Span[UInt8, ...](qi), ) + if crt.pss_sign_with_salt( + Span[UInt8, ...](msg), + Span[UInt8, ...](salt), + SHA256, + SHA256, + Span[mut=True, UInt8, ...](short_rsa_signature), + ): + raise Error("RSA-PSS CRT accepted an undersized signature destination") var sig2 = List[UInt8](unsafe_uninit_length=256) if not crt.pss_sign_with_salt( Span[UInt8, ...](msg), Span[UInt8, ...](salt), SHA256, SHA256, - sig2.unsafe_ptr(), + Span[mut=True, UInt8, ...](sig2), ): raise Error("crt sign") if not rsa_pss_verify( diff --git a/tests/test_wycheproof_p256_ecdh.mojo b/tests/test_wycheproof_p256_ecdh.mojo index 795c061..b4199de 100644 --- a/tests/test_wycheproof_p256_ecdh.mojo +++ b/tests/test_wycheproof_p256_ecdh.mojo @@ -128,7 +128,7 @@ def run_case( # test-harness-only; the p256_ecdh API accepts raw SEC 1. public_key = extract_trailing_sec1_p256_public_key(public_der) var expected = hex_to_bytes(shared_hex) - var actual = StackInlineArray[UInt8, 32](uninitialized=True) + var actual = StackInlineArray[UInt8, 32](fill=0) var got = p256_ecdh( Span[UInt8, ...](private_key), Span[UInt8, ...](public_key), diff --git a/tests/test_wycheproof_p384_ecdh.mojo b/tests/test_wycheproof_p384_ecdh.mojo index b076856..de91302 100644 --- a/tests/test_wycheproof_p384_ecdh.mojo +++ b/tests/test_wycheproof_p384_ecdh.mojo @@ -121,7 +121,7 @@ def run_case( # test-harness-only; the production p384_ecdh API accepts raw SEC 1. public_key = extract_trailing_sec1_p384_public_key(public_der) var expected = hex_to_bytes(shared_hex) - var actual = StackInlineArray[UInt8, 48](uninitialized=True) + var actual = StackInlineArray[UInt8, 48](fill=0) var got = p384_ecdh( Span[UInt8, ...](private_key), Span[UInt8, ...](public_key), diff --git a/tests/test_wycheproof_x25519.mojo b/tests/test_wycheproof_x25519.mojo index 944ddc2..b34a751 100644 --- a/tests/test_wycheproof_x25519.mojo +++ b/tests/test_wycheproof_x25519.mojo @@ -28,7 +28,7 @@ def run_case(tc_id: String, private_hex: String, public_hex: String, shared_hex: var private_key = hex_to_bytes(private_hex) var public_key = hex_to_bytes(public_hex) var expected = hex_to_bytes(shared_hex) - var actual = StackInlineArray[UInt8, 32](uninitialized=True) + var actual = StackInlineArray[UInt8, 32](fill=0) x25519( Span[UInt8, ...](private_key), Span[UInt8, ...](public_key), From 2a73f8154575228798738d716929d5e9bea51d70 Mon Sep 17 00:00:00 2001 From: libalpm64 <159192189+libalpm64@users.noreply.github.com> Date: Thu, 20 Aug 2026 18:57:53 -0400 Subject: [PATCH 4/4] ~ --- src/thistle/__init__.mojo | 4 +++- src/thistle/pbkdf2.mojo | 4 ++++ src/thistle/poly1305.mojo | 28 ++++++++++++++++++++++++++-- src/thistle/rsa.mojo | 19 ++++++++++++++----- src/thistle/sha3.mojo | 10 +++++----- src/thistle/utils.mojo | 2 +- src/thistle/x25519.mojo | 2 -- tests/gen_ed25519_table.mojo | 4 ++-- tests/test_security_boundaries.mojo | 20 ++++++++++++++++++++ tests/test_sha3_capacity_guard.sh | 23 +++++++++++++++++------ tests/test_sha3_rate_guard.mojo | 6 ++++++ 11 files changed, 98 insertions(+), 24 deletions(-) create mode 100644 tests/test_sha3_rate_guard.mojo diff --git a/src/thistle/__init__.mojo b/src/thistle/__init__.mojo index e6c4527..883952f 100644 --- a/src/thistle/__init__.mojo +++ b/src/thistle/__init__.mojo @@ -14,7 +14,9 @@ from .pbkdf2 import pbkdf2_hmac_sha256, pbkdf2_hmac_sha512 from .argon2 import Argon2id, argon2id_hash_string from .aes import AESKey, AESExpandedKey, expand_key_128, expand_key_192, expand_key_256 -from .aes_ni import has_aes_ni, aes_gcm_encrypt, aes_gcm_decrypt, AESGCMContext +# Compatibility exports for low-level kernels; callers must uphold pointer sizes. +from .aes_ni import has_aes_ni, aes_gcm_ctr_kernel, aes_gcm_encrypt, aes_gcm_decrypt, AESGCMContext +from .aes_gpu import aes_gpu_kernel_ecb, aes_gpu_kernel_ctr, aes_gpu_kernel_gcm_ctr from .camellia import CamelliaCipher from .chacha20 import ChaCha20, chacha20_block from .poly1305 import Poly1305, poly1305_mac diff --git a/src/thistle/pbkdf2.mojo b/src/thistle/pbkdf2.mojo index 5d5f411..e5c6538 100644 --- a/src/thistle/pbkdf2.mojo +++ b/src/thistle/pbkdf2.mojo @@ -111,6 +111,8 @@ struct PBKDF2SHA256(Movable): @always_inline def derive(mut self, salt: Span[UInt8, ...], iterations: Int, dklen: Int) raises -> List[UInt8]: + if iterations < 1: + raise Error("PBKDF2-SHA256 iterations must be positive") if dklen < 1 or dklen > PBKDF2_SHA256_MAX_DKLEN: raise Error("PBKDF2-SHA256 dkLen exceeds the RFC 8018 limit") var hLen = 32 @@ -230,6 +232,8 @@ struct PBKDF2SHA512(Movable): @always_inline def derive(mut self, salt: Span[UInt8, ...], iterations: Int, dklen: Int) raises -> List[UInt8]: + if iterations < 1: + raise Error("PBKDF2-SHA512 iterations must be positive") if dklen < 1 or dklen > PBKDF2_SHA512_MAX_DKLEN: raise Error("PBKDF2-SHA512 dkLen exceeds the RFC 8018 limit") var hLen = 64 diff --git a/src/thistle/poly1305.mojo b/src/thistle/poly1305.mojo index d8daad4..af71689 100644 --- a/src/thistle/poly1305.mojo +++ b/src/thistle/poly1305.mojo @@ -39,6 +39,14 @@ struct _RPower(Movable, Copyable, ImplicitlyCopyable): self.r0 = take.r0; self.r1 = take.r1; self.r2 = take.r2 self.s1 = take.s1; self.s2 = take.s2 + @always_inline + def wipe(mut self): + Pointer(to=self.r0).unsafe_store[volatile=True](0, UInt64(0)) + Pointer(to=self.r1).unsafe_store[volatile=True](0, UInt64(0)) + Pointer(to=self.r2).unsafe_store[volatile=True](0, UInt64(0)) + Pointer(to=self.s1).unsafe_store[volatile=True](0, UInt64(0)) + Pointer(to=self.s2).unsafe_store[volatile=True](0, UInt64(0)) + @always_inline def _mul_acc( @@ -125,6 +133,9 @@ struct Poly1305: self.powers4_ready = False self.powers8_ready = False + def __deinit__(deinit self): + self.wipe() + @no_inline def _make_powers4(mut self): self.r2 = Poly1305._rmul(self.r, self.r) @@ -338,12 +349,25 @@ struct Poly1305: self._finalize_into_unchecked(output.unsafe_ptr()) def wipe(mut self): - var p = Pointer(to=self.h0) - p.unsafe_store[volatile=True](0, UInt64(0)) + self.r.wipe() + self.r2.wipe() + self.r3.wipe() + self.r4.wipe() + self.r5.wipe() + self.r6.wipe() + self.r7.wipe() + self.r8.wipe() + Pointer(to=self.h0).unsafe_store[volatile=True](0, UInt64(0)) Pointer(to=self.h1).unsafe_store[volatile=True](0, UInt64(0)) Pointer(to=self.h2).unsafe_store[volatile=True](0, UInt64(0)) Pointer(to=self.pad0).unsafe_store[volatile=True](0, UInt64(0)) Pointer(to=self.pad1).unsafe_store[volatile=True](0, UInt64(0)) + var buf_ptr = self.buf.unsafe_ptr() + for i in range(16): + buf_ptr.unsafe_store[volatile=True](i, UInt8(0)) + self.buf_len = 0 + self.powers4_ready = False + self.powers8_ready = False def poly1305_mac( diff --git a/src/thistle/rsa.mojo b/src/thistle/rsa.mojo index 0d84ca1..47f56d9 100644 --- a/src/thistle/rsa.mojo +++ b/src/thistle/rsa.mojo @@ -734,11 +734,22 @@ struct RsaPublicKey: def _wipe_bn(mut value: StaticTuple[UInt64, _NL], k: Int): - var ptr = Pointer(to=value[0]).unsafe_mut_cast[True]() + # StaticTuple indexing returns an element by value. Take the address of the + # aggregate itself so the volatile stores target the caller's tuple storage. + var ptr = Pointer(to=value).unsafe_bitcast[UInt64]() for i in range(k): ptr.unsafe_store[volatile=True](i, UInt64(0)) +def _wipe_bn_table( + mut table: StaticTuple[StaticTuple[UInt64, _NL], 16], k: Int +): + var ptr = Pointer(to=table).unsafe_bitcast[UInt64]() + for row in range(16): + for i in range(k): + ptr.unsafe_store[volatile=True](row * _NL + i, UInt64(0)) + + struct RsaPrivateKey: var public: RsaPublicKey var d: InlineArray[UInt8, 528] @@ -836,8 +847,7 @@ struct RsaPrivateKey: _wipe_bn(base, k) _wipe_bn(acc, k) _wipe_bn(result, k) - for i in range(16): - _wipe_bn(table[i], k) + _wipe_bn_table(table, k) return valid def pss_sign_with_salt( @@ -943,8 +953,7 @@ def _private_pow( var result = _mont_mul(acc, one, key.n, key.n0, key.k) _wipe_bn(base, key.k) _wipe_bn(acc, key.k) - for i in range(16): - _wipe_bn(table[i], key.k) + _wipe_bn_table(table, key.k) return result diff --git a/src/thistle/sha3.mojo b/src/thistle/sha3.mojo index 650d0fd..bc973e7 100644 --- a/src/thistle/sha3.mojo +++ b/src/thistle/sha3.mojo @@ -343,9 +343,9 @@ struct SHA3Context(Movable): var buffer_len: Int def __init__(out self, rate_bits: Int): - if not (0 < rate_bits <= 1344 and rate_bits % 8 == 0): + if not (0 < rate_bits <= 1344 and rate_bits % 64 == 0): abort( - "SHA-3 rate must be a positive multiple of 8 no larger than 1344 bits" + "SHA-3 rate must be a positive multiple of 64 no larger than 1344 bits" ) self.state = StackBuffer[UInt64, 25](fill=0) self.rate_bytes = rate_bits // 8 @@ -457,7 +457,7 @@ def sha3_final_into(mut ctx: SHA3Context, mut output: StackBuffer[UInt8, ...], o ctx.buffer[ctx.rate_bytes - 1] |= 0x80 sha3_absorb_block(ctx.state.ptr(), ctx.buffer.ptr(), ctx.rate_bytes) - output.set_len_unchecked(output_len_bytes) + output.set_len(output_len_bytes) var offset = 0 while offset < output_len_bytes: @@ -560,7 +560,7 @@ def shake_squeeze_prefix_into(mut ctx: SHA3Context, mut output: StackBuffer[UInt if output_len < 0 or output_len > output.capacity(): abort("SHAKE output length exceeds destination capacity") output.clear() - output.set_len_unchecked(output_len) + output.set_len(output_len) var offset = 0 while offset < output_len: @@ -621,7 +621,7 @@ def shake_final_into(mut ctx: SHA3Context, mut output: StackBuffer[UInt8, ...], output.clear() shake_finalize(ctx) - output.set_len_unchecked(output_len) + output.set_len(output_len) var offset = 0 while offset < output_len: diff --git a/src/thistle/utils.mojo b/src/thistle/utils.mojo index c0cc9e8..51f2c6d 100644 --- a/src/thistle/utils.mojo +++ b/src/thistle/utils.mojo @@ -154,7 +154,7 @@ struct StackBuffer[T: Copyable & Deinitable & Defaultable, N: Int](Movable): self._len = 0 @always_inline - def set_len_unchecked(mut self, new_len: Int): + def set_len(mut self, new_len: Int): if new_len < 0 or new_len > Self.N: abort("StackBuffer length out of bounds") self._len = new_len diff --git a/src/thistle/x25519.mojo b/src/thistle/x25519.mojo index 72743f0..d58ccc4 100644 --- a/src/thistle/x25519.mojo +++ b/src/thistle/x25519.mojo @@ -92,8 +92,6 @@ def x25519_public_key( if len(private_key) != 32: raise Error("X25519 private key must be 32 bytes") var base = StackInlineArray[UInt8, 32](fill=0) - for i in range(32): - base[i] = 0 base[0] = 9 x25519( private_key, diff --git a/tests/gen_ed25519_table.mojo b/tests/gen_ed25519_table.mojo index 2da77fb..7c8ea9f 100644 --- a/tests/gen_ed25519_table.mojo +++ b/tests/gen_ed25519_table.mojo @@ -82,10 +82,10 @@ def main() raises: print(" var p = t.unsafe_ptr()") for j in range(32): print(" p.unsafe_store[alignment=8](" + String(j * 128) + ", _ED25519_BT" + String(j) + ")") - print(" return t") + print(" return t^") print() print("@no_inline") print("def ed25519_b_odd_table() -> InlineArray[UInt64, 128]:") print(" var t = InlineArray[UInt64, 128](fill=0)") print(" t.unsafe_ptr().unsafe_store[alignment=8](0, _ED25519_B_ODD)") - print(" return t") + print(" return t^") diff --git a/tests/test_security_boundaries.mojo b/tests/test_security_boundaries.mojo index 9488163..05d7624 100644 --- a/tests/test_security_boundaries.mojo +++ b/tests/test_security_boundaries.mojo @@ -18,6 +18,8 @@ from thistle.p384 import p384_ecdsa_sign, p384_public_key from thistle.pbkdf2 import ( PBKDF2_SHA256_MAX_DKLEN, PBKDF2_SHA512_MAX_DKLEN, + PBKDF2SHA256, + PBKDF2SHA512, pbkdf2_hmac_sha256, pbkdf2_hmac_sha512, ) @@ -205,6 +207,24 @@ def main() raises: if not rejected: raise Error("PBKDF2-SHA512 accepted an oversized derived key") + var pbkdf256 = PBKDF2SHA256(Span[UInt8, ...](empty)) + rejected = False + try: + _ = pbkdf256.derive(Span[UInt8, ...](empty_salt), 0, 32) + except: + rejected = True + if not rejected: + raise Error("PBKDF2-SHA256 context accepted zero iterations") + + var pbkdf512 = PBKDF2SHA512(Span[UInt8, ...](empty)) + rejected = False + try: + _ = pbkdf512.derive(Span[UInt8, ...](empty_salt), 0, 64) + except: + rejected = True + if not rejected: + raise Error("PBKDF2-SHA512 context accepted zero iterations") + var nonce = List[UInt8](length=12, fill=0) var plaintext = List[UInt8](length=16, fill=0) var ciphertext = List[UInt8](length=15, fill=0) diff --git a/tests/test_sha3_capacity_guard.sh b/tests/test_sha3_capacity_guard.sh index 189ed9d..fa72cbf 100644 --- a/tests/test_sha3_capacity_guard.sh +++ b/tests/test_sha3_capacity_guard.sh @@ -5,18 +5,29 @@ set -uo pipefail guard_dir=$(mktemp -d) trap 'rm -rf -- "$guard_dir"' EXIT -for assert_mode in safe none; do - guard_binary="$guard_dir/sha3-capacity-guard-$assert_mode" - guard_log="$guard_dir/run-$assert_mode.log" +check_guard() { + local assert_mode="$1" + local source_file="$2" + local guard_name="$3" + local expected_message="$4" + local guard_binary="$guard_dir/$guard_name-$assert_mode" + local guard_log="$guard_dir/$guard_name-$assert_mode.log" mojo build -O3 -D ASSERT="$assert_mode" -I src/ \ - tests/test_sha3_capacity_guard.mojo -o "$guard_binary" \ + "$source_file" -o "$guard_binary" \ >/dev/null 2>&1 if "$guard_binary" >"$guard_log" 2>&1; then - echo "SHAKE capacity guard failed with ASSERT=$assert_mode" >&2 + echo "$guard_name failed with ASSERT=$assert_mode" >&2 exit 1 fi - grep -q "SHAKE output length exceeds destination capacity" "$guard_log" + grep -q "$expected_message" "$guard_log" +} + +for assert_mode in safe none; do + check_guard "$assert_mode" tests/test_sha3_capacity_guard.mojo \ + sha3-capacity-guard "SHAKE output length exceeds destination capacity" + check_guard "$assert_mode" tests/test_sha3_rate_guard.mojo \ + sha3-rate-guard "SHA-3 rate must be a positive multiple of 64" done diff --git a/tests/test_sha3_rate_guard.mojo b/tests/test_sha3_rate_guard.mojo new file mode 100644 index 0000000..0cf3de7 --- /dev/null +++ b/tests/test_sha3_rate_guard.mojo @@ -0,0 +1,6 @@ +from thistle.sha3 import SHA3Context + + +def main(): + var invalid = SHA3Context(1000) + _ = invalid