Skip to content

Commit d7c255c

Browse files
committed
Add accessors for channels and nodes
1 parent 635febb commit d7c255c

File tree

2 files changed

+64
-46
lines changed

2 files changed

+64
-46
lines changed

lightning/src/routing/network_graph.rs

Lines changed: 30 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ impl NetworkGraph {
6363
/// Get network addresses by node id
6464
pub fn get_addresses(&self, pubkey: &PublicKey) -> Option<Vec<NetAddress>> {
6565
let network = self.network_map.read().unwrap();
66-
network.nodes.get(pubkey).map(|n| n.addresses.clone())
66+
network.get_nodes().get(pubkey).map(|n| n.addresses.clone())
6767
}
6868

6969
fn remove_channel_in_nodes(nodes: &mut BTreeMap<PublicKey, NodeInfo>, chan: &ChannelInfo, short_channel_id: u64) {
@@ -348,10 +348,28 @@ impl Readable for NodeInfo {
348348
/// Represents the network as nodes and channels between them
349349
#[derive(PartialEq)]
350350
pub struct NetworkMap {
351-
/// Known valid channels
352-
pub channels: BTreeMap<u64, ChannelInfo>,
351+
channels: BTreeMap<u64, ChannelInfo>,
353352
/// A set of known Lightning nodes and their properties
354-
pub nodes: BTreeMap<PublicKey, NodeInfo>,
353+
nodes: BTreeMap<PublicKey, NodeInfo>,
354+
}
355+
356+
impl NetworkMap {
357+
/// Returns a list of known valid channels
358+
pub fn get_channels<'a>(&'a self) -> &'a BTreeMap<u64, ChannelInfo> { &self.channels }
359+
/// Returns a mutable reference to an existing channel
360+
pub fn get_channel_mut<'a>(&'a mut self, chan_id: &u64) -> Option<&mut ChannelInfo> { self.channels.get_mut(chan_id) }
361+
/// Store a valid channel
362+
pub fn add_channel(&mut self, chan_id: u64, channel_info: ChannelInfo) {
363+
self.channels.insert(chan_id, channel_info);
364+
}
365+
/// Returns a list of known nodes
366+
pub fn get_nodes<'a>(&'a self) -> &'a BTreeMap<PublicKey, NodeInfo> { &self.nodes }
367+
/// Returns a mutable reference to an existing node
368+
pub fn get_node_mut<'a>(&'a mut self, node_id: &PublicKey) -> Option<&mut NodeInfo> { self.nodes.get_mut(node_id) }
369+
/// Store a valid node
370+
pub fn add_node(&mut self, node_id: PublicKey, node_info: NodeInfo) {
371+
self.nodes.insert(node_id, node_info);
372+
}
355373
}
356374

