Cleaner msg struct naming.

This commit is contained in:
LeoVasanko committed 2026-10-07 18:40:43 +00:00
1 parent 10b64cbf52
commit 0e1190ffad
2 files changed
+30 -58

No files matched your search

+28 -56
View File
@@ -25,40 +25,17 @@ from .kdf import kdf
# -- Fixed-size record layouts (struct handles packing + length checks) -----
# Hybrid scheme = ML-KEM-768 (pk 1184 / sk-seed 64 / ct 1088) || X25519 (32)
_TAG = struct.Struct("32s") # Bare kdf tag
# Handshake message bodies; box(N) is a chab3 box of N plaintext bytes (N+28)
_MSG1_REGISTER = struct.Struct("1184s32s32s") # e_alice_m | e_alice_x | commitment
_MSG2_REGISTER = struct.Struct(
"32s1088s1244s"
) # e_bob | ee_ct | box(s_bob_m | s_bob_x)
_MSG2_REGISTER_INNER = struct.Struct("1184s32s") # s_bob_m | s_bob_x
_MSG3_REGISTER = struct.Struct("1088s1276s") # es_ct | box(s_alice_m | s_alice_x | r)
_MSG3_REGISTER_INNER = struct.Struct("1184s32s32s") # s_alice_m | s_alice_x | reveal
_MSG4_REGISTER = struct.Struct("1088s32s") # se_ct | tb
_MSG1_RESUME = struct.Struct("1184s32s1088s") # e_alice_m | e_alice_x | sb_ct
_MSG2_RESUME = struct.Struct("32s1088s1088s32s") # e_bob | sa_ct | ee_ct | tb
# msg5 = _TAG (ta); msg3 resumption = _TAG (ta)
REG_MSG1 = struct.Struct("1184s32s32s") # e_alice_m | e_alice_x | commitment
REG_MSG2 = struct.Struct("32s1088s1244s") # e_bob | ee_ct | encrypt_ee(inner)
REG_MSG2_INNER = struct.Struct("1184s32s") # s_bob_m | s_bob_x
REG_MSG3 = struct.Struct("1088s1276s") # es_ct | encrypt_es(inner)
REG_MSG3_INNER = struct.Struct("1184s32s32s") # s_alice_m | s_alice_x | reveal
REG_MSG4 = struct.Struct("1088s32s") # se_ct | tb
# REG_MSG5 is plain tag, no Struct needed
# kdf info strings: boxes are namespaced per flow and message (reg/msg1,
# res/msg2, reg/msg2/inner, ...) so the two exchanges are disjoint; other
# derivations are bare functional labels, shared between the flows when
# the inputs have the same format (a2b/b2a, idroot/*, cred/id).
# --------------------------------------------------------------------------
# Hybrid KEM: ML-KEM-768 || X25519
#
# A static private key is a single 32-byte idroot secret: each scheme's
# key is derived from it with one domain-separated kdf call, so the two
# subkeys are independent PRF outputs and the idroot is the only thing
# held (and the only thing the credential stores). Ephemeral keys are
# generated directly -- nothing of them persists.
#
# Each party ships one X25519 ephemeral public per connection (Alice's
# inside e_alice in msg1, Bob's as e_bob in msg2); every X25519 leg is a
# plain DH against those shipped keys. ML-KEM legs ship as bare ML-KEM
# ciphertexts.
RES_MSG1 = struct.Struct("1184s32s1088s") # e_alice_m | e_alice_x | sb_ct
RES_MSG2 = struct.Struct("32s1088s1088s32s") # e_bob | sa_ct | ee_ct | tb
def _mlkem_private_key(idroot: bytes) -> MLKEM768PrivateKey:
@@ -164,7 +141,7 @@ class Registration:
self._eph_m, self._ephemeral, e_alice_m, e_alice_x = _fresh_ephemeral()
self._reveal = random(32)
commitment = kdf("commit", e_alice_m, e_alice_x, self._reveal)
body = _MSG1_REGISTER.pack(e_alice_m, e_alice_x, commitment)
body = REG_MSG1.pack(e_alice_m, e_alice_x, commitment)
return self._chain.encrypt(body, key=kdf("reg/msg1", self._shared))
def respondA(self, msg2: bytes) -> bytes:
@@ -172,12 +149,12 @@ class Registration:
# proves no key knowledge -- then challenge his static
plain = self._chain.decrypt(msg2, key=kdf("reg/msg2", self._shared))
try:
e_bob, ee_ct, inner = _MSG2_REGISTER.unpack(plain)
e_bob, ee_ct, inner = REG_MSG2.unpack(plain)
# ee is forward secrecy, made to our connection key only
ee = self._eph_m.decapsulate(ee_ct)
ee += self._ephemeral.exchange(X25519PublicKey.from_public_bytes(e_bob))
# Solving ee lets us read Bob's static from the inner box
s_bob_m, s_bob_x = _MSG2_REGISTER_INNER.unpack(
s_bob_m, s_bob_x = REG_MSG2_INNER.unpack(
decrypt(inner, key=kdf("reg/msg2/inner", ee))
)
# We challenge his static: only the real Bob can solve this
@@ -192,7 +169,7 @@ class Registration:
return self._chain.encrypt(
es_ct
+ encrypt(
_MSG3_REGISTER_INNER.pack(self._static_m, self._static_x, self._reveal),
REG_MSG3_INNER.pack(self._static_m, self._static_x, self._reveal),
key=kdf("reg/msg3/inner", es),
),
key=kdf("reg/msg3", self._shared, ee),
@@ -203,7 +180,7 @@ class Registration:
# caller gates the answer (msg5) on the out-of-band comparison.
plain = self._chain.decrypt(msg4, key=kdf("reg/msg4", self._shared, self._ee))
try:
se_ct, tb = _MSG4_REGISTER.unpack(plain)
se_ct, tb = REG_MSG4.unpack(plain)
# Solving se proves to Bob that we hold Alice's static
se = self._mlkem.decapsulate(se_ct)
se += self._x25519.exchange(X25519PublicKey.from_public_bytes(self._e_bob))
@@ -228,7 +205,7 @@ class Registration:
# connection key ships here once, Alice answers every leg with it
self._chain.advance(msg1)
try:
e_alice_m, e_alice_x, commitment = _MSG1_REGISTER.unpack(opened)
e_alice_m, e_alice_x, commitment = REG_MSG1.unpack(opened)
self._ephemeral = X25519PrivateKey.generate()
e_bob = self._ephemeral.public_key().public_bytes_raw()
# Forward secrecy, made to Alice's connection key only
@@ -237,11 +214,11 @@ class Registration:
except (struct.error, ValueError) as e:
raise _Burn(f"bad msg1 contents ({e})") from None
self._commitment = (e_alice_m, e_alice_x, commitment)
body = _MSG2_REGISTER.pack(
body = REG_MSG2.pack(
e_bob,
ee_ct,
encrypt(
_MSG2_REGISTER_INNER.pack(self._static_m, self._static_x),
REG_MSG2_INNER.pack(self._static_m, self._static_x),
key=kdf("reg/msg2/inner", ee),
),
)
@@ -253,13 +230,13 @@ class Registration:
# caller gates the answer (msg4) on the out-of-band comparison.
plain = self._chain.decrypt(msg3, key=kdf("reg/msg3", self._shared, self._ee))
try:
es_ct, inner = _MSG3_REGISTER.unpack(plain)
es_ct, inner = REG_MSG3.unpack(plain)
e_alice_m, e_alice_x, commitment = self._commitment
# Solving es proves to Alice that the static key they encapsulated to was ours
es = self._mlkem.decapsulate(es_ct)
es += self._x25519.exchange(X25519PublicKey.from_public_bytes(e_alice_x))
# This enables us to get Alice's static keys and reveal
s_alice_m, s_alice_x, reveal = _MSG3_REGISTER_INNER.unpack(
s_alice_m, s_alice_x, reveal = REG_MSG3_INNER.unpack(
decrypt(inner, key=kdf("reg/msg3/inner", es))
)
# We send Alice the same test, verifying they have s_alice secrets
@@ -279,7 +256,7 @@ class Registration:
self._se = se
sas = kdf("sas", self._shared, self._ee, reveal, length=4)
msg4 = self._chain.encrypt(
_MSG4_REGISTER.pack(se_ct, kdf("reg/confirm-b", self._root())),
REG_MSG4.pack(se_ct, kdf("reg/confirm-b", self._root())),
key=kdf("reg/msg4", self._shared, self._ee),
)
return msg4, sas
@@ -289,10 +266,9 @@ class Registration:
plain = self._chain.decrypt(
msg5, key=kdf("reg/msg5", self._shared, self._ee, self._es)
)
try:
(ta,) = _TAG.unpack(plain)
except struct.error as e:
raise _Burn(f"bad msg5 contents ({e})") from None
if len(plain) != 32:
raise _Burn("bad msg5 contents (expected a 32-byte tag)")
ta = plain
_check(
hmac.compare_digest(ta, kdf("reg/confirm-a", self._root())),
"initiator tag mismatch",
@@ -331,14 +307,14 @@ class Resumption:
sb, sb_ct = MLKEM768PublicKey.from_public_bytes(peer_m).encapsulate()
sb += self._ephemeral.exchange(X25519PublicKey.from_public_bytes(peer_x))
self._sb = sb
body = _MSG1_RESUME.pack(e_alice_m, e_alice_x, sb_ct)
body = RES_MSG1.pack(e_alice_m, e_alice_x, sb_ct)
return self._chain.encrypt(body, key=kdf("res/msg1", self._shared))
def respondA(self, msg2: bytes) -> bytes:
# Open Bob's reply and solve his three challenges
plain = self._chain.decrypt(msg2, key=kdf("res/msg2", self._shared, self._sb))
try:
e_bob, sa_ct, ee_ct, tb = _MSG2_RESUME.unpack(plain)
e_bob, sa_ct, ee_ct, tb = RES_MSG2.unpack(plain)
# Solving sa proves to Bob that we hold Alice's static
sa = self._mlkem.decapsulate(sa_ct)
sa += self._x25519.exchange(X25519PublicKey.from_public_bytes(e_bob))
@@ -364,7 +340,7 @@ class Resumption:
# Prove our static, challenge Alice's, forward secrecy
self._chain.advance(msg1)
try:
e_alice_m, e_alice_x, sb_ct = _MSG1_RESUME.unpack(opened)
e_alice_m, e_alice_x, sb_ct = RES_MSG1.unpack(opened)
# Solving sb proves to Alice that we hold Bob's static
sb = self._mlkem.decapsulate(sb_ct)
sb += self._x25519.exchange(X25519PublicKey.from_public_bytes(e_alice_x))
@@ -382,18 +358,14 @@ class Resumption:
self._sb = sb
self._sa = sa
self._ee = ee
body = _MSG2_RESUME.pack(e_bob, sa_ct, ee_ct, kdf("res/tag-b", sb, sa, ee))
body = RES_MSG2.pack(e_bob, sa_ct, ee_ct, kdf("res/tag-b", sb, sa, ee))
return self._chain.encrypt(body, key=kdf("res/msg2", self._shared, sb))
def completeB(self, msg3: bytes) -> None:
# Alice answers with just her tag
plain = self._chain.decrypt(
ta = self._chain.decrypt(
msg3, key=kdf("res/msg3", self._shared, self._sb, self._sa, self._ee)
)
try:
(ta,) = _TAG.unpack(plain)
except struct.error as e:
raise _Burn(f"bad msg3 contents ({e})") from None
_check(
hmac.compare_digest(ta, kdf("res/tag-a", self._sb, self._sa, self._ee)),
"initiator tag mismatch",