couchbase_core/memdx/
connection.rs

1/*
2 *
3 *  * Copyright (c) 2025 Couchbase, Inc.
4 *  *
5 *  * Licensed under the Apache License, Version 2.0 (the "License");
6 *  * you may not use this file except in compliance with the License.
7 *  * You may obtain a copy of the License at
8 *  *
9 *  *    http://www.apache.org/licenses/LICENSE-2.0
10 *  *
11 *  * Unless required by applicable law or agreed to in writing, software
12 *  * distributed under the License is distributed on an "AS IS" BASIS,
13 *  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 *  * See the License for the specific language governing permissions and
15 *  * limitations under the License.
16 *
17 */
18
19use crate::memdx::error::Error;
20use crate::memdx::error::Result;
21use crate::tls_config::TlsConfig;
22use socket2::TcpKeepalive;
23use std::fmt::Debug;
24use std::io;
25use std::net::{IpAddr, Ipv4Addr, SocketAddr};
26use std::time::Duration;
27use tokio::io::{AsyncRead, AsyncWrite};
28use tokio::net::TcpStream;
29use tokio::time::{timeout_at, Instant};
30
31use crate::address::Address;
32#[cfg(all(feature = "rustls-tls", not(feature = "native-tls")))]
33use {
34    tokio_rustls::rustls::pki_types::DnsName, tokio_rustls::rustls::pki_types::ServerName,
35    tokio_rustls::TlsConnector,
36};
37
38#[derive(Debug)]
39pub struct ConnectOptions {
40    pub deadline: Instant,
41    pub tcp_keep_alive_time: Duration,
42}
43
44pub trait Stream: Debug + AsyncWrite + AsyncRead + Send + Sync + Unpin + 'static {}
45
46impl Stream for TcpStream {}
47
48#[derive(Debug)]
49#[allow(clippy::large_enum_variant)]
50pub enum ConnectionType {
51    Tcp(TcpConnection),
52    Tls(TlsConnection),
53}
54
55impl ConnectionType {
56    pub fn into_inner(self) -> Box<dyn Stream> {
57        match self {
58            ConnectionType::Tcp(connection) => Box::new(connection.stream),
59            ConnectionType::Tls(connection) => Box::new(connection.stream),
60        }
61    }
62
63    pub fn local_addr(&self) -> &SocketAddr {
64        match self {
65            ConnectionType::Tcp(connection) => &connection.local_addr,
66            ConnectionType::Tls(connection) => &connection.local_addr,
67        }
68    }
69
70    pub fn peer_addr(&self) -> &SocketAddr {
71        match self {
72            ConnectionType::Tcp(connection) => &connection.peer_addr,
73            ConnectionType::Tls(connection) => &connection.peer_addr,
74        }
75    }
76}
77
78#[derive(Debug)]
79pub struct TcpConnection {
80    stream: TcpStream,
81
82    local_addr: SocketAddr,
83    peer_addr: SocketAddr,
84}
85
86impl TcpConnection {
87    async fn tcp_stream(
88        addr: &str,
89        opts: &ConnectOptions,
90    ) -> Result<(TcpStream, SocketAddr, SocketAddr)> {
91        let tcp_socket = timeout_at(opts.deadline, TcpStream::connect(addr))
92            .await
93            .map_err(|e| {
94                Error::new_connection_failed_error(
95                    "failed to connect to server within timeout",
96                    Box::new(io::Error::new(io::ErrorKind::TimedOut, e)),
97                )
98            })?
99            .map_err(|e| {
100                Error::new_connection_failed_error("failed to create tcp stream", Box::new(e))
101            })?;
102
103        let local_addr = tcp_socket
104            .local_addr()
105            .unwrap_or_else(|_e| SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0));
106
107        let peer_addr = tcp_socket
108            .peer_addr()
109            .unwrap_or_else(|_e| SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0));
110
111        // Tokio doesn't expose a keep alive function, but they just call into socket2 for set_linger.
112        socket2::SockRef::from(&tcp_socket)
113            .set_tcp_keepalive(&TcpKeepalive::new().with_time(opts.tcp_keep_alive_time))?;
114
115        tcp_socket.set_nodelay(false).map_err(|e| {
116            Error::new_connection_failed_error("failed to set tcp nodelay", Box::new(e))
117        })?;
118
119        Ok((tcp_socket, local_addr, peer_addr))
120    }
121
122    pub async fn connect(addr: Address, opts: ConnectOptions) -> Result<TcpConnection> {
123        let (stream, local_addr, peer_addr) =
124            TcpConnection::tcp_stream(addr.to_string().as_str(), &opts).await?;
125
126        Ok(TcpConnection {
127            stream,
128            local_addr,
129            peer_addr,
130        })
131    }
132
133    fn local_addr(&self) -> &SocketAddr {
134        &self.local_addr
135    }
136
137    fn peer_addr(&self) -> &SocketAddr {
138        &self.peer_addr
139    }
140}
141
142#[derive(Debug)]
143pub struct TlsConnection {
144    #[cfg(all(feature = "rustls-tls", not(feature = "native-tls")))]
145    stream: tokio_rustls::client::TlsStream<TcpStream>,
146    #[cfg(feature = "native-tls")]
147    stream: tokio_native_tls::TlsStream<TcpStream>,
148
149    local_addr: SocketAddr,
150    peer_addr: SocketAddr,
151}
152
153#[cfg(all(feature = "rustls-tls", not(feature = "native-tls")))]
154impl Stream for tokio_rustls::client::TlsStream<TcpStream> {}
155
156#[cfg(feature = "native-tls")]
157impl Stream for tokio_native_tls::TlsStream<TcpStream> {}
158
159impl TlsConnection {
160    #[cfg(all(feature = "rustls-tls", not(feature = "native-tls")))]
161    pub async fn connect(
162        addr: Address,
163        tls_config: TlsConfig,
164        opts: ConnectOptions,
165    ) -> Result<TlsConnection> {
166        let (tcp_socket, local_addr, peer_addr) =
167            TcpConnection::tcp_stream(addr.to_string().as_str(), &opts).await?;
168
169        let connector = TlsConnector::from(tls_config);
170
171        let server_name = match DnsName::try_from(addr.host) {
172            Ok(name) => ServerName::DnsName(name),
173            Err(_e) => ServerName::IpAddress(tokio_rustls::rustls::pki_types::IpAddr::from(
174                peer_addr.ip(),
175            )),
176        };
177
178        let stream = timeout_at(opts.deadline, connector.connect(server_name, tcp_socket))
179            .await
180            .map_err(|e| {
181                Error::new_connection_failed_error(
182                    "failed to upgrade tcp stream to tls within timeout",
183                    Box::new(io::Error::new(io::ErrorKind::TimedOut, e)),
184                )
185            })?
186            .map_err(|e| {
187                Error::new_connection_failed_error(
188                    "failed to upgrade tcp stream to tls",
189                    Box::new(e),
190                )
191            })?;
192
193        Ok(TlsConnection {
194            stream,
195            local_addr,
196            peer_addr,
197        })
198    }
199
200    #[cfg(feature = "native-tls")]
201    pub async fn connect(
202        addr: Address,
203        tls_config: TlsConfig,
204        opts: ConnectOptions,
205    ) -> Result<TlsConnection> {
206        let (tcp_socket, local_addr, peer_addr) =
207            TcpConnection::tcp_stream(addr.to_string().as_str(), &opts).await?;
208
209        let tls_connector = tokio_native_tls::TlsConnector::from(tls_config);
210
211        let remote_addr = addr.to_string();
212        let stream = timeout_at(
213            opts.deadline,
214            tls_connector.connect(&remote_addr, tcp_socket),
215        )
216        .await
217        .map_err(|e| {
218            Error::new_connection_failed_error(
219                "failed to upgrade tcp stream to tls within timeout",
220                Box::new(io::Error::new(io::ErrorKind::TimedOut, e)),
221            )
222        })?
223        .map_err(|e| {
224            Error::new_connection_failed_error(
225                "failed to upgrade tcp stream to tls",
226                Box::new(io::Error::other(e)),
227            )
228        })?;
229
230        Ok(TlsConnection {
231            stream,
232            local_addr,
233            peer_addr,
234        })
235    }
236
237    fn local_addr(&self) -> &SocketAddr {
238        &self.local_addr
239    }
240
241    fn peer_addr(&self) -> &SocketAddr {
242        &self.peer_addr
243    }
244}
couchbase_core/memdx/connection.rs

couchbase_core/memdx/
connection.rs
