refactor: more general key lookup for process_welcome

fstar/api/MLS.fst

@@ -611,14 +611,13 @@ 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, 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
-      else None
-    )
-    | None -> None
-  ) None in
+  let extract_hpke_sk (x:bytes): result (hpke_private_key bytes) =
+    if not (length x = hpke_private_key_length #bytes) then
+      internal_failure "process_welcome_message: bad HPKE private key length"
+    else
+      return x
+  in
+  let? (group_info, secrets, (_, leaf_decryption_key)) = decrypt_welcome welcome lookup extract_hpke_sk in
   let group_id = group_info.tbs.group_context.group_id in
   let? ratchet_tree = from_option "bad ratchet tree" ((ps_prefix_to_ps_whole #bytes (ps_ratchet_tree_nt tkt)).parse group_info.tbs.extensions) in
   let? treesync_state = (

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, my_init_decryption_key) = 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) (fun x -> return x)) 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

@@ -7,7 +7,6 @@ open MLS.TreeKEM.NetworkTypes
 open MLS.TreeDEM.NetworkTypes
 open MLS.Crypto
 open MLS.Tree
-open MLS.TreeSync.Types
 open MLS.TreeDEM.KeyPackageRef
 open MLS.TreeKEM.KeySchedule
 open MLS.Result
@@ -34,25 +33,26 @@ let welcome_secret_to_nonce #bytes #cb welcome_secret =
 
 #push-options "--ifuel 1"
 val find_my_encrypted_group_secret:
-  #bytes:Type0 -> {|crypto_bytes bytes|} ->
-  (bytes -> option (hpke_private_key bytes)) -> list (encrypted_group_secrets_nt bytes) ->
-  option (hpke_private_key bytes & hpke_ciphertext_nt bytes)
-let rec find_my_encrypted_group_secret #bytes #cb kp_ref_to_hpke_sk l =
+  #bytes:Type0 -> {|crypto_bytes bytes|} -> #a:Type -> 
+  (bytes -> option a) -> list (encrypted_group_secrets_nt bytes) ->
+  option (bytes & a & hpke_ciphertext_nt bytes)
+let rec find_my_encrypted_group_secret #bytes #cb kp_ref_to_kp_secrets l =
   match l with
   | [] -> None
   | h::t -> (
-    match kp_ref_to_hpke_sk h.new_member with
-    | Some sk -> Some (sk, h.encrypted_group_secrets)
-    | None -> find_my_encrypted_group_secret kp_ref_to_hpke_sk t
+    match kp_ref_to_kp_secrets h.new_member with
+    | Some kp_secrets -> Some (h.new_member, kp_secrets, h.encrypted_group_secrets)
+    | None -> find_my_encrypted_group_secret kp_ref_to_kp_secrets t
   )
 #pop-options
 
 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 & 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
+  #bytes:Type0 -> {|crypto_bytes bytes|} -> #a:Type ->
+  welcome_nt bytes -> (bytes -> option a) -> (a -> result (hpke_private_key bytes)) ->
+  result (group_info_nt bytes & group_secrets_nt bytes & (bytes & a))
+let decrypt_welcome #bytes #cb w kp_ref_to_kp_secrets kp_secrets_to_hpke_sk =
+  let? (my_kp_ref, my_kp_secrets, my_hpke_ciphertext) = from_option "decrypt_welcome: can't find my encrypted secret" (find_my_encrypted_group_secret kp_ref_to_kp_secrets w.secrets) in
+  let? my_hpke_sk = kp_secrets_to_hpke_sk my_kp_secrets in
   let? group_secrets = (
     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
@@ -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, my_hpke_sk)
+  return (group_info, group_secrets, (my_kp_ref, my_kp_secrets))
 
 (*** Encrypting a welcome ***)