@@ -16,7 +16,7 @@ use routing::gossip::NetworkUpdate;
1616use  routing:: router:: RouteHop ; 
1717use  util:: chacha20:: { ChaCha20 ,  ChaChaReader } ; 
1818use  util:: errors:: { self ,  APIError } ; 
19- use  util:: ser:: { Readable ,  Writeable ,  LengthCalculatingWriter } ; 
19+ use  util:: ser:: { Readable ,  ReadableArgs ,   Writeable ,  LengthCalculatingWriter } ; 
2020use  util:: logger:: Logger ; 
2121
2222use  bitcoin:: hashes:: { Hash ,  HashEngine } ; 
@@ -636,7 +636,37 @@ pub(super) fn process_onion_failure<T: secp256k1::Signing, L: Deref>(secp_ctx: &
636636	}  else  {  unreachable ! ( ) ;  } 
637637} 
638638
639- /// Data decrypted from the onion payload. 
639+ /// Used in the decoding of inbound payments' and onion messages' routing packets. This enum allows 
640+ /// us to use `decode_next_hop` to return the payloads and next hop packet bytes of both payments 
641+ /// and onion messages. 
642+ enum  Payload  { 
643+ 	/// This payload was for an incoming payment. 
644+  	Payment ( msgs:: OnionHopData ) , 
645+ 	/// This payload was for an incoming onion message. 
646+  	Message ( onion_message:: Payload ) , 
647+ } 
648+ 
649+ enum  NextPacketBytes  { 
650+ 	Payment ( [ u8 ;  20 * 65 ] ) , 
651+ 	Message ( Vec < u8 > ) , 
652+ } 
653+ 
654+ /// Data decrypted from an onion message's onion payload. 
655+ pub ( crate )  enum  MessageHop  { 
656+ 	/// This onion payload was for us, not for forwarding to a next-hop. 
657+  	Receive ( onion_message:: Payload ) , 
658+ 	/// This onion payload needs to be forwarded to a next-hop. 
659+  	Forward  { 
660+ 		/// Onion payload data used in forwarding the onion message. 
661+  		next_hop_data :  onion_message:: Payload , 
662+ 		/// HMAC of the next hop's onion packet. 
663+  		next_hop_hmac :  [ u8 ;  32 ] , 
664+ 		/// Bytes of the onion packet we're forwarding. 
665+  		new_packet_bytes :  Vec < u8 > , 
666+ 	} , 
667+ } 
668+ 
669+ /// Data decrypted from a payment's onion payload. 
640670pub ( crate )  enum  Hop  { 
641671	/// This onion payload was for us, not for forwarding to a next-hop. Contains information for 
642672 	/// verifying the incoming payment. 
@@ -653,6 +683,7 @@ pub(crate) enum Hop {
653683} 
654684
655685/// Error returned when we fail to decode the onion packet. 
686+ #[ derive( Debug ) ]  
656687pub ( crate )  enum  OnionDecodeErr  { 
657688	/// The HMAC of the onion packet did not match the hop data. 
658689 	Malformed  { 
@@ -666,11 +697,44 @@ pub(crate) enum OnionDecodeErr {
666697	} , 
667698} 
668699
669- pub ( crate )  fn  decode_next_hop ( shared_secret :  [ u8 ;  32 ] ,  hop_data :  & [ u8 ] ,  hmac_bytes :  [ u8 ;  32 ] ,  payment_hash :  PaymentHash )  -> Result < Hop ,  OnionDecodeErr >  { 
700+ pub ( crate )  fn  decode_next_message_hop ( shared_secret :  [ u8 ;  32 ] ,  hop_data :  & [ u8 ] ,  hmac_bytes :  [ u8 ;  32 ] ,  encrypted_tlvs_ss :  SharedSecret )  -> Result < MessageHop ,  OnionDecodeErr >  { 
701+ 	match  decode_next_hop ( shared_secret,  hop_data,  hmac_bytes,  None ,  Some ( encrypted_tlvs_ss) )  { 
702+ 		Ok ( ( Payload :: Message ( next_hop_data) ,  None ) )  => Ok ( MessageHop :: Receive ( next_hop_data) ) , 
703+ 		Ok ( ( Payload :: Message ( next_hop_data) ,  Some ( ( next_hop_hmac,  NextPacketBytes :: Message ( new_packet_bytes) ) ) ) )  => { 
704+ 			Ok ( MessageHop :: Forward  { 
705+ 				next_hop_data, 
706+ 				next_hop_hmac, 
707+ 				new_packet_bytes
708+ 			} ) 
709+ 		} , 
710+ 		Err ( e)  => Err ( e) , 
711+ 		_ => unreachable ! ( ) 
712+ 	} 
713+ } 
714+ 
715+ pub ( crate )  fn  decode_next_payment_hop ( shared_secret :  [ u8 ;  32 ] ,  hop_data :  & [ u8 ] ,  hmac_bytes :  [ u8 ;  32 ] ,  payment_hash :  PaymentHash )  -> Result < Hop ,  OnionDecodeErr >  { 
716+ 	match  decode_next_hop ( shared_secret,  hop_data,  hmac_bytes,  Some ( payment_hash) ,  None )  { 
717+ 		Ok ( ( Payload :: Payment ( next_hop_data) ,  None ) )  => Ok ( Hop :: Receive ( next_hop_data) ) , 
718+ 		Ok ( ( Payload :: Payment ( next_hop_data) ,  Some ( ( next_hop_hmac,  NextPacketBytes :: Payment ( new_packet_bytes) ) ) ) )  => { 
719+ 			Ok ( Hop :: Forward  { 
720+ 				next_hop_data, 
721+ 				next_hop_hmac, 
722+ 				new_packet_bytes
723+ 			} ) 
724+ 		} , 
725+ 		Err ( e)  => Err ( e) , 
726+ 		_ => unreachable ! ( ) 
727+ 	} 
728+ } 
729+ 
730+ fn  decode_next_hop ( shared_secret :  [ u8 ;  32 ] ,  hop_data :  & [ u8 ] ,  hmac_bytes :  [ u8 ;  32 ] ,  payment_hash :  Option < PaymentHash > ,  encrypted_tlv_ss :  Option < SharedSecret > )  -> Result < ( Payload ,  Option < ( [ u8 ;  32 ] ,  NextPacketBytes ) > ) ,  OnionDecodeErr >  { 
731+ 	assert ! ( payment_hash. is_some( )  && encrypted_tlv_ss. is_none( )  || payment_hash. is_none( )  && encrypted_tlv_ss. is_some( ) ) ; 
670732	let  ( rho,  mu)  = gen_rho_mu_from_shared_secret ( & shared_secret) ; 
671733	let  mut  hmac = HmacEngine :: < Sha256 > :: new ( & mu) ; 
672734	hmac. input ( hop_data) ; 
673- 	hmac. input ( & payment_hash. 0 [ ..] ) ; 
735+ 	if  let  Some ( payment_hash)  = payment_hash { 
736+ 		hmac. input ( & payment_hash. 0 [ ..] ) ; 
737+ 	} 
674738	if  !fixed_time_eq ( & Hmac :: from_engine ( hmac) . into_inner ( ) ,  & hmac_bytes)  { 
675739		return  Err ( OnionDecodeErr :: Malformed  { 
676740			err_msg :  "HMAC Check failed" , 
@@ -680,7 +744,20 @@ pub(crate) fn decode_next_hop(shared_secret: [u8; 32], hop_data: &[u8], hmac_byt
680744
681745	let  mut  chacha = ChaCha20 :: new ( & rho,  & [ 0u8 ;  8 ] ) ; 
682746	let  mut  chacha_stream = ChaChaReader  {  chacha :  & mut  chacha,  read :  Cursor :: new ( & hop_data[ ..] )  } ; 
683- 	match  <msgs:: OnionHopData  as  Readable >:: read ( & mut  chacha_stream)  { 
747+ 	let  payload_read_res = if  payment_hash. is_some ( )  { 
748+ 		match  <msgs:: OnionHopData  as  Readable >:: read ( & mut  chacha_stream)  { 
749+ 			Ok ( payload)  => Ok ( Payload :: Payment ( payload) ) , 
750+ 			Err ( e)  => Err ( e) 
751+ 		} 
752+ 	}  else  if  encrypted_tlv_ss. is_some ( )  { 
753+ 		match  <onion_message:: Payload  as  ReadableArgs < SharedSecret > >:: read ( & mut  chacha_stream,  encrypted_tlv_ss. unwrap ( ) )  { 
754+ 			Ok ( payload)  => Ok ( Payload :: Message ( payload) ) , 
755+ 			Err ( e)  => Err ( e) 
756+ 		} 
757+ 	}  else  { 
758+ 		unreachable ! ( ) ;  // We already asserted that one of these is `Some` 
759+ 	} ; 
760+ 	match  payload_read_res { 
684761		Err ( err)  => { 
685762			let  error_code = match  err { 
686763				msgs:: DecodeError :: UnknownVersion  => 0x4000  | 1 ,  // unknown realm byte 
@@ -718,10 +795,17 @@ pub(crate) fn decode_next_hop(shared_secret: [u8; 32], hop_data: &[u8], hmac_byt
718795					chacha_stream. read_exact ( & mut  next_bytes) . unwrap ( ) ; 
719796					assert_ne ! ( next_bytes[ ..] ,  [ 0 ;  32 ] [ ..] ) ; 
720797				} 
721- 				return  Ok ( Hop :: Receive ( msg) ) ; 
798+ 				return  Ok ( ( msg,   None ) ) ; 
722799			}  else  { 
723- 				let  mut  new_packet_bytes = [ 0 ;  20 * 65 ] ; 
724- 				let  read_pos = chacha_stream. read ( & mut  new_packet_bytes) . unwrap ( ) ; 
800+ 				let  ( mut  new_packet_bytes,  read_pos)  = if  payment_hash. is_some ( )  { 
801+ 					let  mut  new_packet_bytes = [ 0  as  u8 ;  20 * 65 ] ; 
802+ 					let  read_pos = chacha_stream. read ( & mut  new_packet_bytes) . unwrap ( ) ; 
803+ 					( NextPacketBytes :: Payment ( new_packet_bytes) ,  read_pos) 
804+ 				}  else  { 
805+ 					let  mut  new_packet_bytes = vec ! [ 0  as  u8 ;  hop_data. len( ) ] ; 
806+ 					let  read_pos = chacha_stream. read ( & mut  new_packet_bytes) . unwrap ( ) ; 
807+ 					( NextPacketBytes :: Message ( new_packet_bytes) ,  read_pos) 
808+ 				} ; 
725809				#[ cfg( debug_assertions) ]  
726810				{ 
727811					// Check two things: 
@@ -733,12 +817,11 @@ pub(crate) fn decode_next_hop(shared_secret: [u8; 32], hop_data: &[u8], hmac_byt
733817				} 
734818				// Once we've emptied the set of bytes our peer gave us, encrypt 0 bytes until we 
735819				// fill the onion hop data we'll forward to our next-hop peer. 
736- 				chacha_stream. chacha . process_in_place ( & mut  new_packet_bytes[ read_pos..] ) ; 
737- 				return  Ok ( Hop :: Forward  { 
738- 					next_hop_data :  msg, 
739- 					next_hop_hmac :  hmac, 
740- 					new_packet_bytes, 
741- 				} ) 
820+ 				match  new_packet_bytes { 
821+ 					NextPacketBytes :: Payment ( ref  mut  bytes)  => chacha_stream. chacha . process_in_place ( & mut  bytes[ read_pos..] ) , 
822+ 					NextPacketBytes :: Message ( ref  mut  bytes)  => chacha_stream. chacha . process_in_place ( & mut  bytes[ read_pos..] ) , 
823+ 				} 
824+ 				return  Ok ( ( msg,  Some ( ( hmac,  new_packet_bytes) ) ) ) 
742825			} 
743826		} , 
744827	} 
0 commit comments