core/crypto/mlkem: Initial import

This commit is contained in:
Yawning Angel
2026-04-01 04:21:36 +09:00
parent 8f1067f290
commit d2c29c025e
19 changed files with 2177 additions and 0 deletions

View File

@@ -0,0 +1,84 @@
package benchmark_core_crypto
import "core:log"
import "core:testing"
import "core:text/table"
import "core:time"
import "core:crypto"
import "core:crypto/mlkem"
@(private = "file")
MLKEM_ITERS :: 50000
@(test)
benchmark_crypto_mlkem :: proc(t: ^testing.T) {
if !crypto.HAS_RAND_BYTES {
log.warnf("ML-KEM benchmarks skipped, no system entropy source")
}
tbl: table.Table
table.init(&tbl)
defer table.destroy(&tbl)
table.caption(&tbl, "ML-KEM")
table.aligned_header_of_values(&tbl, .Right, "Parameters", "Keygen", "Encaps", "Decaps")
append_tbl := proc(tbl: ^table.Table, algo_name: string, keygen, encaps, decaps: time.Duration) {
table.aligned_row_of_values(
tbl,
.Right,
algo_name,
table.format(tbl, "%8M", keygen),
table.format(tbl, "%8M", encaps),
table.format(tbl, "%8M", decaps),
)
}
for params in mlkem.Parameters {
if params == .Invalid {
continue
}
param_name := MLKEM_PARAMS_NAMES[params]
decaps_key: mlkem.Decapsulation_Key
start := time.tick_now()
for _ in 0 ..< MLKEM_ITERS {
_ = mlkem.decapsulation_key_generate(&decaps_key, params)
}
keygen := time.tick_since(start) / MLKEM_ITERS
encaps_key := make([]byte, mlkem.ENCAPSULATION_KEY_SIZES[params])
defer delete(encaps_key)
ciphertext := make([]byte, mlkem.CIPHERTEXT_SIZES[params])
defer delete(ciphertext)
mlkem.decapsulation_key_encaps_bytes(&decaps_key, encaps_key)
bob_shared: [mlkem.SHARED_SECRET_SIZE]byte
start = time.tick_now()
for _ in 0 ..< MLKEM_ITERS {
_ = mlkem.encaps(params, encaps_key, bob_shared[:], ciphertext)
}
encaps := time.tick_since(start) / MLKEM_ITERS
alice_shared: [mlkem.SHARED_SECRET_SIZE]byte
start = time.tick_now()
for _ in 0 ..< MLKEM_ITERS {
_ = mlkem.decaps(&decaps_key, ciphertext, alice_shared[:])
}
decaps := time.tick_since(start) / MLKEM_ITERS
append_tbl(&tbl, param_name, keygen, encaps, decaps)
}
log_table(&tbl)
}
@(private="file")
MLKEM_PARAMS_NAMES := [mlkem.Parameters]string {
.Invalid = "invalid",
.ML_KEM_512 = "ML-KEM-512",
.ML_KEM_768 = "ML-KEM-768",
.ML_KEM_1024 = "ML-KEM-1024",
}

View File

@@ -0,0 +1,72 @@
package test_core_crypto
import "core:bytes"
import "core:log"
import "core:testing"
import "core:crypto"
import "core:crypto/mlkem"
@(test)
test_mlkem :: proc(t: ^testing.T) {
if !crypto.HAS_RAND_BYTES {
log.info("rand_bytes not supported - skipping")
return
}
// Test vectors are huge, and are covered by the wycheproof corpus,
// so just test a full key exchange with all supported parameter
// sets.
for params in mlkem.Parameters {
if params == .Invalid {
continue
}
// Alice
decaps_key: mlkem.Decapsulation_Key
if !testing.expectf(
t,
mlkem.decapsulation_key_generate(&decaps_key, params),
"%v: decapsulation_key_generate",
params,
) {
continue
}
defer mlkem.decapsulation_key_clear(&decaps_key)
ek_bytes := make([]byte, mlkem.ENCAPSULATION_KEY_SIZES[params])
defer delete(ek_bytes)
mlkem.decapsulation_key_encaps_bytes(&decaps_key, ek_bytes)
// Bob
bob_shared_secret: [mlkem.SHARED_SECRET_SIZE]byte
ciphertext := make([]byte, mlkem.CIPHERTEXT_SIZES[params])
defer delete(ciphertext)
if !testing.expectf(
t,
mlkem.encaps(params, ek_bytes, bob_shared_secret[:], ciphertext),
"%v: encaps",
params,
) {
continue
}
// Alice
alice_shared_secret: [mlkem.SHARED_SECRET_SIZE]byte
if !testing.expectf(
t,
mlkem.decaps(&decaps_key, ciphertext, alice_shared_secret[:]),
"%v: decaps",
params,
) {
continue
}
testing.expectf(
t,
bytes.equal(alice_shared_secret[:], bob_shared_secret[:]),
"%v: shared secret mismatch",
params,
)
}
}

