1use amplify::DumbDefault;
15use internet2::presentation::sphinx::Hop;
16use p2p::bolt::{
17 ChannelId, HopRealm, Messages, PaymentOnion, PaymentRequest, ShortChannelId,
18};
19use strict_encoding::{strict_deserialize, strict_serialize};
20
21use super::GossipChannelInfo;
22use crate::router::gossip::LocalChannelInfo;
23use crate::router::Router;
24use crate::{extension, router, Extension, RouterExtension};
25
26#[derive(
27 Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug, Display, Error
28)]
29#[display(doc_comments)]
30pub enum Error {}
31
32#[derive(Clone, PartialEq, Eq, Debug, Default)]
33#[derive(StrictEncode, StrictDecode)]
34pub struct RouterState {
35 remote_channels: Vec<GossipChannelInfo>,
36 direct_channels: Vec<LocalChannelInfo>,
37}
38
39impl DumbDefault for RouterState {
40 fn dumb_default() -> Self {
41 RouterState::default()
42 }
43}
44
45#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Display)]
46#[derive(StrictEncode, StrictDecode)]
47#[display(Debug)]
48pub enum GossipExt {
49 MainRouter = 0,
50 DirectRouter = 1,
51 GossipRouter = 2,
52}
53
54impl Default for GossipExt {
55 fn default() -> Self {
56 GossipExt::MainRouter
57 }
58}
59
60impl From<GossipExt> for u16 {
61 fn from(id: GossipExt) -> Self {
62 let mut buf = [0u8; 2];
63 buf.copy_from_slice(
64 &strict_serialize(&id)
65 .expect("Enum in-memory strict encoding can't fail"),
66 );
67 u16::from_be_bytes(buf)
68 }
69}
70
71impl TryFrom<u16> for GossipExt {
72 type Error = strict_encoding::Error;
73
74 fn try_from(value: u16) -> Result<Self, Self::Error> {
75 strict_deserialize(value.to_be_bytes())
76 }
77}
78
79impl extension::Nomenclature for GossipExt {
80 type State = RouterState;
81 type Error = Error;
82 type PeerMessage = lnp2p::bolt::Messages;
83 type UpdateMessage = UpdateMsg;
84 type UpdateRequest = ();
85}
86
87impl router::Nomenclature for GossipExt {
88 type HopPayload = PaymentOnion;
89
90 fn default_extensions() -> Vec<Box<dyn RouterExtension<Self>>> {
91 vec![
92 Box::new(DirectRouter::default()) as Box<dyn RouterExtension<Self>>,
93 Box::new(GossipRouter::default()) as Box<dyn RouterExtension<Self>>,
94 ]
95 }
96
97 fn update_from_peer(
98 _router: &mut Router<Self>,
99 _message: &Messages,
100 ) -> Result<(), Error> {
101 Ok(())
104 }
105}
106
107#[derive(Clone, Eq, PartialEq, Hash, Debug)]
108pub enum UpdateMsg {
109 DirectChannelAdd(LocalChannelInfo),
110 DirectChannelRemove(ChannelId),
111 DirectChannelUpdate {
112 channel_id: ChannelId,
113 local_amount_msat: u64,
114 remote_amount_msat: u64,
115 },
116}
117
118#[derive(Getters, Clone, PartialEq, Eq, Debug, Default)]
121pub struct DirectRouter {
122 channels: Vec<LocalChannelInfo>,
123}
124
125impl DirectRouter {
126 fn add_direct_channel(
127 &mut self,
128 info: LocalChannelInfo,
129 ) -> Option<LocalChannelInfo> {
130 let prev_info = self.remove_direct_channel(info.channel_id);
131 self.channels.push(info);
132 prev_info
133 }
134
135 fn remove_direct_channel(
136 &mut self,
137 channel_id: ChannelId,
138 ) -> Option<LocalChannelInfo> {
139 if let Some((index, _)) = self
140 .channels
141 .iter()
142 .enumerate()
143 .find(|(_, info)| info.channel_id == channel_id)
144 {
145 Some(self.channels.remove(index))
146 } else {
147 None
148 }
149 }
150}
151
152impl Extension<GossipExt> for DirectRouter {
153 fn identity(&self) -> GossipExt {
154 GossipExt::DirectRouter
155 }
156
157 fn update_from_peer(&mut self, message: &Messages) -> Result<(), Error> {
158 #[allow(clippy::match_single_binding)] match message {
160 _ => {} }
169
170 Ok(())
171 }
172
173 fn update_from_local(&mut self, message: &UpdateMsg) -> Result<(), Error> {
174 match message {
175 UpdateMsg::DirectChannelAdd(info) => {
176 self.add_direct_channel(*info);
177 }
178 UpdateMsg::DirectChannelRemove(channel_id) => {
179 self.remove_direct_channel(*channel_id);
180 }
181 UpdateMsg::DirectChannelUpdate {
182 channel_id,
183 local_amount_msat,
184 remote_amount_msat,
185 } => {
186 self.channels.iter_mut().for_each(|ch| {
187 if ch.channel_id == *channel_id {
188 ch.outbound_capacity_msat = *local_amount_msat;
189 ch.inbound_capacity_msat = *remote_amount_msat;
190 };
191 });
192 }
193 }
194 Ok(())
195 }
196
197 fn load_state(&mut self, state: &RouterState) {
198 self.channels = state.direct_channels.clone();
199 }
200
201 fn store_state(&self, state: &mut RouterState) {
202 state.direct_channels = self.channels.clone();
203 }
204}
205
206impl RouterExtension<GossipExt> for DirectRouter {
207 #[inline]
208 fn new() -> Box<dyn RouterExtension<GossipExt>>
209 where
210 Self: Sized,
211 {
212 Box::new(DirectRouter::default())
213 }
214
215 fn build_route(
216 &mut self,
217 payment: PaymentRequest,
218 route: &mut Vec<Hop<PaymentOnion>>,
219 ) {
220 if let Some(channel) = self
221 .channels
222 .iter()
223 .find(|info| info.remote_node == payment.node_id)
224 {
225 if channel.outbound_capacity_msat < payment.amount_msat {
226 return; }
228
229 *route = vec![Hop::with(payment.node_id, PaymentOnion {
230 realm: HopRealm::Legacy(ShortChannelId::default()),
232 amt_to_forward: payment.amount_msat,
233 outgoing_cltv_value: payment.min_final_cltv_expiry,
234 })];
235 }
236 }
237}
238
239#[derive(Getters, Clone, PartialEq, Eq, Debug, Default)]
241pub struct GossipRouter {
242 channels: Vec<GossipChannelInfo>,
243}
244
245impl Extension<GossipExt> for GossipRouter {
246 fn identity(&self) -> GossipExt {
247 GossipExt::GossipRouter
248 }
249
250 fn update_from_local(&mut self, _message: &UpdateMsg) -> Result<(), Error> {
251 Ok(())
253 }
254
255 fn update_from_peer(&mut self, message: &Messages) -> Result<(), Error> {
256 match message {
257 Messages::UpdateFee(_) => {}
259 Messages::ChannelAnnouncement(_) => {}
260 Messages::ChannelUpdate(_) => {}
261 _ => {}
262 }
263 Ok(())
264 }
265
266 fn load_state(&mut self, state: &RouterState) {
267 self.channels = state.remote_channels.clone()
268 }
269
270 fn store_state(&self, state: &mut RouterState) {
271 state.remote_channels = self.channels.clone()
272 }
273}
274
275impl RouterExtension<GossipExt> for GossipRouter {
276 #[inline]
277 fn new() -> Box<dyn RouterExtension<GossipExt>>
278 where
279 Self: Sized,
280 {
281 Box::new(GossipRouter::default())
282 }
283
284 fn build_route(
285 &mut self,
286 _payment: PaymentRequest,
287 _route: &mut Vec<Hop<PaymentOnion>>,
288 ) {
289 }
291}