1use std::{
43 collections::BTreeMap, future::Future, marker::PhantomData, pin::Pin, sync::Arc, time::Duration,
44};
45
46use async_trait::async_trait;
47use bytes::{BufMut, Bytes, BytesMut};
48use log::{debug, error, info, trace, warn};
49use ndn_ndnlp::{FragCount, FragIndex, Fragment, LpPacket, Packet, Sequence};
50use ndn_nfd_mgmt::{make_command, ControlParameters, ControlResponse};
51use ndn_protocol::{
52 signature::{KnownVerifiers, SignMethod, ToVerifier},
53 Data, Interest, Name, SignSettings,
54};
55use ndn_tlv::{NonNegativeInteger, TlvDecode, TlvEncode};
56use tokio::{
57 io::{AsyncRead, AsyncWrite, AsyncWriteExt, BufReader, BufWriter},
58 net::UnixStream,
59 sync::{self, broadcast, mpsc, RwLock},
60};
61use tokio_util::sync::CancellationToken;
62use type_map::concurrent::TypeMap;
63
64use crate::{
65 error::Error,
66 util::add_bytes,
67 verifier::{DataVerifier, InterestVerifier},
68 DataExt, Result, ToName,
69};
70
71#[derive(Debug, Clone)]
72enum Connector {
73 Unix(String),
74}
75
76#[async_trait]
89trait InterestCallbackErased<Context> {
90 async fn run(
91 &self,
92 handler: AppHandler,
93 interest: Interest<Bytes>,
94 context: Context,
95 ) -> std::result::Result<Option<Data<Bytes>>, ()>;
96}
97
98#[async_trait]
99impl<Context, Params, Output, T> InterestCallbackErased<Context> for T
100where
101 Self: Sync,
102 T: InterestCallback<Context, Params = Params, Output = Output>,
103 Context: Send + 'static,
104 Params: TlvDecode + Send,
105 Output: TlvEncode,
106{
107 async fn run(
108 &self,
109 handler: AppHandler,
110 interest: Interest<Bytes>,
111 context: Context,
112 ) -> std::result::Result<Option<Data<Bytes>>, ()> {
113 let has_params = interest.application_parameters().is_some();
118 let input_interest = interest.decode_application_parameters();
119 if has_params && input_interest.application_parameters().is_none() {
120 return Err(());
121 }
122 Ok(
123 InterestCallback::run(self, handler, input_interest, context)
124 .await
125 .map(|x| x.encode_content()),
126 )
127 }
128}
129
130struct InterestCallbackFunction<Input, Context, F> {
131 f: F,
132 _input: PhantomData<fn() -> Input>,
133 _context: PhantomData<fn() -> Context>,
134}
135
136trait IntoInterestCallbackFunction<Input, Context>: Sized {
137 fn into_interest_callback_function(self) -> InterestCallbackFunction<Input, Context, Self> {
138 InterestCallbackFunction {
139 f: self,
140 _input: PhantomData,
141 _context: PhantomData,
142 }
143 }
144}
145
146impl<F, G, Params, Context, Output> IntoInterestCallbackFunction<Params, Context> for F
147where
148 F: Fn(AppHandler, Interest<Params>, Context) -> G,
149 G: Future<Output = Option<Data<Output>>> + Send + 'static,
150{
151}
152
153trait InterestCallback<Context> {
154 type Params;
155 type Output;
156 fn run(
157 &self,
158 handler: AppHandler,
159 interest: Interest<Self::Params>,
160 context: Context,
161 ) -> Pin<Box<dyn Future<Output = Option<Data<Self::Output>>> + Send + 'static>>;
162}
163
164impl<F, G, Context, Params, Output> InterestCallback<Context>
165 for InterestCallbackFunction<Params, Context, F>
166where
167 Params: TlvDecode,
168 Output: TlvEncode,
169 F: Fn(AppHandler, Interest<Params>, Context) -> G,
170 G: Future<Output = Option<Data<Output>>> + Send + 'static,
171{
172 type Params = Params;
173 type Output = Output;
174
175 fn run(
176 &self,
177 handler: AppHandler,
178 interest: Interest<Self::Params>,
179 context: Context,
180 ) -> Pin<Box<(dyn Future<Output = Option<Data<Self::Output>>> + Send)>> {
181 Box::pin((self.f)(handler, interest, context))
182 }
183}
184
185trait OnStartFn<Context> {
186 fn run(
187 &self,
188 handler: AppHandler,
189 context: Context,
190 ) -> Pin<Box<dyn Future<Output = ()> + Send>>;
191}
192
193impl<F, G, Context> OnStartFn<Context> for F
194where
195 F: Send + Sync,
196 F: Fn(AppHandler, Context) -> G,
197 G: Future<Output = ()> + 'static + Send,
198{
199 fn run(
200 &self,
201 handler: AppHandler,
202 context: Context,
203 ) -> Pin<Box<dyn Future<Output = ()> + Send>> {
204 Box::pin(self(handler, context))
205 }
206}
207
208pub struct RouteHandler<Context> {
211 callback: Box<dyn InterestCallbackErased<Context> + Send + Sync>,
212 verifier: Box<dyn InterestVerifier + Send + Sync>,
213}
214
215type InitialisedApp<Signer, Context, Verifiers> =
223 App<Signer, Context, Verifiers, Arc<RwLock<BTreeMap<Name, RouteHandler<Context>>>>>;
224
225type UninitialisedApp<Signer, Context, Verifiers> =
226 App<Signer, Context, Verifiers, BTreeMap<Name, RouteHandler<Context>>>;
227
228pub struct App<Signer, Context, KnownVerifiers, Routes> {
235 routes: Routes,
236 on_start: Option<Box<dyn OnStartFn<Context>>>,
237 connector: Connector,
238 signer: Arc<RwLock<Signer>>,
239 verifier_context: Arc<RwLock<TypeMap>>,
240 context: Context,
241 mtu: usize,
242 known_verifiers: Arc<KnownVerifiers>,
243}
244
245struct InterestToSend<T> {
246 interest: Interest<T>,
247 sign: bool,
248 notifier: tokio::sync::oneshot::Sender<Name>,
249}
250
251#[derive(Clone)]
259pub struct AppHandler {
260 interest_sender: mpsc::Sender<InterestToSend<Bytes>>,
261 in_handler: broadcast::Sender<Packet>,
262 verifier_context: Arc<RwLock<TypeMap>>,
263 known_verifiers: Arc<dyn ToVerifier + Send + Sync + 'static>,
264 shutdown_token: CancellationToken,
265}
266
267impl AppHandler {
268 pub fn shutdown(&self) {
273 self.shutdown_token.cancel();
274 }
275
276 pub async fn express_interest<T>(
280 &mut self,
281 interest: impl std::borrow::BorrowMut<Interest<T>>,
282 verifier: impl DataVerifier,
283 ) -> Result<Data<Bytes>>
284 where
285 T: TlvEncode + TlvDecode + Clone,
286 {
287 self.express_interest_impl(interest, verifier, true).await
288 }
289
290 pub async fn express_interest_unsigned<T>(
298 &mut self,
299 interest: impl std::borrow::BorrowMut<Interest<T>>,
300 verifier: impl DataVerifier,
301 ) -> Result<Data<Bytes>>
302 where
303 T: TlvEncode + TlvDecode + Clone,
304 {
305 self.express_interest_impl(interest, verifier, false).await
306 }
307
308 async fn express_interest_impl<T>(
309 &mut self,
310 mut interest: impl std::borrow::BorrowMut<Interest<T>>,
311 verifier: impl DataVerifier,
312 sign: bool,
313 ) -> Result<Data<Bytes>>
314 where
315 T: TlvEncode + TlvDecode + Clone,
316 {
317 let mut interest = interest
318 .borrow_mut()
319 .clone()
320 .encode_application_parameters();
321 let (notifier_sender, notifier_receiver) = sync::oneshot::channel();
322 self.interest_sender
323 .send(InterestToSend {
324 interest: interest.clone(),
325 sign,
326 notifier: notifier_sender,
327 })
328 .await
329 .map_err(|_| Error::ConnectionClosed)?;
330
331 let signed_name = notifier_receiver
332 .await
333 .unwrap_or_else(|_| interest.name().clone());
334
335 interest.set_name(signed_name.clone());
336
337 let lifetime = interest.interest_lifetime().map(u64::from).unwrap_or(3_000);
338 let wait_for_data = async {
339 let mut in_receiver = self.in_handler.subscribe();
340 while let Ok(packet) = in_receiver.recv().await {
341 match packet {
342 Packet::Data(packet) => {
343 if packet.matches_interest(&interest) {
344 if verifier
345 .verify(
346 &packet,
347 Arc::clone(&self.verifier_context),
348 self.clone(),
349 &*self.known_verifiers,
350 )
351 .await
352 {
353 return Ok(packet);
354 } else {
355 warn!(
356 "Data packet for {} failed verification.",
357 interest.name().to_uri()
358 );
359 return Err(Error::VerificationFailed);
360 }
361 }
362 }
363 Packet::LpPacket(packet) => {
364 if packet.is_nack() {
365 let Some(nack_interest) = packet
366 .fragment()
367 .map(|mut x| Interest::<Bytes>::decode(&mut x).ok())
368 .flatten()
369 else {
370 continue;
371 };
372
373 if nack_interest.name() == &signed_name {
374 debug!(
375 "Received NACK when requesting {}",
376 interest.name().to_uri()
377 );
378 return Err(Error::NackReceived);
379 }
380 }
381 }
382 _ => {}
383 }
384 }
385 Err(Error::ConnectionClosed)
386 };
387 match tokio::time::timeout(Duration::from_millis(lifetime), wait_for_data).await {
388 Ok(x) => x,
389 Err(x) => {
390 debug!("Request for {} timed out", interest.name().to_uri());
391 Err(x.into())
392 }
393 }
394 }
395}
396
397impl<Signer, Context> UninitialisedApp<Signer, Context, KnownVerifiers>
398where
399 Signer: SignMethod + Send + Sync + 'static,
400 Context: Clone + Send + 'static,
401{
402 pub fn new(signer: Signer, context: Context) -> Self {
408 Self {
409 routes: BTreeMap::new(),
410 connector: Connector::Unix("/var/run/nfd/nfd.sock".to_string()),
411 signer: Arc::new(RwLock::new(signer)),
412 on_start: None,
413 verifier_context: Arc::new(RwLock::new(TypeMap::new())),
414 context,
415 mtu: 8800,
416 known_verifiers: Arc::new(KnownVerifiers),
417 }
418 }
419}
420
421impl<Signer, Context, Verifiers> UninitialisedApp<Signer, Context, Verifiers>
422where
423 Signer: SignMethod + Send + Sync + 'static,
424 Context: Clone + Send + 'static,
425 Verifiers: ToVerifier + Send + Sync + 'static,
426{
427 #[allow(private_bounds)]
433 pub fn route<Callback, Verifier, Params, Output, G>(
434 mut self,
435 name: impl ToName,
436 verifier: Verifier,
437 func: Callback,
438 ) -> Self
439 where
440 Callback: IntoInterestCallbackFunction<Params, Context> + Send + Sync + 'static,
441 Callback: Fn(AppHandler, Interest<Params>, Context) -> G,
442 G: Future<Output = Option<Data<Output>>> + Send + 'static,
443 Verifier: InterestVerifier + Send + Sync + 'static,
444 Params: TlvDecode + Send + 'static,
445 Output: TlvEncode + 'static,
446 {
447 self.routes.insert(
448 name.to_name(),
449 RouteHandler {
450 callback: Box::new(func.into_interest_callback_function()),
451 verifier: Box::new(verifier),
452 },
453 );
454 self
455 }
456
457 #[allow(private_bounds)]
462 pub fn on_start<F>(mut self, on_start: F) -> Self
463 where
464 F: OnStartFn<Context> + 'static,
465 {
466 self.on_start = Some(Box::new(on_start));
467 self
468 }
469
470 pub fn mtu(mut self, mtu: usize) -> Self {
473 self.mtu = mtu;
474 self
475 }
476
477 pub async fn start(self) -> Result<()> {
481 self.initialise().start().await
482 }
483
484 fn initialise(self) -> InitialisedApp<Signer, Context, Verifiers> {
485 App {
486 routes: Arc::new(RwLock::new(self.routes)),
487 on_start: self.on_start,
488 connector: self.connector,
489 signer: self.signer,
490 verifier_context: self.verifier_context,
491 context: self.context,
492 mtu: self.mtu,
493 known_verifiers: self.known_verifiers,
494 }
495 }
496}
497
498impl<Signer, Context, Verifiers> InitialisedApp<Signer, Context, Verifiers>
499where
500 Signer: SignMethod + Send + Sync + 'static,
501 Context: Clone + Send + 'static,
502 Verifiers: ToVerifier + Send + Sync + 'static,
503{
504 async fn handle_interest(
516 interest: Interest<Bytes>,
517 routes: Arc<RwLock<BTreeMap<Name, RouteHandler<Context>>>>,
518 verifier_context: Arc<RwLock<TypeMap>>,
519 app_handler: AppHandler,
520 signer: Arc<RwLock<Signer>>,
521 out_sender: mpsc::Sender<Packet>,
522 context: Context,
523 known_verifiers: Arc<Verifiers>,
524 ) -> Result<()> {
525 let routes = routes.read().await;
526 let interest_uri = interest.name().to_uri();
527 trace!("Received interest for {interest_uri}");
528 let mut matching_route_found = false;
529 for (route, route_handler) in routes.iter().rev() {
530 trace!("Checking against route {}", route.to_uri());
531 if interest.name().has_prefix(route) {
532 matching_route_found = true;
533 if !route_handler
534 .verifier
535 .verify(
536 &interest,
537 Arc::clone(&verifier_context),
538 app_handler.clone(),
539 &*known_verifiers,
540 )
541 .await
542 {
543 info!("Verification failed for request for {interest_uri}");
544 let nack = Packet::make_nack(interest);
545 out_sender
546 .send(nack)
547 .await
548 .map_err(|_| Error::ConnectionClosed)?;
549 break;
550 }
551 if let Ok(ret) = route_handler
552 .callback
553 .run(app_handler.clone(), interest.clone(), context.clone()) .await
555 {
556 if let Some(mut ret) = ret {
557 if !ret.is_signed() {
558 let mut signer = signer.write().await;
559 ret.sign(&mut *signer);
560 }
561 out_sender
562 .send(Packet::Data(ret))
563 .await
564 .map_err(|_| Error::ConnectionClosed)?;
565 } else {
566 let nack = Packet::make_nack(interest);
567 out_sender
568 .send(nack)
569 .await
570 .map_err(|_| Error::ConnectionClosed)?;
571 }
572 break;
573 } else {
574 debug!("Interest application parameter decoding failed for {interest_uri}");
575 }
576 }
577 }
578 if !matching_route_found {
579 debug!("Interest for {interest_uri} matches no routes");
580 }
581 Ok(())
582 }
583
584 async fn start(mut self) -> Result<()> {
594 let (reader, writer): (
595 Box<dyn AsyncRead + Unpin + Send>,
596 Box<dyn AsyncWrite + Unpin + Send>,
597 ) = match self.connector.clone() {
598 Connector::Unix(path) => {
599 let sock = UnixStream::connect(path)
600 .await
601 .map_err(|_| Error::ConnectionFailed)?;
602 let (reader, writer) = sock.into_split();
603 (
604 Box::new(BufReader::new(reader)),
605 Box::new(BufWriter::new(writer)),
606 )
607 }
608 };
609
610 let shutdown_token = CancellationToken::new();
611
612 let (out_sender, out_receiver) = mpsc::channel(128);
614 let (in_sender, mut in_receiver) = broadcast::channel(128);
616
617 tokio::spawn(write_thread(
618 writer,
619 out_receiver,
620 self.mtu,
621 shutdown_token.clone(),
622 ));
623 tokio::spawn(read_thread(
624 reader,
625 in_sender.clone(),
626 shutdown_token.clone(),
627 ));
628
629 for (route, _) in self.routes.read().await.iter() {
630 loop {
631 let res = tokio::time::timeout(
632 Duration::from_secs(3),
633 register_route(
634 Arc::clone(&self.signer),
635 in_sender.subscribe(),
636 out_sender.clone(),
637 route.clone(),
638 ),
639 )
640 .await;
641
642 match res {
643 Ok(Ok(())) => break,
644 Ok(Err(x)) => {
645 error!("Route registration had an error: {}", x);
646 return Err(x);
647 }
648 Err(_) => {
649 warn!("Route registration timed out");
650 continue;
651 }
652 }
653 }
654 }
655
656 let (interest_sender, interest_receiver) = mpsc::channel(128);
657
658 tokio::spawn(interest_thread(
659 out_sender.clone(),
660 interest_receiver,
661 Arc::clone(&self.signer),
662 shutdown_token.clone(),
663 ));
664
665 let app_handler = AppHandler {
666 interest_sender,
667 in_handler: in_sender.clone(),
668 verifier_context: Arc::clone(&self.verifier_context),
669 known_verifiers: Arc::<Verifiers>::clone(&self.known_verifiers),
670 shutdown_token: shutdown_token.clone(),
671 };
672
673 if let Some(on_start) = self.on_start.take() {
674 tokio::spawn(on_start.run(app_handler.clone(), self.context.clone()));
675 }
676
677 let mut partial_packet: Vec<Bytes> = Vec::new();
678 let mut partial_count = 0;
679 let mut last_seq = BytesMut::new();
680
681 'main_loop: loop {
682 tokio::select! {
683 _ = shutdown_token.cancelled() => {
684 return Ok(());
685 }
686 received = in_receiver.recv() => {
687 match received {
688 Ok(packet) => match packet {
689 Packet::Interest(interest) => {
690 tokio::spawn(Self::handle_interest(
691 interest.clone(),
692 Arc::clone(&self.routes),
693 Arc::clone(&self.verifier_context),
694 app_handler.clone(),
695 Arc::clone(&self.signer),
696 out_sender.clone(),
697 self.context.clone(),
698 Arc::clone(&self.known_verifiers),
699 ));
700 }
701 Packet::LpPacket(packet) => {
702 for header in packet.other_headers() {
703 if header.is_critical() {
704 continue 'main_loop;
706 }
707 }
708
709 if let Some((frag_idx, frag_cnt)) = packet.frag_info() {
710 let (Some(seq), Some(fragment)) = (packet.seq_num(), packet.fragment())
712 else {
713 partial_packet.clear();
714 partial_count = 0;
715 continue 'main_loop;
716 };
717
718 if frag_idx.as_usize() != partial_packet.len() {
720 partial_packet.clear();
721 partial_count = 0;
722 continue 'main_loop;
723 }
724 if frag_idx.as_usize() == 0 {
726 partial_count = frag_cnt.into();
727 partial_packet.clear();
728 partial_packet.reserve(frag_cnt.as_usize());
729 last_seq = BytesMut::from(&seq[..]);
730 } else if partial_count != frag_cnt.as_u64() {
732 partial_packet.clear();
733 partial_count = 0;
734 continue 'main_loop;
735 } else {
736 add_bytes(&mut last_seq, 1);
737 if last_seq != seq {
739 add_bytes(&mut last_seq, -1);
740 partial_packet.clear();
741 partial_count = 0;
742 continue 'main_loop;
743 }
744 }
745
746 partial_packet.push(fragment);
747 if frag_idx == frag_cnt {
748 let total_size: usize = partial_packet.iter().map(Bytes::len).sum();
749 let mut data = BytesMut::with_capacity(total_size);
750
751 for fragment in &partial_packet {
752 data.put(fragment.clone());
753 }
754 partial_packet.clear();
755 partial_count = 0;
756
757 let packet = Packet::decode(&mut data.freeze());
758 debug!("Reconstituted packet: {packet:#?}");
759 }
760 }
761 }
762 _ => {}
763 },
764 Err(broadcast::error::RecvError::Closed) => return Err(Error::ConnectionClosed),
765 Err(broadcast::error::RecvError::Lagged(n)) => {
766 warn!("Dropped {n} packets in routing handler");
767 }
768 }
769 }
770 }
771 }
772 }
773}
774
775async fn interest_thread(
787 sender: mpsc::Sender<Packet>,
788 mut interest_receiver: mpsc::Receiver<InterestToSend<Bytes>>,
789 signer: Arc<RwLock<impl SignMethod>>,
790 shutdown_token: CancellationToken,
791) -> Result<()> {
792 let _shutdown_guard = shutdown_token.drop_guard();
793 while let Some(mut interest_to_send) = interest_receiver.recv().await {
794 if interest_to_send.sign && !interest_to_send.interest.is_signed() {
795 let mut signer = signer.write().await;
796 interest_to_send
797 .interest
798 .sign(&mut *signer, SignSettings::default()); }
800 let _ = interest_to_send
801 .notifier
802 .send(interest_to_send.interest.name().clone());
803 sender
804 .send(Packet::Interest(interest_to_send.interest))
805 .await
806 .map_err(|_| Error::ConnectionClosed)?;
807 }
808 Err(Error::ConnectionClosed)
809}
810
811async fn write_thread(
822 mut writer: impl AsyncWrite + Unpin,
823 mut receiver: mpsc::Receiver<Packet>,
824 mtu: usize,
825 shutdown_token: CancellationToken,
826) -> Result<()> {
827 let _shutdown_guard = shutdown_token.drop_guard();
828 let mut seq_num = BytesMut::from(&[0; 8][..]);
829
830 while let Some(packet) = receiver.recv().await {
831 let mut data = packet.encode();
832
833 if data.len() > mtu {
834 let header = LpPacket {
835 sequence: Some(Sequence(seq_num.clone().freeze())),
836 frag_index: Some(FragIndex(NonNegativeInteger::U64(0))),
837 frag_count: Some(FragCount(NonNegativeInteger::U64(0))),
838 nack: None,
839 other_headers: Vec::new(),
840 fragment: None,
841 };
842 let header_size = header.size();
843
844 let frag_count = data.len().div_ceil(mtu - header_size);
845
846 for i in 0..frag_count {
847 let frame = LpPacket {
848 sequence: Some(Sequence(seq_num.clone().freeze())),
849 frag_index: Some(FragIndex(NonNegativeInteger::new(i as u64))),
850 frag_count: Some(FragCount(NonNegativeInteger::new(frag_count as u64))),
851 nack: None,
852 other_headers: Vec::new(),
853 fragment: Some(Fragment {
854 data: data.split_to(data.len().min(mtu - header_size)),
855 }),
856 };
857 add_bytes(&mut seq_num, 1);
858 writer.write_all(&frame.encode()).await?;
859 }
860 } else {
861 writer.write_all(&data).await?;
862 }
863 writer.flush().await?;
864 }
865 Err(Error::ConnectionClosed)
866}
867
868async fn read_thread(
877 mut reader: impl AsyncRead + Unpin,
878 sender: broadcast::Sender<Packet>,
879 shutdown_token: CancellationToken,
880) -> Result<()> {
881 let _shutdown_guard = shutdown_token.drop_guard();
882 while let Some(packet) = Packet::from_async_reader(&mut reader).await {
883 sender.send(packet).map_err(|_| Error::ConnectionClosed)?;
884 }
885 Err(Error::ConnectionClosed)
886}
887
888async fn register_route(
897 signer: Arc<RwLock<impl SignMethod>>,
898 mut receiver: broadcast::Receiver<Packet>,
899 sender: mpsc::Sender<Packet>,
900 route: Name,
901) -> Result<()> {
902 info!("Registering route {}", route.to_uri());
903 let control_parameters = ControlParameters::new().set_name(route.clone());
904 let mut interest = make_command("rib", "register", control_parameters).unwrap();
905
906 {
907 let mut signer = signer.write().await;
908 interest.sign(&mut *signer, SignSettings::default());
909 }
910
911 sender
912 .send(Packet::Interest(
913 interest.clone().encode_application_parameters(),
914 ))
915 .await
916 .map_err(|_| Error::ConnectionClosed)?;
917
918 loop {
919 match receiver.recv().await {
920 Ok(Packet::Data(packet)) => {
921 if packet.name().has_prefix(&interest.name()) {
922 let data = packet.decode_content::<ControlResponse<ControlParameters>>();
923 if let Some(content) = data.content() {
924 if content.status_code().as_usize() != 200 {
925 error!("Registering {route} failed with status code {status_code}: {status_text}",
926 route = route.to_uri(),
927 status_code = content.status_code(),
928 status_text = String::from_utf8_lossy(&content.status_text()));
929 } else {
930 info!("Registered route {route}", route = route.to_uri());
931 }
932 return Ok(());
933 }
934 }
935 }
936 Ok(Packet::Interest(_)) => {}
937 Ok(Packet::LpPacket(packet)) => {
938 println!("{:#?}", packet)
939 }
940 Err(broadcast::error::RecvError::Closed) => return Err(Error::ConnectionClosed),
941 Err(broadcast::error::RecvError::Lagged(_)) => {}
942 }
943 }
944}