fix: use KeyPackageRef inside fresh_key_package

fstar/api/MLS.fst

@@ -305,7 +305,7 @@ let fresh_key_package_internal e { identity; signature_key } private_sign_key =
   let kp_tbs = ({
     version = PV_mls10;
     cipher_suite = CS_mls_128_dhkemx25519_chacha20poly1305_sha256_ed25519;
-    init_key = encryption_key;
+    init_key = encryption_key; //TODO: use a different key
     leaf_node;
     extensions = [];
   } <: key_package_tbs_nt bytes tkt) in
@@ -321,8 +321,8 @@ let fresh_key_package_internal e { identity; signature_key } private_sign_key =
 let fresh_key_package e cred private_sign_key =
   let? key_package, leaf_secret = fresh_key_package_internal e cred private_sign_key in
   let key_package_bytes = (ps_prefix_to_ps_whole (ps_key_package_nt _)).serialize key_package in
-  let? hash = hash_leaf_package key_package.tbs.leaf_node in
-  return (key_package_bytes, hash, (leaf_secret <: bytes))
+  let? hash = compute_key_package_ref key_package in
+  return (key_package_bytes, (hash <: bytes), (leaf_secret <: bytes))
 
 let current_epoch s = s.epoch
 
@@ -611,7 +611,7 @@ let find_my_index #l t sign_pk =
 let process_welcome_message w (sign_pk, sign_sk) lookup =
   let (_, welcome_bytes) = w in
   let? welcome = from_option "process_welcome_message: can't parse welcome message" ((ps_prefix_to_ps_whole ps_welcome_nt).parse welcome_bytes) in
-  let? (group_info, secrets) = decrypt_welcome welcome (fun kp_hash ->
+  let? (group_info, secrets, leaf_decryption_key) = decrypt_welcome welcome (fun kp_hash ->
     match lookup kp_hash with
     | Some leaf_secret -> (
       if length leaf_secret = hpke_private_key_length #bytes then Some leaf_secret
@@ -646,17 +646,6 @@ let process_welcome_message w (sign_pk, sign_sk) lookup =
     return ()
   ) in
   let? leaf_index = find_my_index treesync sign_pk in
-  let? leaf_decryption_key = (
-    let opt_my_leaf_package = leaf_at treesync leaf_index in
-    match opt_my_leaf_package with
-    | None -> internal_failure "process_welcome_message: leaf index points to a blank leaf"
-    | Some my_leaf_package -> (
-      let? kp_hash = hash_leaf_package my_leaf_package in
-      match lookup kp_hash with
-      | Some leaf_secret -> mk_hpke_private_key leaf_secret "process_welcome_message" "leaf_decryption_key"
-      | None -> internal_failure "process_welcome_message: decrypt_welcome found my leaf package but not proccess_welcome_message"
-    )
-  ) in
   let opt_path_secret_and_inviter_ind: option (bytes & nat) = match secrets.path_secret with | None -> None | Some {path_secret} -> Some (path_secret, group_info.tbs.signer) in
   let? treekem = treesync_to_treekem treesync in
   assume(leaf_index < pow2 l /\ Some? (leaf_at treekem leaf_index));

fstar/test/MLS.Test.FromExt.Welcome.fst

@@ -47,7 +47,7 @@ let test_welcome_one t =
 
     let kp_ref = extract_result (make_keypackage_ref #bytes (serialize _ key_package)) in
 
-    let (group_info, group_secrets) = extract_result (decrypt_welcome welcome (fun ref -> if ref = kp_ref then Some init_priv else None) None) in
+    let (group_info, group_secrets, my_init_decryption_key) = extract_result (decrypt_welcome welcome (fun ref -> if ref = kp_ref then Some init_priv else None) None) in
 
     if not (extract_result (verify_welcome_group_info (fun _ -> return signer_pub) group_info)) then (
       failwith "test_welcome_one: bad signature"

fstar/treedem/MLS.TreeDEM.Welcome.fst

@@ -50,10 +50,10 @@ let rec find_my_encrypted_group_secret #bytes #cb kp_ref_to_hpke_sk l =
 val decrypt_welcome:
   #bytes:Type0 -> {|crypto_bytes bytes|} ->
   welcome_nt bytes -> (bytes -> option (hpke_private_key bytes)) -> option (l:nat & treesync bytes tkt l 0) ->
-  result (group_info_nt bytes & group_secrets_nt bytes)
+  result (group_info_nt bytes & group_secrets_nt bytes & hpke_private_key bytes)
 let decrypt_welcome #bytes #cb w kp_ref_to_hpke_sk opt_tree =
+  let? (my_hpke_sk, my_hpke_ciphertext) = from_option "decrypt_welcome: can't find my encrypted secret" (find_my_encrypted_group_secret kp_ref_to_hpke_sk w.secrets) in
   let? group_secrets = (
-    let? (my_hpke_sk, my_hpke_ciphertext) = from_option "decrypt_welcome: can't find my encrypted secret" (find_my_encrypted_group_secret kp_ref_to_hpke_sk w.secrets) in
     let? kem_output = mk_hpke_kem_output my_hpke_ciphertext.kem_output "decrypt_welcome" "kem_output" in
     let? group_secrets_bytes = decrypt_with_label my_hpke_sk "Welcome" w.encrypted_group_info kem_output my_hpke_ciphertext.ciphertext in
     let? group_secrets_network = from_option "decrypt_welcome: malformed group secrets" (parse (group_secrets_nt bytes) group_secrets_bytes) in
@@ -68,7 +68,7 @@ let decrypt_welcome #bytes #cb w kp_ref_to_hpke_sk opt_tree =
     return (group_info_network)
   ) in
   //TODO: integrity check, this is where `opt_tree` will be useful
-  return (group_info, group_secrets)
+  return (group_info, group_secrets, my_hpke_sk)
 
 (*** Encrypting a welcome ***)