@@ -25,42 +25,38 @@ use crate::util::ser::{Readable, Writeable};
2525
2626use crate :: io;
2727
28+ use core:: borrow:: Borrow ;
29+
2830#[ allow( unused_imports) ]
2931use crate :: prelude:: * ;
3032
3133// TODO: DRY with onion_utils::construct_onion_keys_callback
32- #[ inline]
33- pub ( crate ) fn construct_keys_callback < ' a , T , I , F > (
34- secp_ctx : & Secp256k1 < T > , unblinded_path : I , destination : Option < Destination > ,
35- session_priv : & SecretKey , mut callback : F
36- ) -> Result < ( ) , secp256k1:: Error >
37- where
38- T : secp256k1:: Signing + secp256k1:: Verification ,
39- I : Iterator < Item =& ' a PublicKey > ,
40- F : FnMut ( PublicKey , SharedSecret , PublicKey , [ u8 ; 32 ] , Option < PublicKey > , Option < Vec < u8 > > ) ,
34+ macro_rules! build_keys_helper {
35+ ( $session_priv: ident, $secp_ctx: ident, $callback: ident) =>
4136{
42- let mut msg_blinding_point_priv = session_priv. clone ( ) ;
43- let mut msg_blinding_point = PublicKey :: from_secret_key ( secp_ctx, & msg_blinding_point_priv) ;
37+ let mut msg_blinding_point_priv = $ session_priv. clone( ) ;
38+ let mut msg_blinding_point = PublicKey :: from_secret_key( $ secp_ctx, & msg_blinding_point_priv) ;
4439 let mut onion_packet_pubkey_priv = msg_blinding_point_priv. clone( ) ;
4540 let mut onion_packet_pubkey = msg_blinding_point. clone( ) ;
4641
4742 macro_rules! build_keys {
48- ( $pk: expr, $blinded: expr, $encrypted_payload: expr) => { {
49- let encrypted_data_ss = SharedSecret :: new( & $pk, & msg_blinding_point_priv) ;
43+ ( $hop: expr, $blinded: expr, $encrypted_payload: expr) => { {
44+ let pk = * $hop. borrow( ) ;
45+ let encrypted_data_ss = SharedSecret :: new( & pk, & msg_blinding_point_priv) ;
5046
51- let blinded_hop_pk = if $blinded { $ pk } else {
47+ let blinded_hop_pk = if $blinded { pk } else {
5248 let hop_pk_blinding_factor = {
5349 let mut hmac = HmacEngine :: <Sha256 >:: new( b"blinded_node_id" ) ;
5450 hmac. input( encrypted_data_ss. as_ref( ) ) ;
5551 Hmac :: from_engine( hmac) . to_byte_array( )
5652 } ;
57- $ pk. mul_tweak( secp_ctx, & Scalar :: from_be_bytes( hop_pk_blinding_factor) . unwrap( ) ) ?
53+ pk. mul_tweak( $ secp_ctx, & Scalar :: from_be_bytes( hop_pk_blinding_factor) . unwrap( ) ) ?
5854 } ;
5955 let onion_packet_ss = SharedSecret :: new( & blinded_hop_pk, & onion_packet_pubkey_priv) ;
6056
6157 let rho = onion_utils:: gen_rho_from_shared_secret( encrypted_data_ss. as_ref( ) ) ;
62- let unblinded_pk_opt = if $blinded { None } else { Some ( $pk ) } ;
63- callback( blinded_hop_pk, onion_packet_ss, onion_packet_pubkey, rho, unblinded_pk_opt , $encrypted_payload) ;
58+ let unblinded_hop_opt = if $blinded { None } else { Some ( $hop ) } ;
59+ $ callback( blinded_hop_pk, onion_packet_ss, onion_packet_pubkey, rho, unblinded_hop_opt , $encrypted_payload) ;
6460 ( encrypted_data_ss, onion_packet_ss)
6561 } }
6662 }
7773 } ;
7874
7975 msg_blinding_point_priv = msg_blinding_point_priv. mul_tweak( & Scalar :: from_be_bytes( msg_blinding_point_blinding_factor) . unwrap( ) ) ?;
80- msg_blinding_point = PublicKey :: from_secret_key( secp_ctx, & msg_blinding_point_priv) ;
76+ msg_blinding_point = PublicKey :: from_secret_key( $ secp_ctx, & msg_blinding_point_priv) ;
8177
8278 let onion_packet_pubkey_blinding_factor = {
8379 let mut sha = Sha256 :: engine( ) ;
@@ -86,45 +82,83 @@ where
8682 Sha256 :: from_engine( sha) . to_byte_array( )
8783 } ;
8884 onion_packet_pubkey_priv = onion_packet_pubkey_priv. mul_tweak( & Scalar :: from_be_bytes( onion_packet_pubkey_blinding_factor) . unwrap( ) ) ?;
89- onion_packet_pubkey = PublicKey :: from_secret_key( secp_ctx, & onion_packet_pubkey_priv) ;
85+ onion_packet_pubkey = PublicKey :: from_secret_key( $ secp_ctx, & onion_packet_pubkey_priv) ;
9086 } ;
9187 }
88+ } }
89+
90+ #[ inline]
91+ pub ( crate ) fn construct_keys_for_onion_message < ' a , T , I , F > (
92+ secp_ctx : & Secp256k1 < T > , unblinded_path : I , destination : Destination , session_priv : & SecretKey ,
93+ mut callback : F ,
94+ ) -> Result < ( ) , secp256k1:: Error >
95+ where
96+ T : secp256k1:: Signing + secp256k1:: Verification ,
97+ I : Iterator < Item =PublicKey > ,
98+ F : FnMut ( PublicKey , SharedSecret , PublicKey , [ u8 ; 32 ] , Option < PublicKey > , Option < Vec < u8 > > ) ,
99+ {
100+ build_keys_helper ! ( session_priv, secp_ctx, callback) ;
92101
93102 for pk in unblinded_path {
94- build_keys_in_loop ! ( * pk, false , None ) ;
103+ build_keys_in_loop ! ( pk, false , None ) ;
95104 }
96- if let Some ( dest) = destination {
97- match dest {
98- Destination :: Node ( pk) => {
99- build_keys ! ( pk, false , None ) ;
100- } ,
101- Destination :: BlindedPath ( BlindedMessagePath ( BlindedPath { blinded_hops, .. } ) ) => {
102- for hop in blinded_hops {
103- build_keys_in_loop ! ( hop. blinded_node_id, true , Some ( hop. encrypted_payload) ) ;
104- }
105- } ,
106- }
105+ match destination {
106+ Destination :: Node ( pk) => {
107+ build_keys ! ( pk, false , None ) ;
108+ } ,
109+ Destination :: BlindedPath ( BlindedMessagePath ( BlindedPath { blinded_hops, .. } ) ) => {
110+ for hop in blinded_hops {
111+ build_keys_in_loop ! ( hop. blinded_node_id, true , Some ( hop. encrypted_payload) ) ;
112+ }
113+ } ,
114+ }
115+ Ok ( ( ) )
116+ }
117+
118+ #[ inline]
119+ pub ( super ) fn construct_keys_for_blinded_path < ' a , T , I , F , H > (
120+ secp_ctx : & Secp256k1 < T > , unblinded_path : I , session_priv : & SecretKey , mut callback : F ,
121+ ) -> Result < ( ) , secp256k1:: Error >
122+ where
123+ T : secp256k1:: Signing + secp256k1:: Verification ,
124+ H : Borrow < PublicKey > ,
125+ I : Iterator < Item =H > ,
126+ F : FnMut ( PublicKey , SharedSecret , PublicKey , [ u8 ; 32 ] , Option < H > , Option < Vec < u8 > > ) ,
127+ {
128+ build_keys_helper ! ( session_priv, secp_ctx, callback) ;
129+
130+ for pk in unblinded_path {
131+ build_keys_in_loop ! ( pk, false , None ) ;
107132 }
108133 Ok ( ( ) )
109134}
110135
111- // Panics if `unblinded_tlvs` length is less than `unblinded_pks` length
112- pub ( crate ) fn construct_blinded_hops < ' a , T , I1 , I2 > (
113- secp_ctx : & Secp256k1 < T > , unblinded_pks : I1 , mut unblinded_tlvs : I2 , session_priv : & SecretKey
136+ struct PublicKeyWithTlvs < W : Writeable > {
137+ pubkey : PublicKey ,
138+ tlvs : W ,
139+ }
140+
141+ impl < W : Writeable > Borrow < PublicKey > for PublicKeyWithTlvs < W > {
142+ fn borrow ( & self ) -> & PublicKey {
143+ & self . pubkey
144+ }
145+ }
146+
147+ pub ( crate ) fn construct_blinded_hops < ' a , T , I , W > (
148+ secp_ctx : & Secp256k1 < T > , unblinded_path : I , session_priv : & SecretKey ,
114149) -> Result < Vec < BlindedHop > , secp256k1:: Error >
115150where
116151 T : secp256k1:: Signing + secp256k1:: Verification ,
117- I1 : Iterator < Item =& ' a PublicKey > ,
118- I2 : Iterator ,
119- I2 :: Item : Writeable
152+ I : Iterator < Item =( PublicKey , W ) > ,
153+ W : Writeable
120154{
121- let mut blinded_hops = Vec :: with_capacity ( unblinded_pks . size_hint ( ) . 0 ) ;
122- construct_keys_callback (
123- secp_ctx, unblinded_pks , None , session_priv,
124- |blinded_node_id, _, _, encrypted_payload_rho, _ , _| {
155+ let mut blinded_hops = Vec :: with_capacity ( unblinded_path . size_hint ( ) . 0 ) ;
156+ construct_keys_for_blinded_path (
157+ secp_ctx, unblinded_path . map ( | ( pubkey , tlvs ) | PublicKeyWithTlvs { pubkey , tlvs } ) , session_priv,
158+ |blinded_node_id, _, _, encrypted_payload_rho, unblinded_hop_data , _| {
125159 blinded_hops. push ( BlindedHop {
126160 blinded_node_id,
127- encrypted_payload : encrypt_payload ( unblinded_tlvs . next ( ) . unwrap ( ) , encrypted_payload_rho) ,
161+ encrypted_payload : encrypt_payload ( unblinded_hop_data . unwrap ( ) . tlvs , encrypted_payload_rho) ,
128162 } ) ;
129163 } ) ?;
130164 Ok ( blinded_hops)
0 commit comments