Rollup merge of #96334 - devnexen:socket_mark, r=dtolnay

socket `set_mark` addition.

to be able to set a marker/id on the socket for network filtering
 (iptables/ipfw here) purpose.

3 files changed, 61 insertions, 0 deletions

diff --git a/library/std/src/os/unix/net/datagram.rs b/library/std/src/os/unix/net/datagram.rs
index 8008acfd1c9..f758f88d0a3 100644
--- a/library/std/src/os/unix/net/datagram.rs
+++ b/library/std/src/os/unix/net/datagram.rs
@@ -838,6 +838,31 @@ impl UnixDatagram {
         self.0.passcred()
     }
 
+    /// Set the id of the socket for network filtering purpose
+    ///
+    #[cfg_attr(
+        any(target_os = "linux", target_os = "freebsd", target_os = "openbsd"),
+        doc = "```no_run"
+    )]
+    #[cfg_attr(
+        not(any(target_os = "linux", target_os = "freebsd", target_os = "openbsd")),
+        doc = "```ignore"
+    )]
+    /// #![feature(unix_set_mark)]
+    /// use std::os::unix::net::UnixDatagram;
+    ///
+    /// fn main() -> std::io::Result<()> {
+    ///     let sock = UnixDatagram::unbound()?;
+    ///     sock.set_mark(32)?;
+    ///     Ok(())
+    /// }
+    /// ```
+    #[cfg(any(doc, target_os = "linux", target_os = "freebsd", target_os = "openbsd",))]
+    #[unstable(feature = "unix_set_mark", issue = "96467")]
+    pub fn set_mark(&self, mark: u32) -> io::Result<()> {
+        self.0.set_mark(mark)
+    }
+
     /// Returns the value of the `SO_ERROR` option.
     ///
     /// # Examples
diff --git a/library/std/src/os/unix/net/stream.rs b/library/std/src/os/unix/net/stream.rs
index cc3a8858793..dff8f6e8567 100644
--- a/library/std/src/os/unix/net/stream.rs
+++ b/library/std/src/os/unix/net/stream.rs
@@ -427,6 +427,31 @@ impl UnixStream {
         self.0.passcred()
     }
 
+    /// Set the id of the socket for network filtering purpose
+    ///
+    #[cfg_attr(
+        any(target_os = "linux", target_os = "freebsd", target_os = "openbsd"),
+        doc = "```no_run"
+    )]
+    #[cfg_attr(
+        not(any(target_os = "linux", target_os = "freebsd", target_os = "openbsd")),
+        doc = "```ignore"
+    )]
+    /// #![feature(unix_set_mark)]
+    /// use std::os::unix::net::UnixStream;
+    ///
+    /// fn main() -> std::io::Result<()> {
+    ///     let sock = UnixStream::connect("/tmp/sock")?;
+    ///     sock.set_mark(32)?;
+    ///     Ok(())
+    /// }
+    /// ```
+    #[cfg(any(doc, target_os = "linux", target_os = "freebsd", target_os = "openbsd",))]
+    #[unstable(feature = "unix_set_mark", issue = "96467")]
+    pub fn set_mark(&self, mark: u32) -> io::Result<()> {
+        self.0.set_mark(mark)
+    }
+
     /// Returns the value of the `SO_ERROR` option.
     ///
     /// # Examples
diff --git a/library/std/src/sys/unix/net.rs b/library/std/src/sys/unix/net.rs
index 85dcfc747af..b84bf8f9264 100644
--- a/library/std/src/sys/unix/net.rs
+++ b/library/std/src/sys/unix/net.rs
@@ -438,6 +438,17 @@ impl Socket {
         self.0.set_nonblocking(nonblocking)
     }
 
+    #[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "openbsd"))]
+    pub fn set_mark(&self, mark: u32) -> io::Result<()> {
+        #[cfg(target_os = "linux")]
+        let option = libc::SO_MARK;
+        #[cfg(target_os = "freebsd")]
+        let option = libc::SO_USER_COOKIE;
+        #[cfg(target_os = "openbsd")]
+        let option = libc::SO_RTABLE;
+        setsockopt(self, libc::SOL_SOCKET, option, mark as libc::c_int)
+    }
+
     pub fn take_error(&self) -> io::Result<Option<io::Error>> {
         let raw: c_int = getsockopt(self, libc::SOL_SOCKET, libc::SO_ERROR)?;
         if raw == 0 { Ok(None) } else { Ok(Some(io::Error::from_raw_os_error(raw as i32))) }