357375
impl Writeable for NetworkMap {
@@ -428,7 +446,7 @@ impl RoutingMessageHandler for NetworkGraph {
428446
secp_verify_sig!(self.secp_ctx, &msg_hash, &msg.signature, &msg.contents.node_id);
429447

430448
let mut network = self.network_map.write().unwrap();
431-
match network.nodes.get_mut(&msg.contents.node_id) {
449+
match network.get_node_mut(&msg.contents.node_id) {
432450
None => Err(LightningError{err: "No existing channels for node_announcement", action: ErrorAction::IgnoreError}),
433451
Some(node) => {
434452
match node.last_update {
@@ -588,7 +606,7 @@ impl RoutingMessageHandler for NetworkGraph {
588606
Self::remove_channel_in_nodes(&mut network.nodes, &chan, *short_channel_id);
589607
}
590608
} else {
591-
if let Some(chan) = network.channels.get_mut(short_channel_id) {
609+
if let Some(chan) = network.get_channel_mut(short_channel_id) {
592610
chan.one_to_two.enabled = false;
593611
chan.two_to_one.enabled = false;
594612
}
@@ -610,7 +628,7 @@ impl RoutingMessageHandler for NetworkGraph {
610628
let chan_enabled = msg.contents.flags & (1 << 1) != (1 << 1);
611629
let chan_was_enabled;
612630

613-
match network.channels.get_mut(&NetworkMap::get_key(msg.contents.short_channel_id, msg.contents.chain_hash)) {
631+
match network.get_channel_mut(&NetworkMap::get_key(msg.contents.short_channel_id, msg.contents.chain_hash)) {
614632
None => return Err(LightningError{err: "Couldn't find channel for update", action: ErrorAction::IgnoreError}),
615633
Some(channel) => {
616634
macro_rules! maybe_update_channel_info {
@@ -646,7 +664,7 @@ impl RoutingMessageHandler for NetworkGraph {
646664
}
647665

648666
if chan_enabled {
649-
let node = network.nodes.get_mut(&dest_node_id).unwrap();
667+
let node = network.get_node_mut(&dest_node_id).unwrap();
650668
let mut base_msat = msg.contents.fee_base_msat;
651669
let mut proportional_millionths = msg.contents.fee_proportional_millionths;
652670
if let Some(fees) = node.lowest_inbound_channel_fees {
@@ -662,7 +680,7 @@ impl RoutingMessageHandler for NetworkGraph {
662680
let mut lowest_inbound_channel_fee_proportional_millionths = u32::max_value();
663681

664682
{
665-
let node = network.nodes.get(&dest_node_id).unwrap();
683+
let node = network.get_nodes().get(&dest_node_id).unwrap();
666684

667685
for chan_id in node.channels.iter() {
668686
let chan = network.channels.get(chan_id).unwrap();
@@ -677,7 +695,7 @@ impl RoutingMessageHandler for NetworkGraph {
677695
}
678696

679697
//TODO: satisfy the borrow-checker without a double-map-lookup :(
680-
let mut_node = network.nodes.get_mut(&dest_node_id).unwrap();
698+
let mut_node = network.get_node_mut(&dest_node_id).unwrap();
681699
if mut_node.channels.len() > 0 {
682700
mut_node.lowest_inbound_channel_fees = Some(RoutingFees {
683701
base_msat: lowest_inbound_channel_fee_base_msat,
@@ -1575,7 +1593,7 @@ mod tests {
15751593
base_msat: 100,
15761594
proportional_millionths: 0,
15771595
};
1578-
network.nodes.insert(node1.clone(), NodeInfo {
1596+
network.add_node(node1.clone(), NodeInfo {
15791597
channels: vec!(NetworkMap::get_key(1, zero_hash.clone()), NetworkMap::get_key(3, zero_hash.clone())),
15801598
lowest_inbound_channel_fees: Some(node_routing_fees),
15811599
features: NodeFeatures::supported(),
@@ -1585,7 +1603,7 @@ mod tests {
15851603
addresses: Vec::new(),
15861604
announcement_message: None,
15871605
});
1588-
network.channels.insert(NetworkMap::get_key(1, zero_hash.clone()), ChannelInfo {
1606+
network.add_channel(NetworkMap::get_key(1, zero_hash.clone()), ChannelInfo {
15891607
features: ChannelFeatures::supported(),
15901608
one_to_two: DirectionalChannelInfo {
15911609
src_node_id: our_id.clone(),

lightning/src/routing/router.rs

Lines changed: 34 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,7 @@ pub fn get_route(our_node_id: &PublicKey, network_graph: &NetworkGraph, target:
191191

192192
let network = network_graph.network_map.read().unwrap();
193193
let mut targets = BinaryHeap::new(); //TODO: Do we care about switching to eg Fibbonaci heap?
194-
let mut dist = HashMap::with_capacity(network.nodes.len());
194+
let mut dist = HashMap::with_capacity(network.get_nodes().len());
195195

196196
let mut first_hop_targets = HashMap::with_capacity(if first_hops.is_some() { first_hops.as_ref().unwrap().len() } else { 0 });
197197
if let Some(hops) = first_hops {
@@ -230,7 +230,7 @@ pub fn get_route(our_node_id: &PublicKey, network_graph: &NetworkGraph, target:
230230
let mut total_fee = $starting_fee_msat as u64;
231231
let hm_entry = dist.entry(&$directional_info.src_node_id);
232232
let old_entry = hm_entry.or_insert_with(|| {
233-
let node = network.nodes.get(&$directional_info.src_node_id).unwrap();
233+
let node = network.get_nodes().get(&$directional_info.src_node_id).unwrap();
234234
let mut fee_base_msat = u32::max_value();
235235
let mut fee_proportional_millionths = u32::max_value();
236236
if let Some(fees) = node.lowest_inbound_channel_fees {
@@ -292,7 +292,7 @@ pub fn get_route(our_node_id: &PublicKey, network_graph: &NetworkGraph, target:
292292

293293
if !$node.features.requires_unknown_bits() {
294294
for chan_id in $node.channels.iter() {
295-
let chan = network.channels.get(chan_id).unwrap();
295+
let chan = network.get_channels().get(chan_id).unwrap();
296296
if !chan.features.requires_unknown_bits() {
297297
if chan.one_to_two.src_node_id == *$node_id {
298298
// ie $node is one, ie next hop in A* is two, via the two_to_one channel
@@ -314,7 +314,7 @@ pub fn get_route(our_node_id: &PublicKey, network_graph: &NetworkGraph, target:
314314
};
315315
}
316316

317-
match network.nodes.get(target) {
317+
match network.get_nodes().get(target) {
318318
None => {},
319319
Some(node) => {
320320
add_entries_to_cheapest_to_target_node!(node, target, 0);
@@ -323,7 +323,7 @@ pub fn get_route(our_node_id: &PublicKey, network_graph: &NetworkGraph, target:
323323

324324
for hop in last_hops.iter() {
325325
if first_hops.is_none() || hop.src_node_id != *our_node_id { // first_hop overrules last_hops
326-
if network.nodes.get(&hop.src_node_id).is_some() {
326+
if network.get_nodes().get(&hop.src_node_id).is_some() {
327327
if first_hops.is_some() {
328328
if let Some(&(ref first_hop, ref features)) = first_hop_targets.get(&hop.src_node_id) {
329329
// Currently there are no channel-context features defined, so we are a
@@ -346,7 +346,7 @@ pub fn get_route(our_node_id: &PublicKey, network_graph: &NetworkGraph, target:
346346
loop {
347347
if let Some(&(_, ref features)) = first_hop_targets.get(&res.last().unwrap().pubkey) {
348348
res.last_mut().unwrap().node_features = NodeFeatures::with_known_relevant_init_flags(&features);
349-
} else if let Some(node) = network.nodes.get(&res.last().unwrap().pubkey) {
349+
} else if let Some(node) = network.get_nodes().get(&res.last().unwrap().pubkey) {
350350
res.last_mut().unwrap().node_features = node.features.clone();
351351
} else {
352352
// We should be able to fill in features for everything except the last
@@ -375,7 +375,7 @@ pub fn get_route(our_node_id: &PublicKey, network_graph: &NetworkGraph, target:
375375
return Ok(route);
376376
}
377377

378-
match network.nodes.get(&pubkey) {
378+
match network.get_nodes().get(&pubkey) {
379379
None => {},
380380
Some(node) => {
381381
add_entries_to_cheapest_to_target_node!(node, &pubkey, lowest_fee_to_node);
@@ -512,7 +512,7 @@ mod tests {
512512
base_msat: 100,
513513
proportional_millionths: 0
514514
};
515-
network.nodes.insert(node1.clone(), NodeInfo {
515+
network.add_node(node1.clone(), NodeInfo {
516516
channels: vec!(NetworkMap::get_key(1, zero_hash.clone()), NetworkMap::get_key(3, zero_hash.clone())),
517517
lowest_inbound_channel_fees: Some(lowest_routing_fees),
518518
features: NodeFeatures::from_le_bytes(id_to_feature_flags!(1)),
@@ -522,7 +522,7 @@ mod tests {
522522
addresses: Vec::new(),
523523
announcement_message: None,
524524
});
525-
network.channels.insert(NetworkMap::get_key(1, zero_hash.clone()), ChannelInfo {
525+
network.add_channel(NetworkMap::get_key(1, zero_hash.clone()), ChannelInfo {
526526
features: ChannelFeatures::from_le_bytes(id_to_feature_flags!(1)),
527527
one_to_two: DirectionalChannelInfo {
528528
src_node_id: our_id.clone(),
@@ -550,7 +550,7 @@ mod tests {
550550
base_msat: 0,
551551
proportional_millionths: 0
552552
};
553-
network.nodes.insert(node2.clone(), NodeInfo {
553+
network.add_node(node2.clone(), NodeInfo {
554554
channels: vec!(NetworkMap::get_key(2, zero_hash.clone()), NetworkMap::get_key(4, zero_hash.clone())),
555555
lowest_inbound_channel_fees: Some(routing_fees.clone()),
556556
features: NodeFeatures::from_le_bytes(id_to_feature_flags!(2)),
@@ -560,7 +560,7 @@ mod tests {
560560
addresses: Vec::new(),
561561
announcement_message: None,
562562
});
563-
network.channels.insert(NetworkMap::get_key(2, zero_hash.clone()), ChannelInfo {
563+
network.add_channel(NetworkMap::get_key(2, zero_hash.clone()), ChannelInfo {
564564
features: ChannelFeatures::from_le_bytes(id_to_feature_flags!(2)),
565565
one_to_two: DirectionalChannelInfo {
566566
src_node_id: our_id.clone(),
@@ -584,7 +584,7 @@ mod tests {
584584
},
585585
announcement_message: None,
586586
});
587-
network.nodes.insert(node8.clone(), NodeInfo {
587+
network.add_node(node8.clone(), NodeInfo {
588588
channels: vec!(NetworkMap::get_key(12, zero_hash.clone()), NetworkMap::get_key(13, zero_hash.clone())),
589589
lowest_inbound_channel_fees: Some(routing_fees.clone()),
590590
features: NodeFeatures::from_le_bytes(id_to_feature_flags!(8)),
@@ -594,7 +594,7 @@ mod tests {
594594
addresses: Vec::new(),
595595
announcement_message: None,
596596
});
597-
network.channels.insert(NetworkMap::get_key(12, zero_hash.clone()), ChannelInfo {
597+
network.add_channel(NetworkMap::get_key(12, zero_hash.clone()), ChannelInfo {
598598
features: ChannelFeatures::from_le_bytes(id_to_feature_flags!(12)),
599599
one_to_two: DirectionalChannelInfo {
600600
src_node_id: our_id.clone(),
@@ -618,7 +618,7 @@ mod tests {
618618
},
619619
announcement_message: None,
620620
});
621-
network.nodes.insert(node3.clone(), NodeInfo {
621+
network.add_node(node3.clone(), NodeInfo {
622622
channels: vec!(
623623
NetworkMap::get_key(3, zero_hash.clone()),
624624
NetworkMap::get_key(4, zero_hash.clone()),
@@ -634,7 +634,7 @@ mod tests {
634634
addresses: Vec::new(),
635635
announcement_message: None,
636636
});
637-
network.channels.insert(NetworkMap::get_key(3, zero_hash.clone()), ChannelInfo {
637+
network.add_channel(NetworkMap::get_key(3, zero_hash.clone()), ChannelInfo {
638638
features: ChannelFeatures::from_le_bytes(id_to_feature_flags!(3)),
639639
one_to_two: DirectionalChannelInfo {
640640
src_node_id: node1.clone(),
@@ -658,7 +658,7 @@ mod tests {
658658
},
659659
announcement_message: None,
660660
});
661-
network.channels.insert(NetworkMap::get_key(4, zero_hash.clone()), ChannelInfo {
661+
network.add_channel(NetworkMap::get_key(4, zero_hash.clone()), ChannelInfo {
662662
features: ChannelFeatures::from_le_bytes(id_to_feature_flags!(4)),
663663
one_to_two: DirectionalChannelInfo {
664664
src_node_id: node2.clone(),
@@ -682,7 +682,7 @@ mod tests {
682682
},
683683
announcement_message: None,
684684
});
685-
network.channels.insert(NetworkMap::get_key(13, zero_hash.clone()), ChannelInfo {
685+
network.add_channel(NetworkMap::get_key(13, zero_hash.clone()), ChannelInfo {
686686
features: ChannelFeatures::from_le_bytes(id_to_feature_flags!(13)),
687687
one_to_two: DirectionalChannelInfo {
688688
src_node_id: node8.clone(),
@@ -706,7 +706,7 @@ mod tests {
706706
},
707707
announcement_message: None,
708708
});
709-
network.nodes.insert(node4.clone(), NodeInfo {
709+
network.add_node(node4.clone(), NodeInfo {
710710
channels: vec!(NetworkMap::get_key(5, zero_hash.clone()), NetworkMap::get_key(11, zero_hash.clone())),
711711
lowest_inbound_channel_fees: Some(routing_fees.clone()),
712712
features: NodeFeatures::from_le_bytes(id_to_feature_flags!(4)),
@@ -716,7 +716,7 @@ mod tests {
716716
addresses: Vec::new(),
717717
announcement_message: None,
718718
});
719-
network.channels.insert(NetworkMap::get_key(5, zero_hash.clone()), ChannelInfo {
719+
network.add_channel(NetworkMap::get_key(5, zero_hash.clone()), ChannelInfo {
720720
features: ChannelFeatures::from_le_bytes(id_to_feature_flags!(5)),
721721
one_to_two: DirectionalChannelInfo {
722722
src_node_id: node3.clone(),
@@ -740,7 +740,7 @@ mod tests {
740740
},
741741
announcement_message: None,
742742
});
743-
network.nodes.insert(node5.clone(), NodeInfo {
743+
network.add_node(node5.clone(), NodeInfo {
744744
channels: vec!(NetworkMap::get_key(6, zero_hash.clone()), NetworkMap::get_key(11, zero_hash.clone())),
745745
lowest_inbound_channel_fees: Some(routing_fees.clone()),
746746
features: NodeFeatures::from_le_bytes(id_to_feature_flags!(5)),
@@ -750,7 +750,7 @@ mod tests {
750750
addresses: Vec::new(),
751751
announcement_message: None,
752752
});
753-
network.channels.insert(NetworkMap::get_key(6, zero_hash.clone()), ChannelInfo {
753+
network.add_channel(NetworkMap::get_key(6, zero_hash.clone()), ChannelInfo {
754754
features: ChannelFeatures::from_le_bytes(id_to_feature_flags!(6)),
755755
one_to_two: DirectionalChannelInfo {
756756
src_node_id: node3.clone(),
@@ -771,7 +771,7 @@ mod tests {
771771
},
772772
announcement_message: None,
773773
});
774-
network.channels.insert(NetworkMap::get_key(11, zero_hash.clone()), ChannelInfo {
774+
network.add_channel(NetworkMap::get_key(11, zero_hash.clone()), ChannelInfo {
775775
features: ChannelFeatures::from_le_bytes(id_to_feature_flags!(11)),
776776
one_to_two: DirectionalChannelInfo {
777777
src_node_id: node5.clone(),
@@ -792,7 +792,7 @@ mod tests {
792792
},
793793
announcement_message: None,
794794
});
795-
network.nodes.insert(node6.clone(), NodeInfo {
795+
network.add_node(node6.clone(), NodeInfo {
796796
channels: vec!(NetworkMap::get_key(7, zero_hash.clone())),
797797
lowest_inbound_channel_fees: Some(routing_fees.clone()),
798798
features: NodeFeatures::from_le_bytes(id_to_feature_flags!(6)),
@@ -802,7 +802,7 @@ mod tests {
802802
addresses: Vec::new(),
803803
announcement_message: None,
804804
});
805-
network.channels.insert(NetworkMap::get_key(7, zero_hash.clone()), ChannelInfo {
805+
network.add_channel(NetworkMap::get_key(7, zero_hash.clone()), ChannelInfo {
806806
features: ChannelFeatures::from_le_bytes(id_to_feature_flags!(7)),
807807
one_to_two: DirectionalChannelInfo {
808808
src_node_id: node3.clone(),
@@ -849,8 +849,8 @@ mod tests {
849849

850850
{ // Disable channels 4 and 12 by requiring unknown feature bits
851851
let mut network = network_graph.network_map.write().unwrap();
852-
network.channels.get_mut(&NetworkMap::get_key(4, zero_hash.clone())).unwrap().features.set_require_unknown_bits();
853-
network.channels.get_mut(&NetworkMap::get_key(12, zero_hash.clone())).unwrap().features.set_require_unknown_bits();
852+
network.get_channel_mut(&NetworkMap::get_key(4, zero_hash.clone())).unwrap().features.set_require_unknown_bits();
853+
network.get_channel_mut(&NetworkMap::get_key(12, zero_hash.clone())).unwrap().features.set_require_unknown_bits();
854854
}
855855

856856
{ // If all the channels require some features we don't understand, route should fail
@@ -891,15 +891,15 @@ mod tests {
891891

892892
{ // Re-enable channels 4 and 12 by wiping the unknown feature bits
893893
let mut network = network_graph.network_map.write().unwrap();
894-
network.channels.get_mut(&NetworkMap::get_key(4, zero_hash.clone())).unwrap().features.clear_require_unknown_bits();
895-
network.channels.get_mut(&NetworkMap::get_key(12, zero_hash.clone())).unwrap().features.clear_require_unknown_bits();
894+
network.get_channel_mut(&NetworkMap::get_key(4, zero_hash.clone())).unwrap().features.clear_require_unknown_bits();
895+
network.get_channel_mut(&NetworkMap::get_key(12, zero_hash.clone())).unwrap().features.clear_require_unknown_bits();
896896
}
897897

898898
{ // Disable nodes 1, 2, and 8 by requiring unknown feature bits
899899
let mut network = network_graph.network_map.write().unwrap();
900-
network.nodes.get_mut(&node1).unwrap().features.set_require_unknown_bits();
901-
network.nodes.get_mut(&node2).unwrap().features.set_require_unknown_bits();
902-
network.nodes.get_mut(&node8).unwrap().features.set_require_unknown_bits();
900+
network.get_node_mut(&node1).unwrap().features.set_require_unknown_bits();
901+
network.get_node_mut(&node2).unwrap().features.set_require_unknown_bits();
902+
network.get_node_mut(&node8).unwrap().features.set_require_unknown_bits();
903903
}
904904

905905
{ // If all nodes require some features we don't understand, route should fail
@@ -940,9 +940,9 @@ mod tests {
940940

941941
{ // Re-enable nodes 1, 2, and 8
942942
let mut network = network_graph.network_map.write().unwrap();
943-
network.nodes.get_mut(&node1).unwrap().features.clear_require_unknown_bits();
944-
network.nodes.get_mut(&node2).unwrap().features.clear_require_unknown_bits();
945-
network.nodes.get_mut(&node8).unwrap().features.clear_require_unknown_bits();
943+
network.get_node_mut(&node1).unwrap().features.clear_require_unknown_bits();
944+
network.get_node_mut(&node2).unwrap().features.clear_require_unknown_bits();
945+
network.get_node_mut(&node8).unwrap().features.clear_require_unknown_bits();
946946
}
947947

948948
// Note that we don't test disabling node 3 and failing to route to it, as we (somewhat

0 commit comments

Comments
 (0)