View File

@@ -35,6 +35,16 @@ import "core:testing"
// - crypto/kmac
// - kmac128_no_customization_test.json
// - kmac256_no_customization_test.json
// - crypto/mlkem
// - mlkem_512_keygen_seed_test.json
// - mlkem_512_encaps_test.json
// - mlkem_512_test.json
// - mlkem_768_keygen_seed_test.json
// - mlkem_768_encaps_test.json
// - mlkem_768_test.json
// - mlkem_1024_keygen_seed_test.json
// - mlkem_1024_encaps_test.json
// - mlkem_1024_test.json
// - crypto/pbkdf2
// - pbkdf2_hmacsha1_test.json
// - pbkdf2_hmacsha224_test.json

View File

@@ -0,0 +1,401 @@
package test_wycheproof
import "core:encoding/hex"
import "core:log"
import "core:mem"
import "core:os"
import "core:testing"
import "core:crypto/_mlkem"
import "core:crypto/mlkem"
import "../common"
@(test)
test_mlkem :: proc(t: ^testing.T) {
arena: mem.Arena
arena_backing := make([]byte, ARENA_SIZE)
defer delete(arena_backing)
mem.arena_init(&arena, arena_backing)
context.allocator = mem.arena_allocator(&arena)
log.debug("mlkem: starting")
files_keygen := []string {
"mlkem_512_keygen_seed_test.json",
"mlkem_768_keygen_seed_test.json",
"mlkem_1024_keygen_seed_test.json",
}
for f in files_keygen {
mem.free_all()
fn, _ := os.join_path([]string{BASE_PATH, f}, context.allocator)
test_vectors: Test_Vectors(Kem_Test_Group)
load_ok := load(&test_vectors, fn)
if !testing.expectf(t, load_ok, "Unable to load {}", f) {
continue
}
testing.expectf(t, test_mlkem_keygen(t, &test_vectors), "ML-KEM KeyGen failed")
}
files_encaps := []string {
"mlkem_512_encaps_test.json",
"mlkem_768_encaps_test.json",
"mlkem_1024_encaps_test.json",
}
for f in files_encaps {
mem.free_all()
fn, _ := os.join_path([]string{BASE_PATH, f}, context.allocator)
test_vectors: Test_Vectors(Kem_Test_Group)
load_ok := load(&test_vectors, fn)
if !testing.expectf(t, load_ok, "Unable to load {}", f) {
continue
}
testing.expectf(t, test_mlkem_encaps(t, &test_vectors), "ML-KEM Encaps failed")
}
files_decaps := []string {
"mlkem_512_test.json",
"mlkem_768_test.json",
"mlkem_1024_test.json",
}
for f in files_decaps {
mem.free_all()
fn, _ := os.join_path([]string{BASE_PATH, f}, context.allocator)
test_vectors: Test_Vectors(Kem_Test_Group)
load_ok := load(&test_vectors, fn)
if !testing.expectf(t, load_ok, "Unable to load {}", f) {
continue
}
testing.expectf(t, test_mlkem_decaps(t, &test_vectors), "ML-KEM Decaps failed")
}
}
test_mlkem_keygen :: proc(t: ^testing.T, test_vectors: ^Test_Vectors(Kem_Test_Group)) -> bool {
params_str := test_vectors.test_groups[0].parameter_set
params := parameter_set_to_params(params_str)
if params == .Invalid {
return false
}
log.debugf("%s: KeyGen starting", params_str)
num_ran, num_passed, num_failed, num_skipped: int
for &test_group, tg_id in test_vectors.test_groups {
for &test_vector in test_group.tests {
num_ran += 1
seed := common.hexbytes_decode(test_vector.seed)
dk: mlkem.Decapsulation_Key
if !testing.expectf(
t,
mlkem.decapsulation_key_set_bytes(&dk, params, seed),
"%s/KeyGen/%d/%d: failed to set decapsulation key from seed",
params_str,
tg_id,
test_vector.tc_id,
test_vector.seed,
) {
num_failed *= 1
continue
}
ek_bytes := make([]byte, mlkem.ENCAPSULATION_KEY_SIZES[params])
mlkem.decapsulation_key_encaps_bytes(&dk, ek_bytes)
ok := common.hexbytes_compare(test_vector.ek, ek_bytes)
if !result_check(test_vector.result, ok) {
x := transmute(string)(hex.encode(ek_bytes))
log.errorf(
"%s/KeyGen/%d/%d: ek: expected %s actual %s",
params_str,
tg_id,
test_vector.tc_id,
test_vector.ek,
x,
)
num_failed += 1
continue
}
dk_bytes := make([]byte, mlkem.DECAPSULATION_KEY_EXPANDED_SIZES[params])
mlkem.decapsulation_key_expanded_bytes(&dk, dk_bytes)
ok = common.hexbytes_compare(test_vector.dk, dk_bytes)
if !result_check(test_vector.result, ok) {
x := transmute(string)(hex.encode(dk_bytes))
log.errorf(
"%s/KeyGen/%d/%d: dk: expected %s actual %s",
tg_id,
params_str,
test_vector.tc_id,
test_vector.dk,
x,
)
num_failed += 1
continue
}
seed_bytes: [mlkem.DECAPSULATION_KEY_SEED_SIZE]byte
mlkem.decapsulation_key_bytes(&dk, seed_bytes[:])
ok = common.hexbytes_compare(test_vector.seed, seed_bytes[:])
if !result_check(test_vector.result, ok) {
x := transmute(string)(hex.encode(seed_bytes[:]))
log.errorf(
"%s/KeyGen/%d/%d: seed: expected %s actual %s",
tg_id,
params_str,
test_vector.tc_id,
test_vector.seed,
x,
)
num_failed += 1
continue
}
num_passed += 1
}
}
assert(num_ran == test_vectors.number_of_tests)
assert(num_passed + num_failed + num_skipped == num_ran)
log.infof(
"%s/KeyGen: ran %d, passed %d, failed %d, skipped %d",
params_str,
num_ran,
num_passed,
num_failed,
num_skipped,
)
return num_failed == 0
}
test_mlkem_encaps :: proc(t: ^testing.T, test_vectors: ^Test_Vectors(Kem_Test_Group)) -> bool {
params_str := test_vectors.test_groups[0].parameter_set
params := parameter_set_to_params(params_str)
if params == .Invalid {
return false
}
log.debugf("%s: Encaps starting", params_str)
num_ran, num_passed, num_failed, num_skipped: int
for &test_group, tg_id in test_vectors.test_groups {
for &test_vector in test_group.tests {
num_ran += 1
ek: mlkem.Encapsulation_Key
ok := mlkem.encapsulation_key_set_bytes(
&ek,
params,
common.hexbytes_decode(test_vector.ek),
)
// The current corpus can only fail if the encapsulation key
// is malformed in some way.
if !result_check(test_vector.result, ok) {
log.errorf(
"%s/Encaps/%d/%d: unexpected set encapsulation key from bytes: %s (%v != %v)",
params_str,
tg_id,
test_vector.tc_id,
test_vector.ek,
test_vector.result,
ok,
)
num_failed += 1
continue
}
if !ok {
num_passed += 1
continue
}
shared_secret: [mlkem.SHARED_SECRET_SIZE]byte
ciphertext := make([]byte, mlkem.CIPHERTEXT_SIZES[params])
_mlkem.kem_encaps_internal(
shared_secret[:],
ciphertext,
&ek,
common.hexbytes_decode(test_vector.m),
)
ok = common.hexbytes_compare(test_vector.c, ciphertext)
if !ok {
x := transmute(string)(hex.encode(ciphertext))
log.errorf(
"%s/Encaps/%d/%d: ciphertext: expected: %s actual: %s",
params_str,
tg_id,
test_vector.tc_id,
test_vector.c,
x,
)
num_failed += 1
continue
}
ok = common.hexbytes_compare(test_vector.k, shared_secret[:])
if !ok {
x := transmute(string)(hex.encode(shared_secret[:]))
log.errorf(
"%s/Encaps/%d/%d: shared_secret: expected: %s actual: %s",
params_str,
tg_id,
test_vector.tc_id,
test_vector.k,
x,
)
num_failed += 1
continue
}
num_passed += 1
}
}
assert(num_ran == test_vectors.number_of_tests)
assert(num_passed + num_failed + num_skipped == num_ran)
log.infof(
"%s/Encaps: ran %d, passed %d, failed %d, skipped %d",
params_str,
num_ran,
num_passed,
num_failed,
num_skipped,
)
return num_failed == 0
}
test_mlkem_decaps :: proc(t: ^testing.T, test_vectors: ^Test_Vectors(Kem_Test_Group)) -> bool {
params_str := test_vectors.test_groups[0].parameter_set
params := parameter_set_to_params(params_str)
if params == .Invalid {
return false
}
log.debugf("%s: Decaps starting", params_str)
num_ran, num_passed, num_failed, num_skipped: int
for &test_group, tg_id in test_vectors.test_groups {
for &test_vector in test_group.tests {
num_ran += 1
// We do not have an API for decaps with raw seed.
seed := common.hexbytes_decode(test_vector.seed)
switch len(seed) {
case mlkem.DECAPSULATION_KEY_SEED_SIZE:
case:
if testing.expectf(
t,
result_is_invalid(test_vector.result),
"%s/Decaps/%d/%d: test vector expects success with invalid seed",
params_str,
tg_id,
test_vector.tc_id,
) {
num_passed += 1
} else {
num_failed += 1
}
continue
}
dk: mlkem.Decapsulation_Key
if !testing.expectf(
t,
mlkem.decapsulation_key_set_bytes(&dk, params, seed),
"%s/Decaps/%d/%d: failed to set decapsulation key from seed",
params_str,
tg_id,
test_vector.tc_id,
test_vector.seed,
) {
num_failed *= 1
continue
}
shared_secret: [mlkem.SHARED_SECRET_SIZE]byte
ok := mlkem.decaps(
&dk,
common.hexbytes_decode(test_vector.c),
shared_secret[:],
)
if !result_check(test_vector.result, ok) {
log.errorf(
"%s/Decaps/%d/%d: unexpected decapsulation failure",
params_str,
tg_id,
test_vector.tc_id,
)
num_failed += 1
continue
}
if !ok {
num_passed += 1
continue
}
ok = common.hexbytes_compare(test_vector.k, shared_secret[:])
if !ok {
x := transmute(string)(hex.encode(shared_secret[:]))
log.errorf(
"%s/Decaps/%d/%d: shared_secret: expected: %s actual: %s",
params_str,
tg_id,
test_vector.tc_id,
test_vector.k,
x,
)
num_failed += 1
continue
}
num_passed += 1
}
}
assert(num_ran == test_vectors.number_of_tests)
assert(num_passed + num_failed + num_skipped == num_ran)
log.infof(
"%s/Decaps: ran %d, passed %d, failed %d, skipped %d",
params_str,
num_ran,
num_passed,
num_failed,
num_skipped,
)
return num_failed == 0
}
@(require_results, private="file")
parameter_set_to_params :: proc(s: string) -> mlkem.Parameters {
switch s {
case "ML-KEM-512":
return .ML_KEM_512
case "ML-KEM-768":
return .ML_KEM_768
case "ML-KEM-1024":
return .ML_KEM_1024
case:
return .Invalid
}
}

View File

@@ -64,6 +64,11 @@ Test_Vectors_Note :: struct {
links: []string `json:"links"`,
}
Test_Group_Source :: struct {
name: string `json:"name"`,
version: string `json:"version"`,
}
Aead_Test_Group :: struct {
iv_size: int `json:"ivSize"`,
key_size: int `json:"keySize"`,
@@ -198,3 +203,23 @@ Pbkdf_Test_Vector :: struct {
result: Result `json:"result"`,
flags: []string `json:"flags"`,
}
Kem_Test_Group :: struct {
type: string `json:"type"`,
source: Test_Group_Source `json:"source"`,
parameter_set: string `json:"parameterSet"`,
tests: []Kem_Test_Vector `json:"tests"`,
}
Kem_Test_Vector :: struct {
tc_id: int `json:"tcId"`,
flags: []string `json:"flags"`,
comment: string `json:"comment"`,
seed: common.Hex_Bytes `json:"seed"`,
m: common.Hex_Bytes `json:"m"`,
ek: common.Hex_Bytes `json:"ek"`,
dk: common.Hex_Bytes `json:"dk"`,
c: common.Hex_Bytes `json:"c"`,
k: common.Hex_Bytes `json:"K"`,
result: Result `json:"result"`,
}