1#[cfg(unix)]
16use crate::protocols::l4::ext::connect_uds;
17use crate::protocols::l4::ext::{
18 connect_error_local_addr, connect_with_attempt as tcp_connect, set_dscp, set_recv_buf,
19 set_tcp_fastopen_connect, ConnectAttempt, ConnectErrorDetails,
20};
21use crate::protocols::l4::socket::SocketAddr;
22use crate::protocols::l4::stream::Stream;
23use crate::protocols::{GetSocketDigest, SocketDigest};
24use crate::upstreams::peer::Peer;
25use async_trait::async_trait;
26use log::debug;
27use pingora_error::{Context, Error, ErrorType::*, OrErr, Result};
28use rand::seq::SliceRandom;
29use std::net::SocketAddr as InetSocketAddr;
30#[cfg(unix)]
31use std::os::unix::io::AsRawFd;
32#[cfg(windows)]
33use std::os::windows::io::AsRawSocket;
34
35#[async_trait]
37pub trait Connect: std::fmt::Debug {
38 async fn connect(&self, addr: &SocketAddr) -> Result<Stream>;
39}
40
41pub trait ConnectErrorExt {
43 fn connect_local_addr(&self) -> Option<InetSocketAddr>;
48}
49
50impl ConnectErrorExt for Error {
51 fn connect_local_addr(&self) -> Option<InetSocketAddr> {
52 connect_error_local_addr(self)
53 }
54}
55
56#[derive(Clone, Debug, Default)]
58pub struct BindTo {
59 pub addr: Option<InetSocketAddr>,
61 port_range: Option<(u16, u16)>,
63 fallback: bool,
65}
66
67impl BindTo {
68 pub fn set_port_range(&mut self, range: Option<(u16, u16)>) -> Result<()> {
74 if range.is_none() && self.port_range.is_none() {
75 return Ok(());
77 }
78
79 match range {
80 None | Some((0, 0)) => self.port_range = Some((0, 0)),
82 Some((low, high)) if low > 0 && low < high => {
84 self.port_range = Some((low, high));
85 }
86 _ => return Error::e_explain(SocketError, "invalid port range: {range}"),
87 }
88 Ok(())
89 }
90
91 pub fn set_fallback(&mut self, fallback: bool) {
93 self.fallback = fallback
94 }
95
96 pub fn port_range(&self) -> Option<(u16, u16)> {
98 self.port_range
99 }
100
101 pub fn will_fallback(&self) -> bool {
103 self.fallback && self.port_range.is_some()
104 }
105}
106
107pub(crate) async fn connect<P>(peer: &P, bind_to: Option<BindTo>) -> Result<Stream>
109where
110 P: Peer + Send + Sync,
111{
112 if peer.get_proxy().is_some() {
113 return proxy_connect(peer)
114 .await
115 .err_context(|| format!("Fail to establish CONNECT proxy: {}", peer));
116 }
117 let peer_addr = peer.address();
118 let mut local_addr = None;
119 let mut stream: Stream =
120 if let Some(custom_l4) = peer.get_peer_options().and_then(|o| o.custom_l4.as_ref()) {
121 custom_l4.connect(peer_addr).await?
122 } else {
123 match peer_addr {
124 SocketAddr::Inet(addr) => {
125 let mut connect_attempt = ConnectAttempt::default();
126 let connect_future = tcp_connect(
127 addr,
128 bind_to.as_ref(),
129 |socket| {
130 #[cfg(unix)]
131 let raw = socket.as_raw_fd();
132 #[cfg(windows)]
133 let raw = socket.as_raw_socket();
134
135 if peer.tcp_fast_open() {
136 set_tcp_fastopen_connect(raw)?;
137 }
138 if let Some(recv_buf) = peer.tcp_recv_buf() {
139 debug!("Setting recv buf size");
140 set_recv_buf(raw, recv_buf)?;
141 }
142 if let Some(dscp) = peer.dscp() {
143 debug!("Setting dscp");
144 set_dscp(raw, dscp)?;
145 }
146
147 if let Some(tweak_hook) = peer
148 .get_peer_options()
149 .and_then(|o| o.upstream_tcp_sock_tweak_hook.clone())
150 {
151 tweak_hook(socket)?;
152 }
153
154 Ok(())
155 },
156 &mut connect_attempt,
157 );
158 let conn_res = match peer.connection_timeout() {
159 Some(t) => match pingora_timeout::timeout(t, connect_future).await {
160 Ok(result) => result,
161 Err(e) => {
162 let context = format!("timeout {t:?} connecting to server {peer}");
163 return Err(match connect_attempt.local_addr() {
164 Some(local_addr) => Error::because(
165 ConnectTimedout,
166 context,
167 ConnectErrorDetails::new(e, Some(local_addr)),
168 ),
169 None => Error::because(ConnectTimedout, context, e),
170 });
171 }
172 },
173 None => connect_future.await,
174 };
175 match conn_res {
176 Ok(socket) => {
177 local_addr = connect_attempt.local_addr();
178 debug!("connected to new server: {}", peer.address());
179 Ok(socket.into())
180 }
181 Err(e) => {
182 let c = format!("Fail to connect to {peer}");
183 match e.etype() {
184 SocketError | BindError => Error::e_because(InternalError, c, e),
185 _ => Err(e.more_context(c)),
186 }
187 }
188 }
189 }
190 #[cfg(unix)]
191 SocketAddr::Unix(addr) => {
192 let connect_future = connect_uds(
193 addr.as_pathname()
194 .expect("non-pathname unix sockets not supported as peer"),
195 );
196 let conn_res = match peer.connection_timeout() {
197 Some(t) => pingora_timeout::timeout(t, connect_future)
198 .await
199 .explain_err(ConnectTimedout, |_| {
200 format!("timeout {t:?} connecting to server {peer}")
201 })?,
202 None => connect_future.await,
203 };
204 match conn_res {
205 Ok(socket) => {
206 debug!("connected to new server: {}", peer.address());
207 Ok(socket.into())
208 }
209 Err(e) => {
210 let c = format!("Fail to connect to {peer}");
211 match e.etype() {
212 SocketError | BindError => Error::e_because(InternalError, c, e),
213 _ => Err(e.more_context(c)),
214 }
215 }
216 }
217 }
218 }?
219 };
220
221 let tracer = peer.get_tracer();
222 if let Some(t) = tracer {
223 t.0.on_connected();
224 stream.tracer = Some(t);
225 }
226
227 if let Some(ka) = peer.tcp_keepalive() {
229 stream.set_keepalive(ka)?;
230 }
231 stream.set_nodelay()?;
232
233 #[cfg(unix)]
234 let digest = SocketDigest::from_raw_fd(stream.as_raw_fd());
235 #[cfg(windows)]
236 let digest = SocketDigest::from_raw_socket(stream.as_raw_socket());
237 digest
238 .peer_addr
239 .set(Some(peer_addr.clone()))
240 .expect("newly created OnceCell must be empty");
241 if let Some(local_addr) = local_addr {
242 digest
243 .local_addr
244 .set(Some(SocketAddr::Inet(local_addr)))
245 .expect("newly created OnceCell must be empty");
246 }
247 stream.set_socket_digest(digest);
248
249 Ok(stream)
250}
251
252pub(crate) fn bind_to_random<P: Peer>(
253 peer: &P,
254 v4_list: &[InetSocketAddr],
255 v6_list: &[InetSocketAddr],
256) -> Option<BindTo> {
257 fn bind_to_ips(ips: &[InetSocketAddr]) -> Option<InetSocketAddr> {
259 match ips.len() {
260 0 => None,
261 1 => Some(ips[0]),
262 _ => {
263 ips.choose(&mut rand::thread_rng()).copied()
265 }
266 }
267 }
268
269 let mut bind_to = peer.get_peer_options().and_then(|o| o.bind_to.clone());
270 if bind_to.as_ref().map(|b| b.addr).is_some() {
271 return bind_to;
273 }
274
275 let addr = match peer.address() {
276 SocketAddr::Inet(sockaddr) => match sockaddr {
277 InetSocketAddr::V4(_) => bind_to_ips(v4_list),
278 InetSocketAddr::V6(_) => bind_to_ips(v6_list),
279 },
280 #[cfg(unix)]
281 SocketAddr::Unix(_) => None,
282 };
283
284 if addr.is_some() {
285 if let Some(bind_to) = bind_to.as_mut() {
286 bind_to.addr = addr;
287 } else {
288 bind_to = Some(BindTo {
289 addr,
290 ..Default::default()
291 });
292 }
293 }
294 bind_to
295}
296
297use crate::protocols::raw_connect;
298
299#[cfg(unix)]
300async fn proxy_connect<P: Peer>(peer: &P) -> Result<Stream> {
301 let proxy = peer.get_proxy().unwrap();
303 let options = peer.get_peer_options().unwrap();
304
305 let mut headers = proxy
307 .headers
308 .iter()
309 .chain(options.extra_proxy_headers.iter());
310
311 let stream: Box<Stream> = Box::new(
313 connect_uds(&proxy.next_hop)
314 .await
315 .or_err_with(ConnectError, || {
316 format!("CONNECT proxy connect() error to {:?}", proxy.next_hop)
317 })?
318 .into(),
319 );
320
321 let req_header = raw_connect::generate_connect_header(&proxy.host, proxy.port, &mut headers)?;
322 let fut = raw_connect::connect(stream, &req_header, peer);
323 let (mut stream, digest) = match peer.connection_timeout() {
324 Some(t) => pingora_timeout::timeout(t, fut)
325 .await
326 .explain_err(ConnectTimedout, |_| "establishing CONNECT proxy")?,
327 None => fut.await,
328 }
329 .map_err(|mut e| {
330 e.retry.decide_reuse(false);
332 e
333 })?;
334 debug!("CONNECT proxy established: {:?}", proxy);
335 stream.set_proxy_digest(digest);
336 let stream = stream.into_any().downcast::<Stream>().unwrap(); Ok(*stream)
338}
339
340#[cfg(windows)]
341async fn proxy_connect<P: Peer>(peer: &P) -> Result<Stream> {
342 panic!("peer proxy not supported on windows")
343}
344
345#[cfg(test)]
346mod tests {
347 use super::*;
348 use crate::upstreams::peer::{BasicPeer, HttpPeer, Proxy};
349 use std::collections::BTreeMap;
350 use std::path::PathBuf;
351 use std::sync::atomic::{AtomicBool, Ordering};
352 use std::sync::Arc;
353 use std::time::Duration;
354
355 #[cfg(target_os = "linux")]
356 async fn wait_for_peer<P>(peer: &P)
357 where
358 P: Peer + Send + Sync,
359 {
360 use pingora_error::ErrorType as E;
361 use std::time::Instant;
362 use tokio::time::sleep;
363
364 let start = Instant::now();
365 let mut res = connect(peer, None).await;
366 let mut delay = Duration::from_millis(5);
367 let max_delay = Duration::from_secs(10);
368
369 while start.elapsed() < max_delay {
370 match &res {
371 Err(e) if e.etype == E::ConnectRefused => {}
372 _ => break,
373 }
374 sleep(delay).await;
375 delay *= 2;
376 res = connect(peer, None).await;
377 }
378 }
379
380 #[tokio::test]
381 async fn test_conn_error_refused() {
382 let peer = BasicPeer::new("127.0.0.1:79"); let new_session = connect(&peer, None).await;
384 assert_eq!(new_session.unwrap_err().etype(), &ConnectRefused)
385 }
386
387 #[tokio::test]
388 async fn test_local_addr_is_cached() {
389 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
390 let peer_addr = listener.local_addr().unwrap().to_string();
391 let peer = BasicPeer::new(&peer_addr);
392 let accept = tokio::spawn(async move { listener.accept().await.unwrap() });
393
394 let stream = connect(&peer, None).await.unwrap();
395 let digest = stream.get_socket_digest().unwrap();
396 let cached_addr = digest
397 .local_addr
398 .get()
399 .expect("local address should be cached during connect")
400 .as_ref()
401 .and_then(SocketAddr::as_inet)
402 .unwrap();
403 let (_, observed_client_addr) = accept.await.unwrap();
404 assert_eq!(*cached_addr, observed_client_addr);
405 }
406
407 #[ignore]
409 #[tokio::test]
410 async fn test_conn_error_no_route() {
411 let peer = BasicPeer::new("[::3]:79"); let new_session = connect(&peer, None).await;
413 assert_eq!(new_session.unwrap_err().etype(), &ConnectNoRoute)
414 }
415
416 #[tokio::test]
417 async fn test_conn_error_addr_not_avail() {
418 let peer = HttpPeer::new("127.0.0.1:121".to_string(), false, "".to_string());
419 let addr = "192.0.2.2:0".parse().ok();
420 let bind_to = BindTo {
421 addr,
422 ..Default::default()
423 };
424 let new_session = connect(&peer, Some(bind_to)).await;
425 assert_eq!(new_session.unwrap_err().etype(), &InternalError)
426 }
427
428 #[tokio::test]
429 async fn test_conn_error_other() {
430 let peer = HttpPeer::new("240.0.0.1:80".to_string(), false, "".to_string()); let addr = "127.0.0.1:0".parse().ok();
432 let bind_to = BindTo {
434 addr,
435 ..Default::default()
436 };
437 let new_session = connect(&peer, Some(bind_to)).await;
438 let error = new_session.unwrap_err();
439 assert!(
441 error.etype() == &ConnectError
442 || error.etype() == &ConnectTimedout
443 || (error.etype() == &InternalError),
445 "{error:?}"
446 )
447 }
448
449 #[tokio::test]
450 async fn test_conn_timeout() {
451 let mut peer = BasicPeer::new("192.0.2.1:79");
455 peer.options.connection_timeout = Some(Duration::from_millis(1));
456 let err = connect(&peer, None).await.unwrap_err();
457 assert!(
458 err.etype() == &ConnectTimedout || err.etype() == &ConnectNoRoute,
459 "unexpected error type: {:?}",
460 err.etype()
461 );
462 if err.etype() == &ConnectTimedout {
463 let local_addr = err
464 .connect_local_addr()
465 .expect("local address should be captured before the timeout");
466 assert_ne!(local_addr.port(), 0);
467 }
468 }
469
470 #[tokio::test]
471 async fn test_tweak_hook() {
472 const INIT_FLAG: bool = false;
473
474 let flag = Arc::new(AtomicBool::new(INIT_FLAG));
475
476 let mut peer = BasicPeer::new("1.1.1.1:80");
477
478 let move_flag = Arc::clone(&flag);
479
480 peer.options.upstream_tcp_sock_tweak_hook = Some(Arc::new(move |_| {
481 move_flag.fetch_not(Ordering::SeqCst);
482 Ok(())
483 }));
484
485 connect(&peer, None).await.unwrap();
486
487 assert_eq!(!INIT_FLAG, flag.load(Ordering::SeqCst));
488 }
489
490 #[tokio::test]
491 async fn test_custom_connect() {
492 #[derive(Debug)]
493 struct MyL4;
494 #[async_trait]
495 impl Connect for MyL4 {
496 async fn connect(&self, _addr: &SocketAddr) -> Result<Stream> {
497 tokio::net::TcpStream::connect("1.1.1.1:80")
498 .await
499 .map(|s| s.into())
500 .or_fail()
501 }
502 }
503 let mut peer = BasicPeer::new("1.1.1.1:79");
505 peer.options.custom_l4 = Some(std::sync::Arc::new(MyL4 {}));
506
507 let new_session = connect(&peer, None).await;
508
509 assert!(new_session.is_ok());
511 }
512
513 #[cfg(unix)]
514 #[tokio::test]
515 async fn test_connect_proxy_fail() {
516 let mut peer = HttpPeer::new("1.1.1.1:80".to_string(), false, "".to_string());
517 let mut path = PathBuf::new();
518 path.push("/tmp/123");
519 peer.proxy = Some(Proxy {
520 next_hop: path.into(),
521 host: "1.1.1.1".into(),
522 port: 80,
523 headers: BTreeMap::new(),
524 });
525 let new_session = connect(&peer, None).await;
526 let e = new_session.unwrap_err();
527 assert_eq!(e.etype(), &ConnectError);
528 assert!(!e.retry());
529 }
530
531 #[cfg(unix)]
532 #[tokio::test(flavor = "multi_thread")]
533 async fn test_connect_proxy_work() {
534 use crate::connectors::test_utils;
535
536 let socket_path = test_utils::unique_uds_path("connect_proxy_work");
537 let (ready_rx, shutdown_tx, server_handle) =
538 test_utils::spawn_mock_uds_server(socket_path.clone(), b"HTTP/1.1 200 OK\r\n\r\n");
539
540 ready_rx.await.unwrap();
542
543 let mut peer = HttpPeer::new("1.1.1.1:80".to_string(), false, "".to_string());
544 let mut path = PathBuf::new();
545 path.push(&socket_path);
546 peer.proxy = Some(Proxy {
547 next_hop: path.into(),
548 host: "1.1.1.1".into(),
549 port: 80,
550 headers: BTreeMap::new(),
551 });
552 let new_session = connect(&peer, None).await;
553 assert!(new_session.is_ok());
554
555 let _ = shutdown_tx.send(());
557 server_handle.await.unwrap();
558 }
559
560 #[cfg(unix)]
561 #[tokio::test(flavor = "multi_thread")]
562 async fn test_connect_proxy_conn_closed() {
563 use crate::connectors::test_utils;
564
565 let socket_path = test_utils::unique_uds_path("connect_proxy_conn_closed");
566 let (ready_rx, shutdown_tx, server_handle) =
567 test_utils::spawn_mock_uds_server_close_immediate(socket_path.clone());
568
569 ready_rx.await.unwrap();
571
572 let mut peer = HttpPeer::new("1.1.1.1:80".to_string(), false, "".to_string());
573 let mut path = PathBuf::new();
574 path.push(&socket_path);
575 peer.proxy = Some(Proxy {
576 next_hop: path.into(),
577 host: "1.1.1.1".into(),
578 port: 80,
579 headers: BTreeMap::new(),
580 });
581 let new_session = connect(&peer, None).await;
582 let err = new_session.unwrap_err();
583 assert_eq!(err.etype(), &ConnectionClosed);
584 assert!(!err.retry());
585
586 let _ = shutdown_tx.send(());
588 server_handle.await.unwrap();
589 }
590
591 #[cfg(target_os = "linux")]
592 #[tokio::test(flavor = "multi_thread")]
593 async fn test_bind_to_port_range_on_connect() {
594 fn get_ip_local_port_range() -> (u16, u16) {
595 let path = "/proc/sys/net/ipv4/ip_local_port_range";
596 let file = std::fs::read_to_string(path).unwrap();
597 let mut parts = file.split_whitespace();
598 (
599 parts.next().unwrap().parse().unwrap(),
600 parts.next().unwrap().parse().unwrap(),
601 )
602 }
603
604 async fn mock_inet_connect_server() -> u16 {
606 use tokio::io::AsyncWriteExt;
607 use tokio::net::TcpListener;
608 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
609
610 let port = listener.local_addr().unwrap().port();
611
612 tokio::spawn(async move {
613 if let Ok((mut stream, _addr)) = listener.accept().await {
614 stream.write_all(b"HTTP/1.1 200 OK\r\n\r\n").await.unwrap();
615 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
617 }
618 });
619
620 port
621 }
622
623 fn in_port_range(session: Stream, lower: u16, upper: u16) -> bool {
624 let digest = session.get_socket_digest();
625 let local_addr = digest
626 .as_ref()
627 .and_then(|s| s.local_addr())
628 .unwrap()
629 .as_inet()
630 .unwrap();
631
632 local_addr.port() >= lower && local_addr.port() <= upper
634 }
635
636 let port = mock_inet_connect_server().await;
637
638 let (low, _) = get_ip_local_port_range();
641 let high = low + 1;
642
643 let peer = HttpPeer::new(format!("127.0.0.1:{port}"), false, "".to_string());
644 let mut bind_to = BindTo {
645 addr: "127.0.0.1:0".parse().ok(),
646 ..Default::default()
647 };
648
649 wait_for_peer(&peer).await;
651
652 bind_to.set_port_range(Some((low, high))).unwrap();
653
654 let mut success_count = 0;
655 let mut address_unavailable_count = 0;
656
657 for _ in 0..10 {
662 match connect(&peer, Some(bind_to.clone())).await {
663 Ok(session) => {
664 assert!(in_port_range(session, low, high));
665 success_count += 1;
666 }
667 Err(e) if format!("{e:?}").contains("AddrNotAvailable") => {
668 address_unavailable_count += 1;
669 }
670 Err(e) => {
671 panic!("Unexpected error {e:?}")
672 }
673 }
674 }
675
676 assert!(address_unavailable_count > 0);
677 assert!(success_count >= (high - low));
678
679 bind_to.set_fallback(true);
681 let session4 = connect(&peer, Some(bind_to.clone())).await.unwrap();
682 assert!(!in_port_range(session4, low, high));
683
684 let low = low + 2;
686 let high = low + 1;
687 let mut bind_to = BindTo::default();
688 bind_to.set_port_range(Some((low, high))).unwrap();
689 let session5 = connect(&peer, Some(bind_to.clone())).await.unwrap();
690 assert!(in_port_range(session5, low, high));
691 }
692
693 #[test]
694 fn test_bind_to_port_ranges() {
695 let addr = "127.0.0.1:0".parse().ok();
696 let mut bind_to = BindTo {
697 addr,
698 ..Default::default()
699 };
700
701 bind_to.set_port_range(None).unwrap();
703 assert!(bind_to.port_range.is_none());
704
705 bind_to.set_port_range(Some((0, 0))).unwrap();
707 assert_eq!(bind_to.port_range, Some((0, 0)));
708
709 bind_to.set_port_range(None).unwrap();
711 assert_eq!(bind_to.port_range, Some((0, 0)));
712
713 assert!(bind_to.set_port_range(Some((2000, 1000))).is_err());
715
716 bind_to.set_port_range(Some((1000, 2000))).unwrap();
718 assert_eq!(bind_to.port_range, Some((1000, 2000)));
719 }
720}