1use core::net::SocketAddr;
13use std::fmt;
14use std::future::Future;
15use std::net::IpAddr;
16use std::net::Ipv6Addr;
17#[cfg(target_os = "linux")]
18use std::os::linux::net::SocketAddrExt;
19use std::os::unix::io::FromRawFd;
20use std::os::unix::io::RawFd;
21use std::panic::Location;
22use std::pin::Pin;
23use std::str::FromStr;
24use std::sync::Arc;
25use std::sync::Mutex;
26use std::sync::atomic::AtomicU8;
27use std::sync::atomic::AtomicUsize;
28use std::sync::atomic::Ordering;
29use std::task::Context;
30use std::task::Poll;
31
32use async_trait::async_trait;
33use enum_as_inner::EnumAsInner;
34use futures::task::AtomicWaker;
35use hyperactor_config::attrs::AttrValue;
36use serde::Deserialize;
37use serde::Serialize;
38use tokio::sync::mpsc;
39use tokio::sync::watch;
40use tokio::time::Instant;
41use tokio_util::sync::CancellationToken;
42
43use crate as hyperactor;
44use crate::RemoteMessage;
45pub(crate) mod local;
46pub(crate) mod net;
47
48pub use net::ServerError;
52pub use net::try_tls_acceptor;
53pub use net::try_tls_connector;
54pub use net::try_tls_pem_bundle;
55
56pub mod duplex {
58 pub use super::net::duplex::DuplexClient;
59 pub use super::net::duplex::DuplexRx;
60 pub use super::net::duplex::DuplexServer;
61 pub use super::net::duplex::DuplexTx;
62 pub use super::net::duplex::dial;
63 pub use super::net::duplex::serve;
64}
65
66#[derive(thiserror::Error, Debug)]
68pub enum ChannelError {
69 #[error("channel closed")]
71 Closed,
72
73 #[error("send: {0}")]
75 Send(#[source] anyhow::Error),
76
77 #[error(transparent)]
79 Client(#[from] net::ClientError),
80
81 #[error("invalid address {0:?}")]
83 InvalidAddress(String),
84
85 #[error(transparent)]
87 Server(#[from] net::ServerError),
88
89 #[error(transparent)]
91 BincodeEncode(#[from] bincode::error::EncodeError),
92
93 #[error(transparent)]
95 BincodeDecode(#[from] bincode::error::DecodeError),
96
97 #[error(transparent)]
99 Data(#[from] wirevalue::Error),
100
101 #[error(transparent)]
103 Other(#[from] anyhow::Error),
104
105 #[error("operation timed out after {0:?}")]
107 Timeout(std::time::Duration),
108}
109
110#[derive(thiserror::Error, Debug, Clone, PartialEq, Eq)]
112pub enum SendErrorReason {
113 #[error(
115 "rejecting oversize frame: len={len} > max={max}. \
116 ack will not arrive before timeout; increase CODEC_MAX_FRAME_LENGTH to allow."
117 )]
118 OversizedFrame {
119 len: usize,
121
122 max: usize,
124 },
125
126 #[error("{0}")]
128 Other(String),
129}
130
131#[derive(thiserror::Error, Debug)]
133#[error("{error} for reason {reason:?}")]
134pub struct SendError<M: RemoteMessage> {
135 #[source]
137 pub error: ChannelError,
138 pub message: M,
140 pub reason: Option<SendErrorReason>,
142}
143
144#[repr(u8)]
145#[derive(Clone, Copy, Debug, PartialEq, Eq)]
146enum CompletionStatus {
147 Pending = 0,
148 Accepted = 1,
149 Rejected = 2,
150}
151
152impl CompletionStatus {
153 fn from_u8(value: u8) -> Self {
154 match value {
155 value if value == Self::Pending as u8 => Self::Pending,
156 value if value == Self::Accepted as u8 => Self::Accepted,
157 value if value == Self::Rejected as u8 => Self::Rejected,
158 _ => panic!("invalid completion state"),
159 }
160 }
161}
162
163struct CompletionState<M: RemoteMessage> {
164 state: AtomicU8,
165 waker: AtomicWaker,
166 rejected: Mutex<Option<Box<SendError<M>>>>,
167}
168
169pub(crate) struct CompletionSender<M: RemoteMessage> {
170 inner: Arc<CompletionState<M>>,
171}
172
173pub struct CompletionReceipt<M: RemoteMessage> {
175 inner: Arc<CompletionState<M>>,
176}
177
178impl<M: RemoteMessage> CompletionSender<M> {
179 fn pair() -> (Self, CompletionReceipt<M>) {
180 let inner = Arc::new(CompletionState {
181 state: AtomicU8::new(CompletionStatus::Pending as u8),
182 waker: AtomicWaker::new(),
183 rejected: Mutex::new(None),
184 });
185
186 (
187 Self {
188 inner: Arc::clone(&inner),
189 },
190 CompletionReceipt { inner },
191 )
192 }
193
194 fn accept(self) {
195 self.inner
196 .state
197 .store(CompletionStatus::Accepted as u8, Ordering::Release);
198 self.inner.waker.wake();
199 }
200
201 fn reject(self, error: SendError<M>) {
202 *self.inner.rejected.lock().unwrap() = Some(Box::new(error));
203 self.inner
204 .state
205 .store(CompletionStatus::Rejected as u8, Ordering::Release);
206 self.inner.waker.wake();
207 }
208}
209
210impl<M: RemoteMessage> Drop for CompletionSender<M> {
211 fn drop(&mut self) {
212 if CompletionStatus::from_u8(self.inner.state.load(Ordering::Acquire))
213 == CompletionStatus::Pending
214 {
215 self.inner
216 .state
217 .store(CompletionStatus::Accepted as u8, Ordering::Release);
218 self.inner.waker.wake();
219 }
220 }
221}
222
223impl<M: RemoteMessage> Future for CompletionReceipt<M> {
224 type Output = Result<(), SendError<M>>;
225
226 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
227 match self.poll_ready() {
228 Poll::Ready(result) => Poll::Ready(result),
229 Poll::Pending => {
230 self.inner.waker.register(cx.waker());
231 self.poll_ready()
232 }
233 }
234 }
235}
236
237impl<M: RemoteMessage> CompletionReceipt<M> {
238 fn poll_ready(&self) -> Poll<Result<(), SendError<M>>> {
239 match CompletionStatus::from_u8(self.inner.state.load(Ordering::Acquire)) {
240 CompletionStatus::Accepted => Poll::Ready(Ok(())),
241 CompletionStatus::Rejected => {
242 let error = self
243 .inner
244 .rejected
245 .lock()
246 .unwrap()
247 .take()
248 .expect("rejected completion should store send error");
249 Poll::Ready(Err(*error))
250 }
251 CompletionStatus::Pending => Poll::Pending,
252 }
253 }
254}
255
256pub(crate) struct CompletionTracker {
258 completed: Arc<AtomicUsize>,
259 completed_notify: Arc<tokio::sync::Notify>,
260}
261
262impl CompletionTracker {
263 pub(crate) fn new(
264 completed: Arc<AtomicUsize>,
265 completed_notify: Arc<tokio::sync::Notify>,
266 ) -> Self {
267 Self {
268 completed,
269 completed_notify,
270 }
271 }
272
273 fn complete(&self) {
274 self.completed.fetch_add(1, Ordering::Relaxed);
275 self.completed_notify.notify_waiters();
276 }
277}
278
279enum CompletionSinkInner<M: RemoteMessage> {
280 Ignore,
281 Receipt(CompletionSender<M>),
282 OnReject(Box<dyn FnOnce(SendError<M>) + Send + Sync>),
283 Tracked {
284 tracker: CompletionTracker,
285 on_reject: Box<dyn FnOnce(SendError<M>) + Send + Sync>,
286 },
287}
288
289pub struct CompletionSink<M: RemoteMessage>(CompletionSinkInner<M>);
291
292impl<M: RemoteMessage> CompletionSink<M> {
293 pub fn ignore() -> Self {
295 Self(CompletionSinkInner::Ignore)
296 }
297
298 pub fn on_reject(f: impl FnOnce(SendError<M>) + Send + Sync + 'static) -> Self {
300 Self(CompletionSinkInner::OnReject(Box::new(f)))
301 }
302
303 pub fn receipt() -> (Self, CompletionReceipt<M>) {
305 let (sender, receipt) = CompletionSender::pair();
306 (Self(CompletionSinkInner::Receipt(sender)), receipt)
307 }
308
309 pub(crate) fn tracked(
311 tracker: CompletionTracker,
312 on_reject: impl FnOnce(SendError<M>) + Send + Sync + 'static,
313 ) -> Self {
314 Self(CompletionSinkInner::Tracked {
315 tracker,
316 on_reject: Box::new(on_reject),
317 })
318 }
319
320 pub fn contramap_rejected<N: RemoteMessage>(
322 self,
323 f: impl FnOnce(SendError<N>) -> Option<SendError<M>> + Send + Sync + 'static,
324 ) -> CompletionSink<N> {
325 match self.0 {
326 CompletionSinkInner::Ignore => CompletionSink::ignore(),
327 CompletionSinkInner::Receipt(sender) => CompletionSink::on_reject(move |error| {
328 if let Some(error) = f(error) {
329 sender.reject(error);
330 } else {
331 sender.accept();
332 }
333 }),
334 CompletionSinkInner::OnReject(on_reject) => CompletionSink::on_reject(move |error| {
335 if let Some(error) = f(error) {
336 on_reject(error);
337 }
338 }),
339 CompletionSinkInner::Tracked { tracker, on_reject } => {
340 CompletionSink::tracked(tracker, move |error| {
341 if let Some(error) = f(error) {
342 on_reject(error);
343 }
344 })
345 }
346 }
347 }
348
349 pub fn accept(self) {
351 match self.0 {
352 CompletionSinkInner::Ignore => {}
353 CompletionSinkInner::Receipt(sender) => sender.accept(),
354 CompletionSinkInner::OnReject(_) => {}
355 CompletionSinkInner::Tracked { tracker, .. } => tracker.complete(),
356 }
357 }
358
359 pub fn reject(self, error: SendError<M>) {
361 match self.0 {
362 CompletionSinkInner::Ignore => {}
363 CompletionSinkInner::Receipt(sender) => sender.reject(error),
364 CompletionSinkInner::OnReject(on_reject) => on_reject(error),
365 CompletionSinkInner::Tracked { tracker, on_reject } => {
366 on_reject(error);
367 tracker.complete();
368 }
369 }
370 }
371}
372
373impl<M: RemoteMessage> From<SendError<M>> for ChannelError {
374 fn from(error: SendError<M>) -> Self {
375 error.error
376 }
377}
378
379#[derive(Debug, Clone, PartialEq)]
384pub enum CloseReason {
385 SequenceMismatch(String),
391 OversizedFrame {
395 size: usize,
397 max: usize,
399 },
400 Other(String),
404}
405
406impl fmt::Display for CloseReason {
407 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
408 match self {
409 Self::SequenceMismatch(s) => write!(f, "stale session: {}", s),
410 Self::OversizedFrame { size, max } => {
411 write!(f, "oversized frame: len={size} > max={max}")
412 }
413 Self::Other(s) => f.write_str(s),
414 }
415 }
416}
417
418#[derive(Debug, Clone, PartialEq, EnumAsInner)]
420pub enum TxStatus {
421 Active,
423 Closed(CloseReason),
425}
426
427#[async_trait]
429pub trait Tx<M: RemoteMessage> {
430 fn do_post(&self, message: M, completion: CompletionSink<M>);
434
435 fn try_post(&self, message: M) -> CompletionReceipt<M> {
438 let (completion, receipt) = CompletionSink::receipt();
439 self.do_post(message, completion);
440 receipt
441 }
442
443 #[hyperactor::instrument_infallible]
445 fn post(&self, message: M) {
446 self.do_post(message, CompletionSink::ignore());
447 }
448
449 async fn send(&self, message: M) -> Result<(), SendError<M>> {
452 self.try_post(message).await
453 }
454
455 fn addr(&self) -> ChannelAddr;
457
458 fn status(&self) -> &watch::Receiver<TxStatus>;
460}
461
462#[async_trait]
464pub trait Rx<M: RemoteMessage> {
465 async fn recv(&mut self) -> Result<M, ChannelError>;
468
469 fn addr(&self) -> ChannelAddr;
471
472 async fn join(self)
476 where
477 Self: Sized;
478}
479
480#[derive(
482 Clone,
483 Debug,
484 PartialEq,
485 Eq,
486 Hash,
487 Serialize,
488 Deserialize,
489 strum::EnumIter,
490 strum::Display,
491 strum::EnumString
492)]
493pub enum TcpMode {
494 Localhost,
496 Hostname,
498}
499
500#[derive(
502 Clone,
503 Debug,
504 PartialEq,
505 Eq,
506 Hash,
507 Serialize,
508 Deserialize,
509 strum::EnumIter,
510 strum::Display,
511 strum::EnumString
512)]
513pub enum TlsMode {
514 IpV6,
516 Hostname,
518 }
520
521#[derive(
523 Clone,
524 Debug,
525 PartialEq,
526 Eq,
527 Hash,
528 Serialize,
529 Deserialize,
530 Ord,
531 PartialOrd
532)]
533pub struct TlsAddr {
534 pub hostname: Hostname,
536 pub port: Port,
538}
539
540impl TlsAddr {
541 pub fn new(hostname: impl Into<Hostname>, port: Port) -> Self {
543 Self {
544 hostname: normalize_host(&hostname.into()),
545 port,
546 }
547 }
548
549 pub fn port(&self) -> Port {
551 self.port
552 }
553
554 pub fn hostname(&self) -> &str {
556 &self.hostname
557 }
558}
559
560impl FromStr for TlsAddr {
561 type Err = anyhow::Error;
562
563 fn from_str(addr: &str) -> Result<Self, Self::Err> {
564 let (hostname, port_str) = addr
565 .rsplit_once(':')
566 .ok_or_else(|| anyhow::anyhow!("invalid TLS address: {}", addr))?;
567 let port = port_str
568 .parse()
569 .map_err(|_| anyhow::anyhow!("invalid TLS address port: {}", port_str))?;
570 Ok(Self::new(hostname, port))
571 }
572}
573
574impl fmt::Display for TlsAddr {
575 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
576 write!(f, "{}:{}", self.hostname, self.port)
577 }
578}
579
580#[derive(
582 Clone,
583 Debug,
584 PartialEq,
585 Eq,
586 Hash,
587 Serialize,
588 Deserialize,
589 typeuri::Named
590)]
591pub enum ChannelTransport {
592 Tcp(TcpMode),
594
595 MetaTls(TlsMode),
597
598 Tls,
600
601 Quic,
603
604 MetaQuic(TlsMode),
606
607 Local,
610
611 Unix,
613}
614
615impl fmt::Display for ChannelTransport {
616 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
617 match self {
618 Self::Tcp(mode) => write!(f, "tcp({:?})", mode),
619 Self::MetaTls(mode) => write!(f, "metatls({:?})", mode),
620 Self::Tls => write!(f, "tls"),
621 Self::Quic => write!(f, "quic"),
622 Self::MetaQuic(mode) => write!(f, "metaquic({:?})", mode),
623 Self::Local => write!(f, "local"),
624 Self::Unix => write!(f, "unix"),
625 }
626 }
627}
628
629impl FromStr for ChannelTransport {
630 type Err = anyhow::Error;
631
632 fn from_str(s: &str) -> Result<Self, Self::Err> {
633 match s {
634 "tcp" => Ok(ChannelTransport::Tcp(TcpMode::Hostname)),
636 s if s.starts_with("tcp(") => {
637 let inner = &s["tcp(".len()..s.len() - 1];
638 let mode = inner.parse()?;
639 Ok(ChannelTransport::Tcp(mode))
640 }
641 "local" => Ok(ChannelTransport::Local),
642 "unix" => Ok(ChannelTransport::Unix),
643 "tls" => Ok(ChannelTransport::Tls),
644 "quic" => Ok(ChannelTransport::Quic),
645 s if s.starts_with("metatls(") && s.ends_with(")") => {
646 let inner = &s["metatls(".len()..s.len() - 1];
647 let mode = inner.parse()?;
648 Ok(ChannelTransport::MetaTls(mode))
649 }
650 s if s.starts_with("metaquic(") && s.ends_with(")") => {
651 let inner = &s["metaquic(".len()..s.len() - 1];
652 let mode = inner.parse()?;
653 Ok(ChannelTransport::MetaQuic(mode))
654 }
655 unknown => Err(anyhow::anyhow!("unknown channel transport: {}", unknown)),
656 }
657 }
658}
659
660impl ChannelTransport {
661 pub fn all() -> [ChannelTransport; 3] {
663 [
664 ChannelTransport::Tcp(TcpMode::Hostname),
667 ChannelTransport::Local,
668 ChannelTransport::Unix,
669 ]
672 }
673
674 pub fn any(&self) -> ChannelAddr {
676 ChannelAddr::any(self.clone())
677 }
678
679 pub fn is_remote(&self) -> bool {
681 match self {
682 ChannelTransport::Tcp(_) => true,
683 ChannelTransport::MetaTls(_) => true,
684 ChannelTransport::Tls => true,
685 ChannelTransport::Quic => true,
686 ChannelTransport::MetaQuic(_) => true,
687 ChannelTransport::Local => false,
688 ChannelTransport::Unix => false,
689 }
690 }
691
692 pub fn is_net(&self) -> bool {
697 match self {
698 ChannelTransport::Tcp(_) => true,
699 ChannelTransport::MetaTls(_) => true,
700 ChannelTransport::Tls => true,
701 ChannelTransport::Quic => false,
702 ChannelTransport::MetaQuic(_) => false,
703 ChannelTransport::Unix => true,
704 ChannelTransport::Local => false,
705 }
706 }
707
708 pub fn is_tls(&self) -> bool {
710 matches!(self, ChannelTransport::Tls | ChannelTransport::MetaTls(_))
711 }
712
713 pub fn supports_duplex(&self) -> bool {
720 match self {
721 ChannelTransport::Tcp(_) => true,
722 ChannelTransport::MetaTls(_) => true,
723 ChannelTransport::Tls => true,
724 ChannelTransport::Quic => false,
726 ChannelTransport::MetaQuic(_) => false,
727 ChannelTransport::Unix => true,
728 ChannelTransport::Local => true,
729 }
730 }
731}
732
733impl AttrValue for ChannelTransport {
734 fn display(&self) -> String {
735 self.to_string()
736 }
737
738 fn parse(s: &str) -> Result<Self, anyhow::Error> {
739 s.parse()
740 }
741}
742
743#[derive(
745 Clone,
746 Debug,
747 PartialEq,
748 Eq,
749 Hash,
750 Serialize,
751 Deserialize,
752 typeuri::Named
753)]
754pub enum BindSpec {
755 Any(ChannelTransport),
757
758 Addr(ChannelAddr),
760}
761
762impl BindSpec {
763 pub fn binding_addr(&self) -> ChannelAddr {
765 match self {
766 BindSpec::Any(transport) => ChannelAddr::any(transport.clone()),
767 BindSpec::Addr(addr) => addr.clone(),
768 }
769 }
770}
771
772impl From<ChannelTransport> for BindSpec {
773 fn from(transport: ChannelTransport) -> Self {
774 BindSpec::Any(transport)
775 }
776}
777
778impl From<ChannelAddr> for BindSpec {
779 fn from(addr: ChannelAddr) -> Self {
780 BindSpec::Addr(addr)
781 }
782}
783
784impl fmt::Display for BindSpec {
785 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
786 match self {
787 Self::Any(transport) => write!(f, "{}", transport),
788 Self::Addr(addr) => write!(f, "{}", addr),
789 }
790 }
791}
792
793impl FromStr for BindSpec {
794 type Err = anyhow::Error;
795
796 fn from_str(s: &str) -> Result<Self, Self::Err> {
797 if let Ok(transport) = ChannelTransport::from_str(s) {
798 Ok(BindSpec::Any(transport))
799 } else if let Ok(addr) = ChannelAddr::from_zmq_url(s) {
800 Ok(BindSpec::Addr(addr))
801 } else if let Ok(addr) = ChannelAddr::from_str(s) {
802 Ok(BindSpec::Addr(addr))
803 } else {
804 Err(anyhow::anyhow!("invalid bind spec: {}", s))
805 }
806 }
807}
808
809impl AttrValue for BindSpec {
810 fn display(&self) -> String {
811 self.to_string()
812 }
813
814 fn parse(s: &str) -> Result<Self, anyhow::Error> {
815 Self::from_str(s)
816 }
817}
818
819pub type Hostname = String;
821
822pub type Port = u16;
824
825#[derive(
849 Clone,
850 Debug,
851 PartialEq,
852 Eq,
853 Ord,
854 PartialOrd,
855 Serialize,
856 Deserialize,
857 Hash,
858 typeuri::Named
859)]
860pub enum ChannelAddr {
861 Tcp(SocketAddr),
864
865 MetaTls(TlsAddr),
868
869 Tls(TlsAddr),
872
873 Quic(TlsAddr),
876
877 MetaQuic(TlsAddr),
880
881 Local(u64),
884
885 Unix(net::unix::SocketAddr),
888
889 Alias {
910 dial_to: Box<ChannelAddr>,
912 bind_to: Box<ChannelAddr>,
914 },
915}
916
917impl From<SocketAddr> for ChannelAddr {
918 fn from(value: SocketAddr) -> Self {
919 Self::Tcp(value)
920 }
921}
922
923impl From<net::unix::SocketAddr> for ChannelAddr {
924 fn from(value: net::unix::SocketAddr) -> Self {
925 Self::Unix(value)
926 }
927}
928
929impl From<std::os::unix::net::SocketAddr> for ChannelAddr {
930 fn from(value: std::os::unix::net::SocketAddr) -> Self {
931 Self::Unix(net::unix::SocketAddr::new(value))
932 }
933}
934
935impl From<tokio::net::unix::SocketAddr> for ChannelAddr {
936 fn from(value: tokio::net::unix::SocketAddr) -> Self {
937 std::os::unix::net::SocketAddr::from(value).into()
938 }
939}
940
941fn find_routable_address(addresses: &[IpAddr]) -> Option<IpAddr> {
943 addresses
944 .iter()
945 .find(|addr| match addr {
946 IpAddr::V6(v6) => !v6.is_unicast_link_local(),
947 IpAddr::V4(v4) => !v4.is_link_local(),
948 })
949 .cloned()
950}
951
952impl ChannelAddr {
953 pub fn any(transport: ChannelTransport) -> Self {
956 match transport {
957 ChannelTransport::Tcp(mode) => {
958 let ip = match mode {
959 TcpMode::Localhost => IpAddr::V6(Ipv6Addr::LOCALHOST),
960 TcpMode::Hostname => {
961 hostname::get()
962 .ok()
963 .and_then(|hostname| {
964 hostname.to_str().and_then(|hostname_str| {
966 dns_lookup::lookup_host(hostname_str)
967 .ok()
968 .and_then(|addresses| find_routable_address(&addresses))
969 })
970 })
971 .unwrap_or(IpAddr::V6(Ipv6Addr::LOCALHOST))
972 }
973 };
974 Self::Tcp(SocketAddr::new(ip, 0))
975 }
976 ChannelTransport::MetaTls(mode) => {
977 let host_address = match mode {
978 TlsMode::Hostname => hostname::get()
979 .ok()
980 .and_then(|hostname| hostname.to_str().map(|s| s.to_string()))
981 .unwrap_or("unknown_host".to_string()),
982 TlsMode::IpV6 => {
983 get_host_ipv6_address().expect("failed to retrieve ipv6 address")
984 }
985 };
986 Self::MetaTls(TlsAddr::new(host_address, 0))
987 }
988 ChannelTransport::MetaQuic(mode) => {
989 let host_address = match mode {
990 TlsMode::Hostname => hostname::get()
991 .ok()
992 .and_then(|hostname| hostname.to_str().map(|s| s.to_string()))
993 .unwrap_or("unknown_host".to_string()),
994 TlsMode::IpV6 => {
995 get_host_ipv6_address().expect("failed to retrieve ipv6 address")
996 }
997 };
998 Self::MetaQuic(TlsAddr::new(host_address, 0))
999 }
1000 ChannelTransport::Local => Self::Local(0),
1001 ChannelTransport::Tls => {
1002 let host_address = hostname::get()
1003 .ok()
1004 .and_then(|hostname| hostname.to_str().map(|s| s.to_string()))
1005 .unwrap_or("localhost".to_string());
1006 Self::Tls(TlsAddr::new(host_address, 0))
1007 }
1008 ChannelTransport::Quic => {
1009 let host_address = hostname::get()
1010 .ok()
1011 .and_then(|hostname| hostname.to_str().map(|s| s.to_string()))
1012 .unwrap_or("localhost".to_string());
1013 Self::Quic(TlsAddr::new(host_address, 0))
1014 }
1015 ChannelTransport::Unix => Self::Unix(net::unix::SocketAddr::from_str("").unwrap()),
1017 }
1018 }
1019
1020 pub fn transport(&self) -> ChannelTransport {
1022 match self {
1023 Self::Tcp(addr) => {
1024 if addr.ip().is_loopback() {
1025 ChannelTransport::Tcp(TcpMode::Localhost)
1026 } else {
1027 ChannelTransport::Tcp(TcpMode::Hostname)
1028 }
1029 }
1030 Self::MetaTls(addr) => match addr.hostname.parse::<IpAddr>() {
1031 Ok(IpAddr::V6(_)) => ChannelTransport::MetaTls(TlsMode::IpV6),
1032 Ok(IpAddr::V4(_)) => ChannelTransport::MetaTls(TlsMode::Hostname),
1033 Err(_) => ChannelTransport::MetaTls(TlsMode::Hostname),
1034 },
1035 Self::Tls(_) => ChannelTransport::Tls,
1036 Self::Quic(_) => ChannelTransport::Quic,
1037 Self::MetaQuic(addr) => match addr.hostname.parse::<IpAddr>() {
1038 Ok(IpAddr::V6(_)) => ChannelTransport::MetaQuic(TlsMode::IpV6),
1039 Ok(IpAddr::V4(_)) => ChannelTransport::MetaQuic(TlsMode::Hostname),
1040 Err(_) => ChannelTransport::MetaQuic(TlsMode::Hostname),
1041 },
1042 Self::Local(_) => ChannelTransport::Local,
1043 Self::Unix(_) => ChannelTransport::Unix,
1044 Self::Alias { bind_to, .. } => bind_to.transport(),
1047 }
1048 }
1049}
1050
1051#[cfg(fbcode_build)]
1052fn get_host_ipv6_address() -> anyhow::Result<String> {
1053 crate::meta::host_ip::host_ipv6_address()
1054}
1055
1056#[cfg(not(fbcode_build))]
1057fn get_host_ipv6_address() -> anyhow::Result<String> {
1058 Ok(local_ip_address::local_ipv6()?.to_string())
1059}
1060
1061impl fmt::Display for ChannelAddr {
1062 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1063 match self {
1064 Self::Tcp(addr) => write!(f, "tcp:{}", addr),
1065 Self::MetaTls(addr) => write!(f, "metatls:{}", addr),
1066 Self::Tls(addr) => write!(f, "tls:{}", addr),
1067 Self::Quic(addr) => write!(f, "quic:{}", addr),
1068 Self::MetaQuic(addr) => write!(f, "metaquic:{}", addr),
1069 Self::Local(index) => write!(f, "local:{}", index),
1070 Self::Unix(addr) => write!(f, "unix:{}", addr),
1071 Self::Alias { dial_to, bind_to } => {
1072 write!(f, "alias:dial_to={};bind_to={}", dial_to, bind_to)
1073 }
1074 }
1075 }
1076}
1077
1078impl FromStr for ChannelAddr {
1079 type Err = anyhow::Error;
1080
1081 fn from_str(addr: &str) -> Result<Self, Self::Err> {
1082 match addr.split_once('!').or_else(|| addr.split_once(':')) {
1083 Some(("local", rest)) => rest
1084 .parse::<u64>()
1085 .map(Self::Local)
1086 .map_err(anyhow::Error::from),
1087 Some(("tcp", rest)) => rest
1088 .parse::<SocketAddr>()
1089 .map(Self::Tcp)
1090 .map_err(anyhow::Error::from),
1091 Some(("metatls", rest)) => net::meta::parse(rest).map_err(|e| e.into()),
1092 Some(("tls", rest)) => net::tls::parse(rest).map_err(|e| e.into()),
1093 Some(("quic", rest)) => TlsAddr::from_str(rest).map(Self::Quic),
1094 Some(("metaquic", rest)) => TlsAddr::from_str(rest).map(Self::MetaQuic),
1095 Some(("unix", rest)) => Ok(Self::Unix(net::unix::SocketAddr::from_str(rest)?)),
1096 Some(("alias", _)) => Err(anyhow::anyhow!(
1097 "detect possible alias address, but we currently do not support \
1098 parsing alias' string representation since we only want to \
1099 support parsing its zmq url format."
1100 )),
1101 Some((r#type, _)) => Err(anyhow::anyhow!("no such channel type: {type}")),
1102 None => Err(anyhow::anyhow!("no channel type specified")),
1103 }
1104 }
1105}
1106
1107pub(crate) fn normalize_host(host: &str) -> String {
1110 let host_clean = host
1113 .strip_prefix('[')
1114 .and_then(|h| h.strip_suffix(']'))
1115 .unwrap_or(host);
1116
1117 if let Ok(ip_addr) = host_clean.parse::<IpAddr>() {
1118 ip_addr.to_string()
1119 } else {
1120 host.to_string()
1121 }
1122}
1123
1124impl ChannelAddr {
1125 pub fn into_dial_addr(self) -> Self {
1132 match self {
1133 Self::Alias { dial_to, .. } => (*dial_to).into_dial_addr(),
1134 addr => addr,
1135 }
1136 }
1137
1138 pub fn from_zmq_url(address: &str) -> Result<Self, anyhow::Error> {
1153 let (addr, _listener) = Self::from_zmq_url_with_listener(address)?;
1154 Ok(addr)
1155 }
1156
1157 pub fn from_zmq_url_with_listener(
1170 address: &str,
1171 ) -> Result<(Self, Option<std::net::TcpListener>), anyhow::Error> {
1172 if let Some(at_pos) = address
1175 .find('@')
1176 .filter(|&pos| address[..pos].starts_with("tcp://"))
1177 {
1178 let dial_to_str = &address[..at_pos];
1179 let bind_to_str = &address[at_pos + 1..];
1180
1181 if !dial_to_str.starts_with("tcp://") {
1183 return Err(anyhow::anyhow!(
1184 "alias format is only supported for TCP addresses, got dial_to: {}",
1185 dial_to_str
1186 ));
1187 }
1188 if !bind_to_str.starts_with("tcp://") {
1189 return Err(anyhow::anyhow!(
1190 "alias format is only supported for TCP addresses, got bind_to: {}",
1191 bind_to_str
1192 ));
1193 }
1194
1195 let dial_to = Self::from_zmq_url(dial_to_str)?;
1196 let bind_to = Self::from_zmq_url(bind_to_str)?;
1197
1198 return Ok((
1199 Self::Alias {
1200 dial_to: Box::new(dial_to),
1201 bind_to: Box::new(bind_to),
1202 },
1203 None,
1204 ));
1205 }
1206
1207 let (scheme, address) = address.split_once("://").ok_or_else(|| {
1209 anyhow::anyhow!("address must be in url form scheme://endppoint {}", address)
1210 })?;
1211
1212 match scheme {
1213 "tcp" => {
1214 let (host, port, listener) = Self::parse_host_port_or_fd(address)?;
1215 let socket_addr = if host == "*" {
1216 SocketAddr::new("::".parse().unwrap(), port)
1217 } else {
1218 Self::resolve_hostname_to_socket_addr(host, port)?
1219 };
1220 Ok((Self::Tcp(socket_addr), listener))
1221 }
1222 "inproc" => {
1223 let port = address.parse::<u64>().map_err(|_| {
1224 anyhow::anyhow!("inproc endpoint must be a valid port number: {}", address)
1225 })?;
1226 Ok((Self::Local(port), None))
1227 }
1228 "ipc" => Ok((Self::Unix(net::unix::SocketAddr::from_str(address)?), None)),
1229 "metatls" | "tls" | "quic" | "metaquic" => {
1230 let (host, port, listener) = Self::parse_host_port_or_fd(address)?;
1231 let hostname = if host == "*" {
1232 std::net::Ipv6Addr::UNSPECIFIED.to_string()
1233 } else {
1234 host.to_string()
1235 };
1236 let addr = match scheme {
1237 "metatls" => Self::MetaTls(TlsAddr::new(hostname, port)),
1238 "metaquic" => Self::MetaQuic(TlsAddr::new(hostname, port)),
1239 "quic" => Self::Quic(TlsAddr::new(hostname, port)),
1240 _ => Self::Tls(TlsAddr::new(hostname, port)),
1241 };
1242 Ok((addr, listener))
1243 }
1244 scheme => Err(anyhow::anyhow!("unsupported ZMQ scheme: {}", scheme)),
1245 }
1246 }
1247
1248 fn parse_host_port_or_fd(
1251 address: &str,
1252 ) -> Result<(&str, u16, Option<std::net::TcpListener>), anyhow::Error> {
1253 let (host, port_str) = address
1254 .rsplit_once(':')
1255 .ok_or_else(|| anyhow::anyhow!("invalid address format: {}", address))?;
1256
1257 if let Some(fd_str) = port_str.strip_prefix("fd") {
1258 let fd_num: RawFd = fd_str
1259 .parse()
1260 .map_err(|_| anyhow::anyhow!("invalid file descriptor number: {}", port_str))?;
1261 let borrowed = unsafe { std::os::unix::io::BorrowedFd::borrow_raw(fd_num) };
1265 nix::sys::socket::listen(&borrowed, nix::sys::socket::Backlog::new(128)?)?;
1266 let std_listener = unsafe { std::net::TcpListener::from_raw_fd(fd_num) };
1268 let local_addr = std_listener.local_addr()?;
1269 Ok((host, local_addr.port(), Some(std_listener)))
1270 } else {
1271 let port: u16 = port_str
1272 .parse()
1273 .map_err(|_| anyhow::anyhow!("invalid port: {}", port_str))?;
1274 Ok((host, port, None))
1275 }
1276 }
1277
1278 pub fn to_zmq_url(&self) -> String {
1280 match self {
1281 Self::Tcp(addr) => format!("tcp://{}", addr),
1282 Self::MetaTls(addr) => format!("metatls://{}:{}", addr.hostname, addr.port),
1283 Self::Tls(addr) => format!("tls://{}:{}", addr.hostname, addr.port),
1284 Self::Quic(addr) => format!("quic://{}:{}", addr.hostname, addr.port),
1285 Self::MetaQuic(addr) => format!("metaquic://{}:{}", addr.hostname, addr.port),
1286 Self::Local(index) => format!("inproc://{}", index),
1287 Self::Unix(addr) => format!("ipc://{}", addr),
1288 Self::Alias { dial_to, bind_to } => {
1289 format!("{}@{}", dial_to.to_zmq_url(), bind_to.to_zmq_url())
1290 }
1291 }
1292 }
1293
1294 fn resolve_hostname_to_socket_addr(host: &str, port: u16) -> Result<SocketAddr, anyhow::Error> {
1296 let host_clean = if host.starts_with('[') && host.ends_with(']') {
1298 &host[1..host.len() - 1]
1299 } else {
1300 host
1301 };
1302
1303 if let Ok(ip_addr) = host_clean.parse::<IpAddr>() {
1305 return Ok(SocketAddr::new(ip_addr, port));
1306 }
1307
1308 use std::net::ToSocketAddrs;
1310 let mut addrs = (host_clean, port)
1311 .to_socket_addrs()
1312 .map_err(|e| anyhow::anyhow!("failed to resolve hostname '{}': {}", host_clean, e))?;
1313
1314 addrs
1315 .next()
1316 .ok_or_else(|| anyhow::anyhow!("no addresses found for hostname '{}'", host_clean))
1317 }
1318}
1319
1320pub struct ChannelTx<M: RemoteMessage> {
1323 sender: mpsc::UnboundedSender<(M, CompletionSink<M>, Instant)>,
1324 dest: ChannelAddr,
1325 status: watch::Receiver<TxStatus>,
1326}
1327
1328impl<M: RemoteMessage> fmt::Debug for ChannelTx<M> {
1329 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1330 f.debug_struct("ChannelTx")
1331 .field("addr", &self.addr())
1332 .finish()
1333 }
1334}
1335
1336#[async_trait]
1337impl<M: RemoteMessage> Tx<M> for ChannelTx<M> {
1338 fn do_post(&self, message: M, completion: CompletionSink<M>) {
1339 tracing::trace!(
1340 name = "post",
1341 dest = %self.dest,
1342 "sending message"
1343 );
1344
1345 if let Err(mpsc::error::SendError((message, completion, _))) =
1346 self.sender.send((message, completion, Instant::now()))
1347 {
1348 let reason = self
1349 .status
1350 .borrow()
1351 .as_closed()
1352 .map(|r| SendErrorReason::Other(r.to_string()));
1353 completion.reject(SendError {
1354 error: ChannelError::Closed,
1355 message,
1356 reason,
1357 });
1358 }
1359 }
1360
1361 fn addr(&self) -> ChannelAddr {
1362 self.dest.clone()
1363 }
1364
1365 fn status(&self) -> &watch::Receiver<TxStatus> {
1366 &self.status
1367 }
1368}
1369
1370pub struct ChannelRx<M: RemoteMessage> {
1372 receiver: mpsc::Receiver<M>,
1373 dest: ChannelAddr,
1374 server: net::ServerHandle,
1375}
1376
1377impl<M: RemoteMessage> fmt::Debug for ChannelRx<M> {
1378 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1379 f.debug_struct("ChannelRx")
1380 .field("addr", &self.addr())
1381 .finish()
1382 }
1383}
1384
1385impl<M: RemoteMessage> ChannelRx<M> {
1386 fn stop(&self, trigger: &str) {
1388 self.server.stop(&format!(
1389 "ChannelRx {trigger}; channel address: {}",
1390 self.dest
1391 ));
1392 }
1393}
1394
1395#[async_trait]
1396impl<M: RemoteMessage> Rx<M> for ChannelRx<M> {
1397 async fn recv(&mut self) -> Result<M, ChannelError> {
1398 tracing::trace!(
1399 name = "recv",
1400 dest = %self.dest,
1401 "receiving message"
1402 );
1403 self.receiver.recv().await.ok_or(ChannelError::Closed)
1404 }
1405
1406 fn addr(&self) -> ChannelAddr {
1407 self.dest.clone()
1408 }
1409
1410 async fn join(mut self) {
1413 self.stop("joined");
1414 let _ = (&mut self.server).await;
1415 }
1417}
1418
1419impl<M: RemoteMessage> Drop for ChannelRx<M> {
1420 fn drop(&mut self) {
1421 self.stop("dropped");
1422 }
1423}
1424
1425#[allow(clippy::result_large_err)] #[track_caller]
1430pub fn dial<M: RemoteMessage>(addr: ChannelAddr) -> Result<ChannelTx<M>, ChannelError> {
1431 let addr = addr.into_dial_addr();
1432 tracing::debug!(name = "dial", caller = %Location::caller(), %addr, "dialing channel {}", addr);
1433 Ok(net::spawn::<M>(net::link(
1434 addr,
1435 net::SessionId::random(),
1436 0,
1437 net::ProtocolKind::Simplex,
1438 )?))
1439}
1440
1441pub mod unordered {
1443 use super::*;
1444
1445 #[track_caller]
1491 pub fn dial<M: RemoteMessage>(
1492 addr: ChannelAddr,
1493 num_streams: usize,
1494 ) -> Result<ChannelTx<M>, ChannelError> {
1495 assert!(num_streams > 0);
1496 let addr = addr.into_dial_addr();
1497 let session_id = net::SessionId::random();
1498 let links: Vec<net::NetLink> = (1..=num_streams)
1499 .map(|i| {
1500 net::link(
1501 addr.clone(),
1502 session_id,
1503 i as u8,
1504 net::ProtocolKind::Simplex,
1505 )
1506 })
1507 .collect::<Result<_, _>>()?;
1508 Ok(net::spawn_unordered::<M>(links))
1509 }
1510
1511 #[track_caller]
1517 pub fn serve<M: RemoteMessage>(
1518 addr: ChannelAddr,
1519 ) -> Result<(ChannelAddr, ChannelRx<M>), ChannelError> {
1520 super::serve(addr)
1521 }
1522
1523 #[track_caller]
1528 pub fn serve_with_listener<M: RemoteMessage>(
1529 addr: ChannelAddr,
1530 listener: Option<std::net::TcpListener>,
1531 ) -> Result<(ChannelAddr, ChannelRx<M>), ChannelError> {
1532 super::serve_with_listener(addr, listener)
1533 }
1534}
1535
1536#[track_caller]
1539pub fn serve<M: RemoteMessage>(
1540 addr: ChannelAddr,
1541) -> Result<(ChannelAddr, ChannelRx<M>), ChannelError> {
1542 serve_with_listener(addr, None)
1543}
1544
1545#[track_caller]
1549pub fn serve_with_listener<M: RemoteMessage>(
1550 addr: ChannelAddr,
1551 listener: Option<std::net::TcpListener>,
1552) -> Result<(ChannelAddr, ChannelRx<M>), ChannelError> {
1553 let caller = Location::caller();
1554 serve_inner(addr, listener).map(|(addr, rx)| {
1555 tracing::debug!(
1556 name = "serve",
1557 %addr,
1558 %caller,
1559 );
1560 (addr, rx)
1561 })
1562}
1563
1564#[track_caller]
1576pub fn serve_mux<M: RemoteMessage, In: RemoteMessage, Out: RemoteMessage>(
1577 addr: ChannelAddr,
1578 prebound_listener: Option<std::net::TcpListener>,
1579) -> Result<MuxServer<M, In, Out>, ChannelError> {
1580 if !addr.transport().is_net() {
1581 return Err(ChannelError::InvalidAddress(format!(
1582 "serve_mux requires a net transport; got {}",
1583 addr
1584 )));
1585 }
1586 let parts = net::mux::serve::<M, In, Out>(addr, prebound_listener)?;
1587 Ok(MuxServer {
1588 addr: parts.addr,
1589 simplex: parts.simplex,
1590 duplex: parts.duplex,
1591 shutdown: MuxShutdown {
1592 join_handle: parts.join_handle,
1593 cancel: parts.cancel,
1594 },
1595 })
1596}
1597
1598pub struct MuxServer<M: RemoteMessage, In: RemoteMessage, Out: RemoteMessage> {
1616 addr: ChannelAddr,
1617 simplex: ChannelRx<M>,
1618 duplex: net::duplex::DuplexServer<In, Out>,
1619 shutdown: MuxShutdown,
1620}
1621
1622impl<M: RemoteMessage, In: RemoteMessage, Out: RemoteMessage> MuxServer<M, In, Out> {
1623 pub fn addr(&self) -> &ChannelAddr {
1625 &self.addr
1626 }
1627
1628 pub fn simplex_mut(&mut self) -> &mut ChannelRx<M> {
1630 &mut self.simplex
1631 }
1632
1633 pub fn duplex_mut(&mut self) -> &mut net::duplex::DuplexServer<In, Out> {
1635 &mut self.duplex
1636 }
1637
1638 pub fn stop(&self, reason: &str) {
1640 self.shutdown.stop(reason);
1641 }
1642
1643 pub fn split(
1650 self,
1651 ) -> (
1652 ChannelAddr,
1653 ChannelRx<M>,
1654 net::duplex::DuplexServer<In, Out>,
1655 MuxShutdown,
1656 ) {
1657 (self.addr, self.simplex, self.duplex, self.shutdown)
1658 }
1659
1660 pub fn serve<SH, DH, DF>(
1687 self,
1688 simplex_handler: SH,
1689 duplex_handler: DH,
1690 ) -> crate::mailbox::MailboxServerHandle
1691 where
1692 SH: FnOnce(ChannelRx<M>) -> crate::mailbox::MailboxServerHandle,
1693 DH: FnOnce(net::duplex::DuplexServer<In, Out>, tokio::sync::watch::Receiver<bool>) -> DF,
1694 DF: std::future::Future<Output = ()> + Send + 'static,
1695 {
1696 let (stopped_tx, mut stopped_rx) = tokio::sync::watch::channel(false);
1697 let duplex_stop = stopped_rx.clone();
1698 let simplex_handle = simplex_handler(self.simplex);
1699 let duplex_task = tokio::spawn(duplex_handler(self.duplex, duplex_stop));
1700 let shutdown = self.shutdown;
1701 let join_handle = tokio::spawn(async move {
1702 let ok = stopped_rx.wait_for(|stopped| *stopped).await.is_ok();
1708 if !ok {
1709 std::future::pending::<()>().await;
1710 }
1711 const REASON: &str = "MuxServer shutdown";
1712 let _ = duplex_task.await;
1719 let _ = simplex_handle.await;
1727 shutdown.stop(REASON);
1731 let _ = shutdown.await;
1732 Ok::<(), crate::mailbox::MailboxServerError>(())
1733 });
1734 crate::mailbox::MailboxServerHandle::from_parts(join_handle, stopped_tx)
1735 }
1736}
1737
1738pub struct MuxShutdown {
1748 join_handle: tokio::task::JoinHandle<Result<(), net::ServerError>>,
1749 cancel: CancellationToken,
1750}
1751
1752impl MuxShutdown {
1753 pub fn stop(&self, reason: &str) {
1757 tracing::info!(
1758 name = "MuxServerStatus",
1759 status = "Stop::Sent",
1760 reason,
1761 "muxed frontend stop signalled",
1762 );
1763 self.cancel.cancel();
1764 }
1765
1766 pub async fn cancelled(&self) {
1770 self.cancel.cancelled().await;
1771 }
1772}
1773
1774impl std::future::Future for MuxShutdown {
1775 type Output =
1776 <tokio::task::JoinHandle<Result<(), net::ServerError>> as std::future::Future>::Output;
1777
1778 fn poll(
1779 mut self: std::pin::Pin<&mut Self>,
1780 cx: &mut std::task::Context<'_>,
1781 ) -> std::task::Poll<Self::Output> {
1782 std::pin::Pin::new(&mut self.join_handle).poll(cx)
1787 }
1788}
1789
1790impl Drop for MuxShutdown {
1791 fn drop(&mut self) {
1792 self.cancel.cancel();
1793 }
1794}
1795
1796fn serve_inner<M: RemoteMessage>(
1797 addr: ChannelAddr,
1798 listener: Option<std::net::TcpListener>,
1799) -> Result<(ChannelAddr, ChannelRx<M>), ChannelError> {
1800 match addr {
1801 ChannelAddr::Unix(_) => {
1802 assert!(
1803 listener.is_none(),
1804 "pre-opened listener not supported for Unix transport"
1805 );
1806 let (addr, rx) = net::server::serve::<M>(addr, listener)?;
1807 Ok((addr, rx))
1808 }
1809 ChannelAddr::Tcp(_)
1810 | ChannelAddr::Local(_)
1811 | ChannelAddr::Tls(_)
1812 | ChannelAddr::MetaTls(_)
1813 | ChannelAddr::Quic(_)
1814 | ChannelAddr::MetaQuic(_)
1815 | ChannelAddr::Alias { .. } => {
1819 let (addr, rx) = net::server::serve::<M>(addr, listener)?;
1820 Ok((addr, rx))
1821 }
1822 }
1823}
1824
1825pub fn serve_local<M: RemoteMessage>() -> (ChannelAddr, ChannelRx<M>) {
1828 serve::<M>(ChannelAddr::Local(0)).expect("fresh local stream port must bind")
1829}
1830
1831pub fn reserve_local_addr() -> ChannelAddr {
1844 ChannelAddr::Local(local::reserve())
1845}
1846
1847#[cfg(test)]
1848mod tests {
1849 use std::assert_matches;
1850 use std::collections::HashSet;
1851 use std::net::IpAddr;
1852 use std::net::Ipv4Addr;
1853 use std::net::Ipv6Addr;
1854 use std::time::Duration;
1855
1856 use rand::RngExt as _;
1857 use rand::distr::Uniform;
1858 use tokio::task::JoinSet;
1859
1860 use super::net::*;
1861 use super::*;
1862 #[test]
1863 fn test_channel_addr() {
1864 let cases_ok = vec![
1865 (
1866 "tcp<DELIM>[::1]:1234",
1867 ChannelAddr::Tcp(SocketAddr::new(
1868 IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)),
1869 1234,
1870 )),
1871 ),
1872 (
1873 "tcp<DELIM>127.0.0.1:8080",
1874 ChannelAddr::Tcp(SocketAddr::new(
1875 IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
1876 8080,
1877 )),
1878 ),
1879 (
1880 "quic<DELIM>example.com:443",
1881 ChannelAddr::Quic(TlsAddr::new("example.com", 443)),
1882 ),
1883 (
1884 "metaquic<DELIM>example.com:443",
1885 ChannelAddr::MetaQuic(TlsAddr::new("example.com", 443)),
1886 ),
1887 #[cfg(target_os = "linux")]
1888 ("local<DELIM>123", ChannelAddr::Local(123)),
1889 (
1890 "unix<DELIM>@yolo",
1891 ChannelAddr::Unix(
1892 unix::SocketAddr::from_abstract_name("yolo")
1893 .expect("can't make socket from abstract name"),
1894 ),
1895 ),
1896 (
1897 "unix<DELIM>/cool/socket-path",
1898 ChannelAddr::Unix(
1899 unix::SocketAddr::from_pathname("/cool/socket-path")
1900 .expect("can't make socket from path"),
1901 ),
1902 ),
1903 ];
1904
1905 for (raw, parsed) in cases_ok {
1906 for delim in ["!", ":"] {
1907 let raw = raw.replace("<DELIM>", delim);
1908 assert_eq!(raw.parse::<ChannelAddr>().unwrap(), parsed);
1909 }
1910 }
1911
1912 let cases_err = vec![
1913 ("tcp:abcdef..123124", "invalid socket address syntax"),
1914 ("xxx:foo", "no such channel type: xxx"),
1915 ("127.0.0.1", "no channel type specified"),
1916 ("local:abc", "invalid digit found in string"),
1917 ];
1918
1919 for (raw, error) in cases_err {
1920 let Err(err) = raw.parse::<ChannelAddr>() else {
1921 panic!("expected error parsing: {}", &raw)
1922 };
1923 assert_eq!(format!("{}", err), error);
1924 }
1925 }
1926
1927 #[test]
1928 fn test_zmq_style_channel_addr() {
1929 assert_eq!(
1931 ChannelAddr::from_zmq_url("tcp://127.0.0.1:8080").unwrap(),
1932 ChannelAddr::Tcp("127.0.0.1:8080".parse().unwrap())
1933 );
1934
1935 assert_eq!(
1937 ChannelAddr::from_zmq_url("tcp://*:5555").unwrap(),
1938 ChannelAddr::Tcp("[::]:5555".parse().unwrap())
1939 );
1940
1941 assert_eq!(
1943 ChannelAddr::from_zmq_url("inproc://12345").unwrap(),
1944 ChannelAddr::Local(12345)
1945 );
1946
1947 assert_eq!(
1949 ChannelAddr::from_zmq_url("ipc:///tmp/my-socket").unwrap(),
1950 ChannelAddr::Unix(unix::SocketAddr::from_pathname("/tmp/my-socket").unwrap())
1951 );
1952
1953 assert_eq!(
1955 ChannelAddr::from_zmq_url("metatls://example.com:443").unwrap(),
1956 ChannelAddr::MetaTls(TlsAddr::new("example.com", 443))
1957 );
1958
1959 assert_eq!(
1961 ChannelAddr::from_zmq_url("metatls://192.168.1.1:443").unwrap(),
1962 ChannelAddr::MetaTls(TlsAddr::new("192.168.1.1", 443))
1963 );
1964
1965 assert_eq!(
1967 ChannelAddr::from_zmq_url("quic://example.com:443").unwrap(),
1968 ChannelAddr::Quic(TlsAddr::new("example.com", 443))
1969 );
1970
1971 assert_eq!(
1973 ChannelAddr::from_zmq_url("quic://*:8443").unwrap(),
1974 ChannelAddr::Quic(TlsAddr::new("::", 8443))
1975 );
1976
1977 assert_eq!(
1979 ChannelAddr::from_zmq_url("metaquic://example.com:443").unwrap(),
1980 ChannelAddr::MetaQuic(TlsAddr::new("example.com", 443))
1981 );
1982
1983 assert_eq!(
1985 ChannelAddr::from_zmq_url("metaquic://*:8443").unwrap(),
1986 ChannelAddr::MetaQuic(TlsAddr::new("::", 8443))
1987 );
1988
1989 assert_eq!(
1991 ChannelAddr::from_zmq_url("metatls://*:8443").unwrap(),
1992 ChannelAddr::MetaTls(TlsAddr::new("::", 8443))
1993 );
1994
1995 let tcp_hostname_result = ChannelAddr::from_zmq_url("tcp://localhost:8080");
1999 assert!(tcp_hostname_result.is_ok());
2000
2001 assert_eq!(
2003 ChannelAddr::from_zmq_url("tcp://[::1]:1234").unwrap(),
2004 ChannelAddr::Tcp("[::1]:1234".parse().unwrap())
2005 );
2006
2007 assert!(ChannelAddr::from_zmq_url("invalid://scheme").is_err());
2009 assert!(ChannelAddr::from_zmq_url("tcp://invalid-port").is_err());
2010 assert!(ChannelAddr::from_zmq_url("metatls://no-port").is_err());
2011 assert!(ChannelAddr::from_zmq_url("inproc://not-a-number").is_err());
2012
2013 assert_eq!(
2015 ChannelAddr::from_zmq_url("metatls://2a03:83e4:5000:c000:56d7:00cf:75ce:144a:443")
2016 .unwrap(),
2017 ChannelAddr::MetaTls(TlsAddr::new("2a03:83e4:5000:c000:56d7:cf:75ce:144a", 443))
2018 );
2019
2020 assert_eq!(
2022 ChannelAddr::from_zmq_url("metatls://2a03:83e4:5000:c000:56d7:00cf:75ce:144a:443")
2023 .unwrap(),
2024 ChannelAddr::from_zmq_url("metatls://2a03:83e4:5000:c000:56d7:cf:75ce:144a:443")
2025 .unwrap(),
2026 );
2027
2028 assert_eq!(
2030 ChannelAddr::from_zmq_url("metatls://[::1]:443").unwrap(),
2031 ChannelAddr::MetaTls(TlsAddr::new("::1", 443))
2032 );
2033
2034 assert_eq!(
2036 ChannelAddr::from_zmq_url("tls://2a03:83e4:5000:c000:56d7:00cf:75ce:144a:443").unwrap(),
2037 ChannelAddr::Tls(TlsAddr::new("2a03:83e4:5000:c000:56d7:cf:75ce:144a", 443))
2038 );
2039 assert_eq!(
2040 ChannelAddr::from_zmq_url("tls://2a03:83e4:5000:c000:56d7:00cf:75ce:144a:443").unwrap(),
2041 ChannelAddr::from_zmq_url("tls://2a03:83e4:5000:c000:56d7:cf:75ce:144a:443").unwrap(),
2042 );
2043 assert_eq!(
2044 ChannelAddr::from_zmq_url("tls://[::1]:443").unwrap(),
2045 ChannelAddr::Tls(TlsAddr::new("::1", 443))
2046 );
2047 }
2048
2049 #[tokio::test]
2050 async fn test_reserved_local_addr_can_be_served() {
2051 let addr = reserve_local_addr();
2052 assert!(dial::<u64>(addr.clone()).is_err());
2053
2054 let (bound_addr, mut rx) = serve::<u64>(addr.clone()).unwrap();
2055 assert_eq!(bound_addr, addr);
2056
2057 let tx = dial::<u64>(addr.clone()).unwrap();
2058 tx.post(123);
2059 assert_eq!(rx.recv().await.unwrap(), 123);
2060 rx.join().await;
2061
2062 let (rebound_addr, _rx) = serve::<u64>(addr.clone()).unwrap();
2063 assert_eq!(rebound_addr, addr);
2064 }
2065
2066 #[test]
2067 fn test_normalize_host() {
2068 assert_eq!(normalize_host("192.168.1.1"), "192.168.1.1");
2070
2071 assert_eq!(normalize_host("example.com"), "example.com");
2073
2074 assert_eq!(
2076 normalize_host("2a03:83e4:5000:c000:56d7:00cf:75ce:144a"),
2077 "2a03:83e4:5000:c000:56d7:cf:75ce:144a"
2078 );
2079
2080 assert_eq!(normalize_host("[::1]"), "::1");
2082
2083 assert!("[::1]".parse::<IpAddr>().is_err());
2087 }
2088
2089 #[test]
2090 fn test_zmq_style_alias_channel_addr() {
2091 let alias_addr = ChannelAddr::from_zmq_url("tcp://127.0.0.1:9000@tcp://[::]:8800").unwrap();
2097 match alias_addr {
2098 ChannelAddr::Alias { dial_to, bind_to } => {
2099 assert_eq!(
2100 *dial_to,
2101 ChannelAddr::Tcp("127.0.0.1:9000".parse().unwrap())
2102 );
2103 assert_eq!(*bind_to, ChannelAddr::Tcp("[::]:8800".parse().unwrap()));
2104 }
2105 _ => panic!("Expected Alias"),
2106 }
2107
2108 let non_alias = ChannelAddr::from_zmq_url("metatls://example.com:443@tcp://127.0.0.1:8080");
2111 assert!(
2112 !matches!(non_alias, Ok(ChannelAddr::Alias { .. })),
2113 "non-tcp left side must not produce Alias"
2114 );
2115
2116 assert!(
2118 ChannelAddr::from_zmq_url("tcp://127.0.0.1:8080@metatls://example.com:443").is_err()
2119 );
2120
2121 assert!(ChannelAddr::from_zmq_url("invalid://scheme@tcp://127.0.0.1:8080").is_err());
2123
2124 assert!(ChannelAddr::from_zmq_url("tcp://127.0.0.1:8080@invalid://scheme").is_err());
2126
2127 assert!(ChannelAddr::from_zmq_url("tcp://host@tcp://127.0.0.1:8080").is_err());
2129
2130 assert!(ChannelAddr::from_zmq_url("tcp://127.0.0.1:8080@tcp://example.com").is_err());
2132 }
2133
2134 #[tokio::test]
2135 async fn test_multiple_connections() {
2136 for addr in ChannelTransport::all().map(ChannelAddr::any) {
2137 let (listen_addr, mut rx) = crate::channel::serve::<u64>(addr).unwrap();
2138
2139 let mut sends: JoinSet<()> = JoinSet::new();
2140 for message in 0u64..100u64 {
2141 let addr = listen_addr.clone();
2142 sends.spawn(async move {
2143 let tx = dial::<u64>(addr).unwrap();
2144 tx.post(message);
2145 });
2146 }
2147
2148 let mut received: HashSet<u64> = HashSet::new();
2149 while received.len() < 100 {
2150 received.insert(rx.recv().await.unwrap());
2151 }
2152
2153 for message in 0u64..100u64 {
2154 assert!(received.contains(&message));
2155 }
2156
2157 loop {
2158 match sends.join_next().await {
2159 Some(Ok(())) => (),
2160 Some(Err(err)) => panic!("{}", err),
2161 None => break,
2162 }
2163 }
2164 }
2165 }
2166
2167 #[tokio::test]
2168 async fn test_server_close() {
2169 for addr in ChannelTransport::all().map(ChannelAddr::any) {
2170 if net::is_net_addr(&addr) {
2171 continue;
2174 }
2175
2176 let (listen_addr, rx) = crate::channel::serve::<u64>(addr).unwrap();
2177
2178 let tx = dial::<u64>(listen_addr).unwrap();
2179 tx.post(123);
2180 drop(rx);
2181
2182 let start = tokio::time::Instant::now();
2187
2188 let result = loop {
2189 let result = tx.try_post(123).await;
2190
2191 if result.is_err() || start.elapsed() > Duration::from_secs(10) {
2192 break result;
2193 }
2194 };
2195 assert_matches!(
2196 result,
2197 Err(SendError {
2198 error: ChannelError::Closed,
2199 message: 123,
2200 reason: None
2201 })
2202 );
2203 }
2204 }
2205
2206 fn addrs() -> Vec<ChannelAddr> {
2207 let rng = rand::rng();
2208 let uniform = Uniform::new_inclusive('a', 'z').unwrap();
2209 vec![
2210 "tcp:[::1]:0".parse().unwrap(),
2211 "local:0".parse().unwrap(),
2212 #[cfg(target_os = "linux")]
2213 "unix:".parse().unwrap(),
2214 #[cfg(target_os = "linux")]
2215 format!(
2216 "unix:@{}",
2217 rng.sample_iter(uniform).take(10).collect::<String>()
2218 )
2219 .parse()
2220 .unwrap(),
2221 ]
2222 }
2223
2224 #[test]
2225 fn test_bind_spec_from_str() {
2226 assert_eq!(
2228 BindSpec::from_str("tcp").unwrap(),
2229 BindSpec::Any(ChannelTransport::Tcp(TcpMode::Hostname))
2230 );
2231 assert_eq!(
2232 BindSpec::from_str("metatls(Hostname)").unwrap(),
2233 BindSpec::Any(ChannelTransport::MetaTls(TlsMode::Hostname))
2234 );
2235
2236 assert_eq!(
2238 BindSpec::from_str("tcp:127.0.0.1:8080").unwrap(),
2239 BindSpec::Addr(ChannelAddr::Tcp("127.0.0.1:8080".parse().unwrap()))
2240 );
2241
2242 assert_eq!(
2244 BindSpec::from_str("tcp://127.0.0.1:9000").unwrap(),
2245 BindSpec::Addr(ChannelAddr::Tcp("127.0.0.1:9000".parse().unwrap()))
2246 );
2247 assert_eq!(
2248 BindSpec::from_str("tcp://127.0.0.1:9000@tcp://[::1]:7200").unwrap(),
2249 BindSpec::Addr(
2250 ChannelAddr::from_zmq_url("tcp://127.0.0.1:9000@tcp://[::1]:7200").unwrap()
2251 )
2252 );
2253
2254 assert!(BindSpec::from_str("invalid_spec").is_err());
2256 assert!(BindSpec::from_str("unknown://scheme").is_err());
2257 assert!(BindSpec::from_str("").is_err());
2258 }
2259
2260 #[tokio::test]
2261 #[cfg_attr(not(fbcode_build), ignore)]
2263 async fn test_dial_serve() {
2264 for addr in addrs() {
2265 let (listen_addr, mut rx) = crate::channel::serve::<i32>(addr).unwrap();
2266 let tx = crate::channel::dial(listen_addr).unwrap();
2267 tx.post(123);
2268 assert_eq!(rx.recv().await.unwrap(), 123);
2269 }
2270 }
2271
2272 #[tokio::test]
2273 #[cfg_attr(not(fbcode_build), ignore)]
2275 async fn test_serve_alias_advertises_dial_to() {
2276 let probe = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
2279 let port = probe.local_addr().unwrap().port();
2280 drop(probe);
2281
2282 let alias =
2286 ChannelAddr::from_zmq_url(&format!("tcp://127.0.0.1:{port}@tcp://0.0.0.0:{port}"))
2287 .unwrap();
2288 assert_matches!(alias, ChannelAddr::Alias { .. });
2289
2290 let (listen_addr, mut rx) = crate::channel::serve::<i32>(alias).unwrap();
2291
2292 assert_eq!(
2298 listen_addr,
2299 ChannelAddr::Tcp(format!("127.0.0.1:{port}").parse().unwrap()),
2300 "serving an alias must advertise dial_to, not the alias itself"
2301 );
2302
2303 let tx = crate::channel::dial(listen_addr).unwrap();
2304 tx.post(123);
2305 assert_eq!(rx.recv().await.unwrap(), 123);
2306 }
2307
2308 #[tokio::test]
2309 #[cfg_attr(not(fbcode_build), ignore)]
2311 async fn test_send() {
2312 let config = hyperactor_config::global::lock();
2313
2314 let _guard1 = config.override_key(
2316 crate::config::MESSAGE_DELIVERY_TIMEOUT,
2317 Duration::from_secs(1),
2318 );
2319 let _guard2 = config.override_key(crate::config::MESSAGE_ACK_EVERY_N_MESSAGES, 1);
2320 for addr in addrs() {
2321 let (listen_addr, mut rx) = crate::channel::serve::<i32>(addr).unwrap();
2322 let tx = crate::channel::dial(listen_addr).unwrap();
2323 tx.send(123).await.unwrap();
2324 assert_eq!(rx.recv().await.unwrap(), 123);
2325
2326 drop(rx);
2327 assert_matches!(
2328 tx.send(123).await.unwrap_err(),
2329 SendError {
2330 error: ChannelError::Closed,
2331 message: 123,
2332 ..
2333 }
2334 );
2335 }
2336 }
2337
2338 #[test]
2339 fn test_find_routable_address_skips_link_local_ipv6() {
2340 let link_local_v6: IpAddr = "fe80::1".parse().unwrap();
2341 let routable_v6: IpAddr = "2001:db8::1".parse().unwrap();
2342 let addrs = vec![link_local_v6, routable_v6];
2343 assert_eq!(find_routable_address(&addrs), Some(routable_v6));
2344 }
2345
2346 #[test]
2347 fn test_find_routable_address_skips_link_local_ipv4() {
2348 let link_local_v4: IpAddr = "169.254.1.1".parse().unwrap();
2349 let routable_v4: IpAddr = "192.168.1.1".parse().unwrap();
2350 let addrs = vec![link_local_v4, routable_v4];
2351 assert_eq!(find_routable_address(&addrs), Some(routable_v4));
2352 }
2353
2354 #[test]
2355 fn test_find_routable_address_returns_none_when_all_link_local() {
2356 let link_local_v6: IpAddr = "fe80::1".parse().unwrap();
2357 let link_local_v4: IpAddr = "169.254.1.1".parse().unwrap();
2358 let addrs = vec![link_local_v6, link_local_v4];
2359 assert_eq!(find_routable_address(&addrs), None);
2360 }
2361
2362 #[test]
2363 fn test_find_routable_address_mixed() {
2364 let link_local_v6: IpAddr = "fe80::1".parse().unwrap();
2365 let link_local_v4: IpAddr = "169.254.0.1".parse().unwrap();
2366 let routable_v4: IpAddr = "10.0.0.1".parse().unwrap();
2367 let routable_v6: IpAddr = "2001:db8::2".parse().unwrap();
2368
2369 let addrs = vec![link_local_v6, link_local_v4, routable_v4, routable_v6];
2371 assert_eq!(find_routable_address(&addrs), Some(routable_v4));
2372 }
2373}