1#![deny(unreachable_pub)]
191#![deny(rustdoc::broken_intra_doc_links)]
192#![deny(rustdoc::private_intra_doc_links)]
193#![deny(rustdoc::invalid_codeblock_attributes)]
194#![deny(rustdoc::invalid_rust_codeblocks)]
195#![cfg_attr(docsrs, feature(doc_cfg))]
196
197use thiserror::Error;
198
199use futures_util::stream::Stream;
200use tokio::io::AsyncWriteExt;
201use tokio::sync::oneshot;
202use tracing::{debug, error};
203
204use core::fmt;
205use std::collections::HashMap;
206use std::collections::VecDeque;
207use std::fmt::Display;
208use std::future::Future;
209use std::iter;
210use std::mem;
211use std::net::SocketAddr;
212use std::option;
213use std::pin::Pin;
214use std::slice;
215use std::str::{self, FromStr};
216use std::sync::atomic::AtomicUsize;
217use std::sync::atomic::Ordering;
218use std::sync::Arc;
219use std::task::{Context, Poll};
220use tokio::io::ErrorKind;
221use tokio::time::{interval, Duration, Interval, MissedTickBehavior};
222use url::{Host, Url};
223
224use bytes::Bytes;
225use serde::{Deserialize, Serialize};
226use serde_repr::{Deserialize_repr, Serialize_repr};
227use tokio::io;
228use tokio::sync::mpsc;
229use tokio::task;
230
231pub type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
232
233const VERSION: &str = env!("CARGO_PKG_VERSION");
234const LANG: &str = "rust";
235const MAX_PENDING_PINGS: usize = 2;
236const MULTIPLEXER_SID: u64 = 0;
237pub(crate) const DEFAULT_SERVER_MAX_PAYLOAD: usize = 1024 * 1024;
238
239pub use tokio_rustls::rustls;
243
244use connection::{Connection, State};
245use connector::{Connector, ConnectorOptions};
246pub use connector::{ReconnectToServer, Server};
247pub use header::{HeaderMap, HeaderName, HeaderValue};
248pub use subject::{Subject, SubjectError, ToSubject};
249
250mod auth;
251pub(crate) mod auth_utils;
252pub mod client;
253pub mod connection;
254mod connector;
255mod options;
256
257pub use auth::Auth;
258pub use client::{
259 Client, PublishError, PublishErrorKind, Request, RequestError, RequestErrorKind,
260 ServerPoolError, ServerPoolErrorKind, SetServerPoolError, SetServerPoolErrorKind, Statistics,
261 SubscribeError, SubscribeErrorKind,
262};
263pub use options::{AuthError, ConnectOptions};
264
265#[cfg(feature = "crypto")]
266#[cfg_attr(docsrs, doc(cfg(feature = "crypto")))]
267mod crypto;
268#[cfg(any(feature = "jetstream", feature = "service", feature = "chrono"))]
269#[cfg_attr(
270 docsrs,
271 doc(cfg(any(feature = "jetstream", feature = "service", feature = "chrono")))
272)]
273pub mod datetime;
274
275pub mod error;
276pub mod header;
277mod id_generator;
278#[cfg(feature = "jetstream")]
279#[cfg_attr(docsrs, doc(cfg(feature = "jetstream")))]
280pub mod jetstream;
281pub mod message;
282#[cfg(feature = "service")]
283#[cfg_attr(docsrs, doc(cfg(feature = "service")))]
284pub mod service;
285pub mod status;
286pub mod subject;
287mod tls;
288
289pub use message::Message;
290pub use status::StatusCode;
291
292#[derive(Debug, Deserialize, Default, Clone, Eq, PartialEq)]
295pub struct ServerInfo {
296 #[serde(default)]
298 pub server_id: String,
299 #[serde(default)]
301 pub server_name: String,
302 #[serde(default)]
304 pub host: String,
305 #[serde(default)]
307 pub port: u16,
308 #[serde(default)]
310 pub version: String,
311 #[serde(default)]
314 pub auth_required: bool,
315 #[serde(default)]
317 pub tls_required: bool,
318 #[serde(default)]
320 pub max_payload: usize,
321 #[serde(default)]
323 pub proto: i8,
324 #[serde(default)]
326 pub client_id: u64,
327 #[serde(default)]
329 pub go: String,
330 #[serde(default)]
332 pub nonce: String,
333 #[serde(default)]
335 pub connect_urls: Vec<String>,
336 #[serde(default)]
338 pub client_ip: String,
339 #[serde(default)]
341 pub headers: bool,
342 #[serde(default, rename = "ldm")]
344 pub lame_duck_mode: bool,
345 #[serde(default)]
347 pub cluster: Option<String>,
348 #[serde(default)]
350 pub domain: Option<String>,
351 #[serde(default)]
353 pub jetstream: bool,
354}
355
356#[derive(Clone, Debug, Eq, PartialEq)]
357pub(crate) enum ServerOp {
358 Ok,
359 Info(Box<ServerInfo>),
360 Ping,
361 Pong,
362 Error(ServerError),
363 Message {
364 sid: u64,
365 subject: Subject,
366 reply: Option<Subject>,
367 payload: Bytes,
368 headers: Option<HeaderMap>,
369 status: Option<StatusCode>,
370 description: Option<String>,
371 length: usize,
372 },
373}
374
375#[deprecated(
379 since = "0.44.0",
380 note = "use `async_nats::message::OutboundMessage` instead"
381)]
382pub type PublishMessage = crate::message::OutboundMessage;
383
384#[derive(Debug)]
386pub(crate) enum Command {
387 Publish(OutboundMessage),
388 Request {
389 subject: Subject,
390 payload: Bytes,
391 respond: Subject,
392 headers: Option<HeaderMap>,
393 sender: oneshot::Sender<Message>,
394 },
395 Subscribe {
396 sid: u64,
397 subject: Subject,
398 queue_group: Option<String>,
399 sender: mpsc::Sender<Message>,
400 },
401 Unsubscribe {
402 sid: u64,
403 max: Option<u64>,
404 },
405 Flush {
406 observer: oneshot::Sender<()>,
407 },
408 Drain {
409 sid: Option<u64>,
410 },
411 Reconnect,
412 SetServerPool {
413 servers: Vec<ServerAddr>,
414 result: oneshot::Sender<Result<(), String>>,
415 },
416 ServerPool {
417 result: oneshot::Sender<Vec<connector::Server>>,
418 },
419}
420
421#[derive(Debug)]
423pub(crate) enum ClientOp {
424 Publish {
425 subject: Subject,
426 payload: Bytes,
427 respond: Option<Subject>,
428 headers: Option<HeaderMap>,
429 },
430 Subscribe {
431 sid: u64,
432 subject: Subject,
433 queue_group: Option<String>,
434 },
435 Unsubscribe {
436 sid: u64,
437 max: Option<u64>,
438 },
439 Ping,
440 Pong,
441 Connect(ConnectInfo),
442}
443
444#[derive(Debug)]
445struct Subscription {
446 subject: Subject,
447 sender: mpsc::Sender<Message>,
448 queue_group: Option<String>,
449 delivered: u64,
450 max: Option<u64>,
451}
452
453#[derive(Debug)]
454struct Multiplexer {
455 subject: Subject,
456 prefix: Subject,
457 senders: HashMap<String, oneshot::Sender<Message>>,
458}
459
460pub(crate) struct ConnectionHandler {
462 connection: Connection,
463 connector: Connector,
464 subscriptions: HashMap<u64, Subscription>,
465 multiplexer: Option<Multiplexer>,
466 pending_pings: usize,
467 info_sender: tokio::sync::watch::Sender<Option<ServerInfo>>,
468 ping_interval: Interval,
469 should_reconnect: bool,
470 flush_observers: Vec<oneshot::Sender<()>>,
471 is_draining: bool,
472 drain_pings: VecDeque<u64>,
473}
474
475impl ConnectionHandler {
476 pub(crate) fn new(
477 connection: Connection,
478 connector: Connector,
479 info_sender: tokio::sync::watch::Sender<Option<ServerInfo>>,
480 ping_period: Duration,
481 ) -> ConnectionHandler {
482 let mut ping_interval = interval(ping_period);
483 ping_interval.set_missed_tick_behavior(MissedTickBehavior::Delay);
484
485 ConnectionHandler {
486 connection,
487 connector,
488 subscriptions: HashMap::new(),
489 multiplexer: None,
490 pending_pings: 0,
491 info_sender,
492 ping_interval,
493 should_reconnect: false,
494 flush_observers: Vec::new(),
495 is_draining: false,
496 drain_pings: VecDeque::new(),
497 }
498 }
499
500 pub(crate) async fn process<'a>(&'a mut self, receiver: &'a mut mpsc::Receiver<Command>) {
501 struct ProcessFut<'a> {
502 handler: &'a mut ConnectionHandler,
503 receiver: &'a mut mpsc::Receiver<Command>,
504 recv_buf: &'a mut Vec<Command>,
505 }
506
507 enum ExitReason {
508 Disconnected(Option<io::Error>),
509 ReconnectRequested,
510 Closed,
511 }
512
513 impl ProcessFut<'_> {
514 const RECV_CHUNK_SIZE: usize = 16;
515
516 #[cold]
517 fn ping(&mut self) -> Poll<ExitReason> {
518 self.handler.pending_pings += 1;
519
520 if self.handler.pending_pings > MAX_PENDING_PINGS {
521 debug!(
522 pending_pings = self.handler.pending_pings,
523 max_pings = MAX_PENDING_PINGS,
524 "disconnecting due to too many pending pings"
525 );
526
527 Poll::Ready(ExitReason::Disconnected(None))
528 } else {
529 self.handler.connection.enqueue_write_op(&ClientOp::Ping);
530
531 Poll::Pending
532 }
533 }
534 }
535
536 impl Future for ProcessFut<'_> {
537 type Output = ExitReason;
538
539 fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
553 while self.handler.ping_interval.poll_tick(cx).is_ready() {
557 if let Poll::Ready(exit) = self.ping() {
558 return Poll::Ready(exit);
559 }
560 }
561
562 loop {
563 match self.handler.connection.poll_read_op(cx) {
564 Poll::Pending => break,
565 Poll::Ready(Ok(Some(server_op))) => {
566 self.handler.handle_server_op(server_op);
567 }
568 Poll::Ready(Ok(None)) => {
569 return Poll::Ready(ExitReason::Disconnected(None))
570 }
571 Poll::Ready(Err(err)) => {
572 return Poll::Ready(ExitReason::Disconnected(Some(err)))
573 }
574 }
575 }
576
577 while let Some(sid) = self.handler.drain_pings.pop_front() {
582 self.handler.subscriptions.remove(&sid);
583 }
584
585 if self.handler.is_draining {
586 return Poll::Ready(ExitReason::Closed);
591 }
592
593 let mut made_progress = true;
599 loop {
600 while !self.handler.connection.is_write_buf_full() {
601 debug_assert!(self.recv_buf.is_empty());
602
603 let Self {
604 recv_buf,
605 handler,
606 receiver,
607 } = &mut *self;
608 match receiver.poll_recv_many(cx, recv_buf, Self::RECV_CHUNK_SIZE) {
609 Poll::Pending => break,
610 Poll::Ready(1..) => {
611 made_progress = true;
612
613 for cmd in recv_buf.drain(..) {
614 handler.handle_command(cmd);
615 }
616 }
617 Poll::Ready(_) => return Poll::Ready(ExitReason::Closed),
619 }
620 }
621
622 if !mem::take(&mut made_progress) {
633 break;
634 }
635
636 match self.handler.connection.poll_write(cx) {
637 Poll::Pending => {
638 break;
640 }
641 Poll::Ready(Ok(())) => {
642 continue;
644 }
645 Poll::Ready(Err(err)) => {
646 return Poll::Ready(ExitReason::Disconnected(Some(err)))
647 }
648 }
649 }
650
651 if let (ShouldFlush::Yes, _) | (ShouldFlush::No, false) = (
652 self.handler.connection.should_flush(),
653 self.handler.flush_observers.is_empty(),
654 ) {
655 match self.handler.connection.poll_flush(cx) {
656 Poll::Pending => {}
657 Poll::Ready(Ok(())) => {
658 for observer in self.handler.flush_observers.drain(..) {
659 let _ = observer.send(());
660 }
661 }
662 Poll::Ready(Err(err)) => {
663 return Poll::Ready(ExitReason::Disconnected(Some(err)))
664 }
665 }
666 }
667
668 if mem::take(&mut self.handler.should_reconnect) {
669 return Poll::Ready(ExitReason::ReconnectRequested);
670 }
671
672 Poll::Pending
673 }
674 }
675
676 let mut recv_buf = Vec::with_capacity(ProcessFut::RECV_CHUNK_SIZE);
677 loop {
678 let process = ProcessFut {
679 handler: self,
680 receiver,
681 recv_buf: &mut recv_buf,
682 };
683 match process.await {
684 ExitReason::Disconnected(err) => {
685 debug!(error = ?err, "disconnected");
686 if self.handle_disconnect().await.is_err() {
687 break;
688 };
689 debug!("reconnected");
690 }
691 ExitReason::Closed => {
692 self.connector.events_tx.try_send(Event::Closed).ok();
694 break;
695 }
696 ExitReason::ReconnectRequested => {
697 debug!("reconnect requested");
698 self.connection.stream.shutdown().await.ok();
700 if self.handle_disconnect().await.is_err() {
701 break;
702 };
703 }
704 }
705 }
706 }
707
708 fn handle_server_op(&mut self, server_op: ServerOp) {
709 self.ping_interval.reset();
710
711 match server_op {
712 ServerOp::Ping => {
713 debug!("received PING");
714 self.connection.enqueue_write_op(&ClientOp::Pong);
715 }
716 ServerOp::Pong => {
717 debug!("received PONG");
718 self.pending_pings = self.pending_pings.saturating_sub(1);
719 }
720 ServerOp::Error(error) => {
721 debug!("received ERROR: {:?}", error);
722 self.connector
723 .events_tx
724 .try_send(Event::ServerError(error))
725 .ok();
726 }
727 ServerOp::Message {
728 sid,
729 subject,
730 reply,
731 payload,
732 headers,
733 status,
734 description,
735 length,
736 } => {
737 debug!("received MESSAGE: sid={}, subject={}", sid, subject);
738 self.connector
739 .connect_stats
740 .in_messages
741 .add(1, Ordering::Relaxed);
742
743 if let Some(subscription) = self.subscriptions.get_mut(&sid) {
744 let message: Message = Message {
745 subject,
746 reply,
747 payload,
748 headers,
749 status,
750 description,
751 length,
752 };
753
754 match subscription.sender.try_send(message) {
757 Ok(_) => {
758 subscription.delivered += 1;
759 if let Some(max) = subscription.max {
763 if subscription.delivered.ge(&max) {
764 debug!("max messages reached for subscription {}", sid);
765 self.subscriptions.remove(&sid);
766 }
767 }
768 }
769 Err(mpsc::error::TrySendError::Full(_)) => {
770 debug!("slow consumer detected for subscription {}", sid);
771 self.connector
772 .events_tx
773 .try_send(Event::SlowConsumer(sid))
774 .ok();
775 }
776 Err(mpsc::error::TrySendError::Closed(_)) => {
777 debug!("subscription {} channel closed", sid);
778 self.subscriptions.remove(&sid);
779 self.connection
780 .enqueue_write_op(&ClientOp::Unsubscribe { sid, max: None });
781 }
782 }
783 } else if sid == MULTIPLEXER_SID {
784 debug!("received message for multiplexer");
785 if let Some(multiplexer) = self.multiplexer.as_mut() {
786 let maybe_token =
787 subject.strip_prefix(multiplexer.prefix.as_ref()).to_owned();
788
789 if let Some(token) = maybe_token {
790 if let Some(sender) = multiplexer.senders.remove(token) {
791 debug!("forwarding message to request with token {}", token);
792 let message = Message {
793 subject,
794 reply,
795 payload,
796 headers,
797 status,
798 description,
799 length,
800 };
801
802 let _ = sender.send(message);
803 }
804 }
805 }
806 }
807 }
808 ServerOp::Info(info) => {
810 debug!("received INFO: server_id={}", info.server_id);
811 if info.lame_duck_mode {
812 debug!("server in lame duck mode");
813 self.connector.events_tx.try_send(Event::LameDuckMode).ok();
814 }
815 }
816
817 _ => {
818 }
820 }
821 }
822
823 fn handle_command(&mut self, command: Command) {
824 match command {
825 Command::Unsubscribe { sid, max } => {
826 if let Some(subscription) = self.subscriptions.get_mut(&sid) {
827 subscription.max = max;
828 match subscription.max {
829 Some(n) => {
830 if subscription.delivered >= n {
831 self.subscriptions.remove(&sid);
832 }
833 }
834 None => {
835 self.subscriptions.remove(&sid);
836 }
837 }
838
839 self.connection
840 .enqueue_write_op(&ClientOp::Unsubscribe { sid, max });
841 }
842 }
843 Command::Flush { observer } => {
844 self.flush_observers.push(observer);
845 }
846 Command::Drain { sid } => {
847 let mut drain_sub = |sid: u64| {
848 self.drain_pings.push_back(sid);
849 self.connection
850 .enqueue_write_op(&ClientOp::Unsubscribe { sid, max: None });
851 };
852
853 if let Some(sid) = sid {
854 if self.subscriptions.get_mut(&sid).is_some() {
855 drain_sub(sid);
856 }
857 } else {
858 self.connector.events_tx.try_send(Event::Draining).ok();
860 self.is_draining = true;
861 for &sid in self.subscriptions.keys() {
862 drain_sub(sid);
863 }
864 }
865 self.connection.enqueue_write_op(&ClientOp::Ping);
866 }
867 Command::Subscribe {
868 sid,
869 subject,
870 queue_group,
871 sender,
872 } => {
873 let subscription = Subscription {
874 sender,
875 delivered: 0,
876 max: None,
877 subject: subject.to_owned(),
878 queue_group: queue_group.to_owned(),
879 };
880
881 self.subscriptions.insert(sid, subscription);
882
883 self.connection.enqueue_write_op(&ClientOp::Subscribe {
884 sid,
885 subject,
886 queue_group,
887 });
888 }
889 Command::Request {
890 subject,
891 payload,
892 respond,
893 headers,
894 sender,
895 } => {
896 let (prefix, token) = respond.rsplit_once('.').expect("malformed request subject");
897
898 let multiplexer = if let Some(multiplexer) = self.multiplexer.as_mut() {
899 multiplexer
900 } else {
901 let prefix = Subject::from(format!("{}.{}.", prefix, id_generator::next()));
902 let subject = Subject::from(format!("{prefix}*"));
903
904 self.connection.enqueue_write_op(&ClientOp::Subscribe {
905 sid: MULTIPLEXER_SID,
906 subject: subject.clone(),
907 queue_group: None,
908 });
909
910 self.multiplexer.insert(Multiplexer {
911 subject,
912 prefix,
913 senders: HashMap::new(),
914 })
915 };
916 self.connector
917 .connect_stats
918 .out_messages
919 .add(1, Ordering::Relaxed);
920
921 multiplexer.senders.insert(token.to_owned(), sender);
922
923 let respond: Subject = format!("{}{}", multiplexer.prefix, token).into();
924
925 let pub_op = ClientOp::Publish {
926 subject,
927 payload,
928 respond: Some(respond),
929 headers,
930 };
931
932 self.connection.enqueue_write_op(&pub_op);
933 }
934
935 Command::Publish(OutboundMessage {
936 subject,
937 payload,
938 reply: respond,
939 headers,
940 }) => {
941 self.connector
942 .connect_stats
943 .out_messages
944 .add(1, Ordering::Relaxed);
945
946 let header_len = headers
947 .as_ref()
948 .map(|headers| headers.len())
949 .unwrap_or_default();
950
951 self.connector.connect_stats.out_bytes.add(
952 (payload.len()
953 + respond.as_ref().map_or_else(|| 0, |r| r.len())
954 + subject.len()
955 + header_len) as u64,
956 Ordering::Relaxed,
957 );
958
959 self.connection.enqueue_write_op(&ClientOp::Publish {
960 subject,
961 payload,
962 respond,
963 headers,
964 });
965 }
966
967 Command::Reconnect => {
968 self.should_reconnect = true;
969 }
970
971 Command::SetServerPool { servers, result } => {
972 let _ = result.send(self.connector.set_server_pool(servers));
973 }
974
975 Command::ServerPool { result } => {
976 let _ = result.send(self.connector.server_pool());
977 }
978 }
979 }
980
981 async fn handle_disconnect(&mut self) -> Result<(), ConnectError> {
982 self.pending_pings = 0;
983 self.connector.events_tx.try_send(Event::Disconnected).ok();
984 self.connector.state_tx.send(State::Disconnected).ok();
985
986 self.handle_reconnect().await
987 }
988
989 async fn handle_reconnect(&mut self) -> Result<(), ConnectError> {
990 let (info, connection) = self.connector.connect().await?;
991 self.connection = connection;
992 let _ = self.info_sender.send(Some(info));
993
994 self.subscriptions
995 .retain(|_, subscription| !subscription.sender.is_closed());
996
997 for (sid, subscription) in &self.subscriptions {
998 self.connection.enqueue_write_op(&ClientOp::Subscribe {
999 sid: *sid,
1000 subject: subscription.subject.to_owned(),
1001 queue_group: subscription.queue_group.to_owned(),
1002 });
1003
1004 if let Some(max) = subscription.max {
1005 self.connection.enqueue_write_op(&ClientOp::Unsubscribe {
1006 sid: *sid,
1007 max: Some(max.saturating_sub(subscription.delivered)),
1008 });
1009 }
1010 }
1011
1012 if let Some(multiplexer) = &self.multiplexer {
1013 self.connection.enqueue_write_op(&ClientOp::Subscribe {
1014 sid: MULTIPLEXER_SID,
1015 subject: multiplexer.subject.to_owned(),
1016 queue_group: None,
1017 });
1018 }
1019 Ok(())
1020 }
1021}
1022
1023pub async fn connect_with_options<A: ToServerAddrs>(
1039 addrs: A,
1040 options: ConnectOptions,
1041) -> Result<Client, ConnectError> {
1042 let ping_period = options.ping_interval;
1043
1044 let (events_tx, mut events_rx) = mpsc::channel(128);
1045 let (state_tx, state_rx) = tokio::sync::watch::channel(State::Pending);
1046 let max_payload = Arc::new(AtomicUsize::new(DEFAULT_SERVER_MAX_PAYLOAD));
1048 let statistics = Arc::new(Statistics::default());
1049
1050 let mut connector = Connector::new(
1051 addrs,
1052 ConnectorOptions {
1053 tls_required: options.tls_required,
1054 certificates: options.certificates,
1055 client_key: options.client_key,
1056 client_cert: options.client_cert,
1057 tls_client_config: options.tls_client_config,
1058 tls_first: options.tls_first,
1059 auth: options.auth,
1060 no_echo: options.no_echo,
1061 connection_timeout: options.connection_timeout,
1062 name: options.name,
1063 ignore_discovered_servers: options.ignore_discovered_servers,
1064 retain_servers_order: options.retain_servers_order,
1065 read_buffer_capacity: options.read_buffer_capacity,
1066 reconnect_delay_callback: options.reconnect_delay_callback,
1067 auth_callback: options.auth_callback,
1068 max_reconnects: options.max_reconnects,
1069 local_address: options.local_address,
1070 reconnect_to_server_callback: options.reconnect_to_server_callback,
1071 },
1072 events_tx,
1073 state_tx,
1074 max_payload.clone(),
1075 statistics.clone(),
1076 )
1077 .map_err(|err| ConnectError::with_source(ConnectErrorKind::ServerParse, err))?;
1078
1079 let mut info = None;
1080 let mut connection = None;
1081 if !options.retry_on_initial_connect {
1082 debug!("retry on initial connect failure is disabled");
1083 let (info_ok, connection_ok) = connector.try_connect().await?;
1084 connection = Some(connection_ok);
1085 info = Some(info_ok);
1086 }
1087
1088 let (info_sender, info_watcher) = tokio::sync::watch::channel(info.clone());
1089 let (sender, mut receiver) = mpsc::channel(options.sender_capacity);
1090
1091 let client = Client::new(
1092 info_watcher,
1093 state_rx,
1094 sender,
1095 options.subscription_capacity,
1096 options.inbox_prefix,
1097 options.request_timeout,
1098 max_payload,
1099 statistics,
1100 options.skip_subject_validation,
1101 );
1102
1103 task::spawn(async move {
1104 while let Some(event) = events_rx.recv().await {
1105 tracing::info!("event: {}", event);
1106 if let Some(event_callback) = &options.event_callback {
1107 event_callback.call(event).await;
1108 }
1109 }
1110 });
1111
1112 task::spawn(async move {
1113 if connection.is_none() && options.retry_on_initial_connect {
1114 let (info, connection_ok) = match connector.connect().await {
1115 Ok((info, connection)) => (info, connection),
1116 Err(err) => {
1117 error!("connection closed: {}", err);
1118 return;
1119 }
1120 };
1121 info_sender.send(Some(info)).ok();
1122 connection = Some(connection_ok);
1123 }
1124 let connection = connection.unwrap();
1125 let mut connection_handler =
1126 ConnectionHandler::new(connection, connector, info_sender, ping_period);
1127 connection_handler.process(&mut receiver).await
1128 });
1129
1130 Ok(client)
1131}
1132
1133#[derive(Debug, Clone, PartialEq, Eq)]
1134pub enum Event {
1135 Connected,
1136 Disconnected,
1137 LameDuckMode,
1138 Draining,
1139 Closed,
1140 SlowConsumer(u64),
1141 ServerError(ServerError),
1142 ClientError(ClientError),
1143}
1144
1145impl fmt::Display for Event {
1146 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1147 match self {
1148 Event::Connected => write!(f, "connected"),
1149 Event::Disconnected => write!(f, "disconnected"),
1150 Event::LameDuckMode => write!(f, "lame duck mode detected"),
1151 Event::Draining => write!(f, "draining"),
1152 Event::Closed => write!(f, "closed"),
1153 Event::SlowConsumer(sid) => write!(f, "slow consumers for subscription {sid}"),
1154 Event::ServerError(err) => write!(f, "server error: {err}"),
1155 Event::ClientError(err) => write!(f, "client error: {err}"),
1156 }
1157 }
1158}
1159
1160pub async fn connect<A: ToServerAddrs>(addrs: A) -> Result<Client, ConnectError> {
1227 connect_with_options(addrs, ConnectOptions::default()).await
1228}
1229
1230#[derive(Debug, Clone, Copy, PartialEq)]
1231pub enum ConnectErrorKind {
1232 ServerParse,
1234 Dns,
1236 Authentication,
1238 AuthorizationViolation,
1240 TimedOut,
1242 Tls,
1244 Io,
1246 MaxReconnects,
1248}
1249
1250impl Display for ConnectErrorKind {
1251 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1252 match self {
1253 Self::ServerParse => write!(f, "failed to parse server or server list"),
1254 Self::Dns => write!(f, "DNS error"),
1255 Self::Authentication => write!(f, "failed signing nonce"),
1256 Self::AuthorizationViolation => write!(f, "authorization violation"),
1257 Self::TimedOut => write!(f, "timed out"),
1258 Self::Tls => write!(f, "TLS error"),
1259 Self::Io => write!(f, "IO error"),
1260 Self::MaxReconnects => write!(f, "reached maximum number of reconnects"),
1261 }
1262 }
1263}
1264
1265pub type ConnectError = error::Error<ConnectErrorKind>;
1268
1269impl From<io::Error> for ConnectError {
1270 fn from(err: io::Error) -> Self {
1271 ConnectError::with_source(ConnectErrorKind::Io, err)
1272 }
1273}
1274
1275#[derive(Debug)]
1289pub struct Subscriber {
1290 sid: u64,
1291 receiver: mpsc::Receiver<Message>,
1292 sender: mpsc::Sender<Command>,
1293}
1294
1295impl Subscriber {
1296 fn new(
1297 sid: u64,
1298 sender: mpsc::Sender<Command>,
1299 receiver: mpsc::Receiver<Message>,
1300 ) -> Subscriber {
1301 Subscriber {
1302 sid,
1303 sender,
1304 receiver,
1305 }
1306 }
1307
1308 pub async fn unsubscribe(&mut self) -> Result<(), UnsubscribeError> {
1323 self.sender
1324 .send(Command::Unsubscribe {
1325 sid: self.sid,
1326 max: None,
1327 })
1328 .await?;
1329 self.receiver.close();
1330 Ok(())
1331 }
1332
1333 pub async fn unsubscribe_after(&mut self, unsub_after: u64) -> Result<(), UnsubscribeError> {
1359 self.sender
1360 .send(Command::Unsubscribe {
1361 sid: self.sid,
1362 max: Some(unsub_after),
1363 })
1364 .await?;
1365 Ok(())
1366 }
1367
1368 pub async fn drain(&mut self) -> Result<(), UnsubscribeError> {
1401 self.sender
1402 .send(Command::Drain {
1403 sid: Some(self.sid),
1404 })
1405 .await?;
1406
1407 Ok(())
1408 }
1409}
1410
1411#[derive(Error, Debug, PartialEq)]
1412#[error("failed to send unsubscribe")]
1413pub struct UnsubscribeError(String);
1414
1415impl From<tokio::sync::mpsc::error::SendError<Command>> for UnsubscribeError {
1416 fn from(err: tokio::sync::mpsc::error::SendError<Command>) -> Self {
1417 UnsubscribeError(err.to_string())
1418 }
1419}
1420
1421impl Drop for Subscriber {
1422 fn drop(&mut self) {
1423 self.receiver.close();
1424 tokio::spawn({
1425 let sender = self.sender.clone();
1426 let sid = self.sid;
1427 async move {
1428 sender
1429 .send(Command::Unsubscribe { sid, max: None })
1430 .await
1431 .ok();
1432 }
1433 });
1434 }
1435}
1436
1437impl Stream for Subscriber {
1438 type Item = Message;
1439
1440 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
1441 self.receiver.poll_recv(cx)
1442 }
1443}
1444
1445#[derive(Clone, Debug, Eq, PartialEq)]
1446pub enum CallbackError {
1447 Client(ClientError),
1448 Server(ServerError),
1449}
1450impl std::fmt::Display for CallbackError {
1451 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1452 match self {
1453 Self::Client(error) => write!(f, "{error}"),
1454 Self::Server(error) => write!(f, "{error}"),
1455 }
1456 }
1457}
1458
1459impl From<ServerError> for CallbackError {
1460 fn from(server_error: ServerError) -> Self {
1461 CallbackError::Server(server_error)
1462 }
1463}
1464
1465impl From<ClientError> for CallbackError {
1466 fn from(client_error: ClientError) -> Self {
1467 CallbackError::Client(client_error)
1468 }
1469}
1470
1471#[derive(Clone, Debug, Eq, PartialEq, Error)]
1472pub enum ServerError {
1473 AuthorizationViolation,
1474 SlowConsumer(u64),
1475 Other(String),
1476}
1477
1478#[derive(Clone, Debug, Eq, PartialEq)]
1479pub enum ClientError {
1480 Other(String),
1481 MaxReconnects,
1482 ServerNotInPool,
1485}
1486impl std::fmt::Display for ClientError {
1487 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1488 match self {
1489 Self::Other(error) => write!(f, "nats: {error}"),
1490 Self::MaxReconnects => write!(f, "nats: max reconnects reached"),
1491 Self::ServerNotInPool => {
1492 write!(f, "nats: reconnect callback returned server not in pool")
1493 }
1494 }
1495 }
1496}
1497
1498impl ServerError {
1499 fn new(error: String) -> ServerError {
1500 match error.to_lowercase().as_str() {
1501 "authorization violation" => ServerError::AuthorizationViolation,
1502 _ => ServerError::Other(error),
1504 }
1505 }
1506}
1507
1508impl std::fmt::Display for ServerError {
1509 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1510 match self {
1511 Self::AuthorizationViolation => write!(f, "nats: authorization violation"),
1512 Self::SlowConsumer(sid) => write!(f, "nats: subscription {sid} is a slow consumer"),
1513 Self::Other(error) => write!(f, "nats: {error}"),
1514 }
1515 }
1516}
1517
1518#[derive(Clone, Debug, Serialize)]
1520pub struct ConnectInfo {
1521 pub verbose: bool,
1523
1524 pub pedantic: bool,
1527
1528 #[serde(rename = "jwt")]
1530 pub user_jwt: Option<String>,
1531
1532 pub nkey: Option<String>,
1534
1535 #[serde(rename = "sig")]
1537 pub signature: Option<String>,
1538
1539 pub name: Option<String>,
1541
1542 pub echo: bool,
1547
1548 pub lang: String,
1550
1551 pub version: String,
1553
1554 pub protocol: Protocol,
1559
1560 pub tls_required: bool,
1562
1563 pub user: Option<String>,
1565
1566 pub pass: Option<String>,
1568
1569 pub auth_token: Option<String>,
1571
1572 pub headers: bool,
1574
1575 pub no_responders: bool,
1577}
1578
1579#[derive(Serialize_repr, Deserialize_repr, PartialEq, Eq, Debug, Clone, Copy)]
1581#[repr(u8)]
1582pub enum Protocol {
1583 Original = 0,
1585 Dynamic = 1,
1587}
1588
1589#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1591pub struct ServerAddr(Url);
1592
1593impl FromStr for ServerAddr {
1594 type Err = io::Error;
1595
1596 fn from_str(input: &str) -> Result<Self, Self::Err> {
1600 let url: Url = if input.contains("://") {
1601 input.parse()
1602 } else {
1603 format!("nats://{input}").parse()
1604 }
1605 .map_err(|e| {
1606 io::Error::new(
1607 ErrorKind::InvalidInput,
1608 format!("NATS server URL is invalid: {e}"),
1609 )
1610 })?;
1611
1612 Self::from_url(url)
1613 }
1614}
1615
1616impl ServerAddr {
1617 pub fn from_url(url: Url) -> io::Result<Self> {
1619 if url.scheme() != "nats"
1620 && url.scheme() != "tls"
1621 && url.scheme() != "ws"
1622 && url.scheme() != "wss"
1623 {
1624 return Err(std::io::Error::new(
1625 ErrorKind::InvalidInput,
1626 format!("invalid scheme for NATS server URL: {}", url.scheme()),
1627 ));
1628 }
1629
1630 Ok(Self(url))
1631 }
1632
1633 pub fn into_inner(self) -> Url {
1635 self.0
1636 }
1637
1638 pub fn tls_required(&self) -> bool {
1640 self.0.scheme() == "tls"
1641 }
1642
1643 pub fn has_user_pass(&self) -> bool {
1645 self.0.username() != ""
1646 }
1647
1648 pub fn scheme(&self) -> &str {
1649 self.0.scheme()
1650 }
1651
1652 pub fn host(&self) -> &str {
1654 match self.0.host() {
1655 Some(Host::Domain(_)) | Some(Host::Ipv4 { .. }) => self.0.host_str().unwrap(),
1656 Some(Host::Ipv6 { .. }) => {
1658 let host = self.0.host_str().unwrap();
1659 &host[1..host.len() - 1]
1660 }
1661 None => "",
1662 }
1663 }
1664
1665 pub fn is_websocket(&self) -> bool {
1666 self.0.scheme() == "ws" || self.0.scheme() == "wss"
1667 }
1668
1669 pub fn port(&self) -> u16 {
1672 self.0.port_or_known_default().unwrap_or(4222)
1673 }
1674
1675 pub fn as_url_str(&self) -> &str {
1677 self.0.as_str()
1678 }
1679
1680 pub fn username(&self) -> Option<&str> {
1682 let user = self.0.username();
1683 if user.is_empty() {
1684 None
1685 } else {
1686 Some(user)
1687 }
1688 }
1689
1690 pub fn password(&self) -> Option<&str> {
1692 self.0.password()
1693 }
1694
1695 pub async fn socket_addrs(&self) -> io::Result<impl Iterator<Item = SocketAddr> + '_> {
1697 tokio::net::lookup_host((self.host(), self.port())).await
1698 }
1699}
1700
1701pub trait ToServerAddrs {
1706 type Iter: Iterator<Item = ServerAddr>;
1709
1710 fn to_server_addrs(&self) -> io::Result<Self::Iter>;
1711}
1712
1713impl ToServerAddrs for ServerAddr {
1714 type Iter = option::IntoIter<ServerAddr>;
1715 fn to_server_addrs(&self) -> io::Result<Self::Iter> {
1716 Ok(Some(self.clone()).into_iter())
1717 }
1718}
1719
1720impl ToServerAddrs for str {
1721 type Iter = option::IntoIter<ServerAddr>;
1722 fn to_server_addrs(&self) -> io::Result<Self::Iter> {
1723 self.parse::<ServerAddr>()
1724 .map(|addr| Some(addr).into_iter())
1725 }
1726}
1727
1728impl ToServerAddrs for String {
1729 type Iter = option::IntoIter<ServerAddr>;
1730 fn to_server_addrs(&self) -> io::Result<Self::Iter> {
1731 (**self).to_server_addrs()
1732 }
1733}
1734
1735impl<T: AsRef<str>> ToServerAddrs for [T] {
1736 type Iter = std::vec::IntoIter<ServerAddr>;
1737 fn to_server_addrs(&self) -> io::Result<Self::Iter> {
1738 self.iter()
1739 .map(AsRef::as_ref)
1740 .map(str::parse)
1741 .collect::<io::Result<_>>()
1742 .map(Vec::into_iter)
1743 }
1744}
1745
1746impl<T: AsRef<str>> ToServerAddrs for Vec<T> {
1747 type Iter = std::vec::IntoIter<ServerAddr>;
1748 fn to_server_addrs(&self) -> io::Result<Self::Iter> {
1749 self.as_slice().to_server_addrs()
1750 }
1751}
1752
1753impl<'a> ToServerAddrs for &'a [ServerAddr] {
1754 type Iter = iter::Cloned<slice::Iter<'a, ServerAddr>>;
1755
1756 fn to_server_addrs(&self) -> io::Result<Self::Iter> {
1757 Ok(self.iter().cloned())
1758 }
1759}
1760
1761impl ToServerAddrs for Vec<ServerAddr> {
1762 type Iter = std::vec::IntoIter<ServerAddr>;
1763
1764 fn to_server_addrs(&self) -> io::Result<Self::Iter> {
1765 Ok(self.clone().into_iter())
1766 }
1767}
1768
1769impl<T: ToServerAddrs + ?Sized> ToServerAddrs for &T {
1770 type Iter = T::Iter;
1771 fn to_server_addrs(&self) -> io::Result<Self::Iter> {
1772 (**self).to_server_addrs()
1773 }
1774}
1775
1776pub(crate) fn is_valid_publish_subject<T: AsRef<str>>(subject: T) -> bool {
1781 let bytes = subject.as_ref().as_bytes();
1782
1783 if bytes.is_empty() {
1784 return false;
1785 }
1786
1787 memchr::memchr3(b' ', b'\r', b'\n', bytes).is_none() && memchr::memchr(b'\t', bytes).is_none()
1788}
1789
1790pub(crate) fn is_valid_subject<T: AsRef<str>>(subject: T) -> bool {
1794 let bytes = subject.as_ref().as_bytes();
1795
1796 if bytes.is_empty() {
1797 return false;
1798 }
1799
1800 bytes[0] != b'.'
1801 && bytes[bytes.len() - 1] != b'.'
1802 && memchr::memmem::find(bytes, b"..").is_none()
1803 && memchr::memchr3(b' ', b'\r', b'\n', bytes).is_none()
1804 && memchr::memchr(b'\t', bytes).is_none()
1805}
1806
1807pub(crate) fn is_valid_queue_group(queue_group: &str) -> bool {
1811 let bytes = queue_group.as_bytes();
1812
1813 if bytes.is_empty() {
1814 return false;
1815 }
1816
1817 memchr::memchr3(b' ', b'\r', b'\n', bytes).is_none() && memchr::memchr(b'\t', bytes).is_none()
1818}
1819
1820#[allow(unused_macros)]
1821macro_rules! from_with_timeout {
1822 ($t:ty, $k:ty, $origin: ty, $origin_kind: ty) => {
1823 impl From<$origin> for $t {
1824 fn from(err: $origin) -> Self {
1825 match err.kind() {
1826 <$origin_kind>::TimedOut => Self::new(<$k>::TimedOut),
1827 _ => Self::with_source(<$k>::Other, err),
1828 }
1829 }
1830 }
1831 };
1832}
1833#[allow(unused_imports)]
1834pub(crate) use from_with_timeout;
1835
1836use crate::connection::ShouldFlush;
1837use crate::message::OutboundMessage;
1838
1839#[cfg(test)]
1840mod tests {
1841 use super::*;
1842
1843 #[test]
1844 fn server_address_ipv6() {
1845 let address = ServerAddr::from_str("nats://[::]").unwrap();
1846 assert_eq!(address.host(), "::")
1847 }
1848
1849 #[test]
1850 fn server_address_ipv4() {
1851 let address = ServerAddr::from_str("nats://127.0.0.1").unwrap();
1852 assert_eq!(address.host(), "127.0.0.1")
1853 }
1854
1855 #[test]
1856 fn server_address_domain() {
1857 let address = ServerAddr::from_str("nats://example.com").unwrap();
1858 assert_eq!(address.host(), "example.com")
1859 }
1860
1861 #[test]
1862 fn to_server_addrs_vec_str() {
1863 let vec = vec!["nats://127.0.0.1", "nats://[::]"];
1864 let mut addrs_iter = vec.to_server_addrs().unwrap();
1865 assert_eq!(addrs_iter.next().unwrap().host(), "127.0.0.1");
1866 assert_eq!(addrs_iter.next().unwrap().host(), "::");
1867 assert_eq!(addrs_iter.next(), None);
1868 }
1869
1870 #[test]
1871 fn to_server_addrs_arr_str() {
1872 let arr = ["nats://127.0.0.1", "nats://[::]"];
1873 let mut addrs_iter = arr.to_server_addrs().unwrap();
1874 assert_eq!(addrs_iter.next().unwrap().host(), "127.0.0.1");
1875 assert_eq!(addrs_iter.next().unwrap().host(), "::");
1876 assert_eq!(addrs_iter.next(), None);
1877 }
1878
1879 #[test]
1880 fn to_server_addrs_vec_string() {
1881 let vec = vec!["nats://127.0.0.1".to_string(), "nats://[::]".to_string()];
1882 let mut addrs_iter = vec.to_server_addrs().unwrap();
1883 assert_eq!(addrs_iter.next().unwrap().host(), "127.0.0.1");
1884 assert_eq!(addrs_iter.next().unwrap().host(), "::");
1885 assert_eq!(addrs_iter.next(), None);
1886 }
1887
1888 #[test]
1889 fn to_server_addrs_arr_string() {
1890 let arr = ["nats://127.0.0.1".to_string(), "nats://[::]".to_string()];
1891 let mut addrs_iter = arr.to_server_addrs().unwrap();
1892 assert_eq!(addrs_iter.next().unwrap().host(), "127.0.0.1");
1893 assert_eq!(addrs_iter.next().unwrap().host(), "::");
1894 assert_eq!(addrs_iter.next(), None);
1895 }
1896
1897 #[test]
1898 fn to_server_ports_arr_string() {
1899 for (arr, expected_port) in [
1900 (
1901 [
1902 "nats://127.0.0.1".to_string(),
1903 "nats://[::]".to_string(),
1904 "tls://127.0.0.1".to_string(),
1905 "tls://[::]".to_string(),
1906 ],
1907 4222,
1908 ),
1909 (
1910 [
1911 "ws://127.0.0.1:80".to_string(),
1912 "ws://[::]:80".to_string(),
1913 "ws://127.0.0.1".to_string(),
1914 "ws://[::]".to_string(),
1915 ],
1916 80,
1917 ),
1918 (
1919 [
1920 "wss://127.0.0.1".to_string(),
1921 "wss://[::]".to_string(),
1922 "wss://127.0.0.1:443".to_string(),
1923 "wss://[::]:443".to_string(),
1924 ],
1925 443,
1926 ),
1927 ] {
1928 let mut addrs_iter = arr.to_server_addrs().unwrap();
1929 assert_eq!(addrs_iter.next().unwrap().port(), expected_port);
1930 }
1931 }
1932}