Add support for setting DSCP on client and server sockets

5 files changed, 50 insertions, 3 deletions

diff --git a/.bleep b/.bleep
index 48a93ee..8439eac 100644
--- a/.bleep
+++ b/.bleep
@@ -1 +1 @@
-bd707d83c6b344fa22ca0e4b61d751acea02f4bc
\ No newline at end of file
+837db6c7ec2d37abf83f9588be99fda00e2012c3
\ No newline at end of file
diff --git a/pingora-core/src/connectors/l4.rs b/pingora-core/src/connectors/l4.rs
index ffbf084..449ea4c 100644
--- a/pingora-core/src/connectors/l4.rs
+++ b/pingora-core/src/connectors/l4.rs
@@ -19,7 +19,7 @@ use std::net::SocketAddr as InetSocketAddr;
 use std::os::unix::io::AsRawFd;
 
 use crate::protocols::l4::ext::{
-    connect_uds, connect_with as tcp_connect, set_recv_buf, set_tcp_fastopen_connect,
+    connect_uds, connect_with as tcp_connect, set_dscp, set_recv_buf, set_tcp_fastopen_connect,
 };
 use crate::protocols::l4::socket::SocketAddr;
 use crate::protocols::l4::stream::Stream;
@@ -47,6 +47,10 @@ where
                     debug!("Setting recv buf size");
                     set_recv_buf(socket.as_raw_fd(), recv_buf)?;
                 }
+                if let Some(dscp) = peer.dscp() {
+                    debug!("Setting dscp");
+                    set_dscp(socket.as_raw_fd(), dscp)?;
+                }
                 Ok(())
             });
             let conn_res = match peer.connection_timeout() {
diff --git a/pingora-core/src/listeners/l4.rs b/pingora-core/src/listeners/l4.rs
index 42adc87..748037a 100644
--- a/pingora-core/src/listeners/l4.rs
+++ b/pingora-core/src/listeners/l4.rs
@@ -25,7 +25,7 @@ use std::os::unix::net::UnixListener as StdUnixListener;
 use std::time::Duration;
 use tokio::net::TcpSocket;
 
-use crate::protocols::l4::ext::set_tcp_fastopen_backlog;
+use crate::protocols::l4::ext::{set_dscp, set_tcp_fastopen_backlog};
 use crate::protocols::l4::listener::Listener;
 pub use crate::protocols::l4::stream::Stream;
 use crate::protocols::TcpKeepalive;
@@ -76,6 +76,9 @@ pub struct TcpSocketOptions {
     /// Enable TCP keepalive on accepted connections.
     /// See the [man page](https://man7.org/linux/man-pages/man7/tcp.7.html) for more information.
     pub tcp_keepalive: Option<TcpKeepalive>,
+    /// Specifies the server should set the following DSCP value on outgoing connections.
+    /// See the [RFC](https://datatracker.ietf.org/doc/html/rfc2474) for more details.
+    pub dscp: Option<u8>,
     // TODO: allow configuring reuseaddr, backlog, etc. from here?
 }
 
@@ -150,6 +153,10 @@ fn apply_tcp_socket_options(sock: &TcpSocket, opt: Option<&TcpSocketOptions>) ->
     if let Some(backlog) = opt.tcp_fastopen {
         set_tcp_fastopen_backlog(sock.as_raw_fd(), backlog)?;
     }
+
+    if let Some(dscp) = opt.dscp {
+        set_dscp(sock.as_raw_fd(), dscp)?;
+    }
     Ok(())
 }
 
@@ -280,6 +287,9 @@ impl ListenerEndpoint {
         if let Some(ka) = op.tcp_keepalive.as_ref() {
             stream.set_keepalive(ka)?;
         }
+        if let Some(dscp) = op.dscp {
+            set_dscp(stream.as_raw_fd(), dscp)?;
+        }
         Ok(())
     }
 
diff --git a/pingora-core/src/protocols/l4/ext.rs b/pingora-core/src/protocols/l4/ext.rs
index f68cdcf..56af522 100644
--- a/pingora-core/src/protocols/l4/ext.rs
+++ b/pingora-core/src/protocols/l4/ext.rs
@@ -276,6 +276,31 @@ pub fn set_tcp_fastopen_backlog(_fd: RawFd, _backlog: usize) -> Result<()> {
 }
 
 #[cfg(target_os = "linux")]
+pub fn set_dscp(fd: RawFd, value: u8) -> Result<()> {
+    use super::socket::SocketAddr;
+    use pingora_error::OkOrErr;
+
+    let sock = SocketAddr::from_raw_fd(fd, false);
+    let addr = sock
+        .as_ref()
+        .and_then(|s| s.as_inet())
+        .or_err(SocketError, "failed to set dscp, invalid IP socket")?;
+
+    if addr.is_ipv6() {
+        set_opt(fd, libc::IPPROTO_IPV6, libc::IPV6_TCLASS, value as c_int)
+            .or_err(SocketError, "failed to set dscp (IPV6_TCLASS)")
+    } else {
+        set_opt(fd, libc::IPPROTO_IP, libc::IP_TOS, value as c_int)
+            .or_err(SocketError, "failed to set dscp (IP_TOS)")
+    }
+}
+
+#[cfg(not(target_os = "linux"))]
+pub fn set_dscp(_fd: RawFd, _value: u8) -> Result<()> {
+    Ok(())
+}
+
+#[cfg(target_os = "linux")]
 pub fn get_socket_cookie(fd: RawFd) -> io::Result<u64> {
     get_opt_sized::<c_ulonglong>(fd, libc::SOL_SOCKET, libc::SO_COOKIE)
 }
diff --git a/pingora-core/src/upstreams/peer.rs b/pingora-core/src/upstreams/peer.rs
index 78e59ed..d0c8125 100644
--- a/pingora-core/src/upstreams/peer.rs
+++ b/pingora-core/src/upstreams/peer.rs
@@ -172,6 +172,12 @@ pub trait Peer: Display + Clone {
         self.get_peer_options().and_then(|o| o.tcp_recv_buf)
     }
 
+    /// The DSCP value that should be applied to the send side of this connection.
+    /// See the [RFC](https://datatracker.ietf.org/doc/html/rfc2474) for more details.
+    fn dscp(&self) -> Option<u8> {
+        self.get_peer_options().and_then(|o| o.dscp)
+    }
+
     /// Whether to enable TCP fast open.
     fn tcp_fast_open(&self) -> bool {
         self.get_peer_options()
@@ -301,6 +307,7 @@ pub struct PeerOptions {
     pub ca: Option<Arc<Box<[X509]>>>,
     pub tcp_keepalive: Option<TcpKeepalive>,
     pub tcp_recv_buf: Option<usize>,
+    pub dscp: Option<u8>,
     pub no_header_eos: bool,
     pub h2_ping_interval: Option<Duration>,
     // how many concurrent h2 stream are allowed in the same connection
@@ -334,6 +341,7 @@ impl PeerOptions {
             ca: None,
             tcp_keepalive: None,
             tcp_recv_buf: None,
+            dscp: None,
             no_header_eos: false,
             h2_ping_interval: None,
             max_h2_streams: 1,