add tweak_new_upstream_tcp_connection hook

Add a new hook which is invoked on a TCP socket before it is connected.

Author: Noah-Kennedy

.bleep

@@ -1 +1 @@
-9df2f3f6e6b919a632b08af4584a1c1a3bcee0fd
+5c08613f7ab19914dc8a45a37590a06a325127bb

Cargo.toml

@@ -27,6 +27,7 @@ tokio = "1"
 async-trait = "0.1.42"
 httparse = "1"
 bytes = "1.0"
+derivative = "2.2.0"
 http = "1.0.0"
 log = "0.4"
 h2 = ">=0.4.6"

pingora-core/Cargo.toml

@@ -35,6 +35,7 @@ bytes = { workspace = true }
 http = { workspace = true }
 log = { workspace = true }
 h2 = { workspace = true }
+derivative.workspace = true
 clap = { version = "3.2.25", features = ["derive"] }
 once_cell = { workspace = true }
 serde = { version = "1.0", features = ["derive"] }
@@ -88,8 +89,8 @@ jemallocator = "0.5"
 
 [features]
 default = []
-openssl = ["pingora-openssl", "openssl_derived",]
-boringssl = ["pingora-boringssl", "openssl_derived",]
+openssl = ["pingora-openssl", "openssl_derived"]
+boringssl = ["pingora-boringssl", "openssl_derived"]
 rustls = ["pingora-rustls", "any_tls", "dep:x509-parser", "ouroboros"]
 patched_http1 = ["pingora-http/patched_http1"]
 openssl_derived = ["any_tls"]

pingora-core/src/connectors/l4.rs

@@ -12,16 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use async_trait::async_trait;
-use log::debug;
-use pingora_error::{Context, Error, ErrorType::*, OrErr, Result};
-use rand::seq::SliceRandom;
-use std::net::SocketAddr as InetSocketAddr;
-#[cfg(unix)]
-use std::os::unix::io::AsRawFd;
-#[cfg(windows)]
-use std::os::windows::io::AsRawSocket;
-
 #[cfg(unix)]
 use crate::protocols::l4::ext::connect_uds;
 use crate::protocols::l4::ext::{
@@ -31,6 +21,15 @@ use crate::protocols::l4::socket::SocketAddr;
 use crate::protocols::l4::stream::Stream;
 use crate::protocols::{GetSocketDigest, SocketDigest};
 use crate::upstreams::peer::Peer;
+use async_trait::async_trait;
+use log::debug;
+use pingora_error::{Context, Error, ErrorType::*, OrErr, Result};
+use rand::seq::SliceRandom;
+use std::net::SocketAddr as InetSocketAddr;
+#[cfg(unix)]
+use std::os::unix::io::AsRawFd;
+#[cfg(windows)]
+use std::os::windows::io::AsRawSocket;
 
 /// The interface to establish a L4 connection
 #[async_trait]
@@ -123,6 +122,14 @@ where
                             debug!("Setting dscp");
                             set_dscp(raw, dscp)?;
                         }
