watermelon_mini/proto/
connector.rs

1use std::io;
2
3use tokio::net::TcpStream;
4use tokio_rustls::{
5    TlsConnector,
6    rustls::{
7        self,
8        pki_types::{InvalidDnsNameError, ServerName},
9    },
10};
11use watermelon_net::{
12    Connection, StreamingConnection, connect_tcp,
13    error::{ConnectionReadError, StreamingReadError},
14    proto_connect,
15};
16#[cfg(feature = "websocket")]
17use watermelon_net::{WebsocketConnection, error::WebsocketReadError};
18#[cfg(feature = "websocket")]
19use watermelon_proto::proto::error::FrameDecoderError;
20use watermelon_proto::{
21    Connect, Host, NonStandardConnect, Protocol, ServerAddr, ServerInfo, Transport,
22    proto::{ServerOp, error::DecoderError},
23};
24
25use crate::{ConnectFlags, ConnectionCompression, util::MaybeConnection};
26
27use super::{
28    authenticator::{AuthenticationError, AuthenticationMethod},
29    connection::ConnectionSecurity,
30};
31
32#[derive(Debug, thiserror::Error)]
33pub enum ConnectError {
34    #[error("connection I/O error")]
35    Io(#[source] io::Error),
36    #[error("TLS handshake failed")]
37    Tls(rustls::Error),
38    #[error("invalid DNS name")]
39    InvalidDnsName(#[source] InvalidDnsNameError),
40    #[error("websocket not supported")]
41    WebsocketUnsupported,
42    #[error("received unexpected server operation during connect")]
43    UnexpectedServerOp,
44    #[error("failed to decode server message")]
45    Decoder(#[source] DecoderError),
46    #[error("authentication failed")]
47    Authentication(#[source] AuthenticationError),
48    #[error("connect handshake failed")]
49    Connect(#[source] watermelon_net::error::ConnectError),
50}
51
52#[expect(clippy::too_many_lines)]
53pub(crate) async fn connect(
54    connector: &TlsConnector,
55    addr: &ServerAddr,
56    client_name: Option<String>,
57    auth_method: Option<&AuthenticationMethod>,
58    flags: ConnectFlags,
59) -> Result<
60    (
61        Connection<
62            ConnectionCompression<ConnectionSecurity<TcpStream>>,
63            ConnectionSecurity<TcpStream>,
64        >,
65        Box<ServerInfo>,
66    ),
67    ConnectError,
68> {
69    let conn = connect_tcp(addr).await.map_err(ConnectError::Io)?;
70    conn.set_nodelay(flags.tcp_nodelay)
71        .map_err(ConnectError::Io)?;
72    let mut conn = ConnectionSecurity::Plain(conn);
73
74    if matches!(addr.protocol(), Protocol::TLS) {
75        let domain = rustls_server_name_from_addr(addr).map_err(ConnectError::InvalidDnsName)?;
76        conn = conn
77            .upgrade_tls(connector, domain.to_owned())
78            .await
79            .map_err(ConnectError::Io)?;
80    }
81
82    let mut conn = match addr.transport() {
83        Transport::TCP => Connection::Streaming(StreamingConnection::new(conn)),
84        #[cfg(feature = "websocket")]
85        Transport::Websocket => {
86            let uri = addr.to_string().parse().unwrap();
87            Connection::Websocket(
88                WebsocketConnection::new(uri, conn)
89                    .await
90                    .map_err(ConnectError::Io)?,
91            )
92        }
93        #[cfg(not(feature = "websocket"))]
94        Transport::Websocket => return Err(ConnectError::WebsocketUnsupported),
95    };
96    let info = match conn.read_next().await {
97        Ok(ServerOp::Info { info }) => info,
98        Ok(_) => return Err(ConnectError::UnexpectedServerOp),
99        Err(ConnectionReadError::Streaming(StreamingReadError::Io(err))) => {
100            return Err(ConnectError::Io(err));
101        }
102        Err(ConnectionReadError::Streaming(StreamingReadError::Decoder(err))) => {
103            return Err(ConnectError::Decoder(err));
104        }
105        #[cfg(feature = "websocket")]
106        Err(ConnectionReadError::Websocket(WebsocketReadError::Io(err))) => {
107            return Err(ConnectError::Io(err));
108        }
109        #[cfg(feature = "websocket")]
110        Err(ConnectionReadError::Websocket(WebsocketReadError::Decoder(
111            FrameDecoderError::Decoder(err),
112        ))) => return Err(ConnectError::Decoder(err)),
113        #[cfg(feature = "websocket")]
114        Err(ConnectionReadError::Websocket(WebsocketReadError::Decoder(
115            FrameDecoderError::IncompleteFrame,
116        ))) => {
117            // The server sent a websocket frame that doesn't contain a
118            // complete NATS operation.
119            return Err(ConnectError::Io(io::Error::new(
120                io::ErrorKind::InvalidData,
121                FrameDecoderError::IncompleteFrame,
122            )));
123        }
124        #[cfg(feature = "websocket")]
125        Err(ConnectionReadError::Websocket(WebsocketReadError::Closed)) => {
126            // Match the streaming connection, which reports an EOF during
127            // the handshake as `UnexpectedEof`.
128            return Err(ConnectError::Io(io::ErrorKind::UnexpectedEof.into()));
129        }
130    };
131
132    let conn = match conn {
133        Connection::Streaming(streaming) => Connection::Streaming(
134            if matches!(
135                (addr.protocol(), info.tls_required),
136                (Protocol::PossiblyPlain, true)
137            ) {
138                let domain =
139                    rustls_server_name_from_addr(addr).map_err(ConnectError::InvalidDnsName)?;
140                StreamingConnection::new(
141                    streaming
142                        .into_inner()
143                        .upgrade_tls(connector, domain.to_owned())
144                        .await
145                        .map_err(ConnectError::Io)?,
146                )
147            } else {
148                streaming
149            },
150        ),
151        Connection::Websocket(websocket) => Connection::Websocket(websocket),
152    };
153
154    let auth;
155    let auth_method = if let Some(auth_method) = auth_method {
156        Some(auth_method)
157    } else if let Some(auth_method) = AuthenticationMethod::try_from_addr(addr) {
158        auth = auth_method;
159        Some(&auth)
160    } else {
161        None
162    };
163
164    #[allow(unused_mut)]
165    let mut non_standard = NonStandardConnect::default();
166    #[cfg(feature = "non-standard-zstd")]
167    if matches!(conn, Connection::Streaming(_)) {
168        non_standard.zstd = flags.zstd_compression_level.is_some() && info.non_standard.zstd;
169    }
170
171    let mut connect = Connect {
172        verbose: true,
173        pedantic: false,
174        require_tls: false,
175        auth_token: None,
176        username: None,
177        password: None,
178        client_name,
179        client_lang: "rust-watermelon",
180        client_version: env!("CARGO_PKG_VERSION"),
181        protocol: 1,
182        echo: flags.echo,
183        signature: None,
184        jwt: None,
185        supports_no_responders: true,
186        supports_headers: true,
187        nkey: None,
188        non_standard,
189    };
190    if let Some(auth_method) = auth_method {
191        auth_method
192            .prepare_for_auth(&info, &mut connect)
193            .map_err(ConnectError::Authentication)?;
194    }
195
196    let mut conn = match conn {
197        Connection::Streaming(streaming) => {
198            Connection::Streaming(streaming.replace_socket(|stream| {
199                MaybeConnection(Some(ConnectionCompression::Plain(stream)))
200            }))
201        }
202        Connection::Websocket(websocket) => Connection::Websocket(websocket),
203    };
204
205    #[cfg(feature = "non-standard-zstd")]
206    let zstd = connect.non_standard.zstd;
207
208    proto_connect(&mut conn, connect, |conn| {
209        #[cfg(feature = "non-standard-zstd")]
210        match conn {
211            Connection::Streaming(streaming) => {
212                if zstd && let Some(zstd_compression_level) = flags.zstd_compression_level {
213                    let stream = streaming.socket_mut().0.take().unwrap();
214                    streaming.socket_mut().0 = Some(stream.upgrade_zstd(zstd_compression_level));
215                }
216            }
217            Connection::Websocket(_websocket) => {}
218        }
219
220        let _ = conn;
221    })
222    .await
223    .map_err(ConnectError::Connect)?;
224
225    let conn = match conn {
226        Connection::Streaming(streaming) => {
227            Connection::Streaming(streaming.replace_socket(|stream| stream.0.unwrap()))
228        }
229        Connection::Websocket(websocket) => Connection::Websocket(websocket),
230    };
231
232    Ok((conn, info))
233}
234
235fn rustls_server_name_from_addr(addr: &ServerAddr) -> Result<ServerName<'_>, InvalidDnsNameError> {
236    match addr.host() {
237        Host::Ip(addr) => Ok(ServerName::IpAddress((*addr).into())),
238        Host::Dns(name) => <_ as AsRef<str>>::as_ref(name).try_into(),
239    }
240}
watermelon_mini/proto/connector.rs

watermelon_mini/proto/
connector.rs
