Remy D. Farley
2 years ago
committed by
B. Blechschmidt
B. Blechschmidt
11 changed files with 615 additions and 48 deletions
@ -0,0 +1,230 @@ |
|||||
|
#![cfg(target_os = "linux")] |
||||
|
|
||||
|
use crate::{error, SocketDomain, SocketProtocol}; |
||||
|
use nix::{ |
||||
|
errno::Errno, |
||||
|
fcntl::{self, FdFlag}, |
||||
|
sys::socket::{cmsg_space, getsockopt, recvmsg, sendmsg, sockopt, ControlMessage, ControlMessageOwned, MsgFlags, SockType}, |
||||
|
}; |
||||
|
use serde::{Deserialize, Serialize}; |
||||
|
use std::{ |
||||
|
io::{ErrorKind, IoSlice, IoSliceMut, Result}, |
||||
|
ops::DerefMut, |
||||
|
os::fd::{AsFd, AsRawFd, FromRawFd, IntoRawFd, OwnedFd, RawFd}, |
||||
|
}; |
||||
|
use tokio::net::{TcpSocket, UdpSocket, UnixDatagram}; |
||||
|
|
||||
|
const REQUEST_BUFFER_SIZE: usize = 64; |
||||
|
|
||||
|
#[derive(Hash, Copy, Clone, Eq, PartialEq, Debug, Serialize, Deserialize)] |
||||
|
struct Request { |
||||
|
protocol: SocketProtocol, |
||||
|
domain: SocketDomain, |
||||
|
number: u32, |
||||
|
} |
||||
|
|
||||
|
#[derive(Hash, Copy, Clone, Eq, PartialEq, Debug, Serialize, Deserialize)] |
||||
|
enum Response { |
||||
|
Ok, |
||||
|
} |
||||
|
|
||||
|
/// Reconstruct socket from raw `fd`
|
||||
|
pub fn reconstruct_socket(fd: RawFd) -> Result<OwnedFd> { |
||||
|
// Check if `fd` is valid
|
||||
|
let fd_flags = fcntl::fcntl(fd, fcntl::F_GETFD)?; |
||||
|
|
||||
|
// `fd` is confirmed to be valid so it should be closed
|
||||
|
let socket = unsafe { OwnedFd::from_raw_fd(fd) }; |
||||
|
|
||||
|
// Insert CLOEXEC flag to the `fd` to prevent further propagation across `execve(2)` calls
|
||||
|
let mut fd_flags = FdFlag::from_bits(fd_flags).ok_or(ErrorKind::Unsupported)?; |
||||
|
if !fd_flags.contains(FdFlag::FD_CLOEXEC) { |
||||
|
fd_flags.insert(FdFlag::FD_CLOEXEC); |
||||
|
fcntl::fcntl(fd, fcntl::F_SETFD(fd_flags))?; |
||||
|
} |
||||
|
|
||||
|
Ok(socket) |
||||
|
} |
||||
|
|
||||
|
/// Reconstruct transfer socket from `fd`
|
||||
|
///
|
||||
|
/// Panics if called outside of tokio runtime
|
||||
|
pub fn reconstruct_transfer_socket(fd: OwnedFd) -> Result<UnixDatagram> { |
||||
|
// Check if socket of type DATAGRAM
|
||||
|
let sock_type = getsockopt(&fd, sockopt::SockType)?; |
||||
|
if !matches!(sock_type, SockType::Datagram) { |
||||
|
return Err(ErrorKind::InvalidInput.into()); |
||||
|
} |
||||
|
|
||||
|
let std_socket: std::os::unix::net::UnixDatagram = fd.into(); |
||||
|
std_socket.set_nonblocking(true)?; |
||||
|
|
||||
|
// Fails if tokio context is absent
|
||||
|
Ok(UnixDatagram::from_std(std_socket).unwrap()) |
||||
|
} |
||||
|
|
||||
|
/// Create pair of interconnected sockets one of which is set to stay open across `execve(2)` calls.
|
||||
|
pub async fn create_transfer_socket_pair() -> std::io::Result<(UnixDatagram, OwnedFd)> { |
||||
|
let (local, remote) = tokio::net::UnixDatagram::pair()?; |
||||
|
|
||||
|
let remote_fd: OwnedFd = remote.into_std().unwrap().into(); |
||||
|
|
||||
|
// Get `remote_fd` flags
|
||||
|
let fd_flags = fcntl::fcntl(remote_fd.as_raw_fd(), fcntl::F_GETFD)?; |
||||
|
|
||||
|
// Remove CLOEXEC flag from the `remote_fd` to allow propagating across `execve(2)`
|
||||
|
let mut fd_flags = FdFlag::from_bits(fd_flags).ok_or(ErrorKind::Unsupported)?; |
||||
|
fd_flags.remove(FdFlag::FD_CLOEXEC); |
||||
|
fcntl::fcntl(remote_fd.as_raw_fd(), fcntl::F_SETFD(fd_flags))?; |
||||
|
|
||||
|
Ok((local, remote_fd)) |
||||
|
} |
||||
|
|
||||
|
pub trait TransferableSocket: Sized { |
||||
|
fn from_fd(fd: OwnedFd) -> Result<Self>; |
||||
|
fn domain() -> SocketProtocol; |
||||
|
} |
||||
|
|
||||
|
impl TransferableSocket for TcpSocket { |
||||
|
fn from_fd(fd: OwnedFd) -> Result<Self> { |
||||
|
// Check if socket is of type STREAM
|
||||
|
let sock_type = getsockopt(&fd, sockopt::SockType)?; |
||||
|
if !matches!(sock_type, SockType::Stream) { |
||||
|
return Err(ErrorKind::InvalidInput.into()); |
||||
|
} |
||||
|
|
||||
|
let std_stream: std::net::TcpStream = fd.into(); |
||||
|
std_stream.set_nonblocking(true)?; |
||||
|
|
||||
|
Ok(TcpSocket::from_std_stream(std_stream)) |
||||
|
} |
||||
|
|
||||
|
fn domain() -> SocketProtocol { |
||||
|
SocketProtocol::Tcp |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
impl TransferableSocket for UdpSocket { |
||||
|
/// Panics if called outside of tokio runtime
|
||||
|
fn from_fd(fd: OwnedFd) -> Result<Self> { |
||||
|
// Check if socket is of type DATAGRAM
|
||||
|
let sock_type = getsockopt(&fd, sockopt::SockType)?; |
||||
|
if !matches!(sock_type, SockType::Datagram) { |
||||
|
return Err(ErrorKind::InvalidInput.into()); |
||||
|
} |
||||
|
|
||||
|
let std_socket: std::net::UdpSocket = fd.into(); |
||||
|
std_socket.set_nonblocking(true)?; |
||||
|
|
||||
|
Ok(UdpSocket::try_from(std_socket).unwrap()) |
||||
|
} |
||||
|
|
||||
|
fn domain() -> SocketProtocol { |
||||
|
SocketProtocol::Udp |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
/// Send [`Request`] to `socket` and return received [`TransferableSocket`]s
|
||||
|
///
|
||||
|
/// Panics if called outside of tokio runtime
|
||||
|
pub async fn request_sockets<S, T>(mut socket: S, domain: SocketDomain, number: u32) -> error::Result<Vec<T>> |
||||
|
where |
||||
|
S: DerefMut<Target = UnixDatagram>, |
||||
|
T: TransferableSocket, |
||||
|
{ |
||||
|
// Borrow socket as mut to prevent multiple simultaneous requests
|
||||
|
let socket = socket.deref_mut(); |
||||
|
|
||||
|
// Send request
|
||||
|
let request = bincode::serialize(&Request { |
||||
|
protocol: T::domain(), |
||||
|
domain, |
||||
|
number, |
||||
|
})?; |
||||
|
|
||||
|
socket.send(&request[..]).await?; |
||||
|
|
||||
|
// Receive response
|
||||
|
loop { |
||||
|
socket.readable().await?; |
||||
|
|
||||
|
let mut buf = [0_u8; REQUEST_BUFFER_SIZE]; |
||||
|
let mut iov = [IoSliceMut::new(&mut buf[..])]; |
||||
|
let mut cmsg = Vec::with_capacity(cmsg_space::<RawFd>() * number as usize); |
||||
|
|
||||
|
let msg = recvmsg::<()>(socket.as_fd().as_raw_fd(), &mut iov, Some(&mut cmsg), MsgFlags::empty()); |
||||
|
|
||||
|
let msg = match msg { |
||||
|
Err(Errno::EAGAIN) => continue, |
||||
|
msg => msg?, |
||||
|
}; |
||||
|
|
||||
|
// Parse response
|
||||
|
let response = &msg.iovs().next().unwrap()[..msg.bytes]; |
||||
|
let response: Response = bincode::deserialize(response)?; |
||||
|
if !matches!(response, Response::Ok) { |
||||
|
return Err("Request for new sockets failed".into()); |
||||
|
} |
||||
|
|
||||
|
// Process received file descriptors
|
||||
|
let mut sockets = Vec::<T>::with_capacity(number as usize); |
||||
|
for cmsg in msg.cmsgs() { |
||||
|
if let ControlMessageOwned::ScmRights(fds) = cmsg { |
||||
|
for fd in fds { |
||||
|
if fd < 0 { |
||||
|
return Err("Received socket is invalid".into()); |
||||
|
} |
||||
|
|
||||
|
let owned_fd = reconstruct_socket(fd)?; |
||||
|
sockets.push(T::from_fd(owned_fd)?); |
||||
|
} |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
return Ok(sockets); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
/// Process [`Request`]s received from `socket`
|
||||
|
///
|
||||
|
/// Panics if called outside of tokio runtime
|
||||
|
pub async fn process_socket_requests(socket: &UnixDatagram) -> error::Result<()> { |
||||
|
loop { |
||||
|
let mut buf = [0_u8; REQUEST_BUFFER_SIZE]; |
||||
|
|
||||
|
let len = socket.recv(&mut buf[..]).await?; |
||||
|
|
||||
|
let request: Request = bincode::deserialize(&buf[..len])?; |
||||
|
|
||||
|
let response = Response::Ok; |
||||
|
let buf = bincode::serialize(&response)?; |
||||
|
|
||||
|
let mut owned_fd_buf: Vec<OwnedFd> = Vec::with_capacity(request.number as usize); |
||||
|
for _ in 0..request.number { |
||||
|
let fd = match request.protocol { |
||||
|
SocketProtocol::Tcp => match request.domain { |
||||
|
SocketDomain::IpV4 => tokio::net::TcpSocket::new_v4(), |
||||
|
SocketDomain::IpV6 => tokio::net::TcpSocket::new_v6(), |
||||
|
} |
||||
|
.map(|s| unsafe { OwnedFd::from_raw_fd(s.into_raw_fd()) }), |
||||
|
SocketProtocol::Udp => match request.domain { |
||||
|
SocketDomain::IpV4 => tokio::net::UdpSocket::bind("0.0.0.0:0").await, |
||||
|
SocketDomain::IpV6 => tokio::net::UdpSocket::bind("[::]:0").await, |
||||
|
} |
||||
|
.map(|s| s.into_std().unwrap().into()), |
||||
|
}; |
||||
|
match fd { |
||||
|
Err(err) => log::warn!("Failed to allocate socket: {err}"), |
||||
|
Ok(fd) => owned_fd_buf.push(fd), |
||||
|
}; |
||||
|
} |
||||
|
|
||||
|
socket.writable().await?; |
||||
|
|
||||
|
let raw_fd_buf: Vec<RawFd> = owned_fd_buf.iter().map(|fd| fd.as_raw_fd()).collect(); |
||||
|
let cmsg = ControlMessage::ScmRights(&raw_fd_buf[..]); |
||||
|
let iov = [IoSlice::new(&buf[..])]; |
||||
|
|
||||
|
sendmsg::<()>(socket.as_raw_fd(), &iov, &[cmsg], MsgFlags::empty(), None)?; |
||||
|
} |
||||
|
} |
||||
Loading…
Reference in new issue