+
+                        if let Some(tweak_hook) = peer
+                            .get_peer_options()
+                            .and_then(|o| o.upstream_tcp_sock_tweak_hook.clone())
+                        {
+                            tweak_hook(socket)?;
+                        }
+
                         Ok(())
                     });
                     let conn_res = match peer.connection_timeout() {
@@ -302,6 +309,8 @@ mod tests {
     use crate::upstreams::peer::{BasicPeer, HttpPeer, Proxy};
     use std::collections::BTreeMap;
     use std::path::PathBuf;
+    use std::sync::atomic::{AtomicBool, Ordering};
+    use std::sync::Arc;
     use tokio::io::AsyncWriteExt;
     #[cfg(unix)]
     use tokio::net::UnixListener;
@@ -358,6 +367,26 @@ mod tests {
         assert_eq!(new_session.unwrap_err().etype(), &ConnectTimedout)
     }
 
+    #[tokio::test]
+    async fn test_tweak_hook() {
+        const INIT_FLAG: bool = false;
+
+        let flag = Arc::new(AtomicBool::new(INIT_FLAG));
+
+        let mut peer = BasicPeer::new("1.1.1.1:80");
+
+        let move_flag = Arc::clone(&flag);
+
+        peer.options.upstream_tcp_sock_tweak_hook = Some(Arc::new(move |_| {
+            move_flag.fetch_xor(true, Ordering::SeqCst);
+            Ok(())
+        }));
+
+        connect(&peer, None).await.unwrap();
+
+        assert_eq!(!INIT_FLAG, flag.load(Ordering::SeqCst));
+    }
+
     #[tokio::test]
     async fn test_custom_connect() {
         #[derive(Debug)]

pingora-core/src/upstreams/peer.rs

@@ -14,7 +14,15 @@
 
 //! Defines where to connect to and how to connect to a remote server
 
+use crate::connectors::{l4::BindTo, L4Connect};
+use crate::protocols::l4::socket::SocketAddr;
+use crate::protocols::tls::CaType;
+#[cfg(unix)]
+use crate::protocols::ConnFdReusable;
+use crate::protocols::TcpKeepalive;
+use crate::utils::tls::{get_organization_unit, CertKey};
 use ahash::AHasher;
+use derivative::Derivative;
 use pingora_error::{
     ErrorType::{InternalError, SocketError},
     OrErr, Result,
@@ -30,14 +38,7 @@ use std::os::windows::io::AsRawSocket;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use std::time::Duration;
-
-use crate::connectors::{l4::BindTo, L4Connect};
-use crate::protocols::l4::socket::SocketAddr;
-use crate::protocols::tls::CaType;
-#[cfg(unix)]
-use crate::protocols::ConnFdReusable;
-use crate::protocols::TcpKeepalive;
-use crate::utils::tls::{get_organization_unit, CertKey};
+use tokio::net::TcpSocket;
 
 pub use crate::protocols::tls::ALPN;
 
@@ -203,6 +204,17 @@ pub trait Peer: Display + Clone {
     fn get_tracer(&self) -> Option<Tracer> {
         None
     }
+
+    /// Returns a hook that should be run before an upstream TCP connection is connected.
+    ///
+    /// This hook can be used to set additional socket options.
+    fn upstream_tcp_sock_tweak_hook(
+        &self,
+    ) -> Option<&Arc<dyn Fn(&TcpSocket) -> Result<()> + Send + Sync + 'static>> {
+        self.get_peer_options()?
+            .upstream_tcp_sock_tweak_hook
+            .as_ref()
+    }
 }
 
 /// A simple TCP or TLS peer without many complicated settings.
@@ -303,7 +315,9 @@ impl Scheme {
 /// The preferences to connect to a remote server
 ///
 /// See [`Peer`] for the meaning of the fields
-#[derive(Clone, Debug)]
+#[non_exhaustive]
+#[derive(Clone, Derivative)]
+#[derivative(Debug)]
 pub struct PeerOptions {
     pub bind_to: Option<BindTo>,
     pub connection_timeout: Option<Duration>,
@@ -335,6 +349,9 @@ pub struct PeerOptions {
     pub tracer: Option<Tracer>,
     // A custom L4 connector to use to establish new L4 connections
     pub custom_l4: Option<Arc<dyn L4Connect + Send + Sync>>,
+    #[derivative(Debug = "ignore")]
+    pub upstream_tcp_sock_tweak_hook:
+        Option<Arc<dyn Fn(&TcpSocket) -> Result<()> + Send + Sync + 'static>>,
 }
 
 impl PeerOptions {
@@ -363,6 +380,7 @@ impl PeerOptions {
             tcp_fast_open: false,
             tracer: None,
             custom_l4: None,
+            upstream_tcp_sock_tweak_hook: None,
         }
     }
 

pingora-load-balancing/Cargo.toml

@@ -30,7 +30,7 @@ tokio = { workspace = true }
 futures = "0"
 log = { workspace = true }
 http = { workspace = true }
-derivative = "2.2.0"
+derivative.workspace = true
 
 [dev-dependencies]