From c0d9a9dd2e1e5291fa3eb94cbe53e76381c139fa Mon Sep 17 00:00:00 2001 From: libalpm64 <159192189+libalpm64@users.noreply.github.com> Date: Wed, 26 Aug 2026 13:28:41 -0400 Subject: [PATCH] ~ --- .gitignore | 1 + src/thistle/__init__.mojo | 34 ++- src/thistle/aes.mojo | 78 +++++-- src/thistle/aes_gpu.mojo | 23 +- src/thistle/aes_ni.mojo | 194 +++++++++++------ src/thistle/argon2.mojo | 298 ++++++++++++++----------- src/thistle/blake2b.mojo | 21 +- src/thistle/blake3.mojo | 68 +++--- src/thistle/camellia.mojo | 79 ++++--- src/thistle/chacha20.mojo | 115 ++++++---- src/thistle/chacha20poly1305.mojo | 141 +++++++----- src/thistle/curve25519.mojo | 71 ++++-- src/thistle/ecdsa_der.mojo | 8 +- src/thistle/ed25519.mojo | 114 +++++++--- src/thistle/fips.mojo | 8 +- src/thistle/kcipher2.mojo | 21 +- src/thistle/ml_dsa.mojo | 189 +++++++++++++--- src/thistle/ml_kem.mojo | 313 +++++++++++++++------------ src/thistle/p256.mojo | 48 ++-- src/thistle/p384.mojo | 48 ++-- src/thistle/pbkdf2.mojo | 17 +- src/thistle/poly1305.mojo | 42 ++-- src/thistle/random.mojo | 14 +- src/thistle/rsa.mojo | 230 +++++++++++--------- src/thistle/sha2.mojo | 62 +++++- src/thistle/sha3.mojo | 42 ++-- src/thistle/sha_ni.mojo | 69 ++++-- src/thistle/utils.mojo | 16 +- src/thistle/weierstrass.mojo | 29 ++- src/thistle/x25519.mojo | 21 +- tests/benchmark.mojo | 111 +++++++--- tests/dudect.mojo | 15 +- tests/test_aes_gpu.mojo | 18 +- tests/test_ml_dsa.mojo | 10 +- tests/test_ml_kem.mojo | 2 +- tests/test_security_boundaries.mojo | 34 +-- tests/test_signing.mojo | 90 +++++--- tests/test_wycheproof_ed25519.mojo | 4 +- tests/test_wycheproof_p256_ecdh.mojo | 12 +- tests/test_wycheproof_p384_ecdh.mojo | 16 +- tests/test_wycheproof_x25519.mojo | 6 +- tests/thistle_test_vectors.mojo | 88 ++++---- 42 files changed, 1753 insertions(+), 1067 deletions(-) diff --git a/.gitignore b/.gitignore index d279157..531ca29 100644 --- a/.gitignore +++ b/.gitignore @@ -13,3 +13,4 @@ build/ kgen.trace.* tests/.DS_Store .DS_Store +/stdlib diff --git a/src/thistle/__init__.mojo b/src/thistle/__init__.mojo index 169426f..7cf489e 100644 --- a/src/thistle/__init__.mojo +++ b/src/thistle/__init__.mojo @@ -13,35 +13,44 @@ 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, AESExpandedKey, expand_key_128, expand_key_192, expand_key_256 +from .aes import ( + AESKey, AESExpandedKey, expand_key_128, expand_key_192, expand_key_256 +) + # 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 .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 from .chacha20poly1305 import ( chacha20_poly1305_encrypt, chacha20_poly1305_decrypt, xchacha20_poly1305_encrypt, xchacha20_poly1305_decrypt, - hchacha20, + hchacha20 ) from .kcipher2 import KCipher2 from .x25519 import x25519, x25519_checked, x25519_public_key, x25519_keygen -from .ed25519 import ed25519_generate_public_key, ed25519_sign, ed25519_verify, Ed25519SigningKey +from .ed25519 import ( + ed25519_generate_public_key, ed25519_sign, ed25519_verify, Ed25519SigningKey +) from .p256 import ( p256_public_key, p256_ecdh, p256_keygen, p256_ecdsa_sign, p256_ecdsa_sign_digest, p256_ecdsa_verify, p256_ecdsa_verify_digest, p256_ecdsa_sign_der, p256_ecdsa_verify_der, - P256_SIZE, P256_POINT_SIZE, P256_SIGNATURE_SIZE, + P256_SIZE, P256_POINT_SIZE, P256_SIGNATURE_SIZE ) from .p384 import ( p384_public_key, p384_ecdh, p384_keygen, p384_ecdsa_sign, p384_ecdsa_sign_digest, p384_ecdsa_verify, p384_ecdsa_verify_digest, p384_ecdsa_sign_der, p384_ecdsa_verify_der, - P384_SIZE, P384_POINT_SIZE, P384_SIGNATURE_SIZE, + P384_SIZE, P384_POINT_SIZE, P384_SIGNATURE_SIZE ) from .rsa import ( RsaPublicKey, RsaPrivateKey, RsaCrtPrivateKey, @@ -54,7 +63,7 @@ from .rsa import ( rsa_pss_crt_sha512_sign, rsa_pkcs1_v15_verify, rsa_pkcs1_v15_sha1_verify, rsa_pkcs1_v15_sha256_verify, rsa_pkcs1_v15_sha384_verify, - rsa_pkcs1_v15_sha512_verify, + rsa_pkcs1_v15_sha512_verify ) from .ml_kem import mlkem512_keygen, mlkem768_keygen, mlkem1024_keygen @@ -64,15 +73,17 @@ from .ml_kem import ( MLKEM512_PUBLICKEYBYTES, MLKEM512_SECRETKEYBYTES, MLKEM512_CIPHERTEXTBYTES, MLKEM768_PUBLICKEYBYTES, MLKEM768_SECRETKEYBYTES, MLKEM768_CIPHERTEXTBYTES, MLKEM1024_PUBLICKEYBYTES, MLKEM1024_SECRETKEYBYTES, MLKEM1024_CIPHERTEXTBYTES, - MLKEM_BYTES, + MLKEM_BYTES ) from .ml_dsa import mldsa44_keygen, mldsa65_keygen, mldsa87_keygen -from .ml_dsa import mldsa_sign, mldsa_sign_hedged, mldsa_sign_deterministic, mldsa_verify +from .ml_dsa import ( + mldsa_sign, mldsa_sign_hedged, mldsa_sign_deterministic, mldsa_verify +) from .ml_dsa import mldsa44_public_key, mldsa65_public_key, mldsa87_public_key from .ml_dsa import ( MLDSA44_PUBLICKEYBYTES, MLDSA44_SECRETKEYBYTES, MLDSA44_BYTES, MLDSA65_PUBLICKEYBYTES, MLDSA65_SECRETKEYBYTES, MLDSA65_BYTES, - MLDSA87_PUBLICKEYBYTES, MLDSA87_SECRETKEYBYTES, MLDSA87_BYTES, + MLDSA87_PUBLICKEYBYTES, MLDSA87_SECRETKEYBYTES, MLDSA87_BYTES ) from .random import random_bytes, random_fill @@ -80,3 +91,4 @@ from .random import random_bytes, random_fill from . import fips comptime VERSION = "1.0.5" +"""Current Thistle package version.""" diff --git a/src/thistle/aes.mojo b/src/thistle/aes.mojo index 948b7da..4c627b5 100644 --- a/src/thistle/aes.mojo +++ b/src/thistle/aes.mojo @@ -1,45 +1,60 @@ -""" -AES CPU implementation -""" +"""Implements AES encryption and key expansion for CPU callers.""" from std.bit import byte_swap from std.memory import unsafe_memset_zero +from std.os import abort from std.utils import StaticTuple from .utils import StackBuffer, volatile_wipe comptime ROUNDS_128: Int = 10 + +@always_inline +def _validate_aes_rounds(rounds: Int): + if rounds != 10 and rounds != 12 and rounds != 14: + abort("AES rounds must be 10, 12, or 14") + + +@always_inline +def _validate_aes_block_count(num_blocks: Int): + if num_blocks < 0: + abort("AES block count cannot be negative") + + @always_inline def _ct_encrypt1( block: Pointer[mut=True, UInt8, _, address_space=_], skey: List[UInt64], - rounds: Int, + rounds: Int ) -> None: var buf = InlineArray[UInt8, 64](fill=0) var bp = buf.unsafe_ptr() for i in range(16): bp[unsafe_offset=i] = block[unsafe_offset=i] - cpu_aes_ct_encrypt4(bp, skey, rounds) + _ct_encrypt_blocks[1](bp, skey.unsafe_ptr(), rounds) for i in range(16): block[unsafe_offset=i] = bp[unsafe_offset=i] + @always_inline def cpu_aes_encrypt( pt_bytes: Pointer[mut=True, UInt8, _, address_space=_], - round_keys: Pointer[mut=True, UInt32, _, 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: Pointer[mut=True, UInt8, _, address_space=_], round_keys: Pointer[mut=True, UInt32, _, address_space=_], - rounds: Int, + rounds: Int ) -> None: var skey = cpu_aes_ct_skey(round_keys, rounds) _ct_encrypt1(pt_bytes, skey, rounds) volatile_wipe(skey.unsafe_ptr(), len(skey)) + @always_inline def cpu_aes_ecb_kernel( input_ptr: Pointer[mut=True, UInt8, _, address_space=_], @@ -48,6 +63,7 @@ def cpu_aes_ecb_kernel( num_blocks: Int, rounds: Int ) -> None: + _validate_aes_block_count(num_blocks) var skey = cpu_aes_ct_skey(round_keys, rounds) var scratch = InlineArray[UInt8, 256](fill=0) var sp = scratch.unsafe_ptr() @@ -58,13 +74,14 @@ def cpu_aes_ecb_kernel( n = 16 for j in range(n * 16): sp[unsafe_offset=j] = input_ptr[unsafe_offset=i * 16 + j] - cpu_aes_ct_encrypt16(sp, skey, rounds) + _ct_encrypt_blocks[4](sp, skey.unsafe_ptr(), rounds) for j in range(n * 16): output_ptr[unsafe_offset=i * 16 + j] = sp[unsafe_offset=j] i += n volatile_wipe(skey.unsafe_ptr(), len(skey)) volatile_wipe(sp, 256) + @always_inline def cpu_aes_cbc_kernel( input_ptr: Pointer[mut=True, UInt8, _, address_space=_], @@ -74,6 +91,7 @@ def cpu_aes_cbc_kernel( iv_ptr: Pointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: + _validate_aes_block_count(num_blocks) var skey = cpu_aes_ct_skey(round_keys, rounds) var prev_block = StaticTuple[UInt8, 16]( iv_ptr[unsafe_offset=0], iv_ptr[unsafe_offset=1], iv_ptr[unsafe_offset=2], iv_ptr[unsafe_offset=3], @@ -98,11 +116,12 @@ def cpu_aes_cbc_kernel( i += 1 volatile_wipe(skey.unsafe_ptr(), len(skey)) + @always_inline def _ctr_write_block( dst: Pointer[mut=True, UInt8, _, address_space=_], nonce_ptr: Pointer[mut=True, UInt8, _, address_space=_], - offset: Int, + offset: Int ) -> None: for j in range(16): dst.unsafe_store(j, nonce_ptr[unsafe_offset=j]) @@ -124,6 +143,7 @@ def cpu_aes_ctr_kernel( nonce_ptr: Pointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: + _validate_aes_block_count(num_blocks) var skey = cpu_aes_ct_skey(round_keys, rounds) var ks = InlineArray[UInt8, 256](fill=0) var kp = ks.unsafe_ptr() @@ -134,7 +154,7 @@ def cpu_aes_ctr_kernel( n = 16 for k in range(16): _ctr_write_block(kp.unsafe_offset(k * 16), nonce_ptr, i + (k if k < n else 0)) - cpu_aes_ct_encrypt16(kp, skey, rounds) + _ct_encrypt_blocks[4](kp, skey.unsafe_ptr(), rounds) for k in range(n): var in_block = input_ptr.unsafe_offset((i + k) * 16) var out_block = output_ptr.unsafe_offset((i + k) * 16) @@ -144,6 +164,7 @@ def cpu_aes_ctr_kernel( volatile_wipe(skey.unsafe_ptr(), len(skey)) volatile_wipe(kp, 256) + @always_inline def cpu_xts_mul_alpha_inplace(tweak_ptr: Pointer[mut=True, UInt8, _, address_space=_]) -> None: var carry = (tweak_ptr.unsafe_load(15) & 0x80) != 0 @@ -154,6 +175,7 @@ def cpu_xts_mul_alpha_inplace(tweak_ptr: Pointer[mut=True, UInt8, _, address_spa t0 = t0 ^ UInt8(0x87) tweak_ptr.unsafe_store(0, t0) + @always_inline def cpu_aes_xts_kernel( input_ptr: Pointer[mut=True, UInt8, _, address_space=_], @@ -164,6 +186,7 @@ def cpu_aes_xts_kernel( tweak_ptr: Pointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: + _validate_aes_block_count(num_blocks) var skey1 = cpu_aes_ct_skey(round_keys1, rounds) var skey2 = cpu_aes_ct_skey(round_keys2, rounds) var tweak = StackBuffer[UInt8, 16]() @@ -194,6 +217,7 @@ def cpu_aes_xts_kernel( volatile_wipe(skey2.unsafe_ptr(), len(skey2)) volatile_wipe(wp, 16) + @always_inline def sub_word(w: UInt32) -> UInt32: var blk = InlineArray[UInt8, 16](fill=0) @@ -220,9 +244,10 @@ comptime RCON: StaticTuple[UInt8, 11] = StaticTuple[UInt8, 11]( 0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36, 0x6c ) + def expand_key_128_into( key_bytes: Pointer[mut=False, UInt8, _, address_space=_], - w: Pointer[mut=True, UInt32, _, address_space=_], + w: Pointer[mut=True, UInt32, _, address_space=_] ) raises -> None: for i in range(4): var key_val: UInt32 = 0 @@ -237,9 +262,10 @@ def expand_key_128_into( temp ^= UInt32(RCON._unsafe_ref(i // 4 - 1)) << 24 w.unsafe_store(i, w.unsafe_load(i - 4) ^ temp) + def expand_key_192_into( key_bytes: Pointer[mut=False, UInt8, _, address_space=_], - w: Pointer[mut=True, UInt32, _, address_space=_], + w: Pointer[mut=True, UInt32, _, address_space=_] ) raises -> None: for i in range(6): var key_val: UInt32 = 0 @@ -254,9 +280,10 @@ def expand_key_192_into( temp ^= UInt32(RCON._unsafe_ref(i // 6 - 1)) << 24 w.unsafe_store(i, w.unsafe_load(i - 6) ^ temp) + def expand_key_256_into( key_bytes: Pointer[mut=False, UInt8, _, address_space=_], - w: Pointer[mut=True, UInt32, _, address_space=_], + w: Pointer[mut=True, UInt32, _, address_space=_] ) raises -> None: for i in range(8): var key_val: UInt32 = 0 @@ -273,6 +300,7 @@ def expand_key_256_into( temp = sub_word(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.""" @@ -329,6 +357,7 @@ def expand_key_256(key: Span[UInt8, ...]) raises -> AESExpandedKey: raise Error("AES-256 keys must contain exactly 32 bytes") return AESExpandedKey(key) + struct AESKey: var _data: StackBuffer[UInt8, 16] var _round_keys: StackBuffer[UInt32, 44] @@ -364,7 +393,7 @@ struct AESKey: @always_inline def _ct_interleave_in[W: Int]( w0: SIMD[DType.uint64, W], w1: SIMD[DType.uint64, W], - w2: SIMD[DType.uint64, W], w3: SIMD[DType.uint64, W], + w2: SIMD[DType.uint64, W], w3: SIMD[DType.uint64, W] ) -> Tuple[SIMD[DType.uint64, W], SIMD[DType.uint64, W]]: var x0 = w0 var x1 = w1 @@ -394,7 +423,7 @@ def _ct_interleave_out[W: Int]( q0: SIMD[DType.uint64, W], q1: SIMD[DType.uint64, W] ) -> Tuple[ SIMD[DType.uint64, W], SIMD[DType.uint64, W], - SIMD[DType.uint64, W], SIMD[DType.uint64, W], + SIMD[DType.uint64, W], SIMD[DType.uint64, W] ]: var x0 = q0 & 0x00FF00FF00FF00FF var x1 = q1 & 0x00FF00FF00FF00FF @@ -424,7 +453,7 @@ def _ct_swapn[W: Int]( var sv = SIMD[DType.uint64, W](s) return ( (x & cl) | ((y & cl) << sv), - ((x & ~cl) >> sv) | (y & ~cl), + ((x & ~cl) >> sv) | (y & ~cl) ) @@ -645,6 +674,7 @@ def _ct_mix_columns[W: Int](mut q: InlineArray[SIMD[DType.uint64, W], 8]): def cpu_aes_ct_skey( round_keys: Pointer[mut=True, UInt32, _, address_space=_], rounds: Int ) -> List[UInt64]: + _validate_aes_rounds(rounds) 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.unsafe_load(r * 4)))) @@ -684,7 +714,7 @@ def _ct_store_le32(p: Pointer[mut=True, UInt8, _, address_space=_], off: Int, w: def _ct_encrypt_blocks[W: Int]( blocks: Pointer[mut=True, UInt8, _, address_space=_], skp: Pointer[mut=False, UInt64, _, address_space=_], - rounds: Int, + rounds: Int ) -> None: var q = InlineArray[SIMD[DType.uint64, W], 8](fill=0) for i in range(4): @@ -730,30 +760,36 @@ def _ct_encrypt_blocks[W: Int]( def cpu_aes_ct_encrypt4( blocks: Pointer[mut=True, UInt8, _, address_space=_], skey: List[UInt64], - rounds: Int, + rounds: Int ) -> None: + _validate_aes_rounds(rounds) + if len(skey) < (rounds + 1) * 8: + abort("AES bitsliced key schedule is too short") _ct_encrypt_blocks[1](blocks, skey.unsafe_ptr(), rounds) def cpu_aes_ct_encrypt16( blocks: Pointer[mut=True, UInt8, _, address_space=_], skey: List[UInt64], - rounds: Int, + rounds: Int ) -> None: + _validate_aes_rounds(rounds) + if len(skey) < (rounds + 1) * 8: + abort("AES bitsliced key schedule is too short") _ct_encrypt_blocks[4](blocks, skey.unsafe_ptr(), rounds) def cpu_aes_ct_encrypt( pt_bytes: Pointer[mut=True, UInt8, _, address_space=_], round_keys: Pointer[mut=True, UInt32, _, address_space=_], - rounds: Int = 10, + 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[unsafe_offset=i] = pt_bytes[unsafe_offset=i] - cpu_aes_ct_encrypt4(bp, skey, rounds) + _ct_encrypt_blocks[1](bp, skey.unsafe_ptr(), rounds) for i in range(16): pt_bytes[unsafe_offset=i] = bp[unsafe_offset=i] volatile_wipe(skey.unsafe_ptr(), len(skey)) diff --git a/src/thistle/aes_gpu.mojo b/src/thistle/aes_gpu.mojo index 9ec6b8e..196117d 100644 --- a/src/thistle/aes_gpu.mojo +++ b/src/thistle/aes_gpu.mojo @@ -1,6 +1,4 @@ -""" -AES-GPU implementation -""" +"""Provides GPU kernels for AES block and counter-mode operations.""" from std.gpu import global_idx from std.memory import stack_allocation @@ -8,6 +6,7 @@ from std.memory.unsafe_pointer import Pointer from .aes import _ct_encrypt_blocks from .aes_ni import _write_gcm_counter + @always_inline def add_counter_offset(counter: Pointer[mut=True, UInt8, _, address_space=_], offset: Int) -> None: var carry = offset @@ -22,22 +21,26 @@ def add_counter_offset(counter: Pointer[mut=True, UInt8, _, address_space=_], of if new_val < old: carry += 1 + @always_inline def _gcm_counter_from_j0( j0: Pointer[mut=True, UInt8, _, address_space=_], block_index: Int, - counter: Pointer[mut=True, UInt8, _, address_space=_], + counter: Pointer[mut=True, UInt8, _, address_space=_] ) -> None: _write_gcm_counter(counter, j0, block_index) + @always_inline def aes_gpu_kernel_ecb( input_data: Pointer[mut=True, UInt8, MutUntrackedOrigin], output_data: Pointer[mut=True, UInt8, MutUntrackedOrigin], skey: Pointer[mut=True, UInt64, MutUntrackedOrigin], n: Int32, - rounds: Int32, + rounds: Int32 ) -> None: + if n <= 0 or (rounds != 10 and rounds != 12 and rounds != 14): + return var tid = global_idx.x var base_block = Int(tid) * 4 var num_blocks = Int(n) @@ -57,6 +60,7 @@ def aes_gpu_kernel_ecb( for j in range(16): output_data[unsafe_offset=blk * 16 + j] = buf[unsafe_offset=k * 16 + j] + @always_inline def aes_gpu_kernel_ctr( input_data: Pointer[mut=True, UInt8, MutUntrackedOrigin], @@ -64,8 +68,10 @@ def aes_gpu_kernel_ctr( skey: Pointer[mut=True, UInt64, MutUntrackedOrigin], n: Int32, nonce: Pointer[mut=True, UInt8, MutUntrackedOrigin], - rounds: Int32, + rounds: Int32 ) -> None: + if n <= 0 or (rounds != 10 and rounds != 12 and rounds != 14): + return var tid = global_idx.x var base_block = Int(tid) * 4 var num_blocks = Int(n) @@ -89,6 +95,7 @@ def aes_gpu_kernel_ctr( for j in range(16): op[unsafe_offset=j] = bp[unsafe_offset=j] ^ buf[unsafe_offset=k * 16 + j] + @always_inline def aes_gpu_kernel_gcm_ctr( input_data: Pointer[mut=True, UInt8, MutUntrackedOrigin], @@ -96,8 +103,10 @@ def aes_gpu_kernel_gcm_ctr( skey: Pointer[mut=True, UInt64, MutUntrackedOrigin], n: Int32, j0: Pointer[mut=True, UInt8, MutUntrackedOrigin], - rounds: Int32, + rounds: Int32 ) -> None: + if n <= 0 or (rounds != 10 and rounds != 12 and rounds != 14): + return var tid = global_idx.x var base_block = Int(tid) * 4 var num_blocks = Int(n) diff --git a/src/thistle/aes_ni.mojo b/src/thistle/aes_ni.mojo index 8daa1e9..83e5edc 100644 --- a/src/thistle/aes_ni.mojo +++ b/src/thistle/aes_ni.mojo @@ -1,27 +1,35 @@ -""" -AES-NI implementation -""" +"""Provides AES hardware acceleration and AES-GCM contexts.""" 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 .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, + _validate_aes_rounds, + _validate_aes_block_count +) from .utils import StackBuffer, load_64be, store_64be, volatile_wipe comptime SIMD16 = SIMD[DType.uint8, 16] comptime SIMD128 = SIMD[DType.uint64, 2] + @always_inline def has_arm_crypto() -> Bool: - return CompilationTarget.has_neon() and not CompilationTarget.is_x86() and ( - CompilationTarget._has_feature["crypto"]() or CompilationTarget._has_feature["aes"]() + return ( + CompilationTarget.has_neon() and not CompilationTarget.is_x86() and ( + CompilationTarget._has_feature["crypto"]() or CompilationTarget._has_feature["aes"]()) ) + @always_inline def has_x86_aes_ni() -> Bool: - return CompilationTarget.is_x86() and CompilationTarget._has_feature["sse"]() and CompilationTarget._has_feature["aes"]() + return ( + CompilationTarget.is_x86() and CompilationTarget._has_feature["sse"]() and CompilationTarget._has_feature["aes"]() + ) + @always_inline def _aese(lhs: SIMD16, rhs: SIMD16) -> SIMD16: @@ -32,6 +40,7 @@ def _aese(lhs: SIMD16, rhs: SIMD16) -> SIMD16: else: return SIMD16(0) + @always_inline def _aesmc(state: SIMD16) -> SIMD16: comptime if has_arm_crypto(): @@ -41,6 +50,7 @@ def _aesmc(state: SIMD16) -> SIMD16: else: return SIMD16(0) + @always_inline def _mm_aesenc_si128(lhs: SIMD128, rhs: SIMD128) -> SIMD128: comptime if has_x86_aes_ni(): @@ -50,6 +60,7 @@ def _mm_aesenc_si128(lhs: SIMD128, rhs: SIMD128) -> SIMD128: else: return SIMD128(0) + @always_inline def _mm_aesenclast_si128(lhs: SIMD128, rhs: SIMD128) -> SIMD128: comptime if has_x86_aes_ni(): @@ -59,20 +70,23 @@ def _mm_aesenclast_si128(lhs: SIMD128, rhs: SIMD128) -> SIMD128: else: return SIMD128(0) + @always_inline 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: Pointer[mut=True, UInt8, _, address_space=_], data: SIMD128) -> None: var bytes: SIMD[DType.uint8, 16] = bitcast[DType.uint8, 16](data) ptr.unsafe_store[width=16, alignment=1](0, bytes) + @always_inline def _write_gcm_counter( counter_ptr: Pointer[mut=True, UInt8, _, address_space=_], j0_ptr: Pointer[mut=True, UInt8, _, address_space=_], - block_index: Int, + block_index: Int ) -> None: if block_index < 0: abort("GCM counter block_index cannot be negative") @@ -92,6 +106,7 @@ def _write_gcm_counter( 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: Pointer[mut=True, UInt32, _, address_space=_]) -> SIMD128: var w0 = round_keys.unsafe_load(idx * 4) @@ -133,6 +148,7 @@ def x86_aes_encrypt_256( x86_aes_encrypt_256_direct(state, round_keys) _mm_storeu_si128(pt, state) + @always_inline def x86_aes_encrypt_128_direct( mut state: SIMD128, @@ -147,6 +163,7 @@ def x86_aes_encrypt_128_direct( state = _mm_aesenc_si128(state, keys[i]) state = _mm_aesenclast_si128(state, keys[10]) + @always_inline def x86_aes_encrypt_192_direct( mut state: SIMD128, @@ -161,6 +178,7 @@ def x86_aes_encrypt_192_direct( state = _mm_aesenc_si128(state, keys[i]) state = _mm_aesenclast_si128(state, keys[12]) + @always_inline def x86_aes_encrypt_256_direct( mut state: SIMD128, @@ -182,9 +200,11 @@ def _arm_load_keys[N: Int]( ) -> InlineArray[SIMD16, N]: 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[ + var raw = ( + (round_keys.unsafe_offset(i * 4)).unsafe_bitcast[UInt8]().unsafe_load[ width=16, alignment=1 ]() + ) keys[i] = raw.shuffle[ 3, 2, 1, 0, 7, 6, 5, 4, 11, 10, 9, 8, 15, 14, 13, 12 ]() @@ -228,17 +248,18 @@ def arm_aes_encrypt_256( 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: 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, + 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 p = input_ptr.unsafe_offset(i * 16) 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) @@ -252,15 +273,15 @@ def _arm_ecb_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 q = output_ptr + i * 16 + var q = output_ptr.unsafe_offset(i * 16) 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).unsafe_load[width=16, alignment=1](0) - (output_ptr + i * 16).unsafe_store[alignment=1](0, _arm_enc_block[NR](x, keys)) + var x = input_ptr.unsafe_offset(i * 16).unsafe_load[width=16, alignment=1](0) + output_ptr.unsafe_offset(i * 16).unsafe_store[alignment=1](0, _arm_enc_block[NR](x, keys)) i += 1 @@ -272,6 +293,8 @@ def arm_aes_ecb_kernel( num_blocks: Int, rounds: Int ) -> None: + _validate_aes_rounds(rounds) + _validate_aes_block_count(num_blocks) if rounds == 10: _arm_ecb_loop[10](input_ptr, output_ptr, round_keys, num_blocks) elif rounds == 12: @@ -286,7 +309,7 @@ def _arm_cbc_loop[NR: Int]( output_ptr: Pointer[mut=True, UInt8, _, address_space=_], round_keys: Pointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - iv_ptr: Pointer[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.unsafe_load[width=16, alignment=1](0) @@ -297,8 +320,8 @@ def _arm_cbc_loop[NR: Int]( var x = src.unsafe_load[width=16, alignment=1](0) ^ prev prev = _arm_enc_block[NR](x, keys) dst.unsafe_store[alignment=1](0, prev) - src += 16 - dst += 16 + src = src.unsafe_offset(16) + dst = dst.unsafe_offset(16) i += 1 @@ -311,6 +334,8 @@ def arm_aes_cbc_kernel( iv_ptr: Pointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: + _validate_aes_rounds(rounds) + _validate_aes_block_count(num_blocks) if rounds == 10: _arm_cbc_loop[10](input_ptr, output_ptr, round_keys, num_blocks, iv_ptr) elif rounds == 12: @@ -318,6 +343,7 @@ def arm_aes_cbc_kernel( else: _arm_cbc_loop[14](input_ptr, output_ptr, round_keys, num_blocks, iv_ptr) + @always_inline def arm_aes_xts_kernel( input_ptr: Pointer[mut=True, UInt8, _, address_space=_], @@ -328,6 +354,8 @@ def arm_aes_xts_kernel( tweak_ptr: Pointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: + _validate_aes_rounds(rounds) + _validate_aes_block_count(num_blocks) 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)) @@ -369,6 +397,7 @@ def arm_aes_xts_kernel( ) i += 1 + @always_inline def x86_aes_ecb_kernel( input_ptr: Pointer[mut=True, UInt8, _, address_space=_], @@ -377,9 +406,11 @@ def x86_aes_ecb_kernel( num_blocks: Int, rounds: Int ) -> None: + _validate_aes_rounds(rounds) + _validate_aes_block_count(num_blocks) var i = 0 while i < num_blocks: - var block = _mm_loadu_si128(input_ptr + i * 16) + var block = _mm_loadu_si128(input_ptr.unsafe_offset(i * 16)) if rounds == 10: x86_aes_encrypt_128_direct(block, round_keys) @@ -388,9 +419,10 @@ def x86_aes_ecb_kernel( else: x86_aes_encrypt_256_direct(block, round_keys) - _mm_storeu_si128(output_ptr + i * 16, block) + _mm_storeu_si128(output_ptr.unsafe_offset(i * 16), block) i += 1 + @always_inline def x86_aes_cbc_kernel( input_ptr: Pointer[mut=True, UInt8, _, address_space=_], @@ -400,11 +432,13 @@ def x86_aes_cbc_kernel( iv_ptr: Pointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: + _validate_aes_rounds(rounds) + _validate_aes_block_count(num_blocks) var prev_block = _mm_loadu_si128(iv_ptr) var i = 0 while i < num_blocks: - var block = _mm_loadu_si128(input_ptr + i * 16) + var block = _mm_loadu_si128(input_ptr.unsafe_offset(i * 16)) block = block ^ prev_block if rounds == 10: @@ -414,10 +448,11 @@ def x86_aes_cbc_kernel( else: x86_aes_encrypt_256_direct(block, round_keys) - _mm_storeu_si128(output_ptr + i * 16, block) + _mm_storeu_si128(output_ptr.unsafe_offset(i * 16), block) prev_block = block i += 1 + @always_inline def _gf_mul2_xts_simd(val: SIMD128) -> SIMD128: var carry_lo_to_hi = val[0] >> 63 @@ -426,6 +461,7 @@ def _gf_mul2_xts_simd(val: SIMD128) -> SIMD128: var shifted_hi = (val[1] << 1) | carry_lo_to_hi return SIMD128(shifted_lo ^ (msb * UInt64(0x87)), shifted_hi) + @always_inline def x86_aes_xts_kernel( input_ptr: Pointer[mut=True, UInt8, _, address_space=_], @@ -436,6 +472,8 @@ def x86_aes_xts_kernel( tweak_ptr: Pointer[mut=True, UInt8, _, address_space=_], rounds: Int ) -> None: + _validate_aes_rounds(rounds) + _validate_aes_block_count(num_blocks) var tweak = _mm_loadu_si128(tweak_ptr) if rounds == 10: @@ -447,7 +485,7 @@ def x86_aes_xts_kernel( var i = 0 while i < num_blocks: - var in_block = _mm_loadu_si128(input_ptr + i * 16) + var in_block = _mm_loadu_si128(input_ptr.unsafe_offset(i * 16)) var xored = in_block ^ tweak if rounds == 10: @@ -458,11 +496,12 @@ def x86_aes_xts_kernel( x86_aes_encrypt_256_direct(xored, round_keys1) var result = xored ^ tweak - _mm_storeu_si128(output_ptr + i * 16, result) + _mm_storeu_si128(output_ptr.unsafe_offset(i * 16), result) tweak = _gf_mul2_xts_simd(tweak) i += 1 + @always_inline def aes_encrypt( pt: Pointer[mut=True, UInt8, _, address_space=_], @@ -494,6 +533,7 @@ def aes_encrypt( else: cpu_aes_encrypt(pt, round_keys, rounds) + @always_inline def aes_gcm_ctr_kernel( input_ptr: Pointer[mut=True, UInt8, _, address_space=_], @@ -564,7 +604,7 @@ def _arm_gcm_ctr_loop[NR: Int]( output_ptr: Pointer[mut=True, UInt8, _, address_space=_], round_keys: Pointer[mut=True, UInt32, _, address_space=_], num_blocks: Int, - j0_ptr: Pointer[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]() @@ -612,7 +652,7 @@ def _arm_gcm_fused_loop[NR: Int]( num_blocks: Int, j0_ptr: Pointer[mut=True, UInt8, _, address_space=_], mut gh: _GHash, - ghash_ciphertext: Bool, + ghash_ciphertext: Bool ) -> None: var keys = _arm_load_keys[NR + 1](round_keys) var y = SIMD128(_bitrev64(gh.y_hi), _bitrev64(gh.y_lo)) @@ -735,6 +775,7 @@ def _soft_gcm_ctr_kernel( volatile_wipe(skey.unsafe_ptr(), len(skey)) volatile_wipe(kp, 256) + @always_inline def has_aes_ni() -> Bool: return has_x86_aes_ni() or has_arm_crypto() @@ -744,6 +785,7 @@ def has_aes_ni() -> Bool: comptime _GCM_MAX_INPUT_BYTES = 68719476704 + @always_inline("nodebug") def _bitrev64(v: UInt64) -> UInt64: return llvm_intrinsic["llvm.bitreverse.i64", UInt64, has_side_effect=False](v) @@ -876,7 +918,8 @@ struct _GHash(Copyable, Movable): self.y_lo = z_lo @always_inline - def update(mut self, data: Pointer[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 @@ -884,7 +927,8 @@ struct _GHash(Copyable, Movable): self._update_soft(data, length) @always_inline - def _update_pmull(mut self, data: Pointer[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 @@ -912,7 +956,8 @@ struct _GHash(Copyable, Movable): self.y_lo = _bitrev64(y[1]) @always_inline - def _update_soft(mut self, data: Pointer[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) @@ -978,7 +1023,7 @@ def _gctr_and_ghash( output_ptr: Pointer[mut=True, UInt8, _, address_space=_], length: Int, mut gh: _GHash, - ghash_ciphertext: Bool, + ghash_ciphertext: Bool ): var j0_buf = InlineArray[UInt8, 16](fill=0) for i in range(16): @@ -1037,7 +1082,7 @@ def _gcm_core_keyed( output_ptr: Pointer[mut=True, UInt8, _, address_space=_], length: Int, mut tag: InlineArray[UInt8, 16], - ghash_ciphertext: Bool, + ghash_ciphertext: Bool ) raises: var j0 = InlineArray[UInt8, 16](fill=0) _derive_j0(gh.h_hi, gh.h_lo, iv, j0) @@ -1100,27 +1145,39 @@ struct AESGCMContext(Copyable, Movable): if n > _GCM_MAX_INPUT_BYTES: raise Error("plaintext too long for AES-GCM") var ciphertext = List[UInt8](unsafe_uninit_length=n) - var pt_ptr = plaintext.unsafe_ptr().unsafe_mut_cast[True]().unsafe_origin_cast[MutAnyOrigin]() + var pt_ptr = ( + plaintext.unsafe_ptr().unsafe_mut_cast[True]().unsafe_origin_cast[MutAnyOrigin]() + ) var tag = InlineArray[UInt8, 16](fill=0) var rk = self._rk.copy() var gh = self._gh0.copy() - _gcm_core_keyed( - rk.unsafe_ptr(), self._rounds, gh, iv, aad, - pt_ptr, ciphertext.unsafe_ptr(), n, tag, - ghash_ciphertext=True, - ) + try: + _gcm_core_keyed( + rk.unsafe_ptr(), + self._rounds, + gh, + iv, + aad, + pt_ptr, + ciphertext.unsafe_ptr(), + n, + tag, + ghash_ciphertext=True, + ) - var tag_out = List[UInt8](capacity=16) - for i in range(16): - tag_out.append(tag[i]) - volatile_wipe(rk.unsafe_ptr(), 60) - volatile_wipe(Pointer(to=gh).unsafe_bitcast[UInt64](), 20) - return (ciphertext^, tag_out^) + var tag_out = List[UInt8](capacity=16) + for i in range(16): + tag_out.append(tag[i]) + return (ciphertext^, tag_out^) + finally: + volatile_wipe(rk.unsafe_ptr(), 60) + volatile_wipe(Pointer(to=gh).unsafe_bitcast[UInt64](), 20) + volatile_wipe(tag.unsafe_ptr(), 16) def decrypt( self, iv: Span[UInt8, ...], ciphertext: Span[UInt8, ...], - aad: Span[UInt8, ...], tag: Span[UInt8, ...], + aad: Span[UInt8, ...], tag: Span[UInt8, ...] ) raises -> Tuple[List[UInt8], Bool]: if len(iv) == 0: raise Error("invalid iv size") @@ -1130,32 +1187,41 @@ struct AESGCMContext(Copyable, Movable): if n > _GCM_MAX_INPUT_BYTES: raise Error("ciphertext too long for AES-GCM") var plaintext = List[UInt8](unsafe_uninit_length=n) - var ct_ptr = ciphertext.unsafe_ptr().unsafe_mut_cast[True]().unsafe_origin_cast[MutAnyOrigin]() + var ct_ptr = ( + ciphertext.unsafe_ptr().unsafe_mut_cast[True]().unsafe_origin_cast[MutAnyOrigin]() + ) var computed_tag = InlineArray[UInt8, 16](fill=0) var rk = self._rk.copy() var gh = self._gh0.copy() - _gcm_core_keyed( - rk.unsafe_ptr(), self._rounds, gh, iv, aad, - ct_ptr, plaintext.unsafe_ptr(), n, computed_tag, - ghash_ciphertext=False, - ) - - var diff = UInt8(0) - for i in range(16): - diff |= computed_tag[i] ^ tag[i] + try: + _gcm_core_keyed( + rk.unsafe_ptr(), + self._rounds, + gh, + iv, + aad, + ct_ptr, + plaintext.unsafe_ptr(), + n, + computed_tag, + ghash_ciphertext=False, + ) - if diff != 0: - var pt_ptr = plaintext.unsafe_ptr() - for i in range(n): - pt_ptr.unsafe_store[volatile=True](i, UInt8(0)) + var diff = UInt8(0) + for i in range(16): + diff |= computed_tag[i] ^ tag[i] + + if diff != 0: + var pt_ptr = plaintext.unsafe_ptr() + for i in range(n): + pt_ptr.unsafe_store[volatile=True](i, UInt8(0)) + return (List[UInt8](), False) + return (plaintext^, True) + finally: volatile_wipe(rk.unsafe_ptr(), 60) volatile_wipe(Pointer(to=gh).unsafe_bitcast[UInt64](), 20) - return (List[UInt8](), False) - - volatile_wipe(rk.unsafe_ptr(), 60) - volatile_wipe(Pointer(to=gh).unsafe_bitcast[UInt64](), 20) - return (plaintext^, True) + volatile_wipe(computed_tag.unsafe_ptr(), 16) def _valid_gcm_key(key: Span[UInt8, ...]) -> Bool: @@ -1164,7 +1230,7 @@ def _valid_gcm_key(key: Span[UInt8, ...]) -> Bool: def aes_gcm_encrypt( key: Span[UInt8, ...], iv: Span[UInt8, ...], - plaintext: Span[UInt8, ...], aad: Span[UInt8, ...], + plaintext: Span[UInt8, ...], aad: Span[UInt8, ...] ) raises -> Tuple[List[UInt8], List[UInt8]]: var ctx = AESGCMContext(key) return ctx.encrypt(iv, plaintext, aad) @@ -1172,7 +1238,7 @@ def aes_gcm_encrypt( def aes_gcm_decrypt( key: Span[UInt8, ...], iv: Span[UInt8, ...], - ciphertext: Span[UInt8, ...], aad: Span[UInt8, ...], tag: Span[UInt8, ...], + ciphertext: Span[UInt8, ...], aad: Span[UInt8, ...], tag: Span[UInt8, ...] ) raises -> Tuple[List[UInt8], Bool]: var ctx = AESGCMContext(key) return ctx.decrypt(iv, ciphertext, aad, tag) diff --git a/src/thistle/argon2.mojo b/src/thistle/argon2.mojo index 1cc6916..9f58fb1 100644 --- a/src/thistle/argon2.mojo +++ b/src/thistle/argon2.mojo @@ -1,16 +1,15 @@ -""" -Argon2id/Argon2d Implementation in Mojo -RFC 9106 -""" +"""Implements Argon2id and Argon2d as specified by RFC 9106.""" from std.collections import List 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 +from .utils import StackBuffer comptime MASK32 = 0xFFFFFFFF + @always_inline def zero_buffer(ptr: Pointer[mut=True, UInt8, _, address_space=_], len: Int): var i = 0 @@ -23,6 +22,7 @@ def zero_buffer(ptr: Pointer[mut=True, UInt8, _, address_space=_], len: Int): ptr.unsafe_store[volatile=True](i, UInt8(0)) i += 1 + @always_inline def zero_buffer_u64(ptr: Pointer[mut=True, UInt64, _, address_space=_], len: Int): var i = 0 @@ -35,20 +35,24 @@ def zero_buffer_u64(ptr: Pointer[mut=True, UInt64, _, address_space=_], len: Int ptr.unsafe_store[volatile=True](i, UInt64(0)) i += 1 + @always_inline def zero_and_free(ptr: Pointer[mut=True, UInt8, _, address_space=_], len: Int): zero_buffer(ptr, len) ptr.unsafe_free() + @always_inline def zero_and_free_u64(ptr: Pointer[mut=True, UInt64, _, address_space=_], len: Int): zero_buffer_u64(ptr, len) ptr.unsafe_free() + @always_inline def f_bla_mka(x: UInt64, y: UInt64) -> UInt64: return x + y + (((x & MASK32) * (y & MASK32)) << UInt64(1)) + @always_inline def gb(a: UInt64, b: UInt64, c: UInt64, d: UInt64) -> Tuple[UInt64, UInt64, UInt64, UInt64]: var a_new = f_bla_mka(a, b) @@ -61,6 +65,7 @@ def gb(a: UInt64, b: UInt64, c: UInt64, d: UInt64) -> Tuple[UInt64, UInt64, UInt b_new = rotate_bits_left[shift=1](b_new ^ c_new) return (a_new, b_new, c_new, d_new) + @always_inline 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]) @@ -84,6 +89,7 @@ def _p_column(base: Int, v: Pointer[mut=True, UInt64, _, address_space=_]): v[unsafe_offset=base + 11] = v11 v[unsafe_offset=base + 15] = v15 + @always_inline 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]) @@ -107,6 +113,7 @@ def _p_diagonal(base: Int, v: Pointer[mut=True, UInt64, _, address_space=_]): v[unsafe_offset=base + 9] = v9 v[unsafe_offset=base + 14] = v14 + struct MemoryPool: var block_buffer: Pointer[UInt64, MutUntrackedOrigin] var temp_buffer: Pointer[UInt64, MutUntrackedOrigin] @@ -129,13 +136,14 @@ struct MemoryPool: def get_temp(self) -> Pointer[UInt64, MutUntrackedOrigin]: return self.temp_buffer + @always_inline def compression_g_with_pool( 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, + pool: MemoryPool ): var block = pool.get_block() var block_xy = pool.get_temp() @@ -201,6 +209,7 @@ def compression_g_with_pool( for i in range(128): out_ptr[unsafe_offset=i] = block[unsafe_offset=i] ^ block_xy[unsafe_offset=i] + @always_inline def store_le32(ptr: Pointer[mut=True, UInt8, _, address_space=_], offset: Int, val: Int): ptr[unsafe_offset=offset + 0] = UInt8(val & 0xFF) @@ -216,8 +225,8 @@ def variable_length_hash_into( 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") + if t_len >= (1 << 32): + raise Error("Argon2 variable-length hash output must fit in 32 bits") var out_ptr = output.unsafe_ptr() @@ -273,15 +282,18 @@ def variable_length_hash_into( 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") + if t_len >= (1 << 32): + raise Error("Argon2 variable-length hash output must fit in 32 bits") var out_buf = alloc(Layout[UInt8](count=t_len)).unsafe_leak() try: variable_length_hash_into( t_len, input, - Span[mut=True, UInt8, ...](unsafe_ptr=out_buf, length=t_len), + Span[mut=True, UInt8, ...](unsafe_ptr=out_buf, length=t_len) ) var result = List[UInt8](capacity=t_len) for i in range(t_len): @@ -290,6 +302,7 @@ def variable_length_hash(t_len: Int, input: Span[UInt8, ...]) raises -> List[UIn finally: zero_and_free(out_buf, t_len) + @always_inline def _argon2_process_lane( memory: Pointer[mut=True, UInt64, _, address_space=_], @@ -303,7 +316,7 @@ def _argon2_process_lane( m_prime_blocks: Int, iterations: Int, type_code: Int, - parallelism: Int, + parallelism: Int ): var addressing_block = alloc(Layout[UInt64](count=128)).unsafe_leak() var z_u64 = alloc(Layout[UInt64](count=128)).unsafe_leak() @@ -364,11 +377,9 @@ def _argon2_process_lane( window_size = slice_idx * segment_length else: if ref_lane == lane: - window_size = ( - q + window_size = q - segment_length + (index % segment_length) - ) else: window_size = q - segment_length @@ -400,7 +411,7 @@ def _argon2_process_lane( p_ptr.unsafe_origin_cast[MutAnyOrigin](), r_ptr.unsafe_origin_cast[MutAnyOrigin](), t > 0, - pool, + pool ) zero_and_free_u64(addressing_block, 128) @@ -408,7 +419,9 @@ def _argon2_process_lane( zero_and_free_u64(zero_u64, 128) zero_and_free_u64(tmp_addr, 128) -def _validate_params(parallelism: Int, tag_length: Int, memory_size_kb: Int, iterations: Int, version: Int) raises: + +def _validate_params(parallelism: Int, tag_length: Int, memory_size_kb: Int, iterations: Int, version: Int +) raises: # RFC 9106 section 3 parameter bounds if parallelism < 1 or parallelism >= (1 << 24): raise Error("Argon2 parallelism must be in [1, 2^24)") @@ -445,7 +458,7 @@ struct Argon2id: tag_length: Int = 32, memory_size_kb: Int = 65536, iterations: Int = 3, - version: Int = 0x13, + version: Int = 0x13 ) raises: _validate_params(parallelism, tag_length, memory_size_kb, iterations, version) self.parallelism = parallelism @@ -469,7 +482,7 @@ struct Argon2id: tag_length: Int = 32, memory_size_kb: Int = 65536, iterations: Int = 3, - version: Int = 0x13, + version: Int = 0x13 ) raises: _validate_params(parallelism, tag_length, memory_size_kb, iterations, version) self.parallelism = parallelism @@ -489,117 +502,156 @@ struct Argon2id: self.ad.append(ad[i]) def hash(self, password: Span[UInt8, ...]) raises -> List[UInt8]: - var h0_ctx = Blake2b(64) - 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) - 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, ...](unsafe_ptr=le_buf, length=4)) - store_le32(le_buf, 0, self.iterations) - h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf, length=4)) - store_le32(le_buf, 0, self.version) - h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf, length=4)) - store_le32(le_buf, 0, self.type_code) - h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf, length=4)) - store_le32(le_buf, 0, len(password)) - 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, ...](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, ...](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, ...](unsafe_ptr=le_buf, length=4)) - h0_ctx.update(Span[UInt8, ...](self.ad)) - zero_and_free(le_buf, 4) + if ( + len(password) >= (1 << 32) + or len(self.salt) >= (1 << 32) + or len(self.secret) >= (1 << 32) + or len(self.ad) >= (1 << 32) + ): + raise Error("Argon2 input lengths must fit in 32 bits") + + # These buffers contain password-derived state. Keeping them on the + # stack avoids per-block heap churn, and the outer finally guarantees + # cleanup on every raised path. + var le_buf = StackBuffer[UInt8, 4](fill=0) + var h0_buf = StackBuffer[UInt8, 64](fill=0) + var h0_input = StackBuffer[UInt8, 72](fill=0) + var b_bytes = StackBuffer[UInt8, 1024](fill=0) + var c_block = StackBuffer[UInt64, 128](fill=0) + var c_bytes = StackBuffer[UInt8, 1024](fill=0) + try: + var h0_ctx = Blake2b(64) + store_le32(le_buf.ptr(), 0, self.parallelism) + h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf.ptr(), length=4)) + store_le32(le_buf.ptr(), 0, self.tag_length) + h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf.ptr(), length=4)) + store_le32(le_buf.ptr(), 0, self.memory_size_kb) + h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf.ptr(), length=4)) + store_le32(le_buf.ptr(), 0, self.iterations) + h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf.ptr(), length=4)) + store_le32(le_buf.ptr(), 0, self.version) + h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf.ptr(), length=4)) + store_le32(le_buf.ptr(), 0, self.type_code) + h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf.ptr(), length=4)) + store_le32(le_buf.ptr(), 0, len(password)) + h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf.ptr(), length=4)) + h0_ctx.update(password) + store_le32(le_buf.ptr(), 0, len(self.salt)) + h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf.ptr(), length=4)) + h0_ctx.update(Span[UInt8, ...](self.salt)) + store_le32(le_buf.ptr(), 0, len(self.secret)) + h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf.ptr(), length=4)) + h0_ctx.update(Span[UInt8, ...](self.secret)) + store_le32(le_buf.ptr(), 0, len(self.ad)) + h0_ctx.update(Span[UInt8, ...](unsafe_ptr=le_buf.ptr(), length=4)) + h0_ctx.update(Span[UInt8, ...](self.ad)) + h0_ctx.finalize_into( + Span[mut=True, UInt8, ...](unsafe_ptr=h0_buf.ptr(), length=64) + ) - var h0_buf = alloc(Layout[UInt8](count=64)).unsafe_leak() - 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 = ( + 4 * self.parallelism * (m_blocks // (4 * self.parallelism)) + ) + if m_prime_blocks < 8 * self.parallelism: + m_prime_blocks = 8 * self.parallelism + var q = m_prime_blocks // self.parallelism + var segment_length = q // 4 + + var memory = alloc(Layout[UInt64](count=m_prime_blocks * 128)).unsafe_leak() + try: + unsafe_memcpy(dest=h0_input.ptr(), src=h0_buf.ptr(), count=64) + for i in range(self.parallelism): + for block_idx in range(2): + store_le32(h0_input.ptr(), 64, block_idx) + store_le32(h0_input.ptr(), 68, i) + variable_length_hash_into( + 1024, + Span[UInt8, ...](unsafe_ptr=h0_input.ptr(), length=72), + Span[mut=True, UInt8, ...]( + unsafe_ptr=b_bytes.ptr(), length=1024 + ), + ) + for k in range(128): + var word = ( + b_bytes.ptr() + .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_buffer(b_bytes.ptr(), 1024) + + var iterations = self.iterations + var type_code = self.type_code + var parallelism = self.parallelism + + for t in range(iterations): + for slice_idx in range(4): + var seg_start = slice_idx * segment_length + var seg_end = (slice_idx + 1) * segment_length + + @always_inline + @__copy_capture( + memory, + seg_start, + seg_end, + segment_length, + q, + m_prime_blocks, + t, + slice_idx, + iterations, + type_code, + parallelism, + ) + @parameter + def process_lane(lane: Int): + _argon2_process_lane( + memory, + lane, + t, + slice_idx, + seg_start, + seg_end, + segment_length, + q, + m_prime_blocks, + iterations, + type_code, + parallelism, + ) + + parallelize[process_lane](parallelism) + + zero_buffer_u64(c_block.ptr(), 128) + for i in range(self.parallelism): + var last_ptr = memory.unsafe_offset(i * q * 128 + (q - 1) * 128) + for k in range(128): + c_block.ptr()[unsafe_offset=k] ^= last_ptr[unsafe_offset=k] - var m_blocks = self.memory_size_kb - var m_prime_blocks = ( - 4 * self.parallelism * (m_blocks // (4 * self.parallelism)) - ) - if m_prime_blocks < 8 * self.parallelism: - m_prime_blocks = 8 * self.parallelism - var q = m_prime_blocks // self.parallelism - var segment_length = q // 4 - - var memory = alloc(Layout[UInt64](count=m_prime_blocks * 128)).unsafe_leak() - - 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): - for block_idx in range(2): - store_le32(h0_input, 64, block_idx) - store_le32(h0_input, 68, i) - - var b_bytes = alloc(Layout[UInt8](count=1024)).unsafe_leak() - 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.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) - - var iterations = self.iterations - var type_code = self.type_code - var parallelism = self.parallelism - - for t in range(iterations): - for slice_idx in range(4): - var seg_start = slice_idx * segment_length - var seg_end = (slice_idx + 1) * segment_length - - @always_inline - @__copy_capture( - memory, seg_start, seg_end, segment_length, - q, m_prime_blocks, t, slice_idx, iterations, - type_code, parallelism, + ( + c_bytes.ptr().unsafe_offset(k * 8) + ).unsafe_bitcast[UInt64]().unsafe_store[ + alignment=1 + ](0, c_block.ptr()[unsafe_offset=k]) + return variable_length_hash( + self.tag_length, + Span[UInt8, ...](unsafe_ptr=c_bytes.ptr(), length=1024), ) - @parameter - def process_lane(lane: Int): - _argon2_process_lane( - memory, lane, t, slice_idx, seg_start, seg_end, - segment_length, q, m_prime_blocks, iterations, - type_code, parallelism, - ) - - parallelize[process_lane](parallelism) - - var c_block = alloc(Layout[UInt64](count=128)).unsafe_leak() - zero_buffer_u64(c_block, 128) - for i in range(self.parallelism): - var last_ptr = memory.unsafe_offset((i * q * 128 + (q - 1) * 128)) - for k in range(128): - c_block[unsafe_offset=k] ^= last_ptr[unsafe_offset=k] - - var c_bytes = alloc(Layout[UInt8](count=1024)).unsafe_leak() - for k in range(128): - (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) - - var result = variable_length_hash(self.tag_length, Span[UInt8, ...](unsafe_ptr=c_bytes, length=1024)) - zero_and_free(c_bytes, 1024) - return result^ + finally: + zero_and_free_u64(memory, m_prime_blocks * 128) + finally: + zero_buffer(le_buf.ptr(), 4) + zero_buffer(h0_buf.ptr(), 64) + zero_buffer(h0_input.ptr(), 72) + zero_buffer(b_bytes.ptr(), 1024) + zero_buffer_u64(c_block.ptr(), 128) + zero_buffer(c_bytes.ptr(), 1024) + def argon2id_hash_string(password: String, salt: String) raises -> String: var p_bytes = password.as_bytes() diff --git a/src/thistle/blake2b.mojo b/src/thistle/blake2b.mojo index 4ae6522..cb0f710 100644 --- a/src/thistle/blake2b.mojo +++ b/src/thistle/blake2b.mojo @@ -1,7 +1,4 @@ -""" -BLAKE2b Implementation in Mojo -RFC 7693 -""" +"""Implements the BLAKE2b hash function specified by RFC 7693.""" from std.collections import List from std.memory import Pointer, unsafe_memcpy, unsafe_memset_zero @@ -15,7 +12,7 @@ comptime BLAKE2B_IV = SIMD[DType.uint64, 8]( 0x510E527FADE682D1, 0x9B05688C2B3E6C1F, 0x1F83D9ABFB41BD6B, - 0x5BE0CD19137E2179, + 0x5BE0CD19137E2179 ) comptime SIGMA = ( @@ -30,7 +27,7 @@ comptime SIGMA = ( SIMD[DType.uint8, 16](6, 15, 14, 9, 11, 3, 0, 8, 12, 2, 13, 7, 1, 4, 10, 5), SIMD[DType.uint8, 16](10, 2, 8, 4, 7, 6, 1, 5, 15, 11, 9, 14, 3, 12, 13, 0), SIMD[DType.uint8, 16](0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15), - SIMD[DType.uint8, 16](14, 10, 4, 8, 9, 15, 13, 6, 1, 12, 0, 2, 11, 7, 5, 3), + SIMD[DType.uint8, 16](14, 10, 4, 8, 9, 15, 13, 6, 1, 12, 0, 2, 11, 7, 5, 3) ) @@ -57,13 +54,14 @@ 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]() +# fmt: off @always_inline def round_fn[r: Int]( mut v0: UInt64, mut v1: UInt64, mut v2: UInt64, mut v3: UInt64, 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: Pointer[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] @@ -78,8 +76,7 @@ def round_fn[r: Int]( v3, v4, v9, v14 = g(v3, v4, v9, v14, _mload(m, Int(s[14])), _mload(m, Int(s[15]))) return (v0, v1, v2, v3, v4, v5, v6, v7, v8, v9, v10, v11, v12, v13, v14, v15) - - +# fmt: on struct Blake2b(Movable): var h: SIMD[DType.uint64, 8] var t_low: UInt64 @@ -157,6 +154,7 @@ struct Blake2b(Movable): if self.t_low < 128: self.t_high += 1 + # fmt: off def compress(mut self, m: Pointer[mut=False, UInt8, _, address_space=_], is_last: Bool): var v0 = self.h[0] var v1 = self.h[1] @@ -203,6 +201,7 @@ struct Blake2b(Movable): self.h[6] ^= v6 ^ v14 self.h[7] ^= v7 ^ v15 + # fmt: on def update(mut self, data: Span[UInt8, ...]): var total = len(data) if total == 0: @@ -248,9 +247,11 @@ struct Blake2b(Movable): self.compress(self._buf_ptr(), True) var h_copy = self.h - var h_bytes = Pointer(to=h_copy).unsafe_bitcast[UInt8]() + var h_copy_ptr = Pointer(to=h_copy).unsafe_mut_cast[True]().unsafe_bitcast[UInt64]() + var h_bytes = h_copy_ptr.unsafe_bitcast[UInt8]() for i in range(self.out_len): output[unsafe_offset=i] = h_bytes[unsafe_offset=i] + volatile_wipe(h_copy_ptr, 8) def finalize_into( mut self, output: Span[mut=True, UInt8, ...] diff --git a/src/thistle/blake3.mojo b/src/thistle/blake3.mojo index 225386d..f576c4f 100644 --- a/src/thistle/blake3.mojo +++ b/src/thistle/blake3.mojo @@ -1,6 +1,4 @@ -""" -BLAKE3 cryptographic hash function -""" +"""Implements the BLAKE3 cryptographic hash function.""" from max.algorithm import parallelize from std.collections import List @@ -17,7 +15,7 @@ comptime IV = SIMD[DType.uint32, 8]( 0x510E527F, 0x9B05688C, 0x1F83D9AB, - 0x5BE0CD19, + 0x5BE0CD19 ) """The BLAKE3 initial chaining value.""" @@ -32,6 +30,7 @@ comptime ROOT = UInt8(1 << 3) comptime CHUNK_LEN = 1024 """The length of a BLAKE3 chunk in bytes.""" + @always_inline def bit_rotr[n: Int, w: Int](v: SIMD[DType.uint32, w]) -> SIMD[DType.uint32, w]: var shift = SIMD[DType.uint32, w](n) @@ -46,7 +45,7 @@ def g_v_half1[ mut b: SIMD[DType.uint32, w], mut c: SIMD[DType.uint32, w], mut d: SIMD[DType.uint32, w], - x: SIMD[DType.uint32, w], + x: SIMD[DType.uint32, w] ): a = a + b + x d = bit_rotr[16, w](d ^ a) @@ -62,13 +61,14 @@ def g_v_half2[ mut b: SIMD[DType.uint32, w], mut c: SIMD[DType.uint32, w], mut d: SIMD[DType.uint32, w], - y: SIMD[DType.uint32, w], + y: SIMD[DType.uint32, w] ): a = a + b + y d = bit_rotr[8, w](d ^ a) c = c + d b = bit_rotr[7, w](b ^ c) + @always_inline def g_v[ w: Int @@ -78,7 +78,7 @@ def g_v[ mut c: SIMD[DType.uint32, w], mut d: SIMD[DType.uint32, w], x: SIMD[DType.uint32, w], - y: SIMD[DType.uint32, w], + y: SIMD[DType.uint32, w] ): g_v_half1[w](a, b, c, d, x) g_v_half2[w](a, b, c, d, y) @@ -90,7 +90,7 @@ def g_idx[ ]( mut v: StackInlineArray[SIMD[DType.uint32, w], 16], x: SIMD[DType.uint32, w], - y: SIMD[DType.uint32, w], + y: SIMD[DType.uint32, w] ): # copy in/out so we never hold two mut refs into the same array var a = v[ai] @@ -103,6 +103,8 @@ def g_idx[ v[ci] = c v[di] = d + +# fmt: off @always_inline def compress_internal[ w: Int @@ -112,19 +114,17 @@ def compress_internal[ counter: UInt64, blen: UInt8, flags: UInt8, - out_ptr: Pointer[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 var v: StackInlineArray[SIMD[DType.uint32, w], 16] = [ cv[0], cv[1], cv[2], cv[3], cv[4], cv[5], cv[6], cv[7], UInt32(0x6A09E667), UInt32(0xBB67AE85), UInt32(0x3C6EF372), UInt32(0xA54FF53A), UInt32(counter & 0xFFFFFFFF), UInt32(counter >> 32), UInt32(blen), - UInt32(flags), + UInt32(flags) ] - # fmt: on @parameter @always_inline @@ -141,12 +141,10 @@ def compress_internal[ @parameter @always_inline def transform(): - # fmt: off m = [ m[2], m[6], m[3], m[10], m[7], m[0], m[4], m[13], - m[1], m[11], m[12], m[5], m[9], m[14], m[15], m[8], + m[1], m[11], m[12], m[5], m[9], m[14], m[15], m[8] ] - # fmt: on round() comptime for _ in range(6): @@ -157,23 +155,22 @@ def compress_internal[ v.unsafe_ptr().unsafe_bitcast[UInt32]().unsafe_load[width=w * 16]() ) + @always_inline def compress_core( cv: SIMD[DType.uint32, 8], block: SIMD[DType.uint32, 16], counter: UInt64, blen: UInt8, - flags: UInt8, + flags: UInt8 ) -> SIMD[DType.uint32, 16]: """Single-block compression returning 16-word output.""" - # fmt: off var m: StackInlineArray[UInt32, 16] = [ block[0], block[1], block[2], block[3], block[4], block[5], block[6], block[7], block[8], block[9], block[10], block[11], - block[12], block[13], block[14], block[15], + block[12], block[13], block[14], block[15] ] - # fmt: on var res = StackInlineArray[SIMD[DType.uint32, 1], 16]( fill=SIMD[DType.uint32, 1](0) @@ -186,6 +183,7 @@ def compress_core( final[i + 8] = res[i + 8][0] ^ cv[i] return final + @always_inline def compress_internal_16way( c: StackInlineArray[SIMD[DType.uint32, 16], 8], @@ -193,23 +191,21 @@ def compress_internal_16way( base_counter: UInt64, blen: UInt8, flags: UInt8, - out_ptr: Pointer[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]( 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15 ) counters_low += UInt32(base_counter & 0xFFFFFFFF) - # fmt: off var v: StackInlineArray[SIMD[DType.uint32, 16], 16] = [ c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7], UInt32(0x6A09E667), UInt32(0xBB67AE85), UInt32(0x3C6EF372), UInt32(0xA54FF53A), counters_low, UInt32(base_counter >> 32), UInt32(blen), - UInt32(flags), + UInt32(flags) ] - # fmt: on @parameter @always_inline @@ -226,12 +222,10 @@ def compress_internal_16way( @parameter @always_inline def transform(): - # fmt: off m = [ m[2], m[6], m[3], m[10], m[7], m[0], m[4], m[13], - m[1], m[11], m[12], m[5], m[9], m[14], m[15], m[8], + m[1], m[11], m[12], m[5], m[9], m[14], m[15], m[8] ] - # fmt: on round() comptime for _ in range(6): @@ -247,6 +241,7 @@ def compress_internal_16way( out_ptr[unsafe_offset=6] = v[6] ^ v[14] out_ptr[unsafe_offset=7] = v[7] ^ v[15] +# fmt: on struct Hasher: var key: SIMD[DType.uint32, 8] var original_key: SIMD[DType.uint32, 8] @@ -304,7 +299,7 @@ struct Hasher: self.chunk_counter, 64, (CHUNK_START if self.blocks_compressed == 0 else UInt8(0)) - | CHUNK_END, + | CHUNK_END ) var chunk_cv = res.slice[8]() self.add_chunk_cv(chunk_cv, self.chunk_counter) @@ -321,7 +316,7 @@ struct Hasher: blk, self.chunk_counter, 64, - (CHUNK_START if self.blocks_compressed == 0 else UInt8(0)), + (CHUNK_START if self.blocks_compressed == 0 else UInt8(0)) ) self.key = res.slice[8]() self.blocks_compressed += 1 @@ -383,7 +378,9 @@ struct Hasher: for i in range(self.buf_len): temp_buf.unsafe_set(i, self.buf[i]) - var blk = temp_buf.unsafe_ptr().unsafe_bitcast[UInt32]().unsafe_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) @@ -439,8 +436,12 @@ struct Hasher: @always_inline def blake3_parallel_hash(input: Span[UInt8, ...], out_len: Int = 32) raises -> List[UInt8]: + if out_len < 0: + raise Error("BLAKE3 output length must be non-negative") var d = input - var total_chunks = (len(d) + CHUNK_LEN - 1) // CHUNK_LEN + var total_chunks = len(d) // CHUNK_LEN + if len(d) % CHUNK_LEN != 0: + total_chunks += 1 if len(d) <= 65536: var h = Hasher() @@ -468,7 +469,7 @@ def blake3_parallel_hash(input: Span[UInt8, ...], out_len: Int = 32) raises -> L # fmt: off var c: StackInlineArray[SIMD[DType.uint32, 16], 8] = [ UInt32(0x6A09E667), UInt32(0xBB67AE85), UInt32(0x3C6EF372), UInt32(0xA54FF53A), - UInt32(0x510E527F), UInt32(0x9B05688C), UInt32(0x1F83D9AB), UInt32(0x5BE0CD19), + UInt32(0x510E527F), UInt32(0x9B05688C), UInt32(0x1F83D9AB), UInt32(0x5BE0CD19) ] # fmt: on @@ -547,7 +548,7 @@ def blake3_parallel_hash(input: Span[UInt8, ...], out_len: Int = 32) raises -> L ma[12].join(mb[12]), ma[13].join(mb[13]), ma[14].join(mb[14]), - ma[15].join(mb[15]), + ma[15].join(mb[15]) ] var res = StackInlineArray[SIMD[DType.uint32, 16], 8]( @@ -559,7 +560,7 @@ def blake3_parallel_hash(input: Span[UInt8, ...], out_len: Int = 32) raises -> L UInt64(base), 64, flags, - res.unsafe_ptr(), + res.unsafe_ptr() ) c[0] = res[0] c[1] = res[1] @@ -598,5 +599,6 @@ def blake3_parallel_hash(input: Span[UInt8, ...], out_len: Int = 32) raises -> L h.update(d[num_full_batches * 64 * 1024 :]) return h.finalize(out_len) + def blake3_hash(input: Span[UInt8, ...], out_len: Int = 32) raises -> List[UInt8]: return blake3_parallel_hash(input, out_len) diff --git a/src/thistle/camellia.mojo b/src/thistle/camellia.mojo index e80cb0d..d486e4d 100644 --- a/src/thistle/camellia.mojo +++ b/src/thistle/camellia.mojo @@ -1,10 +1,9 @@ -""" -Camellia block cipher implementation per RFC 3713 -""" +"""Implements the Camellia block cipher specified by RFC 3713.""" from std.memory import bitcast, Pointer from std.bit import byte_swap, rotate_bits_left from std.collections import InlineArray +from std.os import abort from std.sys import llvm_intrinsic from std.utils import StaticTuple from .aes import _ct_sbox, _ct_ortho, _ctr_write_block @@ -13,11 +12,19 @@ from .aes_ni import ( _aese, _mm_aesenclast_si128, has_arm_crypto, - has_x86_aes_ni, + has_x86_aes_ni ) comptime SIGMA1 = 0xA09E667F3BCC908B + + +@always_inline +def _validate_camellia_block_count(num_blocks: Int): + if num_blocks < 0: + abort("Camellia block count cannot be negative") + + comptime SIGMA2 = 0xB67AE8584CAA73B2 comptime SIGMA3 = 0xC6EF372FE94F82BE comptime SIGMA4 = 0x54FF53A5F1D36F1C @@ -40,12 +47,14 @@ comptime _U8x16 = SIMD[DType.uint8, 16] # y = LO[x & 15] ^ HI[x >> 4], constant LO side only. # s2 = rotl1(s1), s3 = rotr1(s1) fold rotated POST copies; # s4 = s1(rotl1(x)) folds rotl1 into pre matrix. + + def _mk_tbl( cols: StaticTuple[UInt8, 8], add: UInt8, hi: Bool, in_rotl1: Bool, - out_rot: Int, + out_rot: Int ) -> _U8x16: var t = _U8x16(0) for n in range(16): @@ -158,7 +167,7 @@ def _f_planes[W: Int]( left: InlineArray[SIMD[DType.uint64, W], 8], mut right: InlineArray[SIMD[DType.uint64, W], 8], kp: InlineArray[UInt64, 192], - base: Int, + base: Int ): var a = InlineArray[SIMD[DType.uint64, W], 8](fill=0) comptime for k in range(8): @@ -192,11 +201,9 @@ def _f_planes[W: Int]( comptime m23: UInt64 = _LANES_S2 | _LANES_S3 var e = InlineArray[SIMD[DType.uint64, W], 8](fill=0) comptime for k in range(8): - e[k] = ( - (d[k] & ~m23) + e[k] = (d[k] & ~m23) | (d[(k + 7) % 8] & _LANES_S2) | (d[(k + 1) % 8] & _LANES_S3) - ) comptime for k in range(8): var p = e[k] @@ -215,7 +222,7 @@ def _f_planes[W: Int]( def _fl_planes[inv: Bool, W: Int]( mut x: InlineArray[SIMD[DType.uint64, W], 8], kep: InlineArray[UInt64, 48], - base: Int, + base: Int ): comptime if inv: comptime for k in range(8): @@ -245,13 +252,13 @@ def rotl128[n: Int](high: UInt64, low: UInt64) -> SIMD[DType.uint64, 2]: comptime s = UInt64(shift) return SIMD[DType.uint64, 2]( (high << s) | (low >> (UInt64(64) - s)), - (low << s) | (high >> (UInt64(64) - s)), + (low << s) | (high >> (UInt64(64) - s)) ) else: comptime s = UInt64(shift - 64) return SIMD[DType.uint64, 2]( (low << s) | (high >> (UInt64(64) - s)), - (high << s) | (low >> (UInt64(64) - s)), + (high << s) | (low >> (UInt64(64) - s)) ) @@ -272,11 +279,9 @@ def _p_tail(sout0: UInt64) -> UInt64: var orr = ((sout >> 1) & ~(_BIT0_OF_EACH_BYTE << 7)) | ( (sout << 7) & (_BIT0_OF_EACH_BYTE << 7) ) - sout = ( - (sout & ~(_LANES_S2 | _LANES_S3)) + sout = (sout & ~(_LANES_S2 | _LANES_S3)) | (ol & _LANES_S2) | (orr & _LANES_S3) - ) var dw = UInt32(sout & 0xFFFFFFFF) var uw = UInt32(sout >> 32) @@ -353,7 +358,6 @@ def _f_one(x: UInt64, k: UInt64) -> UInt64: return _f_scalar(x, k) - comptime _M_S4V = _U8x16( 0, 0, 0, 0xFF, 0, 0, 0xFF, 0, 0, 0, 0, 0xFF, 0, 0, 0xFF, 0 ) @@ -370,7 +374,7 @@ comptime _P_SRC = StaticTuple[StaticTuple[UInt8, 6], 8]( StaticTuple[UInt8, 6](0, 1, 5, 6, 7, 255), StaticTuple[UInt8, 6](1, 2, 4, 6, 7, 255), StaticTuple[UInt8, 6](2, 3, 4, 5, 7, 255), - StaticTuple[UInt8, 6](0, 3, 4, 5, 6, 255), + StaticTuple[UInt8, 6](0, 3, 4, 5, 6, 255) ) @@ -415,6 +419,7 @@ comptime _FL_SH_B = _mk_fl_idx(8, False, False) comptime _FL_SH2_B = _mk_fl_idx(8, True, False) comptime _FL_DOWN_B = _mk_fl_idx(8, False, True) + @always_inline def _kv_a(k: UInt64) -> _U8x16: return bitcast[DType.uint8, 16](SIMD[DType.uint64, 2](k, 0)) @@ -704,12 +709,10 @@ def _load_half[W: Int]( comptime for e in range(W): comptime for j in range(8): var base = (e * 8 + j) * 16 + off - q[j][e] = ( - bitcast[DType.uint64, 1]( + q[j][e] = bitcast[DType.uint64, 1]( buf.unsafe_load[width=8, alignment=1](base) )[0] ^ kwl - ) _ct_ortho[W](q) return q^ @@ -719,7 +722,7 @@ def _store_half[W: Int]( buf: Pointer[mut=True, UInt8, _, address_space=_], off: Int, mut q: InlineArray[SIMD[DType.uint64, W], 8], - kw: UInt64, + kw: UInt64 ): _ct_ortho(q) var kwl = byte_swap(kw) @@ -730,7 +733,7 @@ def _store_half[W: Int]( base, bitcast[DType.uint8, 8]( SIMD[DType.uint64, 1](q[j][e] ^ kwl) - ), + ) ) @@ -739,7 +742,7 @@ def _six_rounds[forward: Bool, W: Int]( mut a: InlineArray[SIMD[DType.uint64, W], 8], mut b: InlineArray[SIMD[DType.uint64, W], 8], kp: InlineArray[UInt64, 192], - kbase: Int, + kbase: Int ): comptime for r in range(6): comptime off = 8 * (r if forward else 5 - r) @@ -814,6 +817,7 @@ comptime _BITREV4 = StaticTuple[Int, 16]( 0, 8, 4, 12, 2, 10, 6, 14, 1, 9, 5, 13, 3, 11, 7, 15 ) + @always_inline def _transpose16(mut m: InlineArray[_U8x16, 16]): # 16x16 byte matrix transpose setup 4 zip stages. @@ -851,7 +855,7 @@ def _splat_byte(k: UInt64, j: Int) -> _U8x16: def _f_bs( l: InlineArray[_U8x16, 8], mut r: InlineArray[_U8x16, 8], - k: UInt64, + k: UInt64 ): # Byte-sliced F on 16 blocks, k is the pre-byte-swapped subkey # (byte j = t_{j+1}). P-function in CSE form: @@ -913,7 +917,7 @@ def _six_rounds_bs[forward: Bool]( mut a: InlineArray[_U8x16, 8], mut b: InlineArray[_U8x16, 8], cipher: CamelliaCipher, - kbase: Int, + kbase: Int ): comptime for r in range(6): comptime off = r if forward else 5 - r @@ -1089,7 +1093,7 @@ def _camellia_block[encrypt: Bool]( def _camellia_blocks[encrypt: Bool]( cipher: CamelliaCipher, data: Pointer[mut=True, UInt8, _, address_space=_], - num_blocks: Int, + num_blocks: Int ): comptime if _has_hw_sbox(): var i = 0 @@ -1137,16 +1141,18 @@ def camellia_decrypt_block( def camellia_encrypt_blocks( cipher: CamelliaCipher, data: Pointer[mut=True, UInt8, _, address_space=_], - num_blocks: Int, + num_blocks: Int ): + _validate_camellia_block_count(num_blocks) _camellia_blocks[True](cipher, data, num_blocks) def camellia_decrypt_blocks( cipher: CamelliaCipher, data: Pointer[mut=True, UInt8, _, address_space=_], - num_blocks: Int, + num_blocks: Int ): + _validate_camellia_block_count(num_blocks) _camellia_blocks[False](cipher, data, num_blocks) @@ -1155,8 +1161,9 @@ def camellia_cbc_encrypt_kernel( output_ptr: Pointer[mut=True, UInt8, _, address_space=_], cipher: CamelliaCipher, num_blocks: Int, - iv_ptr: Pointer[mut=True, UInt8, _, address_space=_], + iv_ptr: Pointer[mut=True, UInt8, _, address_space=_] ): + _validate_camellia_block_count(num_blocks) var prev = iv_ptr.unsafe_load[width=16, alignment=1](0) for i in range(num_blocks): var x = input_ptr.unsafe_load[width=16, alignment=1](i * 16) ^ prev @@ -1169,8 +1176,9 @@ def camellia_cbc_decrypt_kernel( output_ptr: Pointer[mut=True, UInt8, _, address_space=_], cipher: CamelliaCipher, num_blocks: Int, - iv_ptr: Pointer[mut=True, UInt8, _, address_space=_], + iv_ptr: Pointer[mut=True, UInt8, _, address_space=_] ): + _validate_camellia_block_count(num_blocks) var ct = InlineArray[UInt8, 1024](fill=0) var pt = InlineArray[UInt8, 1024](fill=0) var ctp = ct.unsafe_ptr() @@ -1198,8 +1206,9 @@ def camellia_ctr_kernel( output_ptr: Pointer[mut=True, UInt8, _, address_space=_], cipher: CamelliaCipher, num_blocks: Int, - nonce_ptr: Pointer[mut=True, UInt8, _, address_space=_], + nonce_ptr: Pointer[mut=True, UInt8, _, address_space=_] ): + _validate_camellia_block_count(num_blocks) var ks = InlineArray[UInt8, 512](fill=0) var kp = ks.unsafe_ptr() var i = 0 @@ -1213,7 +1222,7 @@ def camellia_ctr_kernel( output_ptr.unsafe_store[alignment=1]( off, input_ptr.unsafe_load[width=16, alignment=1](off) - ^ kp.unsafe_load[width=16, alignment=1](b * 16), + ^ kp.unsafe_load[width=16, alignment=1](b * 16) ) i += 16 while i < num_blocks: @@ -1228,7 +1237,7 @@ def camellia_ctr_kernel( output_ptr.unsafe_store[alignment=1]( off, input_ptr.unsafe_load[width=16, alignment=1](off) - ^ kp.unsafe_load[width=16, alignment=1](b * 16), + ^ kp.unsafe_load[width=16, alignment=1](b * 16) ) i += n else: @@ -1241,7 +1250,7 @@ def camellia_ctr_kernel( output_ptr.unsafe_store[alignment=1]( off, input_ptr.unsafe_load[width=16, alignment=1](off) - ^ kp.unsafe_load[width=16, alignment=1](b * 16), + ^ kp.unsafe_load[width=16, alignment=1](b * 16) ) i += 32 while i < num_blocks: @@ -1258,6 +1267,6 @@ def camellia_ctr_kernel( output_ptr.unsafe_store[alignment=1]( off, input_ptr.unsafe_load[width=16, alignment=1](off) - ^ kp.unsafe_load[width=16, alignment=1](b * 16), + ^ kp.unsafe_load[width=16, alignment=1](b * 16) ) i += n diff --git a/src/thistle/chacha20.mojo b/src/thistle/chacha20.mojo index 9e72595..a8aa62e 100644 --- a/src/thistle/chacha20.mojo +++ b/src/thistle/chacha20.mojo @@ -1,11 +1,11 @@ -# ChaCha20 stream cipher per RFC 7539. +"""Implements the ChaCha20 stream cipher specified by RFC 7539.""" from std.memory import bitcast from std.memory.unsafe_pointer import Pointer from std.bit import rotate_bits_left comptime CHACHA_CONSTANTS = SIMD[DType.uint32, 4]( - 0x61707865, 0x3320646E, 0x79622D32, 0x6B206574, + 0x61707865, 0x3320646E, 0x79622D32, 0x6B206574 ) @@ -22,7 +22,7 @@ def _rotl[shift: Int, W: Int](x: SIMD[DType.uint32, W]) -> SIMD[DType.uint32, W] h.shuffle[ 1, 0, 3, 2, 5, 4, 7, 6, 9, 8, 11, 10, 13, 12, 15, 14, 17, 16, 19, 18, 21, 20, 23, 22, 25, 24, 27, 26, 29, 28, - 31, 30, + 31, 30 ]() ) else: @@ -40,7 +40,7 @@ def _rotl[shift: Int, W: Int](x: SIMD[DType.uint32, W]) -> SIMD[DType.uint32, W] 19, 16, 17, 18, 23, 20, 21, 22, 27, 24, 25, 26, 31, 28, 29, 30, 35, 32, 33, 34, 39, 36, 37, 38, 43, 40, 41, 42, 47, 44, 45, 46, 51, 48, 49, 50, 55, 52, 53, 54, 59, 56, - 57, 58, 63, 60, 61, 62, + 57, 58, 63, 60, 61, 62 ]() ) else: @@ -48,13 +48,15 @@ def _rotl[shift: Int, W: Int](x: SIMD[DType.uint32, W]) -> SIMD[DType.uint32, W] else: return rotate_bits_left[shift](x) + @always_inline def simd_quarter_round( a: SIMD[DType.uint32, 4], b: SIMD[DType.uint32, 4], c: SIMD[DType.uint32, 4], - d: SIMD[DType.uint32, 4], -) -> Tuple[SIMD[DType.uint32, 4], SIMD[DType.uint32, 4], SIMD[DType.uint32, 4], SIMD[DType.uint32, 4]]: + d: SIMD[DType.uint32, 4] +) -> Tuple[SIMD[DType.uint32, 4], SIMD[DType.uint32, 4], SIMD[DType.uint32, 4], SIMD[DType.uint32, 4] +]: var aa = a var bb = b @@ -85,8 +87,9 @@ def shuffle_for_diagonal( row0: SIMD[DType.uint32, 4], row1: SIMD[DType.uint32, 4], row2: SIMD[DType.uint32, 4], - row3: SIMD[DType.uint32, 4], -) -> Tuple[SIMD[DType.uint32, 4], SIMD[DType.uint32, 4], SIMD[DType.uint32, 4], SIMD[DType.uint32, 4]]: + row3: SIMD[DType.uint32, 4] +) -> Tuple[SIMD[DType.uint32, 4], SIMD[DType.uint32, 4], SIMD[DType.uint32, 4], SIMD[DType.uint32, 4] +]: var a = row0 var b = row1.rotate_left[1]() @@ -101,8 +104,9 @@ def unshuffle_from_diagonal( a: SIMD[DType.uint32, 4], b: SIMD[DType.uint32, 4], c: SIMD[DType.uint32, 4], - d: SIMD[DType.uint32, 4], -) -> Tuple[SIMD[DType.uint32, 4], SIMD[DType.uint32, 4], SIMD[DType.uint32, 4], SIMD[DType.uint32, 4]]: + d: SIMD[DType.uint32, 4] +) -> Tuple[SIMD[DType.uint32, 4], SIMD[DType.uint32, 4], SIMD[DType.uint32, 4], SIMD[DType.uint32, 4] +]: var row0 = a var row1 = b.rotate_right[1]() @@ -117,8 +121,9 @@ def simd_double_round( row0: SIMD[DType.uint32, 4], row1: SIMD[DType.uint32, 4], row2: SIMD[DType.uint32, 4], - row3: SIMD[DType.uint32, 4], -) -> Tuple[SIMD[DType.uint32, 4], SIMD[DType.uint32, 4], SIMD[DType.uint32, 4], SIMD[DType.uint32, 4]]: + row3: SIMD[DType.uint32, 4] +) -> Tuple[SIMD[DType.uint32, 4], SIMD[DType.uint32, 4], SIMD[DType.uint32, 4], SIMD[DType.uint32, 4] +]: var rr0 = row0 var rr1 = row1 @@ -126,23 +131,32 @@ def simd_double_round( var rr3 = row3 var qr = simd_quarter_round(rr0, rr1, rr2, rr3) - rr0 = qr[0]; rr1 = qr[1]; rr2 = qr[2]; rr3 = qr[3] + rr0 = qr[0] + rr1 = qr[1] + rr2 = qr[2] + rr3 = qr[3] var shuffled = shuffle_for_diagonal(rr0, rr1, rr2, rr3) - var diag_a = shuffled[0]; var diag_b = shuffled[1] - var diag_c = shuffled[2]; var diag_d = shuffled[3] + var diag_a = shuffled[0] + var diag_b = shuffled[1] + var diag_c = shuffled[2] + var diag_d = shuffled[3] qr = simd_quarter_round(diag_a, diag_b, diag_c, diag_d) - diag_a = qr[0]; diag_b = qr[1]; diag_c = qr[2]; diag_d = qr[3] + diag_a = qr[0] + diag_b = qr[1] + diag_c = qr[2] + diag_d = qr[3] return unshuffle_from_diagonal(diag_a, diag_b, diag_c, diag_d) + @always_inline def _qr16( mut a: SIMD[DType.uint32, 16], mut b: SIMD[DType.uint32, 16], mut c: SIMD[DType.uint32, 16], - mut d: SIMD[DType.uint32, 16], + mut d: SIMD[DType.uint32, 16] ): a = a + b d = _rotl[16, 16](d ^ a) @@ -159,7 +173,7 @@ def simd_double_round_16x( mut row0: SIMD[DType.uint32, 16], mut row1: SIMD[DType.uint32, 16], mut row2: SIMD[DType.uint32, 16], - mut row3: SIMD[DType.uint32, 16], + mut row3: SIMD[DType.uint32, 16] ): _qr16(row0, row1, row2, row3) @@ -192,7 +206,7 @@ def _quad_block( row1: SIMD[DType.uint32, 16], row2: SIMD[DType.uint32, 16], row3: SIMD[DType.uint32, 16], - blk: Int, + blk: Int ) -> SIMD[DType.uint32, 16]: var out = SIMD[DType.uint32, 16]() if blk == 0: @@ -227,32 +241,32 @@ comptime _CTR_INC4 = SIMD[DType.uint32, 16]( def _quad_rows_init( key: SIMD[DType.uint32, 8], counter: UInt32, - nonce: SIMD[DType.uint32, 4], + nonce: SIMD[DType.uint32, 4] ) -> Tuple[ SIMD[DType.uint32, 16], SIMD[DType.uint32, 16], SIMD[DType.uint32, 16], - SIMD[DType.uint32, 16], + SIMD[DType.uint32, 16] ]: comptime CONST16 = SIMD[DType.uint32, 16]( 0x61707865, 0x3320646E, 0x79622D32, 0x6B206574, 0x61707865, 0x3320646E, 0x79622D32, 0x6B206574, 0x61707865, 0x3320646E, 0x79622D32, 0x6B206574, - 0x61707865, 0x3320646E, 0x79622D32, 0x6B206574, + 0x61707865, 0x3320646E, 0x79622D32, 0x6B206574 ) var row1 = SIMD[DType.uint32, 16]( key[0], key[1], key[2], key[3], key[0], key[1], key[2], key[3], - key[0], key[1], key[2], key[3], key[0], key[1], key[2], key[3], + key[0], key[1], key[2], key[3], key[0], key[1], key[2], key[3] ) var row2 = SIMD[DType.uint32, 16]( key[4], key[5], key[6], key[7], key[4], key[5], key[6], key[7], - key[4], key[5], key[6], key[7], key[4], key[5], key[6], key[7], + key[4], key[5], key[6], key[7], key[4], key[5], key[6], key[7] ) var row3 = SIMD[DType.uint32, 16]( counter, nonce[0], nonce[1], nonce[2], counter + 1, nonce[0], nonce[1], nonce[2], counter + 2, nonce[0], nonce[1], nonce[2], - counter + 3, nonce[0], nonce[1], nonce[2], + counter + 3, nonce[0], nonce[1], nonce[2] ) return Tuple(CONST16, row1, row2, row3) @@ -261,7 +275,7 @@ def _quad_rows_init( def _diag16[back: Bool]( mut r1: SIMD[DType.uint32, 16], mut r2: SIMD[DType.uint32, 16], - mut r3: SIMD[DType.uint32, 16], + mut r3: SIMD[DType.uint32, 16] ): comptime if back: r1 = r1.shuffle[3, 0, 1, 2, 7, 4, 5, 6, 11, 8, 9, 10, 15, 12, 13, 14]() @@ -279,12 +293,12 @@ def chacha20_x12_core_rows( init2: SIMD[DType.uint32, 16], init3a: SIMD[DType.uint32, 16], init3b: SIMD[DType.uint32, 16], - init3c: SIMD[DType.uint32, 16], + init3c: SIMD[DType.uint32, 16] ) -> Tuple[ SIMD[DType.uint32, 16], SIMD[DType.uint32, 16], SIMD[DType.uint32, 16], SIMD[DType.uint32, 16], SIMD[DType.uint32, 16], SIMD[DType.uint32, 16], SIMD[DType.uint32, 16], SIMD[DType.uint32, 16], SIMD[DType.uint32, 16], - SIMD[DType.uint32, 16], SIMD[DType.uint32, 16], SIMD[DType.uint32, 16], + SIMD[DType.uint32, 16], SIMD[DType.uint32, 16], SIMD[DType.uint32, 16] ]: var a0 = init0 var a1 = init1 @@ -338,7 +352,7 @@ def chacha20_x12_core_rows( _quad_block(c0, c1, c2, c3, 0), _quad_block(c0, c1, c2, c3, 1), _quad_block(c0, c1, c2, c3, 2), - _quad_block(c0, c1, c2, c3, 3), + _quad_block(c0, c1, c2, c3, 3) ) @@ -348,7 +362,7 @@ def chacha20_octo_core_rows( init1: SIMD[DType.uint32, 16], init2: SIMD[DType.uint32, 16], init3a: SIMD[DType.uint32, 16], - init3b: SIMD[DType.uint32, 16], + init3b: SIMD[DType.uint32, 16] ) -> Tuple[ SIMD[DType.uint32, 16], SIMD[DType.uint32, 16], @@ -357,7 +371,7 @@ def chacha20_octo_core_rows( SIMD[DType.uint32, 16], SIMD[DType.uint32, 16], SIMD[DType.uint32, 16], - SIMD[DType.uint32, 16], + SIMD[DType.uint32, 16] ]: var a0 = init0 var a1 = init1 @@ -403,7 +417,7 @@ def chacha20_octo_core_rows( _quad_block(b0, b1, b2, b3, 0), _quad_block(b0, b1, b2, b3, 1), _quad_block(b0, b1, b2, b3, 2), - _quad_block(b0, b1, b2, b3, 3), + _quad_block(b0, b1, b2, b3, 3) ) @@ -412,12 +426,12 @@ def chacha20_quad_core_rows( init0: SIMD[DType.uint32, 16], init1: SIMD[DType.uint32, 16], init2: SIMD[DType.uint32, 16], - init3: SIMD[DType.uint32, 16], + init3: SIMD[DType.uint32, 16] ) -> Tuple[ SIMD[DType.uint32, 16], SIMD[DType.uint32, 16], SIMD[DType.uint32, 16], - SIMD[DType.uint32, 16], + SIMD[DType.uint32, 16] ]: var row0 = init0 var row1 = init1 @@ -436,7 +450,7 @@ def chacha20_quad_core_rows( _quad_block(row0, row1, row2, row3, 0), _quad_block(row0, row1, row2, row3, 1), _quad_block(row0, row1, row2, row3, 2), - _quad_block(row0, row1, row2, row3, 3), + _quad_block(row0, row1, row2, row3, 3) ) @@ -456,7 +470,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, 4], + nonce: SIMD[DType.uint32, 4] ) -> SIMD[DType.uint32, 16]: var row0 = CHACHA_CONSTANTS @@ -471,7 +485,10 @@ def chacha20_block_core( comptime for _ in range(10): var dr = simd_double_round(row0, row1, row2, row3) - row0 = dr[0]; row1 = dr[1]; row2 = dr[2]; row3 = dr[3] + row0 = dr[0] + row1 = dr[1] + row2 = dr[2] + row3 = dr[3] row0 = row0 + init0 row1 = row1 + init1 @@ -490,7 +507,7 @@ def chacha20_block_core( def _chacha20_block_scalar( key: SIMD[DType.uint32, 8], counter: UInt32, - nonce: SIMD[DType.uint32, 4], + nonce: SIMD[DType.uint32, 4] ) -> SIMD[DType.uint32, 16]: var x0: UInt32 = 0x61707865 var x1: UInt32 = 0x3320646E @@ -523,7 +540,7 @@ def _chacha20_block_scalar( x0 + 0x61707865, x1 + 0x3320646E, x2 + 0x79622D32, x3 + 0x6B206574, x4 + key[0], x5 + key[1], x6 + key[2], x7 + key[3], x8 + key[4], x9 + key[5], x10 + key[6], x11 + key[7], - x12 + counter, x13 + nonce[0], x14 + nonce[1], x15 + nonce[2], + x12 + counter, x13 + nonce[0], x14 + nonce[1], x15 + nonce[2] ) return result @@ -553,7 +570,7 @@ def _xor_block64( src: Pointer[mut=True, UInt8, _, address_space=_], dst: Pointer[mut=True, UInt8, _, address_space=_], keystream: SIMD[DType.uint32, 16], - offset: Int, + offset: Int ): var ks = bitcast[DType.uint8, 64](keystream) var v = (src.unsafe_offset(offset)).unsafe_load[width=64, alignment=1](0) @@ -570,7 +587,7 @@ struct ChaCha20: out self, key_bytes: SIMD[DType.uint8, 32], nonce_bytes: Span[UInt8, ...], - counter: UInt32 = 1, + counter: UInt32 = 1 ) raises: self.key = bitcast[DType.uint32, 8](key_bytes) self.nonce = _chacha20_nonce_words(nonce_bytes) @@ -588,7 +605,11 @@ struct ChaCha20: Pointer(to=self.exhausted).unsafe_mut_cast[True]().unsafe_bitcast[UInt8]().unsafe_store[volatile=True](0, UInt8(0)) def _check_counter_space(self, data_len: Int) raises: - var blocks_needed = UInt64((data_len + 63) // 64) + if data_len < 0: + raise Error("ChaCha20 input length cannot be negative") + var blocks_needed = UInt64(data_len // 64) + if data_len % 64 != 0: + blocks_needed += 1 if self.exhausted and blocks_needed > 0: raise Error("ChaCha20 counter is exhausted, use a new nonce") var space = UInt64(0x100000000) - UInt64(self.counter) @@ -600,7 +621,7 @@ struct ChaCha20: mut self, src: Pointer[mut=True, UInt8, _, address_space=_], dst: Pointer[mut=True, UInt8, _, address_space=_], - length: Int, + length: Int ) raises: self._check_counter_space(length) var block_idx = 0 @@ -612,7 +633,7 @@ struct ChaCha20: while offset + 768 <= length: var o = chacha20_x12_core_rows( rows[0], rows[1], rows[2], i3, - i3 + _CTR_INC4, i3 + _CTR_INC4 + _CTR_INC4, + i3 + _CTR_INC4, i3 + _CTR_INC4 + _CTR_INC4 ) comptime for j in range(12): _xor_block64(src, dst, o[j], offset + j * 64) @@ -663,20 +684,20 @@ struct ChaCha20: def encrypt_into[origin: Origin[mut=True]]( mut self, plaintext: Span[UInt8, ...], - mut ciphertext: Span[mut=True, UInt8, origin], + mut ciphertext: Span[mut=True, UInt8, origin] ) raises: if len(ciphertext) < len(plaintext): raise Error("ChaCha20 ciphertext buffer too small") self._stream_xor( plaintext.unsafe_ptr().unsafe_mut_cast[True]().unsafe_origin_cast[MutAnyOrigin](), ciphertext.unsafe_ptr().unsafe_origin_cast[MutAnyOrigin](), - len(plaintext), + len(plaintext) ) def decrypt_into[origin: Origin[mut=True]]( mut self, ciphertext: Span[UInt8, ...], - mut plaintext: Span[mut=True, UInt8, origin], + mut plaintext: Span[mut=True, UInt8, origin] ) raises: self.encrypt_into(ciphertext, plaintext) diff --git a/src/thistle/chacha20poly1305.mojo b/src/thistle/chacha20poly1305.mojo index 7942a91..fbd2129 100644 --- a/src/thistle/chacha20poly1305.mojo +++ b/src/thistle/chacha20poly1305.mojo @@ -1,19 +1,21 @@ -""" -ChaCha20-Poly1305 and XChaCha20-Poly1305 AEAD RFC 8439 -""" +"""Implements ChaCha20-Poly1305 and XChaCha20-Poly1305 from RFC 8439.""" from std.memory import bitcast from std.memory.unsafe_pointer import Pointer from std.collections import InlineArray -from .chacha20 import ChaCha20, chacha20_block_core, simd_double_round, CHACHA_CONSTANTS, _chacha20_nonce_words +from .chacha20 import ( + ChaCha20, chacha20_block_core, simd_double_round, CHACHA_CONSTANTS, _chacha20_nonce_words +) from .poly1305 import Poly1305 +from .utils import volatile_wipe + def hchacha20( key: Span[UInt8, ...], input16: Span[UInt8, ...], - output: Span[mut=True, UInt8, ...], + output: Span[mut=True, UInt8, ...] ) raises: - if len(key) < 32 or len(input16) < 16: + 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") @@ -28,7 +30,10 @@ def hchacha20( comptime for _ in range(10): var dr = simd_double_round(row0, row1, row2, row3) - row0 = dr[0]; row1 = dr[1]; row2 = dr[2]; row3 = dr[3] + row0 = dr[0] + row1 = dr[1] + row2 = dr[2] + row3 = dr[3] 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) @@ -38,7 +43,7 @@ def _aead_tag( poly_key: Span[UInt8, ...], aad: Span[UInt8, ...], ciphertext: Span[UInt8, ...], - output: Pointer[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) @@ -64,7 +69,7 @@ def _aead_core[encrypt: Bool]( aad: Span[UInt8, ...], input: Span[UInt8, ...], output: Pointer[mut=True, UInt8, _, address_space=_], - tag: Pointer[mut=True, UInt8, _, address_space=_], + tag: Pointer[mut=True, UInt8, _, address_space=_] ) raises: var key_bytes = key.unsafe_ptr().unsafe_load[width=32, alignment=1](0) var nonce_bytes = InlineArray[UInt8, 12](fill=0) @@ -81,20 +86,35 @@ def _aead_core[encrypt: Bool]( ) var poly_key_span = Span[UInt8, ...](unsafe_ptr=poly_key.unsafe_ptr(), length=32) - 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)) + try: + 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)) - comptime if encrypt: - _aead_tag( - poly_key_span, aad, - Span[UInt8, ...](unsafe_ptr=output, length=len(input)), tag, + comptime if encrypt: + _aead_tag( + poly_key_span, + aad, + Span[UInt8, ...](unsafe_ptr=output, length=len(input)), + tag, + ) + else: + _aead_tag(poly_key_span, aad, input, tag) + finally: + volatile_wipe(poly_key.unsafe_ptr(), 32) + volatile_wipe( + Pointer(to=key_bytes).unsafe_mut_cast[True]().unsafe_bitcast[UInt8](), + 32 + ) + volatile_wipe(Pointer(to=kw).unsafe_mut_cast[True]().unsafe_bitcast[UInt8](), 32) + volatile_wipe( + Pointer(to=block0).unsafe_mut_cast[True]().unsafe_bitcast[UInt8](), + 64 ) - else: - _aead_tag(poly_key_span, aad, input, tag) - var poly_key_ptr = poly_key.unsafe_ptr() - for i in range(32): - poly_key_ptr.unsafe_store[volatile=True](i, UInt8(0)) def chacha20_poly1305_encrypt( @@ -103,7 +123,7 @@ def chacha20_poly1305_encrypt( aad: Span[UInt8, ...], plaintext: Span[UInt8, ...], ciphertext: Span[mut=True, UInt8, ...], - tag: Span[mut=True, UInt8, ...], + tag: Span[mut=True, UInt8, ...] ) raises: if len(key) != 32: raise Error("ChaCha20-Poly1305 key must be 32 bytes") @@ -113,9 +133,7 @@ def chacha20_poly1305_encrypt( 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() - ) + _aead_core[True](key, nonce, aad, plaintext, ciphertext.unsafe_ptr(), tag.unsafe_ptr()) def chacha20_poly1305_decrypt( @@ -124,7 +142,7 @@ def chacha20_poly1305_decrypt( aad: Span[UInt8, ...], ciphertext: Span[UInt8, ...], tag: Span[UInt8, ...], - plaintext: Span[mut=True, UInt8, ...], + plaintext: Span[mut=True, UInt8, ...] ) raises -> Bool: if len(key) != 32: raise Error("ChaCha20-Poly1305 key must be 32 bytes") @@ -144,31 +162,38 @@ def chacha20_poly1305_decrypt( var nw = _chacha20_nonce_words(nonce_span) var block0 = chacha20_block_core(kw, 0, nw) var poly_key = InlineArray[UInt8, 32](fill=0) - poly_key.unsafe_ptr().unsafe_store[alignment=1]( - 0, bitcast[DType.uint8, 64](block0).slice[32]() - ) + poly_key.unsafe_ptr().unsafe_store[alignment=1](0, bitcast[DType.uint8, 64](block0).slice[32]()) var expected = InlineArray[UInt8, 16](fill=0) - _aead_tag( - Span[UInt8, ...](unsafe_ptr=poly_key.unsafe_ptr(), length=32), - aad, ciphertext, expected.unsafe_ptr(), - ) - var diff: UInt8 = 0 - for i in range(16): - diff |= expected[i] ^ tag[i] - if diff != 0: - var poly_key_ptr = poly_key.unsafe_ptr() - for i in range(32): - poly_key_ptr.unsafe_store[volatile=True](i, UInt8(0)) - return False + try: + _aead_tag( + Span[UInt8, ...](unsafe_ptr=poly_key.unsafe_ptr(), length=32), + aad, + ciphertext, + expected.unsafe_ptr() + ) + var diff: UInt8 = 0 + for i in range(16): + diff |= expected[i] ^ tag[i] + if diff != 0: + return False - 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.unsafe_store[volatile=True](i, UInt8(0)) - return True + 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)) + return True + finally: + volatile_wipe(poly_key.unsafe_ptr(), 32) + volatile_wipe(expected.unsafe_ptr(), 16) + volatile_wipe( + Pointer(to=key_bytes).unsafe_mut_cast[True]().unsafe_bitcast[UInt8](), + 32 + ) + volatile_wipe(Pointer(to=kw).unsafe_mut_cast[True]().unsafe_bitcast[UInt8](), 32) + volatile_wipe( + Pointer(to=block0).unsafe_mut_cast[True]().unsafe_bitcast[UInt8](), + 64 + ) def xchacha20_poly1305_encrypt( @@ -177,7 +202,7 @@ def xchacha20_poly1305_encrypt( aad: Span[UInt8, ...], plaintext: Span[UInt8, ...], ciphertext: Span[mut=True, UInt8, ...], - tag: Span[mut=True, UInt8, ...], + tag: Span[mut=True, UInt8, ...] ) raises: if len(key) != 32: raise Error("XChaCha20-Poly1305 key must be 32 bytes") @@ -193,7 +218,10 @@ def xchacha20_poly1305_encrypt( chacha20_poly1305_encrypt( Span[UInt8, ...](unsafe_ptr=sp, length=32), Span[UInt8, ...](unsafe_ptr=sp.unsafe_offset(32), length=12), - aad, plaintext, ciphertext, tag, + aad, + plaintext, + ciphertext, + tag ) finally: for i in range(44): @@ -206,7 +234,7 @@ def xchacha20_poly1305_decrypt( aad: Span[UInt8, ...], ciphertext: Span[UInt8, ...], tag: Span[UInt8, ...], - plaintext: Span[mut=True, UInt8, ...], + plaintext: Span[mut=True, UInt8, ...] ) raises -> Bool: if len(key) != 32: raise Error("XChaCha20-Poly1305 key must be 32 bytes") @@ -221,7 +249,10 @@ def xchacha20_poly1305_decrypt( ok = chacha20_poly1305_decrypt( Span[UInt8, ...](unsafe_ptr=sp, length=32), Span[UInt8, ...](unsafe_ptr=sp.unsafe_offset(32), length=12), - aad, ciphertext, tag, plaintext, + aad, + ciphertext, + tag, + plaintext ) finally: for i in range(44): @@ -229,12 +260,14 @@ def xchacha20_poly1305_decrypt( return ok -def _xchacha_subkey_nonce(key: Span[UInt8, ...], nonce: Span[UInt8, ...]) raises -> InlineArray[UInt8, 44]: +def _xchacha_subkey_nonce( + key: Span[UInt8, ...], nonce: Span[UInt8, ...] +) raises -> InlineArray[UInt8, 44]: var out = InlineArray[UInt8, 44](fill=0) hchacha20( key, Span[UInt8, ...](unsafe_ptr=nonce.unsafe_ptr(), length=16), - Span[mut=True, UInt8, ...](unsafe_ptr=out.unsafe_ptr(), length=32), + Span[mut=True, UInt8, ...](unsafe_ptr=out.unsafe_ptr(), length=32) ) for i in range(8): out[36 + i] = nonce[16 + i] diff --git a/src/thistle/curve25519.mojo b/src/thistle/curve25519.mojo index 671eef4..329ed02 100644 --- a/src/thistle/curve25519.mojo +++ b/src/thistle/curve25519.mojo @@ -1,5 +1,8 @@ +"""Provides finite-field operations used by Curve25519 implementations.""" + from std.builtin.dtype import DType + @always_inline def _u128_shr[shift: Int](x: UInt128) -> UInt128: comptime if shift <= 0: @@ -20,7 +23,8 @@ def _u128_shr[shift: Int](x: UInt128) -> UInt128: else: return UInt128(hi >> UInt64(shift - 64)) -struct FieldElement51(Movable, Copyable, ImplicitlyCopyable): + +struct FieldElement51(Copyable, ImplicitlyCopyable, Movable): var limbs: SIMD[DType.uint64, 8] @always_inline @@ -58,7 +62,7 @@ struct FieldElement51(Movable, Copyable, ImplicitlyCopyable): 2251799813685247, 2251799813685247, 2251799813685247, - 2251799813685247, + 2251799813685247 ) @always_inline @@ -75,10 +79,14 @@ struct FieldElement51(Movable, Copyable, ImplicitlyCopyable): limbs[0] += 19 * q var MASK = UInt64(0x7FFFFFFFFFFFF) - limbs[1] += limbs[0] >> 51; limbs[0] &= MASK - limbs[2] += limbs[1] >> 51; limbs[1] &= MASK - limbs[3] += limbs[2] >> 51; limbs[2] &= MASK - limbs[4] += limbs[3] >> 51; limbs[3] &= MASK + limbs[1] += limbs[0] >> 51 + limbs[0] &= MASK + limbs[2] += limbs[1] >> 51 + limbs[1] &= MASK + limbs[3] += limbs[2] >> 51 + limbs[2] &= MASK + limbs[4] += limbs[3] >> 51 + limbs[3] &= MASK limbs[4] &= MASK var value = Int(limbs[0]) @@ -121,7 +129,7 @@ struct FieldElement51(Movable, Copyable, ImplicitlyCopyable): self.limbs[1] + other.limbs[1], self.limbs[2] + other.limbs[2], self.limbs[3] + other.limbs[3], - self.limbs[4] + other.limbs[4], + self.limbs[4] + other.limbs[4] ) @always_inline @@ -133,11 +141,16 @@ struct FieldElement51(Movable, Copyable, ImplicitlyCopyable): l[3] = (self.limbs[3] + 0x7FFFFFFFFFFFF0) - other.limbs[3] l[4] = (self.limbs[4] + 0x7FFFFFFFFFFFF0) - other.limbs[4] var MASK = UInt64(0x7FFFFFFFFFFFF) - l[1] += l[0] >> 51; l[0] &= MASK - l[2] += l[1] >> 51; l[1] &= MASK - l[3] += l[2] >> 51; l[2] &= MASK - l[4] += l[3] >> 51; l[3] &= MASK - l[0] += (l[4] >> 51) * 19; l[4] &= MASK + l[1] += l[0] >> 51 + l[0] &= MASK + l[2] += l[1] >> 51 + l[1] &= MASK + l[3] += l[2] >> 51 + l[2] &= MASK + l[4] += l[3] >> 51 + l[3] &= MASK + l[0] += (l[4] >> 51) * 19 + l[4] &= MASK return FieldElement51(l) @always_inline @@ -251,7 +264,8 @@ struct FieldElement51(Movable, Copyable, ImplicitlyCopyable): return res @always_inline - def _carry_reduce(self, var c0: UInt128, var c1: UInt128, var c2: UInt128, var c3: UInt128, var c4: UInt128) -> FieldElement51: + def _carry_reduce(self, var c0: UInt128, var c1: UInt128, var c2: UInt128, var c3: UInt128, var c4: UInt128 + ) -> FieldElement51: var MASK = UInt64(0x7FFFFFFFFFFFF) c1 += _u128_shr[51](c0) @@ -277,17 +291,22 @@ struct FieldElement51(Movable, Copyable, ImplicitlyCopyable): var l = limbs var MASK = UInt64(0x7FFFFFFFFFFFF) for _ in range(5): - l[1] += l[0] >> 51; l[0] &= MASK - l[2] += l[1] >> 51; l[1] &= MASK - l[3] += l[2] >> 51; l[2] &= MASK - l[4] += l[3] >> 51; l[3] &= MASK - l[0] += (l[4] >> 51) * 19; l[4] &= MASK + l[1] += l[0] >> 51 + l[0] &= MASK + l[2] += l[1] >> 51 + l[1] &= MASK + l[3] += l[2] >> 51 + l[2] &= MASK + l[4] += l[3] >> 51 + l[3] &= MASK + l[0] += (l[4] >> 51) * 19 + l[4] &= MASK return FieldElement51(l) @staticmethod def from_bytes_span(bytes: Span[UInt8, ...]) raises -> FieldElement51: - if len(bytes) < 32: - raise Error("FieldElement51 input must be at least 32 bytes") + if len(bytes) != 32: + raise Error("FieldElement51 input must be exactly 32 bytes") @always_inline def load8(ptr: Pointer[mut=False, UInt8, _, address_space=_]) -> UInt64: return ptr.unsafe_bitcast[UInt64]().unsafe_load[width=1, alignment=1]() @@ -316,10 +335,14 @@ struct FieldElement51(Movable, Copyable, ImplicitlyCopyable): limbs[0] += 19 * q var MASK = UInt64(0x7FFFFFFFFFFFF) - limbs[1] += limbs[0] >> 51; limbs[0] &= MASK - limbs[2] += limbs[1] >> 51; limbs[1] &= MASK - limbs[3] += limbs[2] >> 51; limbs[2] &= MASK - limbs[4] += limbs[3] >> 51; limbs[3] &= MASK + limbs[1] += limbs[0] >> 51 + limbs[0] &= MASK + limbs[2] += limbs[1] >> 51 + limbs[1] &= MASK + limbs[3] += limbs[2] >> 51 + limbs[2] &= MASK + limbs[4] += limbs[3] >> 51 + limbs[3] &= MASK limbs[4] &= MASK var w0 = limbs[0] | (limbs[1] << 51) diff --git a/src/thistle/ecdsa_der.mojo b/src/thistle/ecdsa_der.mojo index 86bf798..740780f 100644 --- a/src/thistle/ecdsa_der.mojo +++ b/src/thistle/ecdsa_der.mojo @@ -1,10 +1,14 @@ +"""Encodes and decodes ECDSA signatures using the DER format.""" + from std.collections import List def ecdsa_der_encode( signature: Span[UInt8, ...], size: Int ) raises -> List[UInt8]: - if size <= 0 or len(signature) != 2 * size: + # This codec intentionally supports only DER's one-byte length form. A + # 60-byte scalar is the largest one whose two padded INTEGERs can fit. + if size <= 0 or size > 60 or len(signature) != 2 * size: raise Error("invalid ECDSA signature size") var r_start = 0 while r_start < size - 1 and signature[r_start] == 0: @@ -39,7 +43,7 @@ def ecdsa_der_encode( def ecdsa_der_decode(signature: Span[UInt8, ...], size: Int) -> List[UInt8]: - if size <= 0 or len(signature) < 6 or signature[0] != 0x30: + if size <= 0 or size > 60 or len(signature) < 6 or signature[0] != 0x30: return List[UInt8]() if signature[1] >= 0x80 or Int(signature[1]) != len(signature) - 2: return List[UInt8]() diff --git a/src/thistle/ed25519.mojo b/src/thistle/ed25519.mojo index b93e022..5e1b7c4 100644 --- a/src/thistle/ed25519.mojo +++ b/src/thistle/ed25519.mojo @@ -1,6 +1,4 @@ -""" -Ed25519 implementation -""" +"""Implements Ed25519 signing and verification.""" from std.builtin.dtype import DType from std.builtin.simd import SIMD from std.memory import bitcast @@ -15,7 +13,7 @@ comptime L_LIMBS = SIMD[DType.uint64, 8]( 0x000000000014def9, 0x0000000000000000, 0x0000100000000000, - 0, 0, 0, + 0, 0, 0 ) comptime LFACTOR: UInt64 = 0x51da312547e1b @@ -25,7 +23,7 @@ comptime RR_LIMBS = SIMD[DType.uint64, 8]( 0x0005be65cb687604, 0x0003dceec73d217f, 0x000009411b7c309a, - 0, 0, 0, + 0, 0, 0 ) comptime ED25519_D_LIMBS = SIMD[DType.uint64, 8]( @@ -34,7 +32,7 @@ comptime ED25519_D_LIMBS = SIMD[DType.uint64, 8]( 0x005e7a26001c029, 0x00739c663a03cbb, 0x0052036cee2b6ff, - 0, 0, 0, + 0, 0, 0 ) comptime POW2_256_LIMBS = SIMD[DType.uint64, 8]( @@ -43,23 +41,24 @@ comptime POW2_256_LIMBS = SIMD[DType.uint64, 8]( 0x000f5be65cc244cc, 0x000a3dceec73d217, 0x0000099411b7c309, - 0, 0, 0, + 0, 0, 0 ) comptime L_BYTES = SIMD[DType.uint8, 32]( 0xed, 0xd3, 0xf5, 0x5c, 0x1a, 0x63, 0x12, 0x58, 0xd6, 0x9c, 0xf7, 0xa2, 0xde, 0xf9, 0xde, 0x14, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, - 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x10 ) comptime FIELD_P_BYTES = SIMD[DType.uint8, 32]( 0xed, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, - 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x7f, + 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x7f ) + @always_inline def _s_lt_l(s: Span[UInt8, ...]) -> Bool: # RFC 8032 5.1.7 / 8.4: verification requires S < L; this prevents @@ -72,6 +71,7 @@ def _s_lt_l(s: Span[UInt8, ...]) -> Bool: return False return False + @always_inline def _encoded_y_lt_p(y: Span[UInt8, ...]) -> Bool: # RFC 8032 5.1.3: strict decoding rejects y >= p. @@ -87,7 +87,9 @@ def _encoded_y_lt_p(y: Span[UInt8, ...]) -> Bool: gt = 1 return lt == 1 -def _pack_limbs_into(limbs: SIMD[DType.uint64, 8], output: Pointer[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) @@ -97,6 +99,7 @@ def _pack_limbs_into(limbs: SIMD[DType.uint64, 8], output: Pointer[mut=True, UIn for i in range(32): 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().unsafe_bitcast[UInt64]().unsafe_load[width=4, alignment=1]() @@ -110,6 +113,7 @@ def _unpack_limbs(bytes: Span[UInt8, ...]) -> SIMD[DType.uint64, 8]: s[4] = (words[3] >> UInt64(16)) & TOP_MASK return s + 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() @@ -121,27 +125,30 @@ def _from_512_raw(bytes: Span[UInt8, ...]) -> SIMD[DType.uint64, 8]: var res = lo + hi * pow2_256 return res.limbs + @always_inline def ed25519_d() -> FieldElement51: return FieldElement51(ED25519_D_LIMBS) + @always_inline def ed25519_d2() -> FieldElement51: return ed25519_d() * FieldElement51(2, 0, 0, 0, 0) + def ed25519_base_point() -> EdwardsPoint: var X = FieldElement51( 1738742601995546, 1146398526822698, 2070867633025821, - 562264141797630, 587772402128613, + 562264141797630, 587772402128613 ) var Y = FieldElement51( 1801439850948184, 1351079888211148, 450359962737049, - 900719925474099, 1801439850948198, + 900719925474099, 1801439850948198 ) return EdwardsPoint(X, Y, FieldElement51.ONE(), X * Y) -struct Scalar(Movable, Copyable, ImplicitlyCopyable): +struct Scalar(Copyable, ImplicitlyCopyable, Movable): var limbs: SIMD[DType.uint64, 8] @always_inline @@ -262,7 +269,7 @@ struct Scalar(Movable, Copyable, ImplicitlyCopyable): return Scalar(diff) -struct EdwardsPoint(Movable, Copyable, ImplicitlyCopyable): +struct EdwardsPoint(Copyable, ImplicitlyCopyable, Movable): var X: FieldElement51 var Y: FieldElement51 var Z: FieldElement51 @@ -276,18 +283,29 @@ struct EdwardsPoint(Movable, Copyable, ImplicitlyCopyable): self.T = FieldElement51.ZERO() @always_inline - def __init__(out self, X: FieldElement51, Y: FieldElement51, Z: FieldElement51, T: FieldElement51): - self.X = X; self.Y = Y; self.Z = Z; self.T = T + def __init__(out self, X: FieldElement51, Y: FieldElement51, Z: FieldElement51, T: FieldElement51 + ): + self.X = X + self.Y = Y + self.Z = Z + self.T = T @always_inline def __copyinit__(out self, copy: Self): - self.X = copy.X; self.Y = copy.Y; self.Z = copy.Z; self.T = copy.T + self.X = copy.X + self.Y = copy.Y + self.Z = copy.Z + self.T = copy.T @always_inline def __moveinit__(out self, deinit take: Self): - self.X = take.X^; self.Y = take.Y^; self.Z = take.Z^; self.T = take.T^ + self.X = take.X^ + self.Y = take.Y^ + self.Z = take.Z^ + self.T = take.T^ + -struct DecodeResult(Movable, Copyable, ImplicitlyCopyable): +struct DecodeResult(Copyable, ImplicitlyCopyable, Movable): var ok: Bool var p: EdwardsPoint @@ -301,16 +319,20 @@ struct DecodeResult(Movable, Copyable, ImplicitlyCopyable): self.ok = ok self.p = p + def edwards_add(p: EdwardsPoint, q: EdwardsPoint) -> EdwardsPoint: var d2 = ed25519_d2() return _edwards_add_d2(p, q, d2) + def edwards_double(p: EdwardsPoint) -> EdwardsPoint: return _edwards_double_standalone(p) + def edwards_negate(p: EdwardsPoint) -> EdwardsPoint: return EdwardsPoint(FieldElement51.ZERO() - p.X, p.Y, p.Z, FieldElement51.ZERO() - p.T) + @always_inline def _ct_select_fe(a: FieldElement51, b: FieldElement51, choice: UInt8) -> FieldElement51: # constant-time select via mask @@ -320,6 +342,7 @@ def _ct_select_fe(a: FieldElement51, b: FieldElement51, choice: UInt8) -> FieldE limbs[i] = a.limbs[i] ^ (mask & (a.limbs[i] ^ b.limbs[i])) return FieldElement51(limbs) + @no_inline def _edwards_add_d2(p: EdwardsPoint, q: EdwardsPoint, d2: FieldElement51) -> EdwardsPoint: # RFC 8032 5.1.4: complete extended Edwards addition, a = -1. @@ -334,6 +357,7 @@ def _edwards_add_d2(p: EdwardsPoint, q: EdwardsPoint, d2: FieldElement51) -> Edw var H = B + A return EdwardsPoint(E * F, G * H, F * G, E * H) + @no_inline def _edwards_double_standalone(p: EdwardsPoint) -> EdwardsPoint: # RFC 8032 5.1.4: extended Edwards doubling. @@ -366,13 +390,16 @@ def fe_from_bytes(bytes: Span[UInt8, ...]) -> FieldElement51: 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: 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: Pointer[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) @@ -381,6 +408,7 @@ def _edwards_encode_with_zinv(p: EdwardsPoint, z_inv: FieldElement51, output: Po var x_parity = x_bytes[0] & 1 output[unsafe_offset=31] = output[unsafe_offset=31] | (x_parity << 7) + @no_inline def edwards_decode(data: Span[UInt8, ...], strict: Bool = True) -> DecodeResult: # RFC 8032 5.1.3: strict point decoding. @@ -433,6 +461,7 @@ def edwards_decode(data: Span[UInt8, ...], strict: Bool = True) -> DecodeResult: return DecodeResult(False, EdwardsPoint()) + def edwards_decode_verify_compatible(data: Span[UInt8, ...]) -> DecodeResult: # strict RFC decoding only, no ZIP-215 return edwards_decode(data, strict=True) @@ -451,6 +480,7 @@ def _is_small_order(p: EdwardsPoint) -> Bool: diff |= x[i] | yz[i] return diff == 0 + @no_inline def sqrt_ratio_checked(u: FieldElement51, v: FieldElement51) -> Optional[FieldElement51]: # RFC 8032 5.1.3: compute sqrt(u/v) using @@ -482,17 +512,20 @@ def sqrt_ratio_checked(u: FieldElement51, v: FieldElement51) -> Optional[FieldEl if is_zero2: var sqrtm1 = FieldElement51( 1718705420411056, 234908883556509, - 2233514472574048, 2117202627021982, 765476049583133) + 2233514472574048, 2117202627021982, 765476049583133 + ) return Optional[FieldElement51](x * sqrtm1) return None -struct AffineNielsPoint(Movable, Copyable, ImplicitlyCopyable): + +struct AffineNielsPoint(Copyable, ImplicitlyCopyable, Movable): var y_plus_x: FieldElement51 var y_minus_x: FieldElement51 var xy2d: FieldElement51 @always_inline - def __init__(out self, y_plus_x: FieldElement51, y_minus_x: FieldElement51, xy2d: FieldElement51): + def __init__(out self, y_plus_x: FieldElement51, y_minus_x: FieldElement51, xy2d: FieldElement51 + ): self.y_plus_x = y_plus_x self.y_minus_x = y_minus_x self.xy2d = xy2d @@ -510,7 +543,7 @@ struct AffineNielsPoint(Movable, Copyable, ImplicitlyCopyable): self.xy2d = take.xy2d^ -struct ProjectiveNielsPoint(Movable, Copyable, ImplicitlyCopyable): +struct ProjectiveNielsPoint(Copyable, ImplicitlyCopyable, Movable): var Y_plus_X: FieldElement51 var Y_minus_X: FieldElement51 var Z: FieldElement51 @@ -524,7 +557,8 @@ struct ProjectiveNielsPoint(Movable, Copyable, ImplicitlyCopyable): self.T2d = FieldElement51() @always_inline - def __init__(out self, Y_plus_X: FieldElement51, Y_minus_X: FieldElement51, Z: FieldElement51, T2d: FieldElement51): + def __init__(out self, Y_plus_X: FieldElement51, Y_minus_X: FieldElement51, Z: FieldElement51, T2d: FieldElement51 + ): self.Y_plus_X = Y_plus_X self.Y_minus_X = Y_minus_X self.Z = Z @@ -549,6 +583,7 @@ struct ProjectiveNielsPoint(Movable, Copyable, ImplicitlyCopyable): def _to_projective_niels(p: EdwardsPoint, d2: FieldElement51) -> ProjectiveNielsPoint: return ProjectiveNielsPoint(p.Y + p.X, p.Y - p.X, p.Z, p.T * d2) + @always_inline def _add_affine_niels(p: EdwardsPoint, q: AffineNielsPoint) -> EdwardsPoint: var PP = (p.Y + p.X) * q.y_plus_x @@ -561,6 +596,7 @@ def _add_affine_niels(p: EdwardsPoint, q: AffineNielsPoint) -> EdwardsPoint: var T3 = Z2 - Txy return EdwardsPoint(X3 * T3, Y3 * Z3, Z3 * T3, X3 * Y3) + @always_inline def _sub_affine_niels(p: EdwardsPoint, q: AffineNielsPoint) -> EdwardsPoint: var PM = (p.Y + p.X) * q.y_minus_x @@ -573,6 +609,7 @@ def _sub_affine_niels(p: EdwardsPoint, q: AffineNielsPoint) -> EdwardsPoint: var T3 = Z2 + Txy return EdwardsPoint(X3 * T3, Y3 * Z3, Z3 * T3, X3 * Y3) + @always_inline def _add_projective_niels(p: EdwardsPoint, n: ProjectiveNielsPoint) -> EdwardsPoint: var PP = (p.Y + p.X) * n.Y_plus_X @@ -586,6 +623,7 @@ def _add_projective_niels(p: EdwardsPoint, n: ProjectiveNielsPoint) -> EdwardsPo var T3 = ZZ2 - TT2d return EdwardsPoint(X3 * T3, Y3 * Z3, Z3 * T3, X3 * Y3) + @always_inline def _sub_projective_niels(p: EdwardsPoint, n: ProjectiveNielsPoint) -> EdwardsPoint: var PM = (p.Y + p.X) * n.Y_minus_X @@ -599,6 +637,7 @@ def _sub_projective_niels(p: EdwardsPoint, n: ProjectiveNielsPoint) -> EdwardsPo var T3 = ZZ2 + TT2d return EdwardsPoint(X3 * T3, Y3 * Z3, Z3 * T3, X3 * Y3) + @always_inline def _radix16_digits(scalar: Span[UInt8, ...]) -> InlineArray[Int, 64]: var digits = InlineArray[Int, 64](fill=0) @@ -611,6 +650,7 @@ def _radix16_digits(scalar: Span[UInt8, ...]) -> InlineArray[Int, 64]: digits[i + 1] += carry return digits^ + @always_inline def _base_table_lookup(ptr: Pointer[UInt64, _], j: Int, digit: Int) -> AffineNielsPoint: var d = Int64(digit) @@ -636,6 +676,7 @@ def _base_table_lookup(ptr: Pointer[UInt64, _], j: Int, digit: Int) -> AffineNie var xy_sel = _ct_select_fe(xy, xy_neg, UInt8(sm & 1)) return AffineNielsPoint(FieldElement51(yp), FieldElement51(ym), xy_sel) + @no_inline def _mul_base_ct(scalar: Span[UInt8, ...]) -> EdwardsPoint: var table = ed25519_base_table() @@ -655,13 +696,16 @@ def _mul_base_ct(scalar: Span[UInt8, ...]) -> EdwardsPoint: dptr.unsafe_store[volatile=True](i, UInt64(0)) return P + def _naf5(scalar: Span[UInt8, ...]) -> InlineArray[Int, 256]: var naf = InlineArray[Int, 256](fill=0) var words = InlineArray[UInt64, 5](fill=0) words[4] = 0 var ptr = scalar.unsafe_ptr() for w in range(4): - words[w] = (ptr.unsafe_offset(8 * w)).unsafe_bitcast[UInt64]().unsafe_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: @@ -683,15 +727,19 @@ def _naf5(scalar: Span[UInt8, ...]) -> InlineArray[Int, 256]: pos += 5 return naf^ + @always_inline def _b_odd_entry(ptr: Pointer[UInt64, _], k: Int) -> AffineNielsPoint: var base = ptr.unsafe_offset(k * 16) return AffineNielsPoint( - 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]), + 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 def _double_scalar_mult_vartime(a: Span[UInt8, ...], A: EdwardsPoint, b: Span[UInt8, ...]) -> EdwardsPoint: var naf_a = _naf5(a) @@ -724,6 +772,7 @@ def _double_scalar_mult_vartime(a: Span[UInt8, ...], A: EdwardsPoint, b: Span[UI Q = _sub_affine_niels(Q, _b_odd_entry(bptr, (-db - 1) >> 1)) return Q + @no_inline def ed25519_generate_public_key( private_key: Span[UInt8, ...], output: Span[mut=True, UInt8, ...] @@ -752,11 +801,12 @@ def ed25519_generate_public_key( for i in range(32): s_ptr.unsafe_store[volatile=True](i, UInt8(0)) + @no_inline def ed25519_sign( private_key: Span[UInt8, ...], message: Span[UInt8, ...], - output: Span[mut=True, UInt8, ...], + output: Span[mut=True, UInt8, ...] ) raises: # RFC 8032 5.1.6 pure Ed25519: # r = SHA512(prefix || M), R = [r]B, @@ -826,6 +876,7 @@ def ed25519_sign( 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 var _prefix: InlineArray[UInt8, 32] @@ -923,7 +974,8 @@ struct Ed25519SigningKey(Copyable, Movable): @no_inline -def ed25519_verify(public_key: Span[UInt8, ...], message: Span[UInt8, ...], signature: Span[UInt8, ...]) -> Bool: +def ed25519_verify(public_key: Span[UInt8, ...], message: Span[UInt8, ...], signature: Span[UInt8, ...] +) -> Bool: # Uses canonical decoding, S < L, and the uncofactored equation. # Low-order public keys are rejected by library policy. if len(public_key) != 32 or len(signature) != 64: diff --git a/src/thistle/fips.mojo b/src/thistle/fips.mojo index fa600f6..f16b4cc 100644 --- a/src/thistle/fips.mojo +++ b/src/thistle/fips.mojo @@ -1,18 +1,20 @@ +"""Re-exports the FIPS-approved hash and password-based primitives.""" + from .sha2 import ( sha224_hash, sha256_hash, sha384_hash, sha512_hash, SHA256Context, - SHA512Context, + SHA512Context ) from .sha3 import ( sha3_224, sha3_256, sha3_384, - sha3_512, + sha3_512 ) from .pbkdf2 import ( pbkdf2_hmac_sha256, - pbkdf2_hmac_sha512, + pbkdf2_hmac_sha512 ) diff --git a/src/thistle/kcipher2.mojo b/src/thistle/kcipher2.mojo index 77a0330..5097d82 100644 --- a/src/thistle/kcipher2.mojo +++ b/src/thistle/kcipher2.mojo @@ -9,7 +9,7 @@ from .aes_ni import ( _aesmc, _mm_aesenc_si128, has_arm_crypto, - has_x86_aes_ni, + has_x86_aes_ni ) comptime AMUL_BASIS_0 = SIMD[DType.uint32, 4]( @@ -37,6 +37,7 @@ comptime AMUL_BASIS_7 = SIMD[DType.uint32, 4]( 0x4B8AF89E, 0xF85E6A49, 0x59036A01, 0x9C20FEDD ) + @always_inline def _amul4(b: SIMD[DType.uint32, 4]) -> SIMD[DType.uint32, 4]: var zero = SIMD[DType.uint32, 4](0) @@ -206,6 +207,8 @@ def _sbox_planes(mut p: InlineArray[UInt32, 8]): # idx[4*c + r] = 4*((c - r) mod 4) + r. + + @always_inline def sub_k2_x4( w0: UInt32, w1: UInt32, w2: UInt32, w3: UInt32 @@ -284,7 +287,7 @@ def _sub_k2_x4_bitsliced( UInt32(olo & 0xFFFFFFFF), UInt32(olo >> 32), UInt32(ohi & 0xFFFFFFFF), - UInt32(ohi >> 32), + UInt32(ohi >> 32) ) @@ -384,20 +387,16 @@ struct KCipher2: ik[2] = key[2] ik[3] = key[3] - ik[4] = ( - ik[0] + ik[4] = ik[0] ^ sub_k2(((ik[3] << 8) & 0xFFFFFFFF) ^ (ik[3] >> 24)) ^ 0x01000000 - ) ik[5] = ik[1] ^ ik[4] ik[6] = ik[2] ^ ik[5] ik[7] = ik[3] ^ ik[6] - ik[8] = ( - ik[4] + ik[8] = ik[4] ^ sub_k2(((ik[7] << 8) & 0xFFFFFFFF) ^ (ik[7] >> 24)) ^ 0x02000000 - ) ik[9] = ik[5] ^ ik[8] ik[10] = ik[6] ^ ik[9] ik[11] = ik[7] ^ ik[10] @@ -474,13 +473,11 @@ struct KCipher2: var temp2_amul3 = ((self.b8 << 8) & 0xFFFFFF00) ^ am[3] var temp2 = self._select_u32(self.b8, temp2_amul3, (old_a2 >> 31) & 1) - var new_b10 = ( - temp1 + var new_b10 = temp1 ^ self.b1 ^ self.b6 ^ temp2 ^ nlf(self.b10, self.l2, self.l1, old_a0) - ) self.b0 = self.b1 self.b1 = self.b2 @@ -611,7 +608,7 @@ struct KCipher2: UInt8(z >> 32), UInt8(z >> 40), UInt8(z >> 48), - UInt8(z >> 56), + UInt8(z >> 56) ) for j in range(len_data - offset): diff --git a/src/thistle/ml_dsa.mojo b/src/thistle/ml_dsa.mojo index 70cda66..653c5f1 100644 --- a/src/thistle/ml_dsa.mojo +++ b/src/thistle/ml_dsa.mojo @@ -13,7 +13,10 @@ 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 unsafe_memset_zero -from thistle.sha3 import SHA3Context, sha3_update, shake_final, shake128, shake256 +from std.os import abort +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 @@ -75,7 +78,7 @@ comptime ZETAS_TABLE: InlineArray[UInt32, 256] = [ 5341501, 3523897, 3866901, 269760, 2213111, 7404533, 1717735, 472078, 7953734, 1723600, 6577327, 1910376, 6712985, 7276084, 8119771, 4546524, 5441381, 6144432, 7959518, 6094090, 183443, 7403526, 1612842, 4834730, - 7826001, 3919660, 8332111, 7018208, 3937738, 1400424, 7534263, 1976782, + 7826001, 3919660, 8332111, 7018208, 3937738, 1400424, 7534263, 1976782 ] @@ -166,21 +169,116 @@ def params87() -> MLDSAParams: return MLDSAParams(8, 7, 2, 19, 32, 256, 60, 75) +@always_inline +def _valid_params(p: MLDSAParams) -> Bool: + return ( + ( + p.k == 4 + and p.l == 4 + and p.eta == 2 + and p.gamma1_log == 17 + and p.gamma2_denom == 88 + and p.lambda_bits == 128 + and p.tau == 39 + and p.omega == 80 + ) + or ( + p.k == 6 + and p.l == 5 + and p.eta == 4 + and p.gamma1_log == 19 + and p.gamma2_denom == 32 + and p.lambda_bits == 192 + and p.tau == 49 + and p.omega == 55 + ) + or ( + p.k == 8 + and p.l == 7 + and p.eta == 2 + and p.gamma1_log == 19 + and p.gamma2_denom == 32 + and p.lambda_bits == 256 + and p.tau == 60 + and p.omega == 75 + ) + ) + + +@always_inline +def _same_params(a: MLDSAParams, b: MLDSAParams) -> Bool: + return ( + a.k == b.k + and a.l == b.l + and a.eta == b.eta + and a.gamma1_log == b.gamma1_log + and a.gamma2_denom == b.gamma2_denom + and a.lambda_bits == b.lambda_bits + and a.tau == b.tau + and a.omega == b.omega + ) + + +def _require_valid_params(p: MLDSAParams) raises: + if not _valid_params(p): + raise Error("unsupported ML-DSA parameter set") + + def public_key_size(p: MLDSAParams) -> Int: + if not _valid_params(p): + return 0 return 32 + p.k * N * 10 // 8 def signature_size(p: MLDSAParams) -> Int: + if not _valid_params(p): + return 0 return p.lambda_bits // 4 + p.l * N * (p.gamma1_log + 1) // 8 + p.omega + p.k def private_key_size(p: MLDSAParams) -> Int: + if not _valid_params(p): + return 0 var eta_bitlen = 3 if p.eta == 4: eta_bitlen = 4 return 32 + 32 + 64 + p.l * N * eta_bitlen // 8 + p.k * N * eta_bitlen // 8 + p.k * N * 13 // 8 +def _u32_rows_valid(rows: List[List[UInt32]], expected: Int) -> Bool: + if len(rows) != expected: + return False + for i in range(expected): + if len(rows[i]) != N: + return False + return True + + +def _public_key_shape_valid(pub: MLDSAPublicKey) -> Bool: + var p = pub.p.copy() + return ( + _valid_params(p) + and len(pub.raw) == public_key_size(p) + and len(pub.tr) == MLDSA_CRHBYTES + and _u32_rows_valid(pub.a, p.k * p.l) + and _u32_rows_valid(pub.t1_hat, p.k) + ) + + +def _private_key_shape_valid(priv: MLDSAPrivateKey) -> Bool: + var p = priv.p.copy() + return ( + _valid_params(p) + and _same_params(p, priv.pub.p) + and _public_key_shape_valid(priv.pub) + and (len(priv.seed) == 0 or len(priv.seed) == MLDSA_SEEDBYTES) + and len(priv.k_seed) == MLDSA_SEEDBYTES + and _u32_rows_valid(priv.s1, p.l) + and _u32_rows_valid(priv.s2, p.k) + and _u32_rows_valid(priv.t0, p.k) + ) + + def _zero_poly() -> List[UInt32]: var r = List[UInt32](unsafe_uninit_length=N) unsafe_memset_zero(r.unsafe_ptr(), N) @@ -243,6 +341,8 @@ def _append_bytes(mut out: List[UInt8], src: Span[UInt8, ...]): def _append_bytes_stack(mut out: StackBuffer[UInt8, ...], src: Span[UInt8, ...]): + if len(src) > out.remaining(): + abort("ML-DSA stack buffer overflow") for i in range(len(src)): out.push_unchecked(src[i]) @@ -252,6 +352,8 @@ def _ct_bool_to_u32(b: Bool) -> UInt32: return UInt32(Int(b)) # volatile stores so the wipe can't be optimized away + + def _zero_list_u8(mut data: List[UInt8]): var ptr = data.unsafe_ptr() for i in range(len(data)): @@ -408,7 +510,7 @@ def _ntt_mul_into(mut r: List[UInt32], a: List[UInt32], b: List[UInt32]): def _ntt_mul_ptrs( r: Pointer[mut=True, UInt32, _, address_space=_], a: Pointer[mut=False, UInt32, _, address_space=_], - b: Pointer[mut=False, UInt32, _, address_space=_], + b: Pointer[mut=False, UInt32, _, address_space=_] ): var i = 0 while i < N: @@ -521,7 +623,7 @@ def _dsa_inverse_ntt_inplace(mut f: DSAPoly): j + length, _montgomery_reduce_v( zv * (b - a + _U32v(Q)).cast[DType.uint64]() - ), + ) ) j += _VW else: @@ -817,6 +919,7 @@ def _pk_encode(rho: Span[UInt8, ...], t1: List[List[UInt16]], p: MLDSAParams) -> def _pk_decode(pk: Span[UInt8, ...], p: MLDSAParams) raises -> Tuple[List[UInt8], List[List[UInt16]]]: + _require_valid_params(p) if len(pk) != public_key_size(p): raise Error("ML-DSA invalid public key length") var rho = _slice_bytes(pk, 0, 32) @@ -862,6 +965,7 @@ def _compute_t1_hat(t1: List[List[UInt16]], p: MLDSAParams) raises -> List[List[ def new_public_key(pk: Span[UInt8, ...], p: MLDSAParams) raises -> MLDSAPublicKey: + _require_valid_params(p) var decoded = _pk_decode(pk, p) var rho = decoded[0].copy() var t1 = decoded[1].copy() @@ -1029,7 +1133,8 @@ def _dsa_append_w1_encoded_stack(mut out: StackBuffer[UInt8, ...], w: DSAPoly, p out.push_unchecked((b2 >> 4) | (b3 << 2)) -def _append_use_hint_encoded_stack(mut out: StackBuffer[UInt8, ...], w: List[UInt32], h: List[UInt8], p: MLDSAParams): +def _append_use_hint_encoded_stack(mut out: StackBuffer[UInt8, ...], w: List[UInt32], h: List[UInt8], p: MLDSAParams +): if p.gamma2_denom == 32: for i in range(0, N, 2): var b0 = _use_hint_elem(w[i], h[i], p) @@ -1086,7 +1191,8 @@ def _hint_decode(y: Span[UInt8, ...], p: MLDSAParams) raises -> List[List[UInt8] return h^ -def _dsa_sig_encode(ch: InlineArray[UInt8, MLDSA_CRHBYTES], z: DSAPolyVec[MAX_L], h: DSAHintVec, p: MLDSAParams) -> List[UInt8]: +def _dsa_sig_encode(ch: InlineArray[UInt8, MLDSA_CRHBYTES], z: DSAPolyVec[MAX_L], h: DSAHintVec, p: MLDSAParams +) -> List[UInt8]: var sig = List[UInt8](capacity=signature_size(p)) for i in range(p.lambda_bits // 4): sig.append(ch[i]) @@ -1097,6 +1203,7 @@ def _dsa_sig_encode(ch: InlineArray[UInt8, MLDSA_CRHBYTES], z: DSAPolyVec[MAX_L] def _sig_decode(sig: Span[UInt8, ...], p: MLDSAParams) raises -> Tuple[List[UInt8], List[List[UInt32]], List[List[UInt8]]]: + _require_valid_params(p) if len(sig) != signature_size(p): raise Error("ML-DSA invalid signature length") var ch_len = p.lambda_bits // 4 @@ -1134,6 +1241,7 @@ def mldsa_public_key_from_seed(seed: Span[UInt8, ...], p: MLDSAParams) raises -> def mldsa_private_key_from_seed(seed: Span[UInt8, ...], p: MLDSAParams) raises -> MLDSAPrivateKey: + _require_valid_params(p) if len(seed) != 32: raise Error("ML-DSA invalid seed length") var xi = StackBuffer[UInt8, 34]() @@ -1189,6 +1297,7 @@ def mldsa_private_key_from_seed(seed: Span[UInt8, ...], p: MLDSAParams) raises - def mldsa_private_key_from_semiexpanded(sk: Span[UInt8, ...], p: MLDSAParams) raises -> MLDSAPrivateKey: + _require_valid_params(p) if len(sk) != private_key_size(p): raise Error("ML-DSA invalid semi-expanded private key length") var rho = _slice_bytes(sk, 0, 32) @@ -1211,7 +1320,8 @@ def mldsa_private_key_from_semiexpanded(sk: Span[UInt8, ...], p: MLDSAParams) ra var t0_plain = List[List[UInt32]](capacity=p.k) for _ in range(p.k): var length = N * 13 // 8 - t0_plain.append(_bit_unpack_slow(sk.unsafe_subspan(offset=off, length=length), (1 << 12) - 1, 1 << 12)) + t0_plain.append(_bit_unpack_slow(sk.unsafe_subspan(offset=off, length=length), (1 << 12) - 1, 1 << 12 + )) off += length var a = _compute_matrix_a(Span[UInt8, ...](rho), p) @@ -1272,6 +1382,8 @@ def mldsa_private_key_from_semiexpanded(sk: Span[UInt8, ...], p: MLDSAParams) ra def mldsa_sign_external_mu(priv: MLDSAPrivateKey, mu: Span[UInt8, ...], random: Span[UInt8, ...]) raises -> List[UInt8]: + if not _private_key_shape_valid(priv): + raise Error("invalid ML-DSA private key structure") if len(mu) != 64 or len(random) != 32: raise Error("ML-DSA invalid sign input") var p = priv.p.copy() @@ -1438,15 +1550,19 @@ def mldsa_sign_external_mu(priv: MLDSAPrivateKey, mu: Span[UInt8, ...], random: return sig^ -def mldsa_sign(priv: MLDSAPrivateKey, msg: Span[UInt8, ...], context: Span[UInt8, ...], random: Span[UInt8, ...]) raises -> List[UInt8]: +def mldsa_sign(priv: MLDSAPrivateKey, msg: Span[UInt8, ...], context: Span[UInt8, ...], random: Span[UInt8, ...] +) raises -> List[UInt8]: + if not _private_key_shape_valid(priv): + raise Error("invalid ML-DSA private key structure") var mu = _compute_message_hash(Span[UInt8, ...](priv.pub.tr), msg, context) - var sig = mldsa_sign_external_mu(priv, Span[UInt8, ...](mu), random) - _zero_list_u8(mu) - return sig^ + try: + return mldsa_sign_external_mu(priv, Span[UInt8, ...](mu), random) + finally: + _zero_list_u8(mu) def mldsa_verify_external_mu(pub: MLDSAPublicKey, mu: Span[UInt8, ...], sig: Span[UInt8, ...]) raises -> Bool: - if len(mu) != 64: + if not _public_key_shape_valid(pub) or len(mu) != 64: return False var p = pub.p.copy() var beta = p.tau * p.eta @@ -1497,11 +1613,16 @@ def mldsa_verify_external_mu(pub: MLDSAPublicKey, mu: Span[UInt8, ...], sig: Spa _zero_list_u8(computed) return ok -def mldsa_verify(pub: MLDSAPublicKey, msg: Span[UInt8, ...], sig: Span[UInt8, ...], context: Span[UInt8, ...]) raises -> Bool: + +def mldsa_verify(pub: MLDSAPublicKey, msg: Span[UInt8, ...], sig: Span[UInt8, ...], context: Span[UInt8, ...] +) raises -> Bool: + if not _public_key_shape_valid(pub): + return False var mu = _compute_message_hash(Span[UInt8, ...](pub.tr), msg, context) - var ok = mldsa_verify_external_mu(pub, Span[UInt8, ...](mu), sig) - _zero_list_u8(mu) - return ok + try: + return mldsa_verify_external_mu(pub, Span[UInt8, ...](mu), sig) + finally: + _zero_list_u8(mu) def mldsa44_public_key(pk: Span[UInt8, ...]) raises -> MLDSAPublicKey: @@ -1549,10 +1670,12 @@ def _zero_random32() -> List[UInt8]: def mldsa_keygen(p: MLDSAParams) raises -> MLDSAPrivateKey: # FIPS 204 external KeyGen: generate fresh 32-byte xi, then run KeyGen_internal. + _require_valid_params(p) var xi = random_bytes(MLDSA_SEEDBYTES) - var sk = mldsa_private_key_from_seed(Span[UInt8, ...](xi), p) - _zero_list_u8(xi) - return sk^ + try: + return mldsa_private_key_from_seed(Span[UInt8, ...](xi), p) + finally: + _zero_list_u8(xi) def mldsa44_keygen() raises -> MLDSAPrivateKey: @@ -1570,28 +1693,32 @@ def mldsa87_keygen() raises -> MLDSAPrivateKey: def mldsa_sign_hedged(priv: MLDSAPrivateKey, msg: Span[UInt8, ...], context: Span[UInt8, ...]) raises -> List[UInt8]: # FIPS 204 default signing variant: fresh 32-byte rnd. var rnd = random_bytes(MLDSA_RNDBYTES) - var sig = mldsa_sign(priv, msg, context, Span[UInt8, ...](rnd)) - _zero_list_u8(rnd) - return sig^ + try: + return mldsa_sign(priv, msg, context, Span[UInt8, ...](rnd)) + finally: + _zero_list_u8(rnd) def mldsa_sign_deterministic(priv: MLDSAPrivateKey, msg: Span[UInt8, ...], context: Span[UInt8, ...]) raises -> List[UInt8]: # FIPS 204 optional deterministic variant: rnd = {0}^32. var rnd = _zero_random32() - var sig = mldsa_sign(priv, msg, context, Span[UInt8, ...](rnd)) - _zero_list_u8(rnd) - return sig^ + try: + return mldsa_sign(priv, msg, context, Span[UInt8, ...](rnd)) + finally: + _zero_list_u8(rnd) def mldsa_sign_external_mu_hedged(priv: MLDSAPrivateKey, mu: Span[UInt8, ...]) raises -> List[UInt8]: var rnd = random_bytes(MLDSA_RNDBYTES) - var sig = mldsa_sign_external_mu(priv, mu, Span[UInt8, ...](rnd)) - _zero_list_u8(rnd) - return sig^ + try: + return mldsa_sign_external_mu(priv, mu, Span[UInt8, ...](rnd)) + finally: + _zero_list_u8(rnd) def mldsa_sign_external_mu_deterministic(priv: MLDSAPrivateKey, mu: Span[UInt8, ...]) raises -> List[UInt8]: var rnd = _zero_random32() - var sig = mldsa_sign_external_mu(priv, mu, Span[UInt8, ...](rnd)) - _zero_list_u8(rnd) - return sig^ + try: + return mldsa_sign_external_mu(priv, mu, Span[UInt8, ...](rnd)) + finally: + _zero_list_u8(rnd) diff --git a/src/thistle/ml_kem.mojo b/src/thistle/ml_kem.mojo index 27be353..a4a0497 100644 --- a/src/thistle/ml_kem.mojo +++ b/src/thistle/ml_kem.mojo @@ -1,11 +1,10 @@ -""" -ML-KEM primitives implementation in Mojo -""" +"""Implements the ML-KEM primitives specified by FIPS 203.""" from std.collections import List from std.algorithm.functional import vectorize from std.builtin.globals import global_constant from std.sys import simd_width_of +from std.os import abort from thistle.sha3 import ( SHA3Context, sha3_256_into, @@ -16,7 +15,7 @@ from thistle.sha3 import ( shake256_into, shake_advance, shake_finalize, - shake_squeeze_prefix_into, + shake_squeeze_prefix_into ) from thistle.random import random_bytes from thistle.utils import StackBuffer, zero_stack_u8 @@ -86,6 +85,17 @@ comptime MLKEM1024_CIPHERTEXTBYTES = CIPHERTEXTBYTES_1024 comptime MLKEM_BYTES = 32 +@always_inline +def _valid_k(k: Int) -> Bool: + return k == K_512 or k == K_768 or k == K_1024 + + +@always_inline +def _require_valid_k(k: Int): + if not _valid_k(k): + abort("invalid ML-KEM parameter k") + + comptime ZETAS_TABLE: InlineArray[Int16, 128] = [ -1044, -758, -359, -1517, 1493, 1422, 287, 202, -171, 622, 1577, 182, 962, -1202, -1474, 1468, @@ -102,7 +112,7 @@ comptime ZETAS_TABLE: InlineArray[Int16, 128] = [ 817, 1097, 603, 610, 1322, -1285, -1465, 384, -1215, -136, 1218, -1335, -874, 220, -1187, -1659, -1185, -1530, -1278, 794, -1510, -854, -870, 478, - -108, -308, 996, 991, 958, -1460, 1522, 1628, + -108, -308, 996, 991, 958, -1460, 1522, 1628 ] @@ -392,11 +402,9 @@ 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().unsafe_origin_cast[MutAnyOrigin]() - var bp = ( - b.coeffs.unsafe_ptr() + 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.unsafe_store[width=w](i, rp.unsafe_load[width=w](i) + bp.unsafe_load[width=w](i)) @@ -407,11 +415,9 @@ 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().unsafe_origin_cast[MutAnyOrigin]() - var ap = ( - a.coeffs.unsafe_ptr() + 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.unsafe_store[width=w](i, ap.unsafe_load[width=w](i) - rp.unsafe_load[width=w](i)) @@ -436,6 +442,8 @@ def poly_tobytes(mut out: List[UInt8], a: Poly): def poly_tobytes_stack(mut out: StackBuffer[UInt8, ...], a: Poly): + if out.remaining() < POLYBYTES: + abort("ML-KEM serialization destination is too small") for i in range(N // 2): var t0 = _positive_coeff(a.coeffs[2 * i]) var t1 = _positive_coeff(a.coeffs[2 * i + 1]) @@ -445,7 +453,7 @@ def poly_tobytes_stack(mut out: StackBuffer[UInt8, ...], a: Poly): def poly_frombytes(mut r: Poly, a: Span[UInt8, ...]) raises -> Bool: - if len(a) < POLYBYTES: + if len(a) != POLYBYTES: raise Error("ML-KEM poly_frombytes invalid buffer length") var ok = True for i in range(N // 2): @@ -457,11 +465,16 @@ def poly_frombytes(mut r: Poly, a: Span[UInt8, ...]) raises -> Bool: def polyvec_tobytes_stack(mut out: StackBuffer[UInt8, ...], ref a: Polyvec, k: Int): + _require_valid_k(k) + if out.remaining() < k * POLYBYTES: + abort("ML-KEM serialization destination is too small") for i in range(k): poly_tobytes_stack(out, a.vec[i]) def polyvec_frombytes(mut r: Polyvec, a: Span[UInt8, ...], k: Int) raises -> Bool: + if not _valid_k(k): + raise Error("ML-KEM polyvec_frombytes invalid k") if len(a) != k * POLYBYTES: raise Error("ML-KEM polyvec_frombytes invalid buffer length") var ok = True @@ -514,6 +527,11 @@ def polyvec_compress(mut out: List[UInt8], ref a: Polyvec, k: Int) raises: def polyvec_compress_stack(mut out: StackBuffer[UInt8, CIPHERTEXTBYTES_MAX], ref a: Polyvec, k: Int) raises: + var required = polyvec_compressed_byte_size(k) + if required == 0: + raise Error("ML-KEM polyvec_compress invalid k") + if out.remaining() < required: + raise Error("ML-KEM compressed destination is too small") if k == K_512 or k == K_768: for i in range(k): for j in range(N // 4): @@ -679,6 +697,11 @@ def poly_compress(mut out: List[UInt8], ref a: Poly, k: Int) raises: def poly_compress_stack(mut out: StackBuffer[UInt8, CIPHERTEXTBYTES_MAX], ref a: Poly, k: Int) raises: + var required = poly_compressed_bytes(k) + if required == 0: + raise Error("ML-KEM poly_compress invalid k") + if out.remaining() < required: + raise Error("ML-KEM compressed destination is too small") if k == K_512 or k == K_768: for i in range(N // 8): var t = InlineArray[UInt8, 8](fill=0) @@ -773,6 +796,7 @@ def polyvec_compressed_byte_size(k: Int) -> Int: def polyvec_reduce(mut r: Polyvec, k: Int): + _require_valid_k(k) for i in range(k): poly_reduce(r.vec[i]) @@ -784,6 +808,7 @@ def polyvec_reduce_k[k: Int](mut r: Polyvec): def polyvec_ntt(mut r: Polyvec, k: Int): + _require_valid_k(k) for i in range(k): poly_ntt(r.vec[i]) @@ -795,6 +820,7 @@ def polyvec_ntt_k[k: Int](mut r: Polyvec): def polyvec_invntt_tomont(mut r: Polyvec, k: Int): + _require_valid_k(k) for i in range(k): poly_invntt_tomont(r.vec[i]) @@ -806,6 +832,7 @@ def polyvec_invntt_tomont_k[k: Int](mut r: Polyvec): def polyvec_basemul_acc_montgomery(mut r: Poly, ref a: Polyvec, ref b: Polyvec, k: Int): + _require_valid_k(k) var t = Poly() poly_basemul_montgomery(r, a.vec[0], b.vec[0]) for i in range(1, k): @@ -834,7 +861,8 @@ def prf(key: Span[UInt8, ...], iv: UInt8, out_len: Int) raises -> List[UInt8]: 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: +def prf_into(mut out: StackBuffer[UInt8, ...], key: Span[UInt8, ...], iv: UInt8, out_len: Int +) raises: if len(key) != SYMBYTES: raise Error("ML-KEM prf key must be 32 bytes") var input = StackBuffer[UInt8, SYMBYTES + 1]() @@ -857,7 +885,8 @@ def rkprf(key: Span[UInt8, ...], input: Span[UInt8, ...]) raises -> List[UInt8]: 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: +def rkprf_into(mut out: StackBuffer[UInt8, SYMBYTES], key: Span[UInt8, ...], input: Span[UInt8, ...] +) raises: if len(key) != SYMBYTES: raise Error("ML-KEM rkprf key must be 32 bytes") if len(input) > CIPHERTEXTBYTES_MAX: @@ -890,6 +919,8 @@ def xof(seed: Span[UInt8, ...], x: UInt8, y: UInt8, out_len: Int) raises -> List def rej_uniform(mut r: Poly, start: Int, buf: Span[UInt8, ...]) -> Int: + if start < 0 or start > N: + abort("ML-KEM rejection-sampling offset out of bounds") var ctr = start var pos = 0 while ctr < N and pos + 3 <= len(buf): @@ -904,6 +935,7 @@ def rej_uniform(mut r: Poly, start: Int, buf: Span[UInt8, ...]) -> Int: ctr += 1 return ctr - start + def sample_ntt_into(mut out: Poly, seed: Span[UInt8, ...], x: UInt8, y: UInt8) raises: if len(seed) != SYMBYTES: raise Error("ML-KEM xof seed must be 32 bytes") @@ -945,7 +977,10 @@ def poly_getnoise_eta2(mut r: Poly, seed: Span[UInt8, ...], iv: UInt8) raises: 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: +def gen_matrix(mut a: InlineArray[Polyvec, K_MAX], seed: Span[UInt8, ...], transposed: Bool, k: Int +) raises: + if not _valid_k(k): + raise Error("ML-KEM gen_matrix invalid k") for i in range(k): for j in range(k): if transposed: @@ -1015,8 +1050,10 @@ struct DecapsulationKey(Copyable, Movable): 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: - if polyvec_byte_size(k) == 0: +def pack_pk_stack(mut out: StackBuffer[UInt8, ...], ref pk: Polyvec, seed: InlineArray[UInt8, SYMBYTES], k: Int +) -> Bool: + var pk_len = polyvec_byte_size(k) + if pk_len == 0 or out.remaining() < pk_len + SYMBYTES: return False polyvec_tobytes_stack(out, pk, k) for i in range(SYMBYTES): @@ -1036,8 +1073,11 @@ def unpack_pk(mut pk: KPKEEncryptionKey, input: Span[UInt8, ...], k: Int) raises return True -def pack_ciphertext_stack(mut out: StackBuffer[UInt8, CIPHERTEXTBYTES_MAX], ref b: Polyvec, ref v: Poly, k: Int) raises -> Bool: - if polyvec_compressed_byte_size(k) == 0 or poly_compressed_bytes(k) == 0: +def pack_ciphertext_stack(mut out: StackBuffer[UInt8, CIPHERTEXTBYTES_MAX], ref b: Polyvec, ref v: Poly, k: Int +) raises -> Bool: + var b_len = polyvec_compressed_byte_size(k) + var v_len = poly_compressed_bytes(k) + if b_len == 0 or v_len == 0 or out.remaining() < b_len + v_len: return False polyvec_compress_stack(out, b, k) poly_compress_stack(out, v, k) @@ -1054,7 +1094,8 @@ def unpack_ciphertext(mut b: Polyvec, mut v: Poly, c: Span[UInt8, ...], k: Int) return True -def k_pke_keygen(mut ek: KPKEEncryptionKey, mut dk: KPKEDecapsulationKey, d: Span[UInt8, ...], k: Int) raises: +def k_pke_keygen(mut ek: KPKEEncryptionKey, mut dk: KPKEDecapsulationKey, d: Span[UInt8, ...], k: Int +) raises: if len(d) != SYMBYTES: raise Error("ML-KEM K-PKE keygen d must be 32 bytes") if k != K_512 and k != K_768 and k != K_1024: @@ -1157,7 +1198,8 @@ def k_pke_keygen_k[k: Int](mut ek: KPKEEncryptionKey, mut dk: KPKEDecapsulationK ek.k = k -def k_pke_encrypt_into(mut ciphertext: StackBuffer[UInt8, CIPHERTEXTBYTES_MAX], ref ek: KPKEEncryptionKey, m: Span[UInt8, ...], r: Span[UInt8, ...]) raises -> Bool: +def k_pke_encrypt_into(mut ciphertext: StackBuffer[UInt8, CIPHERTEXTBYTES_MAX], ref ek: KPKEEncryptionKey, m: Span[UInt8, ...], r: Span[UInt8, ...] +) raises -> Bool: if len(m) != INDCPA_MSGBYTES or len(r) != SYMBYTES: return False var k = ek.k @@ -1257,7 +1299,8 @@ def k_pke_encrypt_into_k[k: Int](mut ciphertext: StackBuffer[UInt8, CIPHERTEXTBY return pack_ciphertext_stack(ciphertext, b, v, k) -def k_pke_decrypt_msg(mut plaintext: StackBuffer[UInt8, SYMBYTES], ref dk: KPKEDecapsulationKey, c: Span[UInt8, ...]) raises -> Bool: +def k_pke_decrypt_msg(mut plaintext: StackBuffer[UInt8, SYMBYTES], ref dk: KPKEDecapsulationKey, c: Span[UInt8, ...] +) raises -> Bool: var k = dk.k if k != K_512 and k != K_768 and k != K_1024: return False @@ -1342,11 +1385,13 @@ def kem_keygen_internal_k[k: Int](seed: Span[UInt8, ...]) raises -> Decapsulatio return dk^ -def encapsulation_key_bytes_into(mut out: StackBuffer[UInt8, INDCPA_PUBLICKEYBYTES_MAX], ref dk: DecapsulationKey) -> Bool: +def encapsulation_key_bytes_into(mut out: StackBuffer[UInt8, INDCPA_PUBLICKEYBYTES_MAX], ref dk: DecapsulationKey +) -> Bool: out.clear() - var ek_len = polyvec_byte_size(dk.pke_dk.k) + SYMBYTES - if ek_len == 0: + var pv_len = polyvec_byte_size(dk.pke_dk.k) + if pv_len == 0 or dk.ek.pke_ek.k != dk.pke_dk.k: return False + var ek_len = pv_len + SYMBYTES for i in range(ek_len): out.push_unchecked(dk.ek.raw_bytes[i]) return True @@ -1429,7 +1474,6 @@ def _zero_list(mut data: List[UInt8]): ptr.unsafe_store[volatile=True](i, UInt8(0)) - def decapsulation_key_decode(mut dk: DecapsulationKey, input: Span[UInt8, ...], k: Int) raises -> Bool: var dk_len = _decapsulation_key_size(k) var sk_len = polyvec_byte_size(k) @@ -1455,7 +1499,8 @@ def decapsulation_key_decode(mut dk: DecapsulationKey, input: Span[UInt8, ...], return True -def mlkem_keygen_seed_into(mut ek_out: StackBuffer[UInt8, INDCPA_PUBLICKEYBYTES_MAX], mut dk_out: StackBuffer[UInt8, DECAPSKEYBYTES_MAX], seed: Span[UInt8, ...], parameter_set: String) raises -> Bool: +def mlkem_keygen_seed_into(mut ek_out: StackBuffer[UInt8, INDCPA_PUBLICKEYBYTES_MAX], mut dk_out: StackBuffer[UInt8, DECAPSKEYBYTES_MAX], seed: Span[UInt8, ...], parameter_set: String +) raises -> Bool: ek_out.clear() dk_out.clear() var k = _parameter_set_k(parameter_set) @@ -1484,147 +1529,121 @@ def mlkem_keygen_seed_into_k[k: Int](mut ek_out: StackBuffer[UInt8, INDCPA_PUBLI def mlkem_keygen_seed(seed: Span[UInt8, ...], parameter_set: String) raises -> Tuple[List[UInt8], List[UInt8]]: var ek_buf = StackBuffer[UInt8, INDCPA_PUBLICKEYBYTES_MAX]() var dk_buf = StackBuffer[UInt8, DECAPSKEYBYTES_MAX]() - if not mlkem_keygen_seed_into(ek_buf, dk_buf, seed, parameter_set): - return (List[UInt8](), List[UInt8]()) - - var ek = List[UInt8](capacity=ek_buf.len()) - for i in range(ek_buf.len()): - ek.append(ek_buf[i]) - var dk = List[UInt8](capacity=dk_buf.len()) - for i in range(dk_buf.len()): - dk.append(dk_buf[i]) - zero_stack_u8(ek_buf) - zero_stack_u8(dk_buf) - return (ek^, dk^) + try: + if not mlkem_keygen_seed_into(ek_buf, dk_buf, seed, parameter_set): + return (List[UInt8](), List[UInt8]()) + + var ek = List[UInt8](capacity=ek_buf.len()) + for i in range(ek_buf.len()): + ek.append(ek_buf[i]) + var dk = List[UInt8](capacity=dk_buf.len()) + for i in range(dk_buf.len()): + dk.append(dk_buf[i]) + return (ek^, dk^) + finally: + zero_stack_u8(ek_buf) + zero_stack_u8(dk_buf) # FIPS 203 ML-KEM.KeyGen(): fresh d || z, then deterministic expansion. -def mlkem_keygen(parameter_set: String) raises -> Tuple[List[UInt8], List[UInt8]]: + + +def mlkem_keygen(parameter_set: String +) raises -> Tuple[List[UInt8], List[UInt8]]: var seed = random_bytes(2 * SYMBYTES) - var result = mlkem_keygen_seed(Span[UInt8, ...](seed), parameter_set) - _zero_list(seed) - return result^ + try: + return mlkem_keygen_seed(Span[UInt8, ...](seed), parameter_set) + finally: + _zero_list(seed) -def mlkem512_keygen() raises -> Tuple[List[UInt8], List[UInt8]]: +def _mlkem_keygen_random_k[k: Int]() raises -> Tuple[List[UInt8], List[UInt8]]: + comptime assert _valid_k(k), "invalid ML-KEM k" var seed = random_bytes(2 * SYMBYTES) var ek_buf = StackBuffer[UInt8, INDCPA_PUBLICKEYBYTES_MAX]() var dk_buf = StackBuffer[UInt8, DECAPSKEYBYTES_MAX]() - _ = mlkem_keygen_seed_into_k[K_512](ek_buf, dk_buf, Span[UInt8, ...](seed)) - _zero_list(seed) - var ek = List[UInt8](capacity=ek_buf.len()) - for i in range(ek_buf.len()): - ek.append(ek_buf[i]) - var dk = List[UInt8](capacity=dk_buf.len()) - for i in range(dk_buf.len()): - dk.append(dk_buf[i]) - zero_stack_u8(ek_buf) - zero_stack_u8(dk_buf) - return (ek^, dk^) + try: + if not mlkem_keygen_seed_into_k[k](ek_buf, dk_buf, Span[UInt8, ...](seed)): + return (List[UInt8](), List[UInt8]()) + var ek = List[UInt8](capacity=ek_buf.len()) + for i in range(ek_buf.len()): + ek.append(ek_buf[i]) + var dk = List[UInt8](capacity=dk_buf.len()) + for i in range(dk_buf.len()): + dk.append(dk_buf[i]) + return (ek^, dk^) + finally: + _zero_list(seed) + zero_stack_u8(ek_buf) + zero_stack_u8(dk_buf) + + +def mlkem512_keygen() raises -> Tuple[List[UInt8], List[UInt8]]: + return _mlkem_keygen_random_k[K_512]() def mlkem768_keygen() raises -> Tuple[List[UInt8], List[UInt8]]: - var seed = random_bytes(2 * SYMBYTES) - var ek_buf = StackBuffer[UInt8, INDCPA_PUBLICKEYBYTES_MAX]() - var dk_buf = StackBuffer[UInt8, DECAPSKEYBYTES_MAX]() - _ = mlkem_keygen_seed_into_k[K_768](ek_buf, dk_buf, Span[UInt8, ...](seed)) - _zero_list(seed) - var ek = List[UInt8](capacity=ek_buf.len()) - for i in range(ek_buf.len()): - ek.append(ek_buf[i]) - var dk = List[UInt8](capacity=dk_buf.len()) - for i in range(dk_buf.len()): - dk.append(dk_buf[i]) - zero_stack_u8(ek_buf) - zero_stack_u8(dk_buf) - return (ek^, dk^) + return _mlkem_keygen_random_k[K_768]() def mlkem1024_keygen() raises -> Tuple[List[UInt8], List[UInt8]]: - var seed = random_bytes(2 * SYMBYTES) - var ek_buf = StackBuffer[UInt8, INDCPA_PUBLICKEYBYTES_MAX]() - var dk_buf = StackBuffer[UInt8, DECAPSKEYBYTES_MAX]() - _ = mlkem_keygen_seed_into_k[K_1024](ek_buf, dk_buf, Span[UInt8, ...](seed)) - _zero_list(seed) - var ek = List[UInt8](capacity=ek_buf.len()) - for i in range(ek_buf.len()): - ek.append(ek_buf[i]) - var dk = List[UInt8](capacity=dk_buf.len()) - for i in range(dk_buf.len()): - dk.append(dk_buf[i]) - zero_stack_u8(ek_buf) - zero_stack_u8(dk_buf) - return (ek^, dk^) + return _mlkem_keygen_random_k[K_1024]() # FIPS 203 ML-KEM.Encaps(): fresh m, returns (ciphertext, shared_secret, ok). + + def mlkem_encaps(ek_bytes: Span[UInt8, ...], parameter_set: String) raises -> Tuple[List[UInt8], List[UInt8], Bool]: var m = random_bytes(SYMBYTES) - var result = mlkem_encaps_seed(ek_bytes, Span[UInt8, ...](m), parameter_set) - _zero_list(m) - if not result[2]: - return (List[UInt8](), List[UInt8](), False) - return (result[0].copy(), result[1].copy(), True) + try: + var result = mlkem_encaps_seed(ek_bytes, Span[UInt8, ...](m), parameter_set) + if not result[2]: + return (List[UInt8](), List[UInt8](), False) + return (result[0].copy(), result[1].copy(), True) + finally: + _zero_list(m) -def mlkem512_encaps(ek_bytes: Span[UInt8, ...]) raises -> Tuple[List[UInt8], List[UInt8], Bool]: +def _mlkem_encaps_random_k[k: Int](ek_bytes: Span[UInt8, ...]) raises -> Tuple[List[UInt8], List[UInt8], Bool]: + comptime assert _valid_k(k), "invalid ML-KEM k" var m = random_bytes(SYMBYTES) var shared_buf = StackBuffer[UInt8, SYMBYTES]() var ciphertext_buf = StackBuffer[UInt8, CIPHERTEXTBYTES_MAX]() - var ok = mlkem_encaps_seed_into_k[K_512](ciphertext_buf, shared_buf, ek_bytes, Span[UInt8, ...](m)) - _zero_list(m) - if not ok: - return (List[UInt8](), List[UInt8](), False) - var ciphertext = List[UInt8](capacity=ciphertext_buf.len()) - for i in range(ciphertext_buf.len()): - ciphertext.append(ciphertext_buf[i]) - var shared = List[UInt8](capacity=SYMBYTES) - for i in range(SYMBYTES): - shared.append(shared_buf[i]) - zero_stack_u8(shared_buf) - zero_stack_u8(ciphertext_buf) - return (ciphertext^, shared^, True) + try: + if not mlkem_encaps_seed_into_k[k](ciphertext_buf, shared_buf, ek_bytes, Span[UInt8, ...](m)): + return (List[UInt8](), List[UInt8](), False) + var ciphertext = List[UInt8](capacity=ciphertext_buf.len()) + for i in range(ciphertext_buf.len()): + ciphertext.append(ciphertext_buf[i]) + var shared = List[UInt8](capacity=SYMBYTES) + for i in range(SYMBYTES): + shared.append(shared_buf[i]) + return (ciphertext^, shared^, True) + finally: + _zero_list(m) + zero_stack_u8(shared_buf) + zero_stack_u8(ciphertext_buf) + + +def mlkem512_encaps(ek_bytes: Span[UInt8, ...]) raises -> Tuple[List[UInt8], List[UInt8], Bool]: + return _mlkem_encaps_random_k[K_512](ek_bytes) def mlkem768_encaps(ek_bytes: Span[UInt8, ...]) raises -> Tuple[List[UInt8], List[UInt8], Bool]: - var m = random_bytes(SYMBYTES) - var shared_buf = StackBuffer[UInt8, SYMBYTES]() - var ciphertext_buf = StackBuffer[UInt8, CIPHERTEXTBYTES_MAX]() - var ok = mlkem_encaps_seed_into_k[K_768](ciphertext_buf, shared_buf, ek_bytes, Span[UInt8, ...](m)) - _zero_list(m) - if not ok: - return (List[UInt8](), List[UInt8](), False) - var ciphertext = List[UInt8](capacity=ciphertext_buf.len()) - for i in range(ciphertext_buf.len()): - ciphertext.append(ciphertext_buf[i]) - var shared = List[UInt8](capacity=SYMBYTES) - for i in range(SYMBYTES): - shared.append(shared_buf[i]) - zero_stack_u8(shared_buf) - zero_stack_u8(ciphertext_buf) - return (ciphertext^, shared^, True) + return _mlkem_encaps_random_k[K_768](ek_bytes) def mlkem1024_encaps(ek_bytes: Span[UInt8, ...]) raises -> Tuple[List[UInt8], List[UInt8], Bool]: - var m = random_bytes(SYMBYTES) - var shared_buf = StackBuffer[UInt8, SYMBYTES]() - var ciphertext_buf = StackBuffer[UInt8, CIPHERTEXTBYTES_MAX]() - var ok = mlkem_encaps_seed_into_k[K_1024](ciphertext_buf, shared_buf, ek_bytes, Span[UInt8, ...](m)) - _zero_list(m) - if not ok: - return (List[UInt8](), List[UInt8](), False) - var ciphertext = List[UInt8](capacity=ciphertext_buf.len()) - for i in range(ciphertext_buf.len()): - ciphertext.append(ciphertext_buf[i]) - var shared = List[UInt8](capacity=SYMBYTES) - for i in range(SYMBYTES): - shared.append(shared_buf[i]) - zero_stack_u8(shared_buf) - zero_stack_u8(ciphertext_buf) - return (ciphertext^, shared^, True) + return _mlkem_encaps_random_k[K_1024](ek_bytes) -def mlkem_encaps_seed_into(mut ciphertext_out: StackBuffer[UInt8, CIPHERTEXTBYTES_MAX], mut shared_out: StackBuffer[UInt8, SYMBYTES], ek_bytes: Span[UInt8, ...], m: Span[UInt8, ...], parameter_set: String) raises -> Bool: +def mlkem_encaps_seed_into( + mut ciphertext_out: StackBuffer[UInt8, CIPHERTEXTBYTES_MAX], + mut shared_out: StackBuffer[UInt8, SYMBYTES], + ek_bytes: Span[UInt8, ...], + m: Span[UInt8, ...], + parameter_set: String +) raises -> Bool: ciphertext_out.clear() shared_out.clear() var k = _parameter_set_k(parameter_set) @@ -1647,7 +1666,8 @@ 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().unsafe_offset(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 +1699,8 @@ 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().unsafe_offset(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) @@ -1719,7 +1740,8 @@ def mlkem_encaps_seed(ek_bytes: Span[UInt8, ...], m: Span[UInt8, ...], parameter return (result[1].copy(), result[0].copy(), True) -def mlkem_decaps_into(mut shared_out: StackBuffer[UInt8, SYMBYTES], dk_bytes: Span[UInt8, ...], ciphertext: Span[UInt8, ...], parameter_set: String) raises -> Bool: +def mlkem_decaps_into(mut shared_out: StackBuffer[UInt8, SYMBYTES], dk_bytes: Span[UInt8, ...], ciphertext: Span[UInt8, ...], parameter_set: String +) raises -> Bool: shared_out.clear() var k = _parameter_set_k(parameter_set) if len(ciphertext) != _ciphertext_size(k): @@ -1744,7 +1766,8 @@ 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().unsafe_offset(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) @@ -1752,7 +1775,8 @@ 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, ...](unsafe_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)) @@ -1788,7 +1812,8 @@ 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().unsafe_offset(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) @@ -1796,7 +1821,8 @@ 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, ...](unsafe_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)) @@ -1807,7 +1833,8 @@ def mlkem_decaps_into_k[k: Int](mut shared_out: StackBuffer[UInt8, SYMBYTES], dk return True -def mlkem_decaps(dk_bytes: Span[UInt8, ...], ciphertext: Span[UInt8, ...], parameter_set: String) raises -> Tuple[List[UInt8], Bool]: +def mlkem_decaps(dk_bytes: Span[UInt8, ...], ciphertext: Span[UInt8, ...], parameter_set: String +) raises -> Tuple[List[UInt8], Bool]: var shared_buf = StackBuffer[UInt8, SYMBYTES]() if not mlkem_decaps_into(shared_buf, dk_bytes, ciphertext, parameter_set): zero_stack_u8(shared_buf) diff --git a/src/thistle/p256.mojo b/src/thistle/p256.mojo index 31cc74c..7e58e5b 100644 --- a/src/thistle/p256.mojo +++ b/src/thistle/p256.mojo @@ -7,7 +7,9 @@ from .utils import u64_nonzero_choice, u64_zero_choice from .sha2 import sha256_hash from .pbkdf2 import hmac_sha256 from std.utils import StaticTuple -from .weierstrass import Limbs, U256, Point, JacobianPoint, cmp as ws_cmp, sub_raw as ws_sub_raw, add_raw as ws_add_raw, select as ws_select, zero_choice as ws_zero_choice, add_mod as ws_add_mod, sub_mod as ws_sub_mod, from_be as ws_from_be, to_be as ws_to_be, mont_mul as ws_mont_mul, mont_sqr as ws_mont_sqr, to_mont as ws_to_mont, from_mont as ws_from_mont, mul_mod as ws_mul_mod, square_mod as ws_square_mod, is_on_curve as ws_is_on_curve, mul_small_mod as ws_mul_small_mod, jacobian_double_ct as ws_jacobian_double_ct, jacobian_infinity as ws_jacobian_infinity, select_jacobian_ct as ws_select_jacobian_ct, jacobian_add_affine_non_equal_ct as ws_jacobian_add_affine, pow_mod as ws_pow_mod, sqn as ws_sqn, inv_p as ws_inv_p, jacobian_to_affine as ws_jacobian_to_affine, scalar_mult as ws_scalar_mult, base_table_entry as ws_base_table_entry, scalar_mult_base as ws_scalar_mult_base, mod_inv_ct as ws_mod_inv_ct, reduce_mod as ws_reduce_mod, point_add as ws_point_add, rfc6979 as ws_rfc6979 +from .weierstrass import ( + Limbs, U256, Point, JacobianPoint, cmp as ws_cmp, sub_raw as ws_sub_raw, add_raw as ws_add_raw, select as ws_select, zero_choice as ws_zero_choice, add_mod as ws_add_mod, sub_mod as ws_sub_mod, from_be as ws_from_be, to_be as ws_to_be, mont_mul as ws_mont_mul, mont_sqr as ws_mont_sqr, to_mont as ws_to_mont, from_mont as ws_from_mont, mul_mod as ws_mul_mod, square_mod as ws_square_mod, is_on_curve as ws_is_on_curve, mul_small_mod as ws_mul_small_mod, jacobian_double_ct as ws_jacobian_double_ct, jacobian_infinity as ws_jacobian_infinity, select_jacobian_ct as ws_select_jacobian_ct, jacobian_add_affine_non_equal_ct as ws_jacobian_add_affine, pow_mod as ws_pow_mod, sqn as ws_sqn, inv_p as ws_inv_p, jacobian_to_affine as ws_jacobian_to_affine, scalar_mult as ws_scalar_mult, base_table_entry as ws_base_table_entry, scalar_mult_base as ws_scalar_mult_base, mod_inv_ct as ws_mod_inv_ct, reduce_mod as ws_reduce_mod, point_add as ws_point_add, rfc6979 as ws_rfc6979 +) comptime P256_SIZE = 32 comptime P256_POINT_SIZE = 65 @@ -22,7 +24,7 @@ def _p() -> U256: 0xFFFFFFFFFFFFFFFF, 0x00000000FFFFFFFF, 0x0000000000000000, - 0xFFFFFFFF00000001, + 0xFFFFFFFF00000001 ) @@ -31,7 +33,7 @@ def _a() -> U256: 0xFFFFFFFFFFFFFFFC, 0x00000000FFFFFFFF, 0x0000000000000000, - 0xFFFFFFFF00000001, + 0xFFFFFFFF00000001 ) @@ -40,7 +42,7 @@ def _b() -> U256: 0x3BCE3C3E27D2604B, 0x651D06B0CC53B0F6, 0xB3EBBD55769886BC, - 0x5AC635D8AA3A93E7, + 0x5AC635D8AA3A93E7 ) @@ -49,7 +51,7 @@ def _n() -> U256: 0xF3B9CAC2FC632551, 0xBCE6FAADA7179E84, 0xFFFFFFFFFFFFFFFF, - 0xFFFFFFFF00000000, + 0xFFFFFFFF00000000 ) @@ -58,7 +60,7 @@ def _gx() -> U256: 0xF4A13945D898C296, 0x77037D812DEB33A0, 0xF8BCE6E563A440F2, - 0x6B17D1F2E12C4247, + 0x6B17D1F2E12C4247 ) @@ -67,7 +69,7 @@ def _gy() -> U256: 0xCBB6406837BF51F5, 0x2BCE33576B315ECE, 0x8EE7EB4A7C0F9E16, - 0x4FE342E2FE1A7F9B, + 0x4FE342E2FE1A7F9B ) @@ -77,7 +79,7 @@ def _sqrt_exp() -> U256: 0x0000000000000000, 0x0000000040000000, 0x4000000000000000, - 0x3FFFFFFFC0000000, + 0x3FFFFFFFC0000000 ) @@ -87,7 +89,7 @@ def _rr() -> U256: 0x0000000000000003, 0xFFFFFFFBFFFFFFFF, 0xFFFFFFFFFFFFFFFE, - 0x00000004FFFFFFFD, + 0x00000004FFFFFFFD ) @@ -97,7 +99,7 @@ def _one_mont() -> U256: 0x0000000000000001, 0xFFFFFFFF00000000, 0xFFFFFFFFFFFFFFFF, - 0x00000000FFFFFFFE, + 0x00000000FFFFFFFE ) @@ -414,7 +416,7 @@ def p256_public_key( def p256_ecdh( private_key: Span[UInt8, ...], public_key: Span[UInt8, ...], - output: Span[mut=True, UInt8, ...], + output: Span[mut=True, UInt8, ...] ) -> Bool: if len(private_key) != 32 or len(output) < P256_SIZE: return False @@ -445,7 +447,7 @@ def _n_rr() -> U256: 0x83244C95BE79EEA2, 0x4699799C49BD6FA6, 0x2845B2392B6BEC59, - 0x66E12D94F3D95620, + 0x66E12D94F3D95620 ) @@ -454,7 +456,7 @@ def _n_one_mont() -> U256: 0x0C46353D039CDAAF, 0x4319055258E8617B, 0x0000000000000000, - 0x00000000FFFFFFFF, + 0x00000000FFFFFFFF ) @@ -463,7 +465,7 @@ def _n_minus_2() -> U256: 0xF3B9CAC2FC63254F, 0xBCE6FAADA7179E84, 0xFFFFFFFFFFFFFFFF, - 0xFFFFFFFF00000000, + 0xFFFFFFFF00000000 ) @@ -515,12 +517,10 @@ def _rfc6979_p256(private_key: Span[UInt8, ...], digest: Span[UInt8, ...], skip: def p256_ecdsa_sign_digest( private_key: Span[UInt8, ...], digest: Span[UInt8, ...], - signature: Span[mut=True, UInt8, ...], + signature: Span[mut=True, UInt8, ...] ) -> Bool: - if ( - len(private_key) != 32 or len(digest) != 32 - or len(signature) < P256_SIGNATURE_SIZE - ): + 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) @@ -562,7 +562,7 @@ def p256_ecdsa_sign_digest( def p256_ecdsa_sign( private_key: Span[UInt8, ...], message: Span[UInt8, ...], - signature: Span[mut=True, UInt8, ...], + signature: Span[mut=True, UInt8, ...] ) -> Bool: if len(signature) < P256_SIGNATURE_SIZE: return False @@ -584,7 +584,7 @@ def _p256_add_public(a: P256Point, b: P256Point) -> P256Point: def p256_ecdsa_verify_digest( public_key: Span[UInt8, ...], digest: Span[UInt8, ...], - signature: Span[UInt8, ...], + signature: Span[UInt8, ...] ) -> Bool: if len(digest) != 32 or len(signature) != 64: return False @@ -608,7 +608,7 @@ def p256_ecdsa_verify_digest( def p256_ecdsa_verify( public_key: Span[UInt8, ...], message: Span[UInt8, ...], - signature: Span[UInt8, ...], + signature: Span[UInt8, ...] ) -> Bool: var digest = sha256_hash(message) return p256_ecdsa_verify_digest(public_key, Span[UInt8, ...](digest), signature) @@ -628,7 +628,7 @@ def p256_ecdsa_sign_der( def p256_ecdsa_verify_der( public_key: Span[UInt8, ...], message: Span[UInt8, ...], - signature: Span[UInt8, ...], + signature: Span[UInt8, ...] ) -> Bool: from .ecdsa_der import ecdsa_der_decode var raw = ecdsa_der_decode(signature, P256_SIZE) @@ -646,7 +646,7 @@ def p256_keygen() raises -> Tuple[List[UInt8], List[UInt8]]: var public_key = List[UInt8](unsafe_uninit_length=65) if p256_public_key( Span[UInt8, ...](private_key), - Span[mut=True, UInt8, ...](public_key), + Span[mut=True, UInt8, ...](public_key) ): _wipe_u256(d) return (private_key^, public_key^) diff --git a/src/thistle/p384.mojo b/src/thistle/p384.mojo index ed00f88..eeee5c2 100644 --- a/src/thistle/p384.mojo +++ b/src/thistle/p384.mojo @@ -7,7 +7,9 @@ from .utils import u64_nonzero_choice, u64_zero_choice from .sha2 import sha384_hash from .pbkdf2 import hmac_sha384 from std.utils import StaticTuple -from .weierstrass import Limbs, U384, Point, JacobianPoint, cmp as ws_cmp, sub_raw as ws_sub_raw, add_raw as ws_add_raw, select as ws_select, zero_choice as ws_zero_choice, add_mod as ws_add_mod, sub_mod as ws_sub_mod, from_be as ws_from_be, to_be as ws_to_be, mont_mul as ws_mont_mul, mont_sqr as ws_mont_sqr, to_mont as ws_to_mont, from_mont as ws_from_mont, mul_mod as ws_mul_mod, square_mod as ws_square_mod, is_on_curve as ws_is_on_curve, mul_small_mod as ws_mul_small_mod, jacobian_double_ct as ws_jacobian_double_ct, jacobian_infinity as ws_jacobian_infinity, select_jacobian_ct as ws_select_jacobian_ct, jacobian_add_affine_non_equal_ct as ws_jacobian_add_affine, pow_mod as ws_pow_mod, sqn as ws_sqn, inv_p as ws_inv_p, jacobian_to_affine as ws_jacobian_to_affine, scalar_mult as ws_scalar_mult, base_table_entry as ws_base_table_entry, scalar_mult_base as ws_scalar_mult_base, mod_inv_ct as ws_mod_inv_ct, reduce_mod as ws_reduce_mod, point_add as ws_point_add, rfc6979 as ws_rfc6979 +from .weierstrass import ( + Limbs, U384, Point, JacobianPoint, cmp as ws_cmp, sub_raw as ws_sub_raw, add_raw as ws_add_raw, select as ws_select, zero_choice as ws_zero_choice, add_mod as ws_add_mod, sub_mod as ws_sub_mod, from_be as ws_from_be, to_be as ws_to_be, mont_mul as ws_mont_mul, mont_sqr as ws_mont_sqr, to_mont as ws_to_mont, from_mont as ws_from_mont, mul_mod as ws_mul_mod, square_mod as ws_square_mod, is_on_curve as ws_is_on_curve, mul_small_mod as ws_mul_small_mod, jacobian_double_ct as ws_jacobian_double_ct, jacobian_infinity as ws_jacobian_infinity, select_jacobian_ct as ws_select_jacobian_ct, jacobian_add_affine_non_equal_ct as ws_jacobian_add_affine, pow_mod as ws_pow_mod, sqn as ws_sqn, inv_p as ws_inv_p, jacobian_to_affine as ws_jacobian_to_affine, scalar_mult as ws_scalar_mult, base_table_entry as ws_base_table_entry, scalar_mult_base as ws_scalar_mult_base, mod_inv_ct as ws_mod_inv_ct, reduce_mod as ws_reduce_mod, point_add as ws_point_add, rfc6979 as ws_rfc6979 +) comptime P384_SIZE = 48 comptime P384_POINT_SIZE = 97 @@ -24,7 +26,7 @@ def _p() -> U384: 0xFFFFFFFFFFFFFFFE, 0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF, - 0xFFFFFFFFFFFFFFFF, + 0xFFFFFFFFFFFFFFFF ) @@ -35,7 +37,7 @@ def _a() -> U384: 0xFFFFFFFFFFFFFFFE, 0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF, - 0xFFFFFFFFFFFFFFFF, + 0xFFFFFFFFFFFFFFFF ) @@ -46,7 +48,7 @@ def _b() -> U384: 0x0314088F5013875A, 0x181D9C6EFE814112, 0x988E056BE3F82D19, - 0xB3312FA7E23EE7E4, + 0xB3312FA7E23EE7E4 ) @@ -57,7 +59,7 @@ def _n() -> U384: 0xC7634D81F4372DDF, 0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF, - 0xFFFFFFFFFFFFFFFF, + 0xFFFFFFFFFFFFFFFF ) @@ -68,7 +70,7 @@ def _gx() -> U384: 0x59F741E082542A38, 0x6E1D3B628BA79B98, 0x8EB1C71EF320AD74, - 0xAA87CA22BE8B0537, + 0xAA87CA22BE8B0537 ) @@ -79,7 +81,7 @@ def _gy() -> U384: 0xE9DA3113B5F0B8C0, 0xF8F41DBD289A147C, 0x5D9E98BF9292DC29, - 0x3617DE4A96262C6F, + 0x3617DE4A96262C6F ) @@ -91,7 +93,7 @@ def _sqrt_exp() -> U384: 0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF, - 0x3FFFFFFFFFFFFFFF, + 0x3FFFFFFFFFFFFFFF ) @@ -102,7 +104,7 @@ def _rr() -> U384: 0xFFFFFFFE00000000, 0x0000000200000000, 0x0000000000000001, - 0x0000000000000000, + 0x0000000000000000 ) @@ -113,7 +115,7 @@ def _one_mont() -> U384: 0x0000000000000001, 0x0000000000000000, 0x0000000000000000, - 0x0000000000000000, + 0x0000000000000000 ) @@ -419,7 +421,7 @@ def p384_public_key( def p384_ecdh( private_key: Span[UInt8, ...], public_key: Span[UInt8, ...], - output: Span[mut=True, UInt8, ...], + output: Span[mut=True, UInt8, ...] ) -> Bool: if len(private_key) != 48 or len(output) < P384_SIZE: return False @@ -452,7 +454,7 @@ def _n_rr() -> U384: 0xBC3E483AFCB82947, 0xD40D49174AAB1CC5, 0x3FB05B7A28266895, - 0x0C84EE012B39BF21, + 0x0C84EE012B39BF21 ) @@ -463,7 +465,7 @@ def _n_one_mont() -> U384: 0x389CB27E0BC8D220, 0, 0, - 0, + 0 ) @@ -474,7 +476,7 @@ def _n_minus_2() -> U384: 0xC7634D81F4372DDF, 0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF, - 0xFFFFFFFFFFFFFFFF, + 0xFFFFFFFFFFFFFFFF ) @@ -526,12 +528,10 @@ def _rfc6979_p384(private_key: Span[UInt8, ...], digest: Span[UInt8, ...], skip: def p384_ecdsa_sign_digest( private_key: Span[UInt8, ...], digest: Span[UInt8, ...], - signature: Span[mut=True, UInt8, ...], + signature: Span[mut=True, UInt8, ...] ) -> Bool: - if ( - len(private_key) != 48 or len(digest) != 48 - or len(signature) < P384_SIGNATURE_SIZE - ): + 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) @@ -573,7 +573,7 @@ def p384_ecdsa_sign_digest( def p384_ecdsa_sign( private_key: Span[UInt8, ...], message: Span[UInt8, ...], - signature: Span[mut=True, UInt8, ...], + signature: Span[mut=True, UInt8, ...] ) -> Bool: if len(signature) < P384_SIGNATURE_SIZE: return False @@ -595,7 +595,7 @@ def _p384_add_public(a: P384Point, b: P384Point) -> P384Point: def p384_ecdsa_verify_digest( public_key: Span[UInt8, ...], digest: Span[UInt8, ...], - signature: Span[UInt8, ...], + signature: Span[UInt8, ...] ) -> Bool: if len(digest) != 48 or len(signature) != 96: return False @@ -619,7 +619,7 @@ def p384_ecdsa_verify_digest( def p384_ecdsa_verify( public_key: Span[UInt8, ...], message: Span[UInt8, ...], - signature: Span[UInt8, ...], + signature: Span[UInt8, ...] ) -> Bool: var digest = sha384_hash(message) return p384_ecdsa_verify_digest(public_key, Span[UInt8, ...](digest), signature) @@ -639,7 +639,7 @@ def p384_ecdsa_sign_der( def p384_ecdsa_verify_der( public_key: Span[UInt8, ...], message: Span[UInt8, ...], - signature: Span[UInt8, ...], + signature: Span[UInt8, ...] ) -> Bool: from .ecdsa_der import ecdsa_der_decode var raw = ecdsa_der_decode(signature, P384_SIZE) @@ -657,7 +657,7 @@ def p384_keygen() raises -> Tuple[List[UInt8], List[UInt8]]: var public_key = List[UInt8](unsafe_uninit_length=97) if p384_public_key( Span[UInt8, ...](private_key), - Span[mut=True, UInt8, ...](public_key), + Span[mut=True, UInt8, ...](public_key) ): _wipe_u384(d) return (private_key^, public_key^) diff --git a/src/thistle/pbkdf2.mojo b/src/thistle/pbkdf2.mojo index 44a7849..6c19a67 100644 --- a/src/thistle/pbkdf2.mojo +++ b/src/thistle/pbkdf2.mojo @@ -1,7 +1,4 @@ -""" -PBKDF2 (Password-Based Key Derivation Function 2) Implementation in Mojo -SP 800-132 / FIPS 140-2 / RFC 8018 -""" +"""Implements PBKDF2 and HMAC-SHA-2 from RFC 8018.""" from std.collections import List from std.memory import unsafe_memcpy, Pointer from .utils import StackBuffer @@ -14,7 +11,7 @@ from .sha2 import ( sha256_final_to_buffer, sha512_update, sha512_final_to_buffer, - sha512_final_with_len, + sha512_final_with_len ) comptime PBKDF2_SHA256_MAX_DKLEN: Int = 0xFFFFFFFF * 32 @@ -27,8 +24,10 @@ def _secure_zero(ptr: Pointer[mut=True, UInt8, _, address_space=_], count: Int): for i in range(count): ptr.unsafe_store[volatile=True](i, UInt8(0)) + @always_inline -def _xor_block[WIDTH: Int](dst: Pointer[mut=True, UInt8, _, address_space=_], src: Pointer[mut=True, UInt8, _, address_space=_]): +def _xor_block[WIDTH: Int](dst: Pointer[mut=True, UInt8, _, address_space=_], src: Pointer[mut=True, UInt8, _, address_space=_] +): var d = dst.unsafe_bitcast[UInt64]().unsafe_load[width=WIDTH, alignment=1]() var s = src.unsafe_bitcast[UInt64]().unsafe_load[width=WIDTH, alignment=1]() dst.unsafe_bitcast[UInt64]().unsafe_store[width=WIDTH, alignment=1](0, d ^ s) @@ -73,6 +72,7 @@ def _pbkdf2_derive[H: HMACer](mut h: H, salt: Span[UInt8, ...], iterations: Int, _secure_zero(input_block.unsafe_ptr(), 64) return derived_key^ + struct PBKDF2SHA256(HMACer): comptime BLOCK = 64 comptime HASH = 32 @@ -151,6 +151,7 @@ struct PBKDF2SHA256(HMACer): def derive(mut self, salt: Span[UInt8, ...], iterations: Int, dklen: Int) raises -> List[UInt8]: return _pbkdf2_derive(self, salt, iterations, dklen) + def pbkdf2_hmac_sha256( password: Span[UInt8, ...], salt: Span[UInt8, ...], iterations: Int, dkLen: Int ) raises -> List[UInt8]: @@ -163,6 +164,7 @@ def pbkdf2_hmac_sha256( var ctx = PBKDF2SHA256(password) return ctx.derive(salt, iterations, dkLen) + struct PBKDF2SHA512(HMACer): comptime BLOCK = 128 comptime HASH = 64 @@ -241,6 +243,7 @@ struct PBKDF2SHA512(HMACer): def derive(mut self, salt: Span[UInt8, ...], iterations: Int, dklen: Int) raises -> List[UInt8]: return _pbkdf2_derive(self, salt, iterations, dklen) + def pbkdf2_hmac_sha512( password: Span[UInt8, ...], salt: Span[UInt8, ...], iterations: Int, dkLen: Int ) raises -> List[UInt8]: @@ -253,6 +256,7 @@ def pbkdf2_hmac_sha512( var ctx = PBKDF2SHA512(password) return ctx.derive(salt, iterations, dkLen) + def hmac_sha256(key: Span[UInt8, ...], data: Span[UInt8, ...]) -> List[UInt8]: var ctx = PBKDF2SHA256(key) ctx.hmac(data) @@ -261,6 +265,7 @@ def hmac_sha256(key: Span[UInt8, ...], data: Span[UInt8, ...]) -> List[UInt8]: result.append(ctx.u_block[i]) return result^ + def hmac_sha512(key: Span[UInt8, ...], data: Span[UInt8, ...]) -> List[UInt8]: var ctx = PBKDF2SHA512(key) ctx.hmac(data) diff --git a/src/thistle/poly1305.mojo b/src/thistle/poly1305.mojo index af71689..b2e3e0d 100644 --- a/src/thistle/poly1305.mojo +++ b/src/thistle/poly1305.mojo @@ -1,6 +1,4 @@ -""" -Poly1305 per RFC 8439 -""" +"""Implements the Poly1305 one-time authenticator from RFC 8439.""" from std.memory import Pointer from std.collections import InlineArray @@ -14,7 +12,7 @@ def _le64(ptr: Pointer[mut=False, UInt8, _, address_space=_], offset: Int) -> UI return (ptr.unsafe_offset(offset)).unsafe_bitcast[UInt64]().unsafe_load[width=1, alignment=1]() -struct _RPower(Movable, Copyable, ImplicitlyCopyable): +struct _RPower(Copyable, ImplicitlyCopyable, Movable): var r0: UInt64 var r1: UInt64 var r2: UInt64 @@ -31,13 +29,19 @@ struct _RPower(Movable, Copyable, ImplicitlyCopyable): @always_inline def __copyinit__(out self, copy: Self): - self.r0 = copy.r0; self.r1 = copy.r1; self.r2 = copy.r2 - self.s1 = copy.s1; self.s2 = copy.s2 + self.r0 = copy.r0 + self.r1 = copy.r1 + self.r2 = copy.r2 + self.s1 = copy.s1 + self.s2 = copy.s2 @always_inline def __moveinit__(out self, deinit take: Self): - self.r0 = take.r0; self.r1 = take.r1; self.r2 = take.r2 - self.s1 = take.s1; self.s2 = take.s2 + 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): @@ -51,7 +55,7 @@ struct _RPower(Movable, Copyable, ImplicitlyCopyable): @always_inline def _mul_acc( h0: UInt64, h1: UInt64, h2: UInt64, r: _RPower, - mut d0: UInt128, mut d1: UInt128, mut d2: UInt128, + mut d0: UInt128, mut d1: UInt128, mut d2: UInt128 ): d0 += UInt128(h0) * UInt128(r.r0) + UInt128(h1) * UInt128(r.s2) + UInt128(h2) * UInt128(r.s1) d1 += UInt128(h0) * UInt128(r.r1) + UInt128(h1) * UInt128(r.r0) + UInt128(h2) * UInt128(r.s2) @@ -59,7 +63,8 @@ def _mul_acc( @always_inline -def _reduce(mut h0: UInt64, mut h1: UInt64, mut h2: UInt64, d0: UInt128, d1: UInt128, d2: UInt128): +def _reduce(mut h0: UInt64, mut h1: UInt64, mut h2: UInt64, d0: UInt128, d1: UInt128, d2: UInt128 +): var c = d0 >> 44 h0 = d0.cast[DType.uint64]() & _M44 var e1 = d1 + c @@ -75,14 +80,15 @@ def _reduce(mut h0: UInt64, mut h1: UInt64, mut h2: UInt64, d0: UInt128, d1: UIn @always_inline -def _limbs_at(ptr: Pointer[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]( t0 & _M44, ((t0 >> 44) | (t1 << 20)) & _M44, ((t1 >> 24) & _M42) | hibit, - 0, + 0 ) @@ -106,7 +112,7 @@ struct Poly1305: var powers8_ready: Bool def __init__(out self, key: Span[UInt8, ...]) raises: - if len(key) < 32: + if len(key) != 32: raise Error("Poly1305 key must be 32 bytes") var kp = key.unsafe_ptr() var t0 = _le64(kp, 0) @@ -114,7 +120,7 @@ struct Poly1305: self.r = _RPower( t0 & 0xFFC0FFFFFFF, ((t0 >> 44) | (t1 << 20)) & 0xFFFFFC0FFFF, - (t1 >> 24) & 0x00FFFFFFC0F, + (t1 >> 24) & 0x00FFFFFFC0F ) self.pad0 = _le64(kp, 16) self.pad1 = _le64(kp, 24) @@ -186,7 +192,8 @@ struct Poly1305: _reduce(self.h0, self.h1, self.h2, d0, d1, d2) @no_inline - def _blocks8(mut self, ptr: Pointer[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 @@ -218,7 +225,8 @@ struct Poly1305: self.h2 = h2 @no_inline - def _blocks4(mut self, ptr: Pointer[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 @@ -373,7 +381,7 @@ struct Poly1305: def poly1305_mac( key: Span[UInt8, ...], message: Span[UInt8, ...], - output: Span[mut=True, UInt8, ...], + output: Span[mut=True, UInt8, ...] ) raises: if len(output) < 16: raise Error("Poly1305 output needs at least 16 writable bytes") diff --git a/src/thistle/random.mojo b/src/thistle/random.mojo index 5115f47..6a3d9fc 100644 --- a/src/thistle/random.mojo +++ b/src/thistle/random.mojo @@ -1,3 +1,5 @@ +"""Provides operating-system-backed cryptographically secure randomness.""" + from std.collections import List from std.memory import Pointer from std.sys import CompilationTarget, inlined_assembly @@ -12,7 +14,7 @@ def _getrandom_linux_x86(buf: Pointer[mut=True, UInt8, _, address_space=_], leng "syscall", Int64, constraints="={rax},0,{rdi},{rsi},{rdx},~{rcx},~{r11},~{memory}", - has_side_effect=True, + has_side_effect=True ](Int64(318), Int64(Int(buf)), Int64(length), Int64(0)) ) @@ -32,7 +34,7 @@ def _getrandom_linux_arm(buf: Pointer[mut=True, UInt8, _, address_space=_], leng """, Int64, constraints="=&{x0},{x9},{x10},~{x1},~{x2},~{x8},~{memory}", - has_side_effect=True, + has_side_effect=True ](Int64(Int(buf)), Int64(length)) ) @@ -56,7 +58,7 @@ def _getentropy_macos_arm(buf: Pointer[mut=True, UInt8, _, address_space=_], len """, Int64, constraints="=&{x0},{x9},{x10},~{x1},~{x16},~{cc},~{memory}", - has_side_effect=True, + has_side_effect=True ](Int64(Int(buf)), Int64(length)) ) @@ -71,6 +73,8 @@ def _fill_linux_x86(buf: Pointer[mut=True, UInt8, _, address_space=_], length: I raise Error("getrandom syscall failed") if ret == 0: raise Error("getrandom returned zero bytes") + if ret > length - offset: + raise Error("getrandom returned too many bytes") offset += ret @@ -84,6 +88,8 @@ def _fill_linux_arm(buf: Pointer[mut=True, UInt8, _, address_space=_], length: I raise Error("getrandom syscall failed") if ret == 0: raise Error("getrandom returned zero bytes") + if ret > length - offset: + raise Error("getrandom returned too many bytes") offset += ret @@ -93,6 +99,8 @@ def _fill_macos_arm(buf: Pointer[mut=True, UInt8, _, address_space=_], length: I var chunk = min(256, length - offset) var ret = _getentropy_macos_arm(buf.unsafe_offset(offset), chunk) if ret < 0: + if ret == -4: # EINTR + continue raise Error("getentropy syscall failed") if ret != 0: raise Error("getentropy returned unexpected value") diff --git a/src/thistle/rsa.mojo b/src/thistle/rsa.mojo index 285dd7e..37aec4c 100644 --- a/src/thistle/rsa.mojo +++ b/src/thistle/rsa.mojo @@ -1,6 +1,4 @@ -""" -RSASSA-PSS signature verification per RFC 8017 -""" +"""Implements RSA-PSS signatures and PKCS#1 v1.5 verification.""" from std.collections import List, InlineArray from std.memory import Pointer @@ -10,7 +8,7 @@ from .random import random_bytes from .sha2 import ( sha224_hash, sha256_hash, sha384_hash, sha512_hash, SHA256Context, sha256_update, sha256_final_to_buffer, - SHA512Context, sha512_update, sha512_final_to_buffer, + SHA512Context, sha512_update, sha512_final_to_buffer ) comptime _NL = 66 @@ -121,7 +119,8 @@ def _hash_len(alg: Int) raises -> Int: raise Error("unsupported hash for RSA-PSS") -def _hash_into(alg: Int, data: Span[UInt8, ...], output: Pointer[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) @@ -158,7 +157,7 @@ def _mgf1( seed: Pointer[UInt8, _], seed_len: Int, mask_len: Int, - output: Pointer[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) @@ -175,7 +174,7 @@ def _mgf1( _ = _hash_into( alg, Span[UInt8, ...](unsafe_ptr=block.unsafe_ptr(), length=seed_len + 4), - digest.unsafe_ptr(), + digest.unsafe_ptr() ) var take = mask_len - done if take > h_len: @@ -192,7 +191,7 @@ def _emsa_pss_encode( sha: Int, mgf_sha: Int, em_bits: Int, - output: Pointer[mut=True, UInt8, _, address_space=_], + output: Pointer[mut=True, UInt8, _, address_space=_] ) raises -> Bool: var h_len = _hash_len(sha) _ = _hash_len(mgf_sha) @@ -211,7 +210,7 @@ def _emsa_pss_encode( _ = _hash_into( sha, Span[UInt8, ...](unsafe_ptr=mprime.unsafe_ptr(), length=8 + h_len + len(salt)), - h.unsafe_ptr(), + h.unsafe_ptr() ) var db_len = em_len - h_len - 1 @@ -282,16 +281,14 @@ def _bn_ge(a: StaticTuple[UInt64, _NL], b: StaticTuple[UInt64, _NL], k: Int) -> def _bn_ge_ct( a: StaticTuple[UInt64, _NL], b: StaticTuple[UInt64, _NL], - k: Int, + k: Int ) -> Bool: var borrow: UInt64 = 0 for i in range(k): - var d = ( - (UInt128(1) << 64) + var d = (UInt128(1) << 64) + UInt128(a[i]) - UInt128(b[i]) - UInt128(borrow) - ) borrow = 1 - (d >> 64).cast[DType.uint64]() return borrow == 0 @@ -323,7 +320,7 @@ def _bn_select( a: StaticTuple[UInt64, _NL], b: StaticTuple[UInt64, _NL], choice: UInt64, - k: Int, + k: Int ) -> StaticTuple[UInt64, _NL]: var out = a var mask = UInt64(0) - (choice & UInt64(1)) @@ -355,7 +352,7 @@ def _mont_mul_k[K: Int]( a: StaticTuple[UInt64, _NL], b: StaticTuple[UInt64, _NL], n: StaticTuple[UInt64, _NL], - n0: UInt64, + n0: UInt64 ) -> StaticTuple[UInt64, _NL]: var t = StaticTuple[UInt64, _NL]() comptime for z in range(K): @@ -416,7 +413,7 @@ def _mont_mul_any( b: StaticTuple[UInt64, _NL], n: StaticTuple[UInt64, _NL], n0: UInt64, - k: Int, + k: Int ) -> StaticTuple[UInt64, _NL]: var t = _bn_zero() var t_hi: UInt64 = 0 @@ -445,7 +442,7 @@ def _mont_mul_any( def _mont_sqr_k[K: Int]( a: StaticTuple[UInt64, _NL], n: StaticTuple[UInt64, _NL], - n0: UInt64, + n0: UInt64 ) -> StaticTuple[UInt64, _NL]: var t = InlineArray[UInt64, 2 * _NL + 2](fill=0) comptime for z in range(2 * K + 1): @@ -501,7 +498,7 @@ def _mont_sqr( a: StaticTuple[UInt64, _NL], n: StaticTuple[UInt64, _NL], n0: UInt64, - k: Int, + k: Int ) -> StaticTuple[UInt64, _NL]: if k == 16: return _mont_sqr_k[16](a, n, n0) @@ -520,7 +517,7 @@ def _mont_mul( b: StaticTuple[UInt64, _NL], n: StaticTuple[UInt64, _NL], n0: UInt64, - k: Int, + k: Int ) -> StaticTuple[UInt64, _NL]: if k == 16: return _mont_mul_k[16](a, b, n, n0) @@ -621,7 +618,7 @@ struct RsaPublicKey: def _public_op( self, sig: Span[UInt8, ...], - output: Pointer[mut=True, UInt8, _, address_space=_], + output: Pointer[mut=True, UInt8, _, address_space=_] ) raises -> Bool: var nb = self.nb var k = self.k @@ -665,7 +662,7 @@ struct RsaPublicKey: signature: Span[UInt8, ...], sha: Int, mgf_sha: Int, - salt_len: Int, + salt_len: Int ) raises -> Bool: var h_len = _hash_len(sha) _ = _hash_len(mgf_sha) @@ -722,7 +719,7 @@ struct RsaPublicKey: _ = _hash_into( sha, Span[UInt8, ...](unsafe_ptr=mprime.unsafe_ptr(), length=8 + h_len + salt_len), - h2.unsafe_ptr(), + h2.unsafe_ptr() ) var diff: UInt8 = 0 @@ -756,7 +753,7 @@ struct RsaPrivateKey: out self, modulus: Span[UInt8, ...], exponent: Span[UInt8, ...], - private_exponent: Span[UInt8, ...], + private_exponent: Span[UInt8, ...] ) raises: self.public = RsaPublicKey(modulus, exponent) if self.public.mod_bits < 2048: @@ -787,7 +784,7 @@ struct RsaPrivateKey: def _private_op( self, encoded: Span[UInt8, ...], - signature: Pointer[mut=True, UInt8, _, address_space=_], + signature: Pointer[mut=True, UInt8, _, address_space=_] ) raises -> Bool: var nb = self.public.nb var k = self.public.k @@ -854,33 +851,40 @@ struct RsaPrivateKey: salt: Span[UInt8, ...], sha: Int, mgf_sha: Int, - signature: Span[mut=True, UInt8, ...], + 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) - if not _emsa_pss_encode( - message, salt, sha, mgf_sha, em_bits, - encoded.unsafe_ptr().unsafe_offset((self.public.nb - em_len)), - ): - return False - var ok = self._private_op( - 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): - ep.unsafe_store[volatile=True](i, UInt8(0)) - return ok + try: + if not _emsa_pss_encode( + message, + salt, + sha, + mgf_sha, + em_bits, + encoded.unsafe_ptr().unsafe_offset(self.public.nb - em_len), + ): + return False + return self._private_op( + Span[UInt8, ...]( + unsafe_ptr=encoded.unsafe_ptr(), length=self.public.nb + ), + signature.unsafe_ptr(), + ) + finally: + var ep = encoded.unsafe_ptr() + for i in range(self.public.nb): + ep.unsafe_store[volatile=True](i, UInt8(0)) def pss_sign( self, message: Span[UInt8, ...], sha: Int, mgf_sha: Int, - salt_len: Int, + salt_len: Int ) raises -> List[UInt8]: if salt_len < 0: raise Error("RSA-PSS salt length must be non-negative") @@ -890,17 +894,21 @@ struct RsaPrivateKey: if em_len < h_len + 2 or salt_len > em_len - h_len - 2: raise Error("RSA-PSS salt is too long for the modulus") 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, - Span[mut=True, UInt8, ...](signature), - ) - var salt_ptr = salt.unsafe_ptr() - for i in range(len(salt)): - salt_ptr.unsafe_store[volatile=True](i, UInt8(0)) - if not ok: - raise Error("RSA-PSS signing failed") - return signature^ + try: + var signature = List[UInt8](unsafe_uninit_length=self.public.nb) + if not self.pss_sign_with_salt( + message, + Span[UInt8, ...](salt), + sha, + mgf_sha, + Span[mut=True, UInt8, ...](signature) + ): + raise Error("RSA-PSS signing failed") + return signature^ + finally: + var salt_ptr = salt.unsafe_ptr() + for i in range(len(salt)): + salt_ptr.unsafe_store[volatile=True](i, UInt8(0)) def _bn_reduce_bytes( @@ -924,7 +932,7 @@ def _bn_reduce_bytes( def _private_pow( key: RsaPublicKey, exponent: InlineArray[UInt8, 528], - input: StaticTuple[UInt64, _NL], + input: StaticTuple[UInt64, _NL] ) -> StaticTuple[UInt64, _NL]: var base = _mont_mul(input, key.r2, key.n, key.n0, key.k) var table = StaticTuple[StaticTuple[UInt64, _NL], 16]() @@ -957,7 +965,7 @@ def _private_pow( def _bn_mul_parts( a: StaticTuple[UInt64, _NL], a_len: Int, - b: StaticTuple[UInt64, _NL], b_len: Int, + b: StaticTuple[UInt64, _NL], b_len: Int ) -> StaticTuple[UInt64, _NL]: var out = _bn_zero() for i in range(a_len): @@ -992,7 +1000,7 @@ struct RsaCrtPrivateKey: modulus: Span[UInt8, ...], exponent: Span[UInt8, ...], prime1: Span[UInt8, ...], prime2: Span[UInt8, ...], exponent1: Span[UInt8, ...], exponent2: Span[UInt8, ...], - coefficient: Span[UInt8, ...], + coefficient: Span[UInt8, ...] ) raises: self.public = RsaPublicKey(modulus, exponent) if self.public.mod_bits < 2048: @@ -1061,7 +1069,7 @@ struct RsaCrtPrivateKey: ) var check = _mont_mul( _mont_mul(q_mod_p, self.p.r2, self.p.n, self.p.n0, self.p.k), - self.qinv, self.p.n, self.p.n0, self.p.k, + self.qinv, self.p.n, self.p.n0, self.p.k ) var coefficient_diff = check[0] ^ UInt64(1) for i in range(1, self.p.k): @@ -1093,7 +1101,7 @@ struct RsaCrtPrivateKey: def _private_op( self, encoded: Span[UInt8, ...], - signature: Pointer[mut=True, UInt8, _, address_space=_], + signature: Pointer[mut=True, UInt8, _, address_space=_] ) raises -> Bool: if len(encoded) != self.public.nb: return False @@ -1119,7 +1127,7 @@ struct RsaCrtPrivateKey: h = _bn_select(h, h_plus_p, borrow, self.p.k) h = _mont_mul( _mont_mul(h, self.p.r2, self.p.n, self.p.n0, self.p.k), - self.qinv, self.p.n, self.p.n0, self.p.k, + self.qinv, self.p.n, self.p.n0, self.p.k ) var result = _bn_mul_parts(self.q.n, self.q.k, h, self.p.k) @@ -1163,33 +1171,40 @@ struct RsaCrtPrivateKey: def pss_sign_with_salt( self, message: Span[UInt8, ...], salt: Span[UInt8, ...], sha: Int, mgf_sha: Int, - signature: Span[mut=True, UInt8, ...], + 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) - if not _emsa_pss_encode( - message, salt, sha, mgf_sha, em_bits, - encoded.unsafe_ptr().unsafe_offset((self.public.nb - em_len)), - ): - return False - var ok = self._private_op( - 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): - ep.unsafe_store[volatile=True](i, UInt8(0)) - return ok + try: + if not _emsa_pss_encode( + message, + salt, + sha, + mgf_sha, + em_bits, + encoded.unsafe_ptr().unsafe_offset(self.public.nb - em_len), + ): + return False + return self._private_op( + Span[UInt8, ...]( + unsafe_ptr=encoded.unsafe_ptr(), length=self.public.nb + ), + signature.unsafe_ptr(), + ) + finally: + var ep = encoded.unsafe_ptr() + for i in range(self.public.nb): + ep.unsafe_store[volatile=True](i, UInt8(0)) def pss_sign( self, message: Span[UInt8, ...], sha: Int, mgf_sha: Int, - salt_len: Int, + salt_len: Int ) raises -> List[UInt8]: if salt_len < 0: raise Error("RSA-PSS salt length must be non-negative") @@ -1199,23 +1214,28 @@ struct RsaCrtPrivateKey: if em_len < h_len + 2 or salt_len > em_len - h_len - 2: raise Error("RSA-PSS salt is too long for the modulus") 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, - Span[mut=True, UInt8, ...](signature), - ) - var salt_ptr = salt.unsafe_ptr() - for i in range(len(salt)): - salt_ptr.unsafe_store[volatile=True](i, UInt8(0)) - if not ok: - raise Error("RSA-PSS signing failed") - return signature^ + try: + var signature = List[UInt8](unsafe_uninit_length=self.public.nb) + if not self.pss_sign_with_salt( + message, + Span[UInt8, ...](salt), + sha, + mgf_sha, + Span[mut=True, UInt8, ...](signature) + ): + raise Error("RSA-PSS signing failed") + return signature^ + finally: + var salt_ptr = salt.unsafe_ptr() + for i in range(len(salt)): + salt_ptr.unsafe_store[volatile=True](i, UInt8(0)) + def _pkcs1_v15_verify( key: RsaPublicKey, message: Span[UInt8, ...], signature: Span[UInt8, ...], - sha: Int, + sha: Int ) raises -> Bool: var h_len = _hash_len(sha) var prefix_len = _digest_info_prefix_len(sha) @@ -1252,7 +1272,7 @@ def rsa_pss_verify( signature: Span[UInt8, ...], sha: Int, mgf_sha: Int, - salt_len: Int, + salt_len: Int ) raises -> Bool: var key: RsaPublicKey try: @@ -1269,7 +1289,7 @@ def rsa_pss_sign_with_salt( message: Span[UInt8, ...], salt: Span[UInt8, ...], sha: Int, - mgf_sha: Int, + mgf_sha: Int ) raises -> List[UInt8]: var key = RsaPrivateKey(modulus, exponent, private_exponent) var signature = List[UInt8](unsafe_uninit_length=key.public.nb) @@ -1287,7 +1307,7 @@ def rsa_pss_sign( message: Span[UInt8, ...], sha: Int, mgf_sha: Int, - salt_len: Int, + salt_len: Int ) raises -> List[UInt8]: var key = RsaPrivateKey(modulus, exponent, private_exponent) return key.pss_sign(message, sha, mgf_sha, salt_len) @@ -1298,7 +1318,7 @@ def rsa_pss_crt_sign_with_salt( prime1: Span[UInt8, ...], prime2: Span[UInt8, ...], exponent1: Span[UInt8, ...], exponent2: Span[UInt8, ...], coefficient: Span[UInt8, ...], message: Span[UInt8, ...], - salt: Span[UInt8, ...], sha: Int, mgf_sha: Int, + salt: Span[UInt8, ...], sha: Int, mgf_sha: Int ) raises -> List[UInt8]: var key = RsaCrtPrivateKey( modulus, exponent, prime1, prime2, exponent1, exponent2, coefficient @@ -1316,7 +1336,7 @@ def rsa_pss_crt_sign( prime1: Span[UInt8, ...], prime2: Span[UInt8, ...], exponent1: Span[UInt8, ...], exponent2: Span[UInt8, ...], coefficient: Span[UInt8, ...], message: Span[UInt8, ...], - sha: Int, mgf_sha: Int, salt_len: Int, + sha: Int, mgf_sha: Int, salt_len: Int ) raises -> List[UInt8]: var key = RsaCrtPrivateKey( modulus, exponent, prime1, prime2, exponent1, exponent2, coefficient @@ -1326,21 +1346,21 @@ def rsa_pss_crt_sign( def rsa_pss_sha256_sign( modulus: Span[UInt8, ...], exponent: Span[UInt8, ...], - private_exponent: Span[UInt8, ...], message: Span[UInt8, ...], + private_exponent: Span[UInt8, ...], message: Span[UInt8, ...] ) raises -> List[UInt8]: return rsa_pss_sign(modulus, exponent, private_exponent, message, SHA256, SHA256, 32) def rsa_pss_sha384_sign( modulus: Span[UInt8, ...], exponent: Span[UInt8, ...], - private_exponent: Span[UInt8, ...], message: Span[UInt8, ...], + private_exponent: Span[UInt8, ...], message: Span[UInt8, ...] ) raises -> List[UInt8]: return rsa_pss_sign(modulus, exponent, private_exponent, message, SHA384, SHA384, 48) def rsa_pss_sha512_sign( modulus: Span[UInt8, ...], exponent: Span[UInt8, ...], - private_exponent: Span[UInt8, ...], message: Span[UInt8, ...], + private_exponent: Span[UInt8, ...], message: Span[UInt8, ...] ) raises -> List[UInt8]: return rsa_pss_sign(modulus, exponent, private_exponent, message, SHA512, SHA512, 64) @@ -1349,11 +1369,11 @@ def rsa_pss_crt_sha256_sign( modulus: Span[UInt8, ...], exponent: Span[UInt8, ...], prime1: Span[UInt8, ...], prime2: Span[UInt8, ...], exponent1: Span[UInt8, ...], exponent2: Span[UInt8, ...], - coefficient: Span[UInt8, ...], message: Span[UInt8, ...], + coefficient: Span[UInt8, ...], message: Span[UInt8, ...] ) raises -> List[UInt8]: return rsa_pss_crt_sign( modulus, exponent, prime1, prime2, exponent1, exponent2, - coefficient, message, SHA256, SHA256, 32, + coefficient, message, SHA256, SHA256, 32 ) @@ -1361,11 +1381,11 @@ def rsa_pss_crt_sha384_sign( modulus: Span[UInt8, ...], exponent: Span[UInt8, ...], prime1: Span[UInt8, ...], prime2: Span[UInt8, ...], exponent1: Span[UInt8, ...], exponent2: Span[UInt8, ...], - coefficient: Span[UInt8, ...], message: Span[UInt8, ...], + coefficient: Span[UInt8, ...], message: Span[UInt8, ...] ) raises -> List[UInt8]: return rsa_pss_crt_sign( modulus, exponent, prime1, prime2, exponent1, exponent2, - coefficient, message, SHA384, SHA384, 48, + coefficient, message, SHA384, SHA384, 48 ) @@ -1373,11 +1393,11 @@ def rsa_pss_crt_sha512_sign( modulus: Span[UInt8, ...], exponent: Span[UInt8, ...], prime1: Span[UInt8, ...], prime2: Span[UInt8, ...], exponent1: Span[UInt8, ...], exponent2: Span[UInt8, ...], - coefficient: Span[UInt8, ...], message: Span[UInt8, ...], + coefficient: Span[UInt8, ...], message: Span[UInt8, ...] ) raises -> List[UInt8]: return rsa_pss_crt_sign( modulus, exponent, prime1, prime2, exponent1, exponent2, - coefficient, message, SHA512, SHA512, 64, + coefficient, message, SHA512, SHA512, 64 ) @@ -1385,7 +1405,7 @@ def rsa_pss_sha256_verify( modulus: Span[UInt8, ...], exponent: Span[UInt8, ...], message: Span[UInt8, ...], - signature: Span[UInt8, ...], + signature: Span[UInt8, ...] ) raises -> Bool: return rsa_pss_verify(modulus, exponent, message, signature, SHA256, SHA256, 32) @@ -1394,7 +1414,7 @@ def rsa_pss_sha384_verify( modulus: Span[UInt8, ...], exponent: Span[UInt8, ...], message: Span[UInt8, ...], - signature: Span[UInt8, ...], + signature: Span[UInt8, ...] ) raises -> Bool: return rsa_pss_verify(modulus, exponent, message, signature, SHA384, SHA384, 48) @@ -1403,7 +1423,7 @@ def rsa_pss_sha512_verify( modulus: Span[UInt8, ...], exponent: Span[UInt8, ...], message: Span[UInt8, ...], - signature: Span[UInt8, ...], + signature: Span[UInt8, ...] ) raises -> Bool: return rsa_pss_verify(modulus, exponent, message, signature, SHA512, SHA512, 64) @@ -1413,7 +1433,7 @@ def rsa_pkcs1_v15_verify( exponent: Span[UInt8, ...], message: Span[UInt8, ...], signature: Span[UInt8, ...], - sha: Int, + sha: Int ) raises -> Bool: var key: RsaPublicKey try: @@ -1425,27 +1445,27 @@ def rsa_pkcs1_v15_verify( def rsa_pkcs1_v15_sha1_verify( modulus: Span[UInt8, ...], exponent: Span[UInt8, ...], - message: Span[UInt8, ...], signature: Span[UInt8, ...], + message: Span[UInt8, ...], signature: Span[UInt8, ...] ) raises -> Bool: return rsa_pkcs1_v15_verify(modulus, exponent, message, signature, SHA1) def rsa_pkcs1_v15_sha256_verify( modulus: Span[UInt8, ...], exponent: Span[UInt8, ...], - message: Span[UInt8, ...], signature: Span[UInt8, ...], + message: Span[UInt8, ...], signature: Span[UInt8, ...] ) raises -> Bool: return rsa_pkcs1_v15_verify(modulus, exponent, message, signature, SHA256) def rsa_pkcs1_v15_sha384_verify( modulus: Span[UInt8, ...], exponent: Span[UInt8, ...], - message: Span[UInt8, ...], signature: Span[UInt8, ...], + message: Span[UInt8, ...], signature: Span[UInt8, ...] ) raises -> Bool: return rsa_pkcs1_v15_verify(modulus, exponent, message, signature, SHA384) def rsa_pkcs1_v15_sha512_verify( modulus: Span[UInt8, ...], exponent: Span[UInt8, ...], - message: Span[UInt8, ...], signature: Span[UInt8, ...], + message: Span[UInt8, ...], signature: Span[UInt8, ...] ) raises -> Bool: return rsa_pkcs1_v15_verify(modulus, exponent, message, signature, SHA512) diff --git a/src/thistle/sha2.mojo b/src/thistle/sha2.mojo index 632e7a2..1c2c32e 100644 --- a/src/thistle/sha2.mojo +++ b/src/thistle/sha2.mojo @@ -1,7 +1,4 @@ -""" -SHA-2 (SHA-256/SHA-512) Implementation in Mojo -RFC 6234 / FIPS 180-4 / CAVP validated -""" +"""Implements SHA-224, SHA-256, SHA-384, and SHA-512.""" from std.collections import List from std.memory import Pointer, unsafe_memcpy, unsafe_memset_zero @@ -10,6 +7,7 @@ from std.builtin.simd import SIMD from std.builtin.dtype import DType from std.sys import CompilationTarget from std.utils import StaticTuple +from std.os import abort from .sha_ni import sha512ni_transform_blocks, sha256ni_transform_blocks from .utils import bytes_to_hex, string_to_bytes, load_32be, load_64be @@ -26,7 +24,7 @@ comptime SHA256_K = SIMD[DType.uint32, 64]( 0xD6990624, 0xF40E3585, 0x106AA070, 0x19A4C116, 0x1E376C08, 0x2748774C, 0x34B0BCB5, 0x391C0CB3, 0x4ED8AA4A, 0x5B9CCA4F, 0x682E6FF3, 0x748F82EE, 0x78A5636F, 0x84C87814, 0x8CC70208, - 0x90BEFFFA, 0xA4506CEB, 0xBEF9A3F7, 0xC67178F2, + 0x90BEFFFA, 0xA4506CEB, 0xBEF9A3F7, 0xC67178F2 ) comptime SHA512_K = StaticTuple[UInt64, 80]( @@ -45,9 +43,10 @@ comptime SHA512_K = StaticTuple[UInt64, 80]( 0x90BEFFFA23631E28, 0xA4506CEBDE82BDE9, 0xBEF9A3F7B2C67915, 0xC67178F2E372532B, 0xCA273ECEEA26619C, 0xD186B8C721C0C207, 0xEADA7DD6CDE0EB1E, 0xF57D4F7FEE6ED178, 0x06F067AA72176FBA, 0x0A637DC5A2C898A6, 0x113F9804BEF90DAE, 0x1B710B35131C471B, 0x28DB77F523047D84, 0x32CAAB7B40C72493, 0x3C9EBE0A15C9BEBC, - 0x431D67C49C100D4C, 0x4CC5D4BECB3E42B6, 0x597F299CFC657E2A, 0x5FCB6FAB3AD6FAEC, 0x6C44198C4A475817, + 0x431D67C49C100D4C, 0x4CC5D4BECB3E42B6, 0x597F299CFC657E2A, 0x5FCB6FAB3AD6FAEC, 0x6C44198C4A475817 ) + @always_inline def ch32(x: UInt32, y: UInt32, z: UInt32) -> UInt32: return (x & y) ^ ((~x) & z) @@ -116,7 +115,7 @@ comptime SHA256_IV = SIMD[DType.uint32, 8]( 0x510E527F, 0x9B05688C, 0x1F83D9AB, - 0x5BE0CD19, + 0x5BE0CD19 ) comptime SHA224_IV = SIMD[DType.uint32, 8]( @@ -127,7 +126,7 @@ comptime SHA224_IV = SIMD[DType.uint32, 8]( 0xFFC00B31, 0x68581511, 0x64F98FA7, - 0xBEFA4FA4, + 0xBEFA4FA4 ) comptime SHA512_IV = SIMD[DType.uint64, 8]( @@ -138,7 +137,7 @@ comptime SHA512_IV = SIMD[DType.uint64, 8]( 0x510E527FADE682D1, 0x9B05688C2B3E6C1F, 0x1F83D9ABFB41BD6B, - 0x5BE0CD19137E2179, + 0x5BE0CD19137E2179 ) comptime SHA384_IV = SIMD[DType.uint64, 8]( @@ -149,9 +148,10 @@ comptime SHA384_IV = SIMD[DType.uint64, 8]( 0x67332667FFC00B31, 0x8EB44A8768581511, 0xDB0C2E0D64F98FA7, - 0x47B5481DBEFA4FA4, + 0x47B5481DBEFA4FA4 ) + struct SHA256Context(Movable): var state: SIMD[DType.uint32, 8] var count: UInt64 @@ -198,11 +198,12 @@ struct SHA256Context(Movable): 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: Pointer[mut=False, UInt8, _, address_space=_], - nblocks: Int, + 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"]()): sha256ni_transform_blocks(state, data, nblocks) @@ -326,6 +327,7 @@ def sha256_final(mut ctx: SHA256Context) -> List[UInt8]: sha256_final_to_buffer(ctx, output.unsafe_ptr()) return output^ + def sha256_hash(data: Span[UInt8, ...]) -> List[UInt8]: var ctx = SHA256Context() sha256_update(ctx, data) @@ -373,6 +375,7 @@ def sha256_final_to_buffer(mut ctx: SHA256Context, output: Pointer[mut=True, UIn for i in range(8): (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] var count_high: UInt64 @@ -444,7 +447,7 @@ struct SHA512Context(Movable): def sha512_transform_blocks( mut state: SIMD[DType.uint64, 8], data: Pointer[mut=False, UInt8, _, address_space=_], - nblocks: Int, + nblocks: Int ): comptime if CompilationTarget.has_neon() and CompilationTarget._has_feature["sha3"]() and not CompilationTarget.is_x86(): sha512ni_transform_blocks(state, data, nblocks) @@ -578,6 +581,7 @@ def sha512_final(mut ctx: SHA512Context) -> List[UInt8]: sha512_final_to_buffer(ctx, output.unsafe_ptr()) return output^ + def sha512_hash(data: Span[UInt8, ...]) -> List[UInt8]: var ctx = SHA512Context(SHA512_IV) sha512_update(ctx, data) @@ -630,7 +634,10 @@ def sha512_final_to_buffer(mut ctx: SHA512Context, output: Pointer[mut=True, UIn for i in range(8): (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]: + if output_len < 0 or output_len > 32: + abort("SHA-256 output length must be between 0 and 32 bytes") var bit_count = ctx.count + UInt64(ctx.buffer_len) * 8 _sha256_pad(ctx, bit_count) return _sha256_output(ctx, output_len) @@ -657,6 +664,10 @@ def padding_bit(n: Int) -> UInt8: def sha256_final_partial( mut ctx: SHA256Context, final_octet: UInt8, final_bit_count: Int, output_len: Int ) -> List[UInt8]: + if final_bit_count < 0 or final_bit_count > 7: + abort("SHA-256 final bit count must be between 0 and 7") + if output_len < 0 or output_len > 32: + abort("SHA-256 output length must be between 0 and 32 bytes") var bit_count = ctx.count + UInt64(ctx.buffer_len) * 8 + UInt64(final_bit_count) var pad = (final_octet & high_bits_mask(final_bit_count)) | padding_bit(final_bit_count) _sha256_pad_byte(ctx, pad, bit_count) @@ -664,9 +675,14 @@ def sha256_final_partial( def sha224_hash_bits(data: Span[UInt8, ...], bit_len: Int) -> List[UInt8]: + if bit_len < 0: + abort("SHA-224 bit length cannot be negative") + var available_bytes = len(data) var ctx = SHA256Context(SHA224_IV) var full_bytes = bit_len // 8 var rem_bits = bit_len % 8 + if full_bytes > available_bytes or (full_bytes == available_bytes and rem_bits != 0): + abort("SHA-224 bit length exceeds input") if full_bytes > 0: sha256_update(ctx, data[0:full_bytes]) if rem_bits == 0: @@ -675,9 +691,14 @@ def sha224_hash_bits(data: Span[UInt8, ...], bit_len: Int) -> List[UInt8]: def sha256_hash_bits(data: Span[UInt8, ...], bit_len: Int) -> List[UInt8]: + if bit_len < 0: + abort("SHA-256 bit length cannot be negative") + var available_bytes = len(data) var ctx = SHA256Context() var full_bytes = bit_len // 8 var rem_bits = bit_len % 8 + if full_bytes > available_bytes or (full_bytes == available_bytes and rem_bits != 0): + abort("SHA-256 bit length exceeds input") if full_bytes > 0: sha256_update(ctx, data[0:full_bytes]) if rem_bits == 0: @@ -686,6 +707,8 @@ def sha256_hash_bits(data: Span[UInt8, ...], bit_len: Int) -> List[UInt8]: def sha512_final_with_len(mut ctx: SHA512Context, output_len: Int) -> List[UInt8]: + if output_len < 0 or output_len > 64: + abort("SHA-512 output length must be between 0 and 64 bytes") var low = ctx.count_low + UInt64(ctx.buffer_len) * 8 var high = ctx.count_high if low < ctx.count_low: @@ -703,6 +726,10 @@ def sha384_hash(data: Span[UInt8, ...]) -> List[UInt8]: def sha512_final_partial( mut ctx: SHA512Context, final_octet: UInt8, final_bit_count: Int, output_len: Int ) -> List[UInt8]: + if final_bit_count < 0 or final_bit_count > 7: + abort("SHA-512 final bit count must be between 0 and 7") + if output_len < 0 or output_len > 64: + abort("SHA-512 output length must be between 0 and 64 bytes") var low = ctx.count_low + UInt64(ctx.buffer_len) * 8 + UInt64(final_bit_count) var high = ctx.count_high if low < ctx.count_low: @@ -713,9 +740,14 @@ def sha512_final_partial( def sha384_hash_bits(data: Span[UInt8, ...], bit_len: Int) -> List[UInt8]: + if bit_len < 0: + abort("SHA-384 bit length cannot be negative") + var available_bytes = len(data) var ctx = SHA512Context(SHA384_IV) var full_bytes = bit_len // 8 var rem_bits = bit_len % 8 + if full_bytes > available_bytes or (full_bytes == available_bytes and rem_bits != 0): + abort("SHA-384 bit length exceeds input") if full_bytes > 0: sha512_update(ctx, data[0:full_bytes]) if rem_bits == 0: @@ -724,15 +756,21 @@ def sha384_hash_bits(data: Span[UInt8, ...], bit_len: Int) -> List[UInt8]: def sha512_hash_bits(data: Span[UInt8, ...], bit_len: Int) -> List[UInt8]: + if bit_len < 0: + abort("SHA-512 bit length cannot be negative") + var available_bytes = len(data) var ctx = SHA512Context(SHA512_IV) var full_bytes = bit_len // 8 var rem_bits = bit_len % 8 + if full_bytes > available_bytes or (full_bytes == available_bytes and rem_bits != 0): + abort("SHA-512 bit length exceeds input") if full_bytes > 0: sha512_update(ctx, data[0:full_bytes]) if rem_bits == 0: return sha512_final(ctx) return sha512_final_partial(ctx, data[full_bytes], rem_bits, 64) + def sha256_hash_string(s: String) -> String: var data = string_to_bytes(s) var hash = sha256_hash(Span[UInt8, ...](data)) diff --git a/src/thistle/sha3.mojo b/src/thistle/sha3.mojo index bc973e7..b128bd6 100644 --- a/src/thistle/sha3.mojo +++ b/src/thistle/sha3.mojo @@ -1,7 +1,4 @@ -""" -SHA-3 (Keccak) + shake128/shake256 Implementation in Mojo -FIPS 202 -""" +"""Implements SHA-3 and SHAKE as specified by FIPS 202.""" from std.collections import List from std.memory import Pointer, unsafe_memcpy, unsafe_memset_zero @@ -25,9 +22,10 @@ comptime KECCAK_RC = StaticTuple[UInt64, 24]( 0x8000000000008002, 0x8000000000000080, 0x000000000000800A, 0x800000008000000A, 0x8000000080008081, 0x8000000000008080, - 0x0000000080000001, 0x8000000080008008, + 0x0000000080000001, 0x8000000080008008 ) + @always_inline def rotl64[n: Int](x: UInt64) -> UInt64: return rotate_bits_left[n](x) @@ -368,10 +366,13 @@ struct SHA3Context(Movable): @always_inline -def sha3_absorb_block(state: Pointer[mut=True, UInt64, _, address_space=_], block: Pointer[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[unsafe_offset=i] ^= block.unsafe_offset(i * 8).unsafe_bitcast[UInt64]().unsafe_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) @@ -431,7 +432,7 @@ def sha3_final(mut ctx: SHA3Context, output_len_bytes: Int) -> List[UInt8]: unsafe_memcpy( dest=output.unsafe_ptr().unsafe_offset(offset), src=ctx.state.ptr().unsafe_bitcast[UInt8](), - count=limit, + count=limit ) offset += limit @@ -442,7 +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): +def sha3_final_into(mut ctx: SHA3Context, mut output: StackBuffer[UInt8, ...], output_len_bytes: Int +): if output_len_bytes < 0 or output_len_bytes > output.capacity(): abort("SHA-3 output length exceeds destination capacity") output.clear() @@ -468,7 +470,7 @@ def sha3_final_into(mut ctx: SHA3Context, mut output: StackBuffer[UInt8, ...], o unsafe_memcpy( dest=output.ptr().unsafe_offset(offset), src=ctx.state.ptr().unsafe_bitcast[UInt8](), - count=limit, + count=limit ) offset += limit @@ -484,7 +486,8 @@ def sha3_hash(rate_bits: Int, data: Span[UInt8, ...], output_len: Int) -> List[U @always_inline -def sha3_hash_into(mut output: StackBuffer[UInt8, ...], rate_bits: Int, data: Span[UInt8, ...], output_len: Int): +def sha3_hash_into(mut output: StackBuffer[UInt8, ...], rate_bits: Int, data: Span[UInt8, ...], output_len: Int +): var ctx = SHA3Context(rate_bits) sha3_update(ctx, data) sha3_final_into(ctx, output, output_len) @@ -571,7 +574,7 @@ def shake_squeeze_prefix_into(mut ctx: SHA3Context, mut output: StackBuffer[UInt unsafe_memcpy( dest=output.ptr().unsafe_offset(offset), src=ctx.state.ptr().unsafe_bitcast[UInt8](), - count=limit, + count=limit ) offset += limit @@ -584,7 +587,6 @@ def shake_advance(mut ctx: SHA3Context): keccak_f1600(ctx.state.ptr()) - @always_inline def shake_final(mut ctx: SHA3Context, output_len: Int) -> List[UInt8]: if output_len < 0: @@ -604,7 +606,7 @@ def shake_final(mut ctx: SHA3Context, output_len: Int) -> List[UInt8]: unsafe_memcpy( dest=output.unsafe_ptr().unsafe_offset(offset), src=ctx.state.ptr().unsafe_bitcast[UInt8](), - count=limit, + count=limit ) offset += limit @@ -632,7 +634,7 @@ def shake_final_into(mut ctx: SHA3Context, mut output: StackBuffer[UInt8, ...], unsafe_memcpy( dest=output.ptr().unsafe_offset(offset), src=ctx.state.ptr().unsafe_bitcast[UInt8](), - count=limit, + count=limit ) offset += limit @@ -648,13 +650,13 @@ def shake_hash(rate_bits: Int, data: Span[UInt8, ...], output_len: Int) -> List[ @always_inline -def shake_hash_into(mut output: StackBuffer[UInt8, ...], rate_bits: Int, data: Span[UInt8, ...], output_len: Int): +def shake_hash_into(mut output: StackBuffer[UInt8, ...], rate_bits: Int, data: Span[UInt8, ...], output_len: Int +): var ctx = SHA3Context(rate_bits) sha3_update(ctx, data) shake_final_into(ctx, output, output_len) - def shake128(data: Span[UInt8, ...], output_len_bytes: Int) -> List[UInt8]: return shake_hash(1344, data, output_len_bytes) @@ -664,10 +666,12 @@ def shake256(data: Span[UInt8, ...], output_len_bytes: Int) -> List[UInt8]: @always_inline -def shake128_into(mut output: StackBuffer[UInt8, ...], data: Span[UInt8, ...], output_len_bytes: Int): +def shake128_into(mut output: StackBuffer[UInt8, ...], data: Span[UInt8, ...], output_len_bytes: Int +): shake_hash_into(output, 1344, data, output_len_bytes) @always_inline -def shake256_into(mut output: StackBuffer[UInt8, ...], data: Span[UInt8, ...], output_len_bytes: Int): +def shake256_into(mut output: StackBuffer[UInt8, ...], data: Span[UInt8, ...], output_len_bytes: Int +): shake_hash_into(output, 1088, data, output_len_bytes) diff --git a/src/thistle/sha_ni.mojo b/src/thistle/sha_ni.mojo index 503edac..e3615bc 100644 --- a/src/thistle/sha_ni.mojo +++ b/src/thistle/sha_ni.mojo @@ -1,8 +1,10 @@ -""" -SHA-NI implementation In Mojo. -""" +"""Provides hardware-accelerated SHA-2 transforms.""" -from std.sys import llvm_intrinsic, CompilationTarget, prefetch, PrefetchOptions +from std.sys import ( + llvm_intrinsic, + inlined_assembly, + CompilationTarget, prefetch, PrefetchOptions +) from std.memory import Pointer, bitcast from .utils import StackBuffer from std.builtin.simd import SIMD @@ -34,18 +36,22 @@ comptime SHA256_K = SIMD[DType.uint32, 64]( 0xD6990624, 0xF40E3585, 0x106AA070, 0x19A4C116, 0x1E376C08, 0x2748774C, 0x34B0BCB5, 0x391C0CB3, 0x4ED8AA4A, 0x5B9CCA4F, 0x682E6FF3, 0x748F82EE, 0x78A5636F, 0x84C87814, 0x8CC70208, - 0x90BEFFFA, 0xA4506CEB, 0xBEF9A3F7, 0xC67178F2, + 0x90BEFFFA, 0xA4506CEB, 0xBEF9A3F7, 0xC67178F2 ) @always_inline def has_x86_sha_ni() -> Bool: - return CompilationTarget.is_x86() and CompilationTarget._has_feature["sse"]() and CompilationTarget._has_feature["sha"]() + return ( + CompilationTarget.is_x86() and CompilationTarget._has_feature["sse"]() and CompilationTarget._has_feature["sha"]() + ) @always_inline def has_arm_sha2() -> Bool: - return CompilationTarget.has_neon() and CompilationTarget._has_feature["sha2"]() and not CompilationTarget.is_x86() + return ( + CompilationTarget.has_neon() and CompilationTarget._has_feature["sha2"]() and not CompilationTarget.is_x86() + ) @always_inline("nodebug") @@ -152,7 +158,7 @@ def _sha256ni_transform_arm(state: SIMD[DType.uint32, 8], block: Span[UInt8, ... w[i & 3] = _arm_sha256su1( _arm_sha256su0(w[i & 3], w[(i + 1) & 3]), w[(i + 2) & 3], - w[(i + 3) & 3], + w[(i + 3) & 3] ) var tmp = st0 st0 = _arm_sha256h(st0, st1, wk) @@ -266,6 +272,7 @@ def sha256ni_hash(data: Span[UInt8, ...]) -> List[UInt8]: return output^ + def has_sha_ni() -> Bool: return has_x86_sha_ni() or has_arm_sha2() @@ -273,7 +280,7 @@ def has_sha_ni() -> Bool: def sha256ni_transform_blocks( mut state: SIMD[DType.uint32, 8], data: Pointer[mut=False, UInt8, _, address_space=_], - nblocks: Int, + nblocks: Int ): comptime if CompilationTarget.has_neon() and CompilationTarget._has_feature["sha2"]() and not CompilationTarget.is_x86(): var st0 = SIMD128(state[0], state[1], state[2], state[3]) @@ -298,7 +305,7 @@ def sha256ni_transform_blocks( w[i & 3] = _arm_sha256su1( _arm_sha256su0(w[i & 3], w[(i + 1) & 3]), w[(i + 2) & 3], - w[(i + 3) & 3], + w[(i + 3) & 3] ) var tmp = st0 st0 = _arm_sha256h(st0, st1, wk) @@ -319,6 +326,40 @@ def sha256ni_transform_blocks( comptime SIMD64x2 = SIMD[DType.uint64, 2] + +@always_inline("nodebug") +def _sha512_dit_begin() -> Int64: + """Enable ARM data-independent timing and return the prior DIT state.""" + comptime if CompilationTarget.is_macos() and CompilationTarget.has_neon(): + return inlined_assembly[ + """ + mrs x0, DIT + msr DIT, #1 + sb + """, + Int64, + constraints="={x0},~{memory}", + has_side_effect=True + ]() + else: + return 0 + + +@always_inline("nodebug") +def _sha512_dit_restore(previous: Int64): + """Restore the ARM DIT state saved by `_sha512_dit_begin`.""" + comptime if CompilationTarget.is_macos() and CompilationTarget.has_neon(): + _ = inlined_assembly[ + """ + msr DIT, x9 + sb + """, + Int64, + constraints="={x0},0,{x9},~{memory}", + has_side_effect=True + ](Int64(0), previous) + + comptime SHA512NI_K = StaticTuple[UInt64, 80]( 0x428A2F98D728AE22, 0x7137449123EF65CD, 0xB5C0FBCFEC4D3B2F, 0xE9B5DBA58189DBBC, 0x3956C25BF348B538, 0x59F111F1B605D019, 0x923F82A4AF194F9B, 0xAB1C5ED5DA6D8118, 0xD807AA98A3030242, 0x12835B0145706FBE, @@ -335,7 +376,7 @@ comptime SHA512NI_K = StaticTuple[UInt64, 80]( 0x90BEFFFA23631E28, 0xA4506CEBDE82BDE9, 0xBEF9A3F7B2C67915, 0xC67178F2E372532B, 0xCA273ECEEA26619C, 0xD186B8C721C0C207, 0xEADA7DD6CDE0EB1E, 0xF57D4F7FEE6ED178, 0x06F067AA72176FBA, 0x0A637DC5A2C898A6, 0x113F9804BEF90DAE, 0x1B710B35131C471B, 0x28DB77F523047D84, 0x32CAAB7B40C72493, 0x3C9EBE0A15C9BEBC, - 0x431D67C49C100D4C, 0x4CC5D4BECB3E42B6, 0x597F299CFC657E2A, 0x5FCB6FAB3AD6FAEC, 0x6C44198C4A475817, + 0x431D67C49C100D4C, 0x4CC5D4BECB3E42B6, 0x597F299CFC657E2A, 0x5FCB6FAB3AD6FAEC, 0x6C44198C4A475817 ) @@ -396,8 +437,9 @@ def prefetch_next_block512(ptr: Pointer[mut=False, UInt8, _, address_space=_]): def sha512ni_transform_blocks( mut state: SIMD[DType.uint64, 8], data: Pointer[mut=False, UInt8, _, address_space=_], - nblocks: Int, + nblocks: Int ): + var previous_dit = _sha512_dit_begin() var ab = SIMD64x2(state[0], state[1]) var cd = SIMD64x2(state[2], state[3]) var ef = SIMD64x2(state[4], state[5]) @@ -427,7 +469,7 @@ def sha512ni_transform_blocks( w[t & 7] = _sha512su1( _sha512su0(w[t & 7], w[(t + 1) & 7]), w[(t + 7) & 7], - _ext1(w[(t + 4) & 7], w[(t + 5) & 7]), + _ext1(w[(t + 4) & 7], w[(t + 5) & 7]) ) v[d] = _sha512h2(intermed, v[(d + 2) & 3], v[(d + 1) & 3]) v[(d + 2) & 3] += intermed @@ -440,3 +482,4 @@ def sha512ni_transform_blocks( state = SIMD[DType.uint64, 8]( ab[0], ab[1], cd[0], cd[1], ef[0], ef[1], gh[0], gh[1] ) + _sha512_dit_restore(previous_dit) diff --git a/src/thistle/utils.mojo b/src/thistle/utils.mojo index 64182b8..3617c35 100644 --- a/src/thistle/utils.mojo +++ b/src/thistle/utils.mojo @@ -1,3 +1,5 @@ +"""Provides bounded stack buffers and constant-time utility operations.""" + from std.bit import byte_swap from std.os import abort from std.memory import Pointer @@ -18,13 +20,13 @@ struct StackInlineArray[ElementType: Copyable & Deinitable, size: Int](Copyable) @always_inline def __init__[ - origin: MutOrigin, + origin: MutOrigin ]( out self, *, var storage: VariadicList[ elt_is_mutable=True, origin=origin, Self.ElementType, is_owned=True - ], + ] ): if len(storage) != Self.size: abort("StackInlineArray storage length must match its size") @@ -99,7 +101,9 @@ struct StackBuffer[T: Copyable & Deinitable & Defaultable, N: Int](Movable): @always_inline def __init__(out self): - comptime assert Self.T.__del__is_trivial, "StackBuffer requires trivially destructible types (UInt8, UInt32, UInt64, etc)" + comptime assert ( + Self.T.__del__is_trivial + ), "StackBuffer requires trivially destructible types (UInt8, UInt32, UInt64, etc)" self._data = InlineArray[Self.T, Self.N](fill=Self.T()) self._len = 0 @@ -266,10 +270,8 @@ def nibble_to_hex_char(nibble: UInt8) -> UInt8: @always_inline 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", - ) + if len < 0 or len > Int.MAX // 2: + abort("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[unsafe_offset=i] diff --git a/src/thistle/weierstrass.mojo b/src/thistle/weierstrass.mojo index 4980d45..e744630 100644 --- a/src/thistle/weierstrass.mojo +++ b/src/thistle/weierstrass.mojo @@ -4,6 +4,7 @@ Generic Weierstrass limb operations for P-256 (N=4) and P-384 (N=6). from std.utils import StaticTuple from std.memory import Pointer +from std.os import abort from .utils import u64_nonzero_choice, u64_zero_choice, volatile_wipe from .pbkdf2 import hmac_sha256, hmac_sha384 @@ -31,7 +32,8 @@ struct Limbs[N: Int](Copyable, ImplicitlyCopyable, Movable): self.limbs[2] = l2 self.limbs[3] = l3 - def __init__(out self, l0: UInt64, l1: UInt64, l2: UInt64, l3: UInt64, l4: UInt64, l5: UInt64) where Self.N == 6: + def __init__(out self, l0: UInt64, l1: UInt64, l2: UInt64, l3: UInt64, l4: UInt64, l5: UInt64 + ) where Self.N == 6: self.limbs = StaticTuple[UInt64, Self.N]() comptime for i in range(Self.N): self.limbs[i] = 0 @@ -140,7 +142,7 @@ def sub_mod[N: Int](a: Limbs[N], b: Limbs[N], m: Limbs[N]) -> Limbs[N]: def from_be[N: Int](bytes: Span[UInt8, ...]) -> Limbs[N]: var out = Limbs[N].zero() comptime for i in range(N): - var off = (N * 8 - 8 - i * 8) + var off = N * 8 - 8 - i * 8 var v: UInt64 = 0 for k in range(8): v = (v << 8) | UInt64(bytes[off + k]) @@ -184,7 +186,8 @@ struct JacobianPoint[N: Int](Copyable, ImplicitlyCopyable, Movable): self.z = Limbs[Self.N].zero() self.infinity = True - def __init__(out self, x: Limbs[Self.N], y: Limbs[Self.N], z: Limbs[Self.N], infinity: Bool): + def __init__(out self, x: Limbs[Self.N], y: Limbs[Self.N], z: Limbs[Self.N], infinity: Bool + ): self.x = x self.y = y self.z = z @@ -198,7 +201,7 @@ def _p256_final_sub( acc2: UInt64, acc3: UInt64, acc4: UInt64, - p: Limbs[4], + p: Limbs[4] ) -> Limbs[4]: var out = Limbs[4](acc0, acc1, acc2, acc3) var d, borrow = sub_raw(out, p) @@ -312,19 +315,15 @@ def _p256_mont_sqr(a: Limbs[4], p: Limbs[4]) -> Limbs[4]: var acc1 = UInt64(p10 & _MASK64) var s = (p20 & _MASK64) + (p10 >> UInt128(64)) var acc2 = UInt64(s & _MASK64) - s = ( - (p30 & _MASK64) + s = (p30 & _MASK64) + (p20 >> UInt128(64)) + (p21 & _MASK64) + (s >> UInt128(64)) - ) var acc3 = UInt64(s & _MASK64) - s = ( - (p30 >> UInt128(64)) + s = (p30 >> UInt128(64)) + (p21 >> UInt128(64)) + (p31 & _MASK64) + (s >> UInt128(64)) - ) var acc4 = UInt64(s & _MASK64) s = (p31 >> UInt128(64)) + (p32 & _MASK64) + (s >> UInt128(64)) var acc5 = UInt64(s & _MASK64) @@ -694,7 +693,8 @@ def base_table_entry[N: Int](tptr: Pointer[UInt64, _], j: Int, d: UInt64) -> Poi @always_inline -def scalar_mult_base[N: Int, N0: UInt64](tptr: Pointer[UInt64, _], k: Limbs[N], mod: Limbs[N], rr: Limbs[N], one_mont: Limbs[N]) -> Point[N]: +def scalar_mult_base[N: Int, N0: UInt64](tptr: Pointer[UInt64, _], k: Limbs[N], mod: Limbs[N], rr: Limbs[N], one_mont: Limbs[N] +) -> Point[N]: var acc = jacobian_infinity(one_mont) for i in range(1, N * 16, 2): var d = (k.limbs[i >> 4] >> UInt64(4 * (i & 15))) & UInt64(0xF) @@ -722,6 +722,10 @@ def _hmac[N: Int](key: Span[UInt8, ...], data: Span[UInt8, ...]) -> List[UInt8]: def rfc6979[N: Int](private_key: Span[UInt8, ...], digest: Span[UInt8, ...], skip: Int, n: Limbs[N]) -> Limbs[N]: + if len(private_key) != N * 8 or len(digest) != N * 8: + abort("RFC 6979 inputs must match the curve scalar size") + if skip < 0: + abort("RFC 6979 skip count cannot be negative") var h1 = reduce_mod(from_be[N](digest), n) var h1_bytes = InlineArray[UInt8, N * 8](fill=0) to_be(h1, h1_bytes.unsafe_ptr()) @@ -796,7 +800,8 @@ def rfc6979[N: Int](private_key: Span[UInt8, ...], digest: Span[UInt8, ...], ski @always_inline -def jacobian_add_affine_non_equal_ct[N: Int, N0: UInt64](p: JacobianPoint[N], q: Point[N], mod: Limbs[N], rr: Limbs[N], one_mont: Limbs[N]) -> JacobianPoint[N]: +def jacobian_add_affine_non_equal_ct[N: Int, N0: UInt64](p: JacobianPoint[N], q: Point[N], mod: Limbs[N], rr: Limbs[N], one_mont: Limbs[N] +) -> JacobianPoint[N]: var z1z1 = mont_sqr[N, N0](p.z, mod) var u2 = mont_mul[N, N0](q.x, z1z1, mod) var s2 = mont_mul[N, N0](q.y, mont_mul[N, N0](p.z, z1z1, mod), mod) diff --git a/src/thistle/x25519.mojo b/src/thistle/x25519.mojo index 9599c40..99e1d3c 100644 --- a/src/thistle/x25519.mojo +++ b/src/thistle/x25519.mojo @@ -1,12 +1,11 @@ -""" -X25519 implementation -""" +"""Implements X25519 key agreement.""" from .curve25519 import FieldElement51 from .utils import StackInlineArray from .random import random_bytes from std.collections import List + @always_inline def _cswap_fe(swap: UInt64, mut a: FieldElement51, mut b: FieldElement51): var mask = UInt64(0) - swap @@ -15,26 +14,28 @@ def _cswap_fe(swap: UInt64, mut a: FieldElement51, mut b: FieldElement51): a.limbs[i] = a.limbs[i] ^ dummy b.limbs[i] = b.limbs[i] ^ dummy + @always_inline def _cswap_pair( swap: UInt64, mut x_2: FieldElement51, mut x_3: FieldElement51, mut z_2: FieldElement51, - mut z_3: FieldElement51, + mut z_3: FieldElement51 ): _cswap_fe(swap, x_2, x_3) _cswap_fe(swap, z_2, z_3) + @no_inline def x25519( scalar_in: Span[UInt8, ...], point: Span[UInt8, ...], - output: Span[mut=True, UInt8, ...], + output: Span[mut=True, UInt8, ...] ) raises: - if len(scalar_in) < 32: + if len(scalar_in) != 32: raise Error("X25519 scalar must be 32 bytes") - if len(point) < 32: + 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") @@ -89,7 +90,7 @@ def x25519( def x25519_checked( scalar_in: Span[UInt8, ...], point: Span[UInt8, ...], - output: Span[mut=True, UInt8, ...], + output: Span[mut=True, UInt8, ...] ) raises: x25519(scalar_in, point, output) var out_ptr = output.unsafe_ptr() @@ -110,7 +111,7 @@ def x25519_public_key( x25519( private_key, Span[UInt8, ...](unsafe_ptr=base.unsafe_ptr(), length=32), - output, + output ) @@ -119,6 +120,6 @@ def x25519_keygen() raises -> Tuple[List[UInt8], List[UInt8]]: var public_key = List[UInt8](unsafe_uninit_length=32) x25519_public_key( Span[UInt8, ...](private_key), - Span[mut=True, UInt8, ...](public_key), + Span[mut=True, UInt8, ...](public_key) ) return (private_key^, public_key^) diff --git a/tests/benchmark.mojo b/tests/benchmark.mojo index 1f4ca87..c162b60 100644 --- a/tests/benchmark.mojo +++ b/tests/benchmark.mojo @@ -9,15 +9,21 @@ from std.memory import Layout, alloc from thistle.argon2 import Argon2id from thistle.blake2b import Blake2b from thistle.blake3 import blake3_parallel_hash -from thistle.camellia import CamelliaCipher, camellia_encrypt_blocks, camellia_ctr_kernel +from thistle.camellia import ( + CamelliaCipher, camellia_encrypt_blocks, camellia_ctr_kernel +) from thistle.chacha20 import ChaCha20 from thistle.kcipher2 import KCipher2 from thistle.sha2 import sha256_hash, sha512_hash from thistle.sha_ni import sha256ni_hash, has_sha_ni from thistle.sha3 import sha3_256 -from thistle.aes import AESKey, cpu_aes_ct_encrypt16, cpu_aes_ct_skey, ROUNDS_128, expand_key_128 +from thistle.aes import ( + AESKey, cpu_aes_ct_encrypt16, cpu_aes_ct_skey, ROUNDS_128, expand_key_128 +) from thistle.x25519 import x25519 -from thistle.ed25519 import ed25519_sign, ed25519_verify, ed25519_generate_public_key +from thistle.ed25519 import ( + ed25519_sign, ed25519_verify, ed25519_generate_public_key +) from thistle.p256 import p256_ecdsa_sign from thistle.p384 import p384_public_key, p384_ecdsa_sign from thistle.utils import StackInlineArray @@ -33,6 +39,7 @@ comptime TEST_CT: StaticTuple[UInt8, 16] = StaticTuple[UInt8, 16]( 0x3a, 0xd7, 0x7b, 0xb4, 0x0d, 0x7a, 0x36, 0x60, 0xa8, 0x9e, 0xca, 0xf3, 0x24, 0x66, 0xef, 0x97 ) + def generate_data(length: Int) -> List[UInt8]: var data = List[UInt8](capacity=length) for i in range(length): @@ -58,7 +65,9 @@ def benchmark_x25519(duration_secs: Float64) raises -> String: count += 1 var duration = perf_counter() - start var ops = Float64(count) / duration - return "x25519 | throughput: " + String(ops) + " ops/s, ops: " + String(count) + ", time: " + String(duration) + "s" + return ( + "x25519 | throughput: " + String(ops) + " ops/s, ops: " + String(count) + ", time: " + String(duration) + "s" + ) def benchmark_p384(duration_secs: Float64) -> String: @@ -74,7 +83,9 @@ def benchmark_p384(duration_secs: Float64) -> String: count += 1 var duration = perf_counter() - start var ops = Float64(count) / duration - return "p384-public-key | throughput: " + String(ops) + " ops/s, ops: " + String(count) + ", time: " + String(duration) + "s" + return ( + "p384-public-key | throughput: " + String(ops) + " ops/s, ops: " + String(count) + ", time: " + String(duration) + "s" + ) def benchmark_ecdsa(duration_secs: Float64) -> String: @@ -91,7 +102,7 @@ def benchmark_ecdsa(duration_secs: Float64) -> String: _ = p256_ecdsa_sign( Span[UInt8, ...](p256_key), msg, - Span[mut=True, UInt8, ...](unsafe_ptr=p256_sig.unsafe_ptr(), length=64), + Span[mut=True, UInt8, ...](unsafe_ptr=p256_sig.unsafe_ptr(), length=64) ) p256_count += 1 var p256_time = perf_counter() - start @@ -102,7 +113,7 @@ def benchmark_ecdsa(duration_secs: Float64) -> String: _ = p384_ecdsa_sign( Span[UInt8, ...](p384_key), msg, - Span[mut=True, UInt8, ...](unsafe_ptr=p384_sig.unsafe_ptr(), length=96), + Span[mut=True, UInt8, ...](unsafe_ptr=p384_sig.unsafe_ptr(), length=96) ) p384_count += 1 var p384_time = perf_counter() - start @@ -156,8 +167,12 @@ 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) + " 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" + 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 @@ -175,7 +190,9 @@ 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) + " mb/s, hashes: " + String(count) + ", time: " + String(duration) + "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: @@ -192,7 +209,9 @@ 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) + " mb/s, hashes: " + String(count) + ", time: " + String(duration) + "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: @@ -207,7 +226,9 @@ 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) + " mb/s, hashes: " + String(count) + ", time: " + String(duration) + "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: @@ -222,7 +243,9 @@ 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) + " mb/s, hashes: " + String(count) + ", time: " + String(duration) + "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: @@ -238,7 +261,9 @@ 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) + " mb/s, hashes: " + String(count) + ", time: " + String(duration) + "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: @@ -253,7 +278,9 @@ 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) + " mb/s, hashes: " + String(count) + ", time: " + String(duration) + "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: @@ -281,7 +308,9 @@ def benchmark_camellia(data_size: Int, duration_secs: Float64) raises -> String: 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" + return ( + "camellia | throughput: " + String(mbps) + " mb/s, blocks: " + String(count) + ", time: " + String(duration) + "s" + ) def benchmark_camellia_ctr(duration_secs: Float64) raises -> String: @@ -312,7 +341,9 @@ def benchmark_camellia_ctr(duration_secs: Float64) raises -> String: 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" + 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: @@ -340,7 +371,9 @@ 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) + " mb/s, encrypts: " + String(count) + ", time: " + String(duration) + "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: @@ -363,7 +396,9 @@ 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) + " mb/s, encrypts: " + String(count) + ", time: " + String(duration) + "s" + return ( + "kcipher2 | throughput: " + String(mbps) + " mb/s, encrypts: " + String(count) + ", time: " + String(duration) + "s" + ) def benchmark_argon2(duration_secs: Float64) raises -> String: @@ -381,7 +416,9 @@ 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) + " h/s, hashes: " + String(count) + ", time: " + String(duration) + "s" + return ( + "argon2id | throughput: " + String(hps) + " h/s, hashes: " + String(count) + ", time: " + String(duration) + "s" + ) def benchmark_aes_cpu(duration_secs: Float64) raises -> String: @@ -406,7 +443,9 @@ def benchmark_aes_cpu(duration_secs: Float64) raises -> String: 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" + return ( + "aes-128-cpu | throughput: " + String(mbps) + " mb/s, blocks: " + String(count) + ", time: " + String(duration) + "s" + ) def benchmark_aes_gpu_ecb() raises -> String: @@ -452,7 +491,7 @@ def benchmark_aes_gpu_ecb() raises -> String: Int32(num_blocks), Int32(10), grid_dim=grid_dim, - block_dim=block_dim, + block_dim=block_dim ) ctx.synchronize() @@ -466,7 +505,7 @@ def benchmark_aes_gpu_ecb() raises -> String: Int32(num_blocks), Int32(10), grid_dim=grid_dim, - block_dim=block_dim, + block_dim=block_dim ) ctx.synchronize() var end = perf_counter() @@ -479,7 +518,9 @@ def benchmark_aes_gpu_ecb() raises -> String: output_host.unsafe_free() key_ptr.unsafe_free() - return "aes-128-gpu-ecb | throughput: " + String(gbps) + " 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: @@ -531,7 +572,7 @@ def benchmark_aes_gpu_ctr() raises -> String: nonce_buffer, Int32(10), grid_dim=grid_dim, - block_dim=block_dim, + block_dim=block_dim ) ctx.synchronize() @@ -546,7 +587,7 @@ def benchmark_aes_gpu_ctr() raises -> String: nonce_buffer, Int32(10), grid_dim=grid_dim, - block_dim=block_dim, + block_dim=block_dim ) ctx.synchronize() var end = perf_counter() @@ -560,10 +601,9 @@ def benchmark_aes_gpu_ctr() raises -> String: nonce_host.unsafe_free() key_ptr.unsafe_free() - return "aes-128-gpu-ctr | throughput: " + String(gbps) + " gb/s, iterations: " + String(iterations) - - - + return ( + "aes-128-gpu-ctr | throughput: " + String(gbps) + " gb/s, iterations: " + String(iterations) + ) def benchmark_aes_gpu_gcm() raises -> String: @@ -616,7 +656,7 @@ def benchmark_aes_gpu_gcm() raises -> String: nonce_buffer, Int32(10), grid_dim=grid_dim, - block_dim=block_dim, + block_dim=block_dim ) ctx.synchronize() @@ -631,7 +671,7 @@ def benchmark_aes_gpu_gcm() raises -> String: nonce_buffer, Int32(10), grid_dim=grid_dim, - block_dim=block_dim, + block_dim=block_dim ) ctx.synchronize() var end = perf_counter() @@ -645,10 +685,9 @@ def benchmark_aes_gpu_gcm() raises -> String: nonce_host.unsafe_free() key_ptr.unsafe_free() - return "aes-128-gpu-gcm | throughput: " + String(gbps) + " gb/s, iterations: " + String(iterations) - - - + return ( + "aes-128-gpu-gcm | throughput: " + String(gbps) + " gb/s, iterations: " + String(iterations) + ) def main() raises: diff --git a/tests/dudect.mojo b/tests/dudect.mojo index a22c819..c165007 100644 --- a/tests/dudect.mojo +++ b/tests/dudect.mojo @@ -25,6 +25,7 @@ comptime N_ASYM = 2_000 comptime BATCH = 16 comptime T_THRESHOLD = 4.5 + struct Rng: var s: UInt64 @@ -148,6 +149,8 @@ def _classes(n: Int, mut rng: Rng) -> List[Int]: return cls^ # This is intentionally leaked + + @no_inline def _leaky(secret: UInt64) -> UInt64: var acc = secret @@ -219,7 +222,7 @@ def run_kcipher2(mut rng: Rng) raises -> Bool: for i in range(len(cls)): var key = SIMD[DType.uint32, 4]( UInt32(keys[2 * i] & 0xFFFFFFFF), UInt32(keys[2 * i] >> 32), - UInt32(keys[2 * i + 1] & 0xFFFFFFFF), UInt32(keys[2 * i + 1] >> 32), + UInt32(keys[2 * i + 1] & 0xFFFFFFFF), UInt32(keys[2 * i + 1] >> 32) ) var t0 = perf_counter_ns() for _ in range(8): @@ -403,7 +406,7 @@ def run_x25519(mut rng: Rng) raises -> Bool: x25519( Span[UInt8, ...](sc), Span[UInt8, ...](base), - Span[mut=True, UInt8, ...](out), + Span[mut=True, UInt8, ...](out) ) times.append(Float64(perf_counter_ns() - t0)) sink ^= out[0] @@ -434,7 +437,7 @@ def run_ed25519(mut rng: Rng) raises -> Bool: var t0 = perf_counter_ns() ed25519_sign( Span[UInt8, ...](sk), Span[UInt8, ...](msg), - Span[mut=True, UInt8, ...](sig), + Span[mut=True, UInt8, ...](sig) ) times.append(Float64(perf_counter_ns() - t0)) sink ^= sig[0] @@ -495,7 +498,7 @@ def run_p256_sign(mut rng: Rng) -> Bool: var ok = p256_ecdsa_sign( Span[UInt8, ...](private_key), Span[UInt8, ...](message), - Span[mut=True, UInt8, ...](signature), + Span[mut=True, UInt8, ...](signature) ) times.append(Float64(perf_counter_ns() - t0)) sink ^= signature[0] ^ (UInt8(1) if ok else UInt8(0)) @@ -516,7 +519,7 @@ def run_p384_sign(mut rng: Rng) -> Bool: var ok = p384_ecdsa_sign( Span[UInt8, ...](private_key), Span[UInt8, ...](message), - Span[mut=True, UInt8, ...](signature), + Span[mut=True, UInt8, ...](signature) ) times.append(Float64(perf_counter_ns() - t0)) sink ^= signature[0] ^ (UInt8(1) if ok else UInt8(0)) @@ -557,7 +560,7 @@ def run_mlkem_decaps(mut rng: Rng) raises -> Bool: def main() raises: print( "dudect harness: fast", N_FAST, "batch", BATCH, "| asym", N_ASYM, - "| |t| threshold", T_THRESHOLD, + "| |t| threshold", T_THRESHOLD ) print("") var rng = Rng(0x1234567890ABCDEF) diff --git a/tests/test_aes_gpu.mojo b/tests/test_aes_gpu.mojo index e5a8171..38aeaed 100644 --- a/tests/test_aes_gpu.mojo +++ b/tests/test_aes_gpu.mojo @@ -4,7 +4,9 @@ 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, AESExpandedKey -from thistle.aes_gpu import aes_gpu_kernel_ecb, aes_gpu_kernel_ctr, aes_gpu_kernel_gcm_ctr +from thistle.aes_gpu import ( + aes_gpu_kernel_ecb, aes_gpu_kernel_ctr, aes_gpu_kernel_gcm_ctr +) from max.gpu.host import DeviceContext from std.memory.unsafe_pointer import Pointer from std.memory import Layout, alloc @@ -15,6 +17,7 @@ def byte_to_hex(b: UInt8) -> String: var lo = Int(b & 0xF) return chr(48 + hi if hi < 10 else 87 + hi) + chr(48 + lo if lo < 10 else 87 + lo) + @fieldwise_init struct TestResult(Copyable, Movable): var passed: Int @@ -116,7 +119,7 @@ def test_aes_gpu_basic(json_data: PythonObject, py: PythonObject) raises -> Test Int32(4), Int32(rounds), grid_dim=grid_dim, - block_dim=block_dim, + block_dim=block_dim ) ctx.synchronize() _ = skey_buffer @@ -211,7 +214,7 @@ def test_mode_gpu(json_data: PythonObject, mode: String) raises -> TestResult: Int32(n_blocks), Int32(rounds), grid_dim=grid_dim, - block_dim=block_dim, + block_dim=block_dim ) elif "CTR" in mode: var iv_hex = String(tv.get("iv", PythonObject())) @@ -232,7 +235,7 @@ def test_mode_gpu(json_data: PythonObject, mode: String) raises -> TestResult: nonce_buffer, Int32(rounds), grid_dim=grid_dim, - block_dim=block_dim, + block_dim=block_dim ) elif "GCM" in mode: var nonce_hex = String(tv.get("nonce", PythonObject())) @@ -257,7 +260,7 @@ def test_mode_gpu(json_data: PythonObject, mode: String) raises -> TestResult: nonce_buffer, Int32(rounds), grid_dim=grid_dim, - block_dim=block_dim, + block_dim=block_dim ) else: ctx.enqueue_function[aes_gpu_kernel_ecb]( @@ -267,7 +270,7 @@ def test_mode_gpu(json_data: PythonObject, mode: String) raises -> TestResult: Int32(n_blocks), Int32(rounds), grid_dim=grid_dim, - block_dim=block_dim, + block_dim=block_dim ) ctx.synchronize() @@ -340,7 +343,8 @@ def main() raises: var json_data = load_json("tests/vectors/aes_test_vectors.json", py) var modes = ["AES-128-ECB", "AES-192-ECB", "AES-256-ECB", "AES-128-CTR", "AES-192-CTR", "AES-256-CTR", - "AES-128-GCM", "AES-192-GCM", "AES-256-GCM"] + "AES-128-GCM", "AES-192-GCM", "AES-256-GCM" + ] for mode in modes: print("Loading " + mode + " vectors...") diff --git a/tests/test_ml_dsa.mojo b/tests/test_ml_dsa.mojo index bd0bb3f..afa9405 100644 --- a/tests/test_ml_dsa.mojo +++ b/tests/test_ml_dsa.mojo @@ -22,9 +22,10 @@ from thistle.ml_dsa import ( private_key_size, public_key_size, signature_size, - new_public_key, + new_public_key ) + def hex_to_bytes(s: String) -> List[UInt8]: var r = List[UInt8]() var b = s.as_bytes() @@ -350,7 +351,6 @@ def run_nist_sigver_file(prompt_path: String, expected_path: String, py: PythonO return passed, failed - def run_external_api_smoke_tests() raises -> Tuple[Int, Int]: var passed = 0 var failed = 0 @@ -370,14 +370,16 @@ def run_external_api_smoke_tests() raises -> Tuple[Int, Int]: passed += 1 var sig_h = mldsa_sign_hedged(priv44, Span[UInt8, ...](msg), Span[UInt8, ...](ctx)) - if not mldsa_verify(priv44.pub, Span[UInt8, ...](msg), Span[UInt8, ...](sig_h), Span[UInt8, ...](ctx)): + if not mldsa_verify(priv44.pub, Span[UInt8, ...](msg), Span[UInt8, ...](sig_h), Span[UInt8, ...](ctx) + ): print("external API: ML-DSA-44 hedged signature did not verify") failed += 1 else: passed += 1 var sig_d = mldsa_sign_deterministic(priv44, Span[UInt8, ...](msg), Span[UInt8, ...](ctx)) - if not mldsa_verify(priv44.pub, Span[UInt8, ...](msg), Span[UInt8, ...](sig_d), Span[UInt8, ...](ctx)): + if not mldsa_verify(priv44.pub, Span[UInt8, ...](msg), Span[UInt8, ...](sig_d), Span[UInt8, ...](ctx) + ): print("external API: ML-DSA-44 deterministic signature did not verify") failed += 1 else: diff --git a/tests/test_ml_kem.mojo b/tests/test_ml_kem.mojo index e9233d1..116826a 100644 --- a/tests/test_ml_kem.mojo +++ b/tests/test_ml_kem.mojo @@ -29,7 +29,7 @@ from thistle.ml_kem import ( poly_ntt, poly_reduce, poly_tobytes, - poly_tomsg, + poly_tomsg ) diff --git a/tests/test_security_boundaries.mojo b/tests/test_security_boundaries.mojo index 49920fb..efd3174 100644 --- a/tests/test_security_boundaries.mojo +++ b/tests/test_security_boundaries.mojo @@ -5,13 +5,13 @@ from thistle.argon2 import variable_length_hash_into from thistle.blake2b import Blake2b from thistle.chacha20poly1305 import ( chacha20_poly1305_encrypt, - hchacha20, + hchacha20 ) from thistle.chacha20 import ChaCha20 from thistle.ed25519 import ( Ed25519SigningKey, ed25519_generate_public_key, - ed25519_sign, + ed25519_sign ) from thistle.p256 import p256_ecdsa_sign, p256_public_key from thistle.p384 import p384_ecdsa_sign, p384_public_key @@ -21,7 +21,7 @@ from thistle.pbkdf2 import ( PBKDF2SHA256, PBKDF2SHA512, pbkdf2_hmac_sha256, - pbkdf2_hmac_sha512, + pbkdf2_hmac_sha512 ) from thistle.poly1305 import Poly1305 from thistle.x25519 import x25519, x25519_checked @@ -45,7 +45,7 @@ def main() raises: variable_length_hash_into( 5, Span[UInt8, ...](empty), - Span[mut=True, UInt8, ...](argon_output), + Span[mut=True, UInt8, ...](argon_output) ) except: rejected = True @@ -83,7 +83,7 @@ def main() raises: x25519( Span[UInt8, ...](key32), Span[UInt8, ...](point32), - Span[mut=True, UInt8, ...](output31), + Span[mut=True, UInt8, ...](output31) ) except: rejected = True @@ -97,7 +97,7 @@ def main() raises: x25519_checked( Span[UInt8, ...](key32), Span[UInt8, ...](zero_point), - Span[mut=True, UInt8, ...](x25519_output), + Span[mut=True, UInt8, ...](x25519_output) ) except: rejected = True @@ -110,7 +110,7 @@ def main() raises: hchacha20( Span[UInt8, ...](key32), Span[UInt8, ...](input16), - Span[mut=True, UInt8, ...](output31), + Span[mut=True, UInt8, ...](output31) ) except: rejected = True @@ -179,7 +179,7 @@ def main() raises: try: ed25519_generate_public_key( Span[UInt8, ...](ed_private), - Span[mut=True, UInt8, ...](ed_public_short), + Span[mut=True, UInt8, ...](ed_public_short) ) except: rejected = True @@ -192,7 +192,7 @@ def main() raises: ed25519_sign( Span[UInt8, ...](ed_private), Span[UInt8, ...](empty), - Span[mut=True, UInt8, ...](ed_signature_short), + Span[mut=True, UInt8, ...](ed_signature_short) ) except: rejected = True @@ -208,7 +208,7 @@ def main() raises: try: ed_key.sign( Span[UInt8, ...](empty), - Span[mut=True, UInt8, ...](ed_signature_short), + Span[mut=True, UInt8, ...](ed_signature_short) ) except: rejected = True @@ -222,7 +222,7 @@ def main() raises: Span[UInt8, ...](empty), Span[UInt8, ...](empty_salt), 1, - PBKDF2_SHA256_MAX_DKLEN + 1, + PBKDF2_SHA256_MAX_DKLEN + 1 ) except: rejected = True @@ -235,7 +235,7 @@ def main() raises: Span[UInt8, ...](empty), Span[UInt8, ...](empty_salt), 1, - PBKDF2_SHA512_MAX_DKLEN + 1, + PBKDF2_SHA512_MAX_DKLEN + 1 ) except: rejected = True @@ -272,7 +272,7 @@ def main() raises: Span[UInt8, ...](empty), Span[UInt8, ...](plaintext), Span[mut=True, UInt8, ...](ciphertext), - Span[mut=True, UInt8, ...](tag), + Span[mut=True, UInt8, ...](tag) ) except: rejected = True @@ -287,12 +287,12 @@ def main() raises: var p384_output = List[UInt8](length=96, fill=0) if p256_public_key( Span[UInt8, ...](p256_private), - Span[mut=True, UInt8, ...](p256_output), + 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), + Span[mut=True, UInt8, ...](p384_output) ): raise Error("P-384 public-key API accepted an undersized destination") @@ -301,13 +301,13 @@ def main() raises: if p256_ecdsa_sign( Span[UInt8, ...](p256_private), Span[UInt8, ...](plaintext), - Span[mut=True, UInt8, ...](short_signature), + 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), + Span[mut=True, UInt8, ...](short_signature384) ): raise Error("P-384 signing accepted an undersized destination") diff --git a/tests/test_signing.mojo b/tests/test_signing.mojo index 6902520..59564f4 100644 --- a/tests/test_signing.mojo +++ b/tests/test_signing.mojo @@ -6,7 +6,7 @@ from thistle.p256 import ( p256_ecdsa_sign_der, p256_ecdsa_verify_der, p256_keygen, - p256_public_key, + p256_public_key ) from thistle.p384 import ( p384_ecdsa_sign, @@ -14,7 +14,7 @@ from thistle.p384 import ( p384_ecdsa_sign_der, p384_ecdsa_verify_der, p384_keygen, - p384_public_key, + p384_public_key ) from thistle.x25519 import x25519_keygen, x25519_public_key from thistle.pbkdf2 import hmac_sha384 @@ -28,9 +28,10 @@ from thistle.rsa import ( SHA224, SHA256, SHA384, - SHA512, + SHA512 ) + def hex_bytes(s: String) -> List[UInt8]: var out = List[UInt8]() var data = s.as_bytes() @@ -102,7 +103,7 @@ def test_p256() raises: if not p256_ecdsa_sign( Span[UInt8, ...](private_key), Span[UInt8, ...](message), - Span[mut=True, UInt8, ...](signature), + Span[mut=True, UInt8, ...](signature) ): raise Error("P-256 signing failed") if not equal(signature, expected): @@ -110,14 +111,14 @@ def test_p256() raises: if not p256_ecdsa_verify( Span[UInt8, ...](public_key), Span[UInt8, ...](message), - Span[UInt8, ...](signature), + Span[UInt8, ...](signature) ): raise Error("P-256 verification failed") signature[0] ^= 1 if p256_ecdsa_verify( Span[UInt8, ...](public_key), Span[UInt8, ...](message), - Span[UInt8, ...](signature), + Span[UInt8, ...](signature) ): raise Error("P-256 accepted a changed signature") var der = p256_ecdsa_sign_der( @@ -126,7 +127,7 @@ def test_p256() raises: if not p256_ecdsa_verify_der( Span[UInt8, ...](public_key), Span[UInt8, ...](message), - Span[UInt8, ...](der), + Span[UInt8, ...](der) ): raise Error("P-256 DER signature failed") @@ -158,7 +159,7 @@ def test_p384() raises: if not p384_ecdsa_sign( Span[UInt8, ...](private_key), Span[UInt8, ...](message), - Span[mut=True, UInt8, ...](signature), + Span[mut=True, UInt8, ...](signature) ): raise Error("P-384 signing failed") if not equal(signature, expected): @@ -166,14 +167,14 @@ def test_p384() raises: if not p384_ecdsa_verify( Span[UInt8, ...](public_key), Span[UInt8, ...](message), - Span[UInt8, ...](signature), + Span[UInt8, ...](signature) ): raise Error("P-384 verification failed") signature[0] ^= 1 if p384_ecdsa_verify( Span[UInt8, ...](public_key), Span[UInt8, ...](message), - Span[UInt8, ...](signature), + Span[UInt8, ...](signature) ): raise Error("P-384 accepted a changed signature") var der = p384_ecdsa_sign_der( @@ -182,7 +183,7 @@ def test_p384() raises: if not p384_ecdsa_verify_der( Span[UInt8, ...](public_key), Span[UInt8, ...](message), - Span[UInt8, ...](der), + Span[UInt8, ...](der) ): raise Error("P-384 DER signature failed") @@ -194,7 +195,7 @@ def test_keygen() raises: var p256_check = List[UInt8](unsafe_uninit_length=65) if not p256_public_key( Span[UInt8, ...](p256_private), - Span[mut=True, UInt8, ...](p256_check), + Span[mut=True, UInt8, ...](p256_check) ) or not equal(p256_public, p256_check): raise Error("P-256 key generation failed") @@ -204,7 +205,7 @@ def test_keygen() raises: var p384_check = List[UInt8](unsafe_uninit_length=97) if not p384_public_key( Span[UInt8, ...](p384_private), - Span[mut=True, UInt8, ...](p384_check), + Span[mut=True, UInt8, ...](p384_check) ) or not equal(p384_public, p384_check): raise Error("P-384 key generation failed") @@ -238,7 +239,7 @@ def test_rsa_pss_signing() raises: Span[UInt8, ...](salt), SHA256, SHA256, - Span[mut=True, UInt8, ...](short_rsa_signature), + Span[mut=True, UInt8, ...](short_rsa_signature) ): raise Error("RSA-PSS accepted an undersized signature destination") var sig = rsa_pss_sign_with_salt( @@ -248,7 +249,7 @@ def test_rsa_pss_signing() raises: Span[UInt8, ...](msg), Span[UInt8, ...](salt), SHA256, - SHA256, + SHA256 ) if not rsa_pss_verify( Span[UInt8, ...](n), @@ -257,7 +258,7 @@ def test_rsa_pss_signing() raises: Span[UInt8, ...](sig), SHA256, SHA256, - 32, + 32 ): raise Error("verify") var p = hex_bytes( @@ -282,14 +283,14 @@ def test_rsa_pss_signing() raises: Span[UInt8, ...](q), Span[UInt8, ...](dp), Span[UInt8, ...](dq), - Span[UInt8, ...](qi), + Span[UInt8, ...](qi) ) if crt.pss_sign_with_salt( Span[UInt8, ...](msg), Span[UInt8, ...](salt), SHA256, SHA256, - Span[mut=True, UInt8, ...](short_rsa_signature), + 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) @@ -298,7 +299,7 @@ def test_rsa_pss_signing() raises: Span[UInt8, ...](salt), SHA256, SHA256, - Span[mut=True, UInt8, ...](sig2), + Span[mut=True, UInt8, ...](sig2) ): raise Error("crt sign") if not rsa_pss_verify( @@ -308,7 +309,7 @@ def test_rsa_pss_signing() raises: Span[UInt8, ...](sig2), SHA256, SHA256, - 32, + 32 ): raise Error("crt verify") if not equal(sig, sig2): @@ -316,7 +317,7 @@ def test_rsa_pss_signing() raises: var full_key = RsaPrivateKey( Span[UInt8, ...](n), Span[UInt8, ...](e), - Span[UInt8, ...](d), + Span[UInt8, ...](d) ) var oversized_salt_rejected = False try: @@ -324,7 +325,7 @@ def test_rsa_pss_signing() raises: Span[UInt8, ...](msg), SHA256, SHA256, - 1_000_000_000, + 1_000_000_000 ) except: oversized_salt_rejected = True @@ -336,7 +337,7 @@ def test_rsa_pss_signing() raises: _ = RsaPrivateKey( Span[UInt8, ...](n), Span[UInt8, ...](e), - Span[UInt8, ...](bad_private_exponent), + Span[UInt8, ...](bad_private_exponent) ) except: invalid_private_rejected = True @@ -350,16 +351,17 @@ def test_rsa_pss_signing() raises: Span[UInt8, ...](sig2), SHA256, SHA256, - 32, + 32 ): raise Error("changed signature accepted") print("RSA-PSS signing tests passed") + def run_ecdsa_file( json: PythonObject, builtins: PythonObject, path: String, - curve_size: Int, + curve_size: Int ) raises -> Bool: var file = builtins.open(path, "r") var root = json.load(file) @@ -376,13 +378,13 @@ def run_ecdsa_file( valid = p256_ecdsa_verify_der( Span[UInt8, ...](public_key), Span[UInt8, ...](message), - Span[UInt8, ...](signature), + Span[UInt8, ...](signature) ) else: valid = p384_ecdsa_verify_der( Span[UInt8, ...](public_key), Span[UInt8, ...](message), - Span[UInt8, ...](signature), + Span[UInt8, ...](signature) ) var result = String(test["result"]) if ( @@ -405,20 +407,21 @@ def test_ecdsa_wycheproof() raises: json, builtins, "tests/Wycheproof/ecdsa_secp256r1_sha256_test.json", - 32, + 32 ) ok = ( run_ecdsa_file( json, builtins, "tests/Wycheproof/ecdsa_secp384r1_sha384_test.json", - 48, + 48 ) and ok ) if not ok: raise Error("Wycheproof ECDSA failures") + def test_rsa_pkcs1_wycheproof() raises: var json = Python.import_module("json") var builtins = Python.import_module("builtins") @@ -439,7 +442,7 @@ def test_rsa_pkcs1_wycheproof() raises: Span[UInt8, ...](modulus), Span[UInt8, ...](exponent), Span[UInt8, ...](message), - Span[UInt8, ...](signature), + Span[UInt8, ...](signature) ) var result = String(test["result"]) if ( @@ -455,6 +458,7 @@ def test_rsa_pkcs1_wycheproof() raises: if failed != 0: raise Error("Wycheproof RSA PKCS#1 v1.5 failures") + def sha_id(name: String) raises -> Int: if name == "SHA-1": return SHA1 @@ -491,7 +495,7 @@ def run_pss_file(py: PythonObject, builtins: PythonObject, path: String) raises got = rsa_pss_verify( Span[UInt8, ...](n), Span[UInt8, ...](e), Span[UInt8, ...](msg), Span[UInt8, ...](sig), - sha, mgf_sha, s_len, + sha, mgf_sha, s_len ) except: got = False @@ -515,14 +519,32 @@ def test_rsa_pss_wycheproof() raises: var builtins = Python.import_module("builtins") var all_ok = True - all_ok = run_pss_file(py, builtins, "tests/Wycheproof/rsa_pss_2048_sha256_mgf1_32_test.json") and all_ok - all_ok = run_pss_file(py, builtins, "tests/Wycheproof/rsa_pss_3072_sha256_mgf1_32_test.json") and all_ok - all_ok = run_pss_file(py, builtins, "tests/Wycheproof/rsa_pss_4096_sha512_mgf1_64_test.json") and all_ok + all_ok = ( + run_pss_file(py, builtins, "tests/Wycheproof/rsa_pss_2048_sha256_mgf1_32_test.json") + and all_ok + ) + all_ok = ( + run_pss_file( + py, + builtins, + "tests/Wycheproof/rsa_pss_3072_sha256_mgf1_32_test.json" + ) + and all_ok + ) + all_ok = ( + run_pss_file( + py, + builtins, + "tests/Wycheproof/rsa_pss_4096_sha512_mgf1_64_test.json" + ) + and all_ok + ) all_ok = run_pss_file(py, builtins, "tests/Wycheproof/rsa_pss_misc_test.json") and all_ok all_ok = run_pss_file(py, builtins, "tests/Wycheproof/rsa_pss_misc_params_test.json") and all_ok if not all_ok: raise Error("Wycheproof RSA-PSS failures") + def main() raises: test_hmac_sha384() test_p256() diff --git a/tests/test_wycheproof_ed25519.mojo b/tests/test_wycheproof_ed25519.mojo index 34074e5..17adceb 100644 --- a/tests/test_wycheproof_ed25519.mojo +++ b/tests/test_wycheproof_ed25519.mojo @@ -6,6 +6,7 @@ from std.collections import List from std.python import Python from thistle.ed25519 import ed25519_verify + def hex_to_bytes(s: String) -> List[UInt8]: var r = List[UInt8]() var b = s.as_bytes() @@ -19,7 +20,8 @@ def hex_to_bytes(s: String) -> List[UInt8]: return r^ -def run_case(tc_id: String, pk_hex: String, msg_hex: String, sig_hex: String, expected_valid: Bool) -> Bool: +def run_case(tc_id: String, pk_hex: String, msg_hex: String, sig_hex: String, expected_valid: Bool +) -> Bool: var pk = hex_to_bytes(pk_hex) var msg = hex_to_bytes(msg_hex) var sig = hex_to_bytes(sig_hex) diff --git a/tests/test_wycheproof_p256_ecdh.mojo b/tests/test_wycheproof_p256_ecdh.mojo index b4199de..cee8cb3 100644 --- a/tests/test_wycheproof_p256_ecdh.mojo +++ b/tests/test_wycheproof_p256_ecdh.mojo @@ -19,6 +19,7 @@ def hex_to_bytes(s: String) -> List[UInt8]: r.append((hi << 4) | lo) return r^ + def extract_p256_public_key(public_der: List[UInt8]) -> List[UInt8]: # RFC 5480 section 2.1 identifies EC public keys with id-ecPublicKey plus # a namedCurve OID. SEC 2 v2.0 appendix A assigns secp256r1: @@ -63,7 +64,7 @@ def extract_p256_public_key(public_der: List[UInt8]) -> List[UInt8]: def extract_trailing_sec1_p256_public_key( - public_der: List[UInt8], + public_der: List[UInt8] ) -> List[UInt8]: if len(public_der) >= 65 and public_der[len(public_der) - 65] == 0x04: var out = List[UInt8](capacity=65) @@ -115,7 +116,7 @@ def run_case( private_hex: String, public_hex: String, shared_hex: String, - result: String, + result: String ) -> Bool: var is_valid = result == "valid" var is_acceptable = result == "acceptable" @@ -132,13 +133,13 @@ def run_case( var got = p256_ecdh( Span[UInt8, ...](private_key), Span[UInt8, ...](public_key), - Span[mut=True, UInt8, ...](unsafe_ptr=actual.unsafe_ptr(), length=32), + Span[mut=True, UInt8, ...](unsafe_ptr=actual.unsafe_ptr(), length=32) ) if is_valid and not got: print( "Test ", tc_id, - " validity mismatch: valid vector was rejected", + " validity mismatch: valid vector was rejected" ) return False if (not is_valid and not is_acceptable) and got: @@ -149,6 +150,7 @@ def run_case( return False return True + def main() raises: print("Wycheproof P-256 ECDH") var py = Python.import_module("json") @@ -169,7 +171,7 @@ def main() raises: String(t["private"]), String(t["public"]), String(t["shared"]), - result, + result ): ok_count += 1 else: diff --git a/tests/test_wycheproof_p384_ecdh.mojo b/tests/test_wycheproof_p384_ecdh.mojo index de91302..a8a89bc 100644 --- a/tests/test_wycheproof_p384_ecdh.mojo +++ b/tests/test_wycheproof_p384_ecdh.mojo @@ -19,6 +19,7 @@ def hex_to_bytes(s: String) -> List[UInt8]: r.append((hi << 4) | lo) return r^ + def extract_p384_public_key(public_der: List[UInt8]) -> List[UInt8]: # RFC 5480 section 2.1 identifies EC public keys with id-ecPublicKey plus # a namedCurve OID. SEC 2 v2.0 appendix A assigns secp384r1: @@ -58,8 +59,9 @@ def extract_p384_public_key(public_der: List[UInt8]) -> List[UInt8]: return out^ return List[UInt8]() + def extract_trailing_sec1_p384_public_key( - public_der: List[UInt8], + public_der: List[UInt8] ) -> List[UInt8]: if len(public_der) >= 97 and public_der[len(public_der) - 97] == 0x04: var out = List[UInt8](capacity=97) @@ -76,6 +78,7 @@ def extract_trailing_sec1_p384_public_key( return out^ return List[UInt8]() + def normalize_p384_private_key(var private_key: List[UInt8]) -> List[UInt8]: # DER INTEGER values may include a leading 00 octet to keep the integer # positive. SEC 1 scalar input to p384_ecdh is the fixed 48-octet value. @@ -95,6 +98,7 @@ def normalize_p384_private_key(var private_key: List[UInt8]) -> List[UInt8]: return out^ return List[UInt8]() + def matches48( actual: StackInlineArray[UInt8, 48], expected: List[UInt8] ) -> Bool: @@ -103,12 +107,13 @@ def matches48( return False return True + def run_case( tc_id: String, private_hex: String, public_hex: String, shared_hex: String, - result: String, + result: String ) -> Bool: var is_valid = result == "valid" var is_acceptable = result == "acceptable" @@ -125,13 +130,13 @@ def run_case( var got = p384_ecdh( Span[UInt8, ...](private_key), Span[UInt8, ...](public_key), - Span[mut=True, UInt8, ...](unsafe_ptr=actual.unsafe_ptr(), length=48), + Span[mut=True, UInt8, ...](unsafe_ptr=actual.unsafe_ptr(), length=48) ) if is_valid and not got: print( "Test ", tc_id, - " validity mismatch: valid vector was rejected", + " validity mismatch: valid vector was rejected" ) return False if (not is_valid and not is_acceptable) and got: @@ -142,6 +147,7 @@ def run_case( return False return True + def main() raises: print("Wycheproof P-384 ECDH") var py = Python.import_module("json") @@ -162,7 +168,7 @@ def main() raises: String(t["private"]), String(t["public"]), String(t["shared"]), - result, + result ): ok_count += 1 else: diff --git a/tests/test_wycheproof_x25519.mojo b/tests/test_wycheproof_x25519.mojo index b34a751..300b669 100644 --- a/tests/test_wycheproof_x25519.mojo +++ b/tests/test_wycheproof_x25519.mojo @@ -6,6 +6,7 @@ from std.python import Python from thistle.x25519 import x25519 from thistle.utils import StackInlineArray + def hex_to_bytes(s: String) -> List[UInt8]: var r = List[UInt8]() var b = s.as_bytes() @@ -18,12 +19,14 @@ def hex_to_bytes(s: String) -> List[UInt8]: r.append((hi << 4) | lo) return r^ + def matches32(actual: StackInlineArray[UInt8, 32], expected: List[UInt8]) -> Bool: for i in range(32): if actual[i] != expected[i]: return False return True + def run_case(tc_id: String, private_hex: String, public_hex: String, shared_hex: String) raises -> Bool: var private_key = hex_to_bytes(private_hex) var public_key = hex_to_bytes(public_hex) @@ -34,13 +37,14 @@ def run_case(tc_id: String, private_hex: String, public_hex: String, shared_hex: 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 return True + def main() raises: print("Wycheproof X25519") var py = Python.import_module("json") diff --git a/tests/thistle_test_vectors.mojo b/tests/thistle_test_vectors.mojo index 0877c5c..6cbb7d0 100644 --- a/tests/thistle_test_vectors.mojo +++ b/tests/thistle_test_vectors.mojo @@ -7,7 +7,7 @@ from thistle.sha2 import ( string_to_bytes, sha224_hash_bits, sha256_hash_bits, - sha384_hash_bits, + sha384_hash_bits ) from thistle.argon2 import Argon2id from thistle.blake2b import Blake2b @@ -20,22 +20,24 @@ from thistle.camellia import ( camellia_decrypt_blocks, camellia_cbc_encrypt_kernel, camellia_cbc_decrypt_kernel, - camellia_ctr_kernel, + camellia_ctr_kernel ) from thistle.chacha20 import ChaCha20 from thistle.chacha20poly1305 import ( chacha20_poly1305_encrypt, chacha20_poly1305_decrypt, xchacha20_poly1305_encrypt, - xchacha20_poly1305_decrypt, + xchacha20_poly1305_decrypt ) 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, AESExpandedKey, + 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.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 @@ -138,7 +140,7 @@ def test_argon2(data: PythonObject, py: PythonObject) raises -> TestResult: tag_length=Int(py=v["tag_length"]), memory_size_kb=Int(py=v["memory_size_kb"]), iterations=Int(py=v["iterations"]), - version=Int(py=v["version"]), + version=Int(py=v["version"]) ) var got = bytes_to_hex(argon2.hash(Span[UInt8, ...](pass_bytes))) var expected = String(v["hash"]) @@ -183,7 +185,7 @@ def test_blake3(data: PythonObject, py: PythonObject) raises -> TestResult: var got = bytes_to_hex( blake3_parallel_hash( Span[UInt8, ...](generate_blake3_input(input_len)), - expected.byte_length() // 2, + expected.byte_length() // 2 ) ) if got == expected: @@ -300,7 +302,7 @@ def test_camellia(data: PythonObject, py: PythonObject) raises -> TestResult: ctr_out.unsafe_ptr().unsafe_origin_cast[MutAnyOrigin](), cipher, ctr_nb, - nonce.unsafe_ptr(), + nonce.unsafe_ptr() ) for i2 in range(ctr_nb * 16): if ctr_out[i2] != msg[i2]: @@ -337,7 +339,7 @@ def test_camellia(data: PythonObject, py: PythonObject) raises -> TestResult: cbc_out.unsafe_ptr().unsafe_origin_cast[MutAnyOrigin](), cipher, cbc_nb, - nonce.unsafe_ptr(), + nonce.unsafe_ptr() ) for i2 in range(cbc_nb * 16): if cbc_out[i2] != cbc_msg[i2]: @@ -360,7 +362,7 @@ def test_chacha20(data: PythonObject, py: PythonObject) raises -> TestResult: var cipher = ChaCha20( list_to_simd32(hex_to_bytes(String(v["key"]))), Span[UInt8, ...](nonce), - UInt32(Int(py=v["counter"])), + UInt32(Int(py=v["counter"])) ) var pt_bytes = hex_to_bytes(String(v["plaintext"])) var expected_ct = hex_to_bytes(String(v["ciphertext"])) @@ -454,7 +456,7 @@ def _test_pbkdf2_sha256( password, Span[UInt8, ...](salt), Int(py=v["iterations"]), - Int(py=v["dklen"]), + Int(py=v["dklen"]) ) ) var expected = String(v["derived_key"]) @@ -482,7 +484,7 @@ def _test_pbkdf2_sha512( password, Span[UInt8, ...](salt), Int(py=v["iterations"]), - Int(py=v["dklen"]), + Int(py=v["dklen"]) ) ) var expected = String(v["derived_key"]) @@ -512,9 +514,7 @@ def _test_sha224(data: PythonObject, py: PythonObject) raises -> TestResult: for i in range(Int(py=data.__len__())): var v = data[i] var bit_len = Int(py=v["len"]) - var msg = ( - hex_to_bytes(String(v["msg"])) if bit_len > 0 else List[UInt8]() - ) + var msg = hex_to_bytes(String(v["msg"])) if bit_len > 0 else List[UInt8]() var got = bytes_to_hex(sha224_hash_bits(Span[UInt8, ...](msg), bit_len)) var expected = String(v["md"]) if got == expected: @@ -538,9 +538,7 @@ def _test_sha256(data: PythonObject, py: PythonObject) raises -> TestResult: for i in range(Int(py=data.__len__())): var v = data[i] var bit_len = Int(py=v["len"]) - var msg = ( - hex_to_bytes(String(v["msg"])) if bit_len > 0 else List[UInt8]() - ) + var msg = hex_to_bytes(String(v["msg"])) if bit_len > 0 else List[UInt8]() var got = bytes_to_hex(sha256_hash_bits(Span[UInt8, ...](msg), bit_len)) var expected = String(v["md"]) if got == expected: @@ -564,9 +562,7 @@ def _test_sha384(data: PythonObject, py: PythonObject) raises -> TestResult: for i in range(Int(py=data.__len__())): var v = data[i] var bit_len = Int(py=v["len"]) - var msg = ( - hex_to_bytes(String(v["msg"])) if bit_len > 0 else List[UInt8]() - ) + var msg = hex_to_bytes(String(v["msg"])) if bit_len > 0 else List[UInt8]() var got = bytes_to_hex(sha384_hash_bits(Span[UInt8, ...](msg), bit_len)) var expected = String(v["md"]) if got == expected: @@ -678,9 +674,7 @@ def test_sha_ni(data: PythonObject, py: PythonObject) raises -> TestResult: for i in range(Int(py=sha256_data.__len__())): var v = sha256_data[i] var bit_len = Int(py=v["len"]) - var msg = ( - hex_to_bytes(String(v["msg"])) if bit_len > 0 else List[UInt8]() - ) + var msg = hex_to_bytes(String(v["msg"])) if bit_len > 0 else List[UInt8]() var got = bytes_to_hex(sha256ni_hash(Span[UInt8, ...](msg))) var expected = String(v["md"]) if got == expected: @@ -718,12 +712,12 @@ def test_aes_gcm(data: PythonObject, py: PythonObject) raises -> TestResult: try: var enc = aes_gcm_encrypt( Span[UInt8, ...](key), Span[UInt8, ...](iv), - Span[UInt8, ...](msg), Span[UInt8, ...](aad), + Span[UInt8, ...](msg), Span[UInt8, ...](aad) ) var dec = aes_gcm_decrypt( Span[UInt8, ...](key), Span[UInt8, ...](iv), Span[UInt8, ...](ct), Span[UInt8, ...](aad), - Span[UInt8, ...](tag), + Span[UInt8, ...](tag) ) ok = ( bytes_to_hex(enc[0]) == bytes_to_hex(ct) @@ -738,7 +732,7 @@ def test_aes_gcm(data: PythonObject, py: PythonObject) raises -> TestResult: var dec = aes_gcm_decrypt( Span[UInt8, ...](key), Span[UInt8, ...](iv), Span[UInt8, ...](ct), Span[UInt8, ...](aad), - Span[UInt8, ...](tag), + Span[UInt8, ...](tag) ) ok = (not dec[1]) and len(dec[0]) == 0 except: @@ -785,14 +779,14 @@ def test_chacha20_poly1305(data: PythonObject, py: PythonObject, xchacha: Bool) Span[UInt8, ...](key), Span[UInt8, ...](iv), Span[UInt8, ...](aad), Span[UInt8, ...](msg), Span[mut=True, UInt8, ...](out_ct), - Span[mut=True, UInt8, ...](out_tag), + Span[mut=True, UInt8, ...](out_tag) ) else: chacha20_poly1305_encrypt( Span[UInt8, ...](key), Span[UInt8, ...](iv), Span[UInt8, ...](aad), Span[UInt8, ...](msg), Span[mut=True, UInt8, ...](out_ct), - Span[mut=True, UInt8, ...](out_tag), + Span[mut=True, UInt8, ...](out_tag) ) enc_matches = len(ct) == n and len(tag) == 16 for i in range(n): @@ -810,13 +804,13 @@ def test_chacha20_poly1305(data: PythonObject, py: PythonObject, xchacha: Bool) dec_ok = xchacha20_poly1305_decrypt( Span[UInt8, ...](key), Span[UInt8, ...](iv), Span[UInt8, ...](aad), Span[UInt8, ...](ct), Span[UInt8, ...](tag), - Span[mut=True, UInt8, ...](out_pt), + 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), - Span[mut=True, UInt8, ...](out_pt), + Span[mut=True, UInt8, ...](out_pt) ) if dec_ok: for i in range(len(ct)): @@ -863,7 +857,7 @@ def test_aes_cpu_modes(data: PythonObject, py: PythonObject) raises -> TestResul var aad = List[UInt8]() var enc = aes_gcm_encrypt( Span[UInt8, ...](key), Span[UInt8, ...](nonce), - Span[UInt8, ...](pt), Span[UInt8, ...](aad), + Span[UInt8, ...](pt), Span[UInt8, ...](aad) ) ok = bytes_to_hex(enc[0]) == bytes_to_hex(ct_exp) if len(tag_exp) == 16: @@ -974,7 +968,7 @@ def main() raises: test_argon2(load_json("tests/vectors/argon2.json", py), py), tp, tf, - af, + af ) except e: print("Argon2 [error] " + String(e)) @@ -988,7 +982,7 @@ def main() raises: test_blake2b(load_json("tests/vectors/blake2b.json", py), py), tp, tf, - af, + af ) except e: print("BLAKE2b [error] " + String(e)) @@ -1002,7 +996,7 @@ def main() raises: test_blake3(load_json("tests/vectors/blake3.json", py), py), tp, tf, - af, + af ) except e: print("BLAKE3 [error] " + String(e)) @@ -1016,7 +1010,7 @@ def main() raises: test_camellia(load_json("tests/vectors/camellia.json", py), py), tp, tf, - af, + af ) except e: print("Camellia [error] " + String(e)) @@ -1030,7 +1024,7 @@ def main() raises: test_chacha20(load_json("tests/vectors/chacha20.json", py), py), tp, tf, - af, + af ) except e: print("ChaCha20 [error] " + String(e)) @@ -1044,7 +1038,7 @@ def main() raises: test_kcipher2(load_json("tests/vectors/kcipher2.json", py), py), tp, tf, - af, + af ) except e: print("KCipher2 [error] " + String(e)) @@ -1058,7 +1052,7 @@ def main() raises: test_pbkdf2(load_json("tests/vectors/pbkdf2.json", py), py), tp, tf, - af, + af ) except e: print("PBKDF2 [error] " + String(e)) @@ -1072,7 +1066,7 @@ def main() raises: test_sha(load_json("tests/vectors/sha.json", py), py), tp, tf, - af, + af ) except e: print("SHA [error] " + String(e)) @@ -1086,7 +1080,7 @@ def main() raises: test_aes_cpu(load_json("tests/vectors/aes.json", py), py), tp, tf, - af, + af ) except e: print("AES-CPU [error] " + String(e)) @@ -1100,7 +1094,7 @@ def main() raises: test_aes_ni(load_json("tests/vectors/aes.json", py), py), tp, tf, - af, + af ) except e: print("AES-NI [error] " + String(e)) @@ -1114,7 +1108,7 @@ def main() raises: test_aes_cpu_modes(load_json("tests/vectors/aes_test_vectors.json", py), py), tp, tf, - af, + af ) except e: print("AES-CPU-Modes [error] " + String(e)) @@ -1130,7 +1124,7 @@ def main() raises: ), tp, tf, - af, + af ) except e: print("ChaCha20-Poly1305 [error] " + String(e)) @@ -1146,7 +1140,7 @@ def main() raises: ), tp, tf, - af, + af ) except e: print("XChaCha20-Poly1305 [error] " + String(e)) @@ -1160,7 +1154,7 @@ def main() raises: test_aes_gcm(load_json("tests/Wycheproof/aes_gcm_test.json", py), py), tp, tf, - af, + af ) except e: print("AES-GCM [error] " + String(e)) @@ -1174,7 +1168,7 @@ def main() raises: test_sha_ni(load_json("tests/vectors/sha.json", py), py), tp, tf, - af, + af ) except e: print("SHA256-NI [error] " + String(e))