//! Wayland socket manipulation
use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd};
use nix::{
sys::{socket, uio},
Result as NixResult,
};
use crate::wire::{ArgumentType, Message, MessageParseError, MessageWriteError};
/// Maximum number of FD that can be sent in a single socket message
pub const MAX_FDS_OUT: usize = 28;
/// Maximum number of bytes that can be sent in a single socket message
pub const MAX_BYTES_OUT: usize = 4096;
/*
* Socket
*/
/// A wayland socket
pub struct Socket {
fd: RawFd,
}
impl Socket {
/// Send a single message to the socket
///
/// A single socket message can contain several wayland messages
///
/// The `fds` slice should not be longer than `MAX_FDS_OUT`, and the `bytes`
/// slice should not be longer than `MAX_BYTES_OUT` otherwise the receiving
/// end may lose some data.
pub fn send_msg(&self, bytes: &[u8], fds: &[RawFd]) -> NixResult<()> {
let iov = [uio::IoVec::from_slice(bytes)];
if !fds.is_empty() {
let cmsgs = [socket::ControlMessage::ScmRights(fds)];
socket::sendmsg(self.fd, &iov, &cmsgs, socket::MsgFlags::MSG_DONTWAIT, None)?;
} else {
socket::sendmsg(self.fd, &iov, &[], socket::MsgFlags::MSG_DONTWAIT, None)?;
};
Ok(())
}
/// Receive a single message from the socket
///
/// Return the number of bytes received and the number of Fds received.
///
/// Errors with `WouldBlock` is no message is available.
///
/// A single socket message can contain several wayland messages.
///
/// The `buffer` slice should be at least `MAX_BYTES_OUT` long and the `fds`
/// slice `MAX_FDS_OUT` long, otherwise some data of the received message may
/// be lost.
pub fn rcv_msg(&self, buffer: &mut [u8], fds: &mut [RawFd]) -> NixResult<(usize, usize)> {
let mut cmsg = cmsg_space!([RawFd; MAX_FDS_OUT]);
let iov = [uio::IoVec::from_mut_slice(buffer)];
let msg = socket::recvmsg(self.fd, &iov[..], Some(&mut cmsg), socket::MsgFlags::MSG_DONTWAIT)?;
let mut fd_count = 0;
let received_fds = msg.cmsgs().flat_map(|cmsg| match cmsg {
socket::ControlMessageOwned::ScmRights(s) => s,
_ => Vec::new(),
});
for (fd, place) in received_fds.zip(fds.iter_mut()) {
fd_count += 1;
*place = fd;
}
Ok((msg.bytes, fd_count))
}
}
impl FromRawFd for Socket {
unsafe fn from_raw_fd(fd: RawFd) -> Socket {
Socket { fd }
}
}
impl AsRawFd for Socket {
fn as_raw_fd(&self) -> RawFd {
self.fd
}
}
impl IntoRawFd for Socket {
fn into_raw_fd(self) -> RawFd {
self.fd
}
}
impl Drop for Socket {
fn drop(&mut self) {
let _ = ::nix::unistd::close(self.fd);
}
}
/*
* BufferedSocket
*/
/// An adapter around a raw Socket that directly handles buffering and
/// conversion from/to wayland messages
pub struct BufferedSocket {
socket: Socket,
in_data: Buffer<u32>,
in_fds: Buffer<RawFd>,
out_data: Buffer<u32>,
out_fds: Buffer<RawFd>,
}
impl BufferedSocket {
/// Wrap a Socket into a Buffered Socket
pub fn new(socket: Socket) -> BufferedSocket {
BufferedSocket {
socket,
in_data: Buffer::new(2 * MAX_BYTES_OUT / 4), // Incoming buffers are twice as big in order to be
in_fds: Buffer::new(2 * MAX_FDS_OUT), // able to store leftover data if needed
out_data: Buffer::new(MAX_BYTES_OUT / 4),
out_fds: Buffer::new(MAX_FDS_OUT),
}
}
/// Get direct access to the underlying socket
pub fn get_socket(&mut self) -> &mut Socket {
&mut self.socket
}
/// Retrieve ownership of the underlying Socket
///
/// Any leftover content in the internal buffers will be lost
pub fn into_socket(self) -> Socket {
self.socket
}
/// Flush the contents of the outgoing buffer into the socket
pub fn flush(&mut self) -> NixResult<()> {
{
let words = self.out_data.get_contents();
let bytes = unsafe { ::std::slice::from_raw_parts(words.as_ptr() as *const u8, words.len() * 4) };
let fds = self.out_fds.get_contents();
self.socket.send_msg(bytes, fds)?;
for &fd in fds {
// once the fds are sent, we can close them
let _ = ::nix::unistd::close(fd);
}
}
self.out_data.clear();
self.out_fds.clear();
Ok(())
}
// internal method
//
// attempts to write a message in the internal out buffers,
// returns true if successful
//
// if false is returned, it means there is not enough space
// in the buffer
fn attempt_write_message(&mut self, msg: &Message) -> NixResult<bool> {
match msg.write_to_buffers(
self.out_data.get_writable_storage(),
self.out_fds.get_writable_storage(),
) {
Ok((bytes_out, fds_out)) => {
self.out_data.advance(bytes_out);
self.out_fds.advance(fds_out);
Ok(true)
}
Err(MessageWriteError::BufferTooSmall) => Ok(false),
Err(MessageWriteError::DupFdFailed(e)) => Err(e),
}
}
/// Write a message to the outgoing buffer
///
/// This method may flush the internal buffer if necessary (if it is full).
///
/// If the message is too big to fit in the buffer, the error `Error::Sys(E2BIG)`
/// will be returned.
pub fn write_message(&mut self, msg: &Message) -> NixResult<()> {
if !self.attempt_write_message(msg)? {
// the attempt failed, there is not enough space in the buffer
// we need to flush it
self.flush()?;
if !self.attempt_write_message(msg)? {
// If this fails again, this means the message is too big
// to be transmitted at all
return Err(::nix::Error::Sys(::nix::errno::Errno::E2BIG));
}
}
Ok(())
}
/// Try to fill the incoming buffers of this socket, to prepare
/// a new round of parsing.
pub fn fill_incoming_buffers(&mut self) -> NixResult<()> {
// clear the buffers if they have no content
if !self.in_data.has_content() {
self.in_data.clear();
}
if !self.in_fds.has_content() {
self.in_fds.clear();
}
// receive a message
let (in_bytes, in_fds) = {
let words = self.in_data.get_writable_storage();
let bytes =
unsafe { ::std::slice::from_raw_parts_mut(words.as_ptr() as *mut u8, words.len() * 4) };
let fds = self.in_fds.get_writable_storage();
self.socket.rcv_msg(bytes, fds)?
};
if in_bytes == 0 {
// the other end of the socket was closed
return Err(::nix::Error::Sys(::nix::errno::Errno::EPIPE));
}
// advance the storage
self.in_data
.advance(in_bytes / 4 + if in_bytes % 4 > 0 { 1 } else { 0 });
self.in_fds.advance(in_fds);
Ok(())
}
/// Read and deserialize a single message from the incoming buffers socket
///
/// This method requires one closure that given an object id and an opcode,
/// must provide the signature of the associated request/event, in the form of
/// a `&'static [ArgumentType]`. If it returns `None`, meaning that
/// the couple object/opcode does not exist, an error will be returned.
///
/// There are 3 possibilities of return value:
///
/// - `Ok(Ok(msg))`: no error occurred, this is the message
/// - `Ok(Err(e))`: either a malformed message was encountered or we need more data,
/// in the latter case you need to try calling `fill_incoming_buffers()`.
/// - `Err(e)`: an I/O error occurred reading from the socked, details are in `e`
/// (this can be a "wouldblock" error, which just means that no message is available
/// to read)
pub fn read_one_message<F>(&mut self, mut signature: F) -> Result<Message, MessageParseError>
where
F: FnMut(u32, u16) -> Option<&'static [ArgumentType]>,
{
let (msg, read_data, read_fd) = {
let data = self.in_data.get_contents();
let fds = self.in_fds.get_contents();
if data.len() < 2 {
return Err(MessageParseError::MissingData);
}
let object_id = data[0];
let opcode = (data[1] & 0x0000_FFFF) as u16;
if let Some(sig) = signature(object_id, opcode) {
match Message::from_raw(data, sig, fds) {
Ok((msg, rest_data, rest_fds)) => {
(msg, data.len() - rest_data.len(), fds.len() - rest_fds.len())
}
// TODO: gracefully handle wayland messages split across unix messages ?
Err(e) => return Err(e),
}
} else {
// no signature found ?
return Err(MessageParseError::Malformed);
}
};
self.in_data.offset(read_data);
self.in_fds.offset(read_fd);
Ok(msg)
}
/// Read and deserialize messages from the socket
///
/// This method requires two closures:
///
/// - The first one, given an object id and an opcode, must provide
/// the signature of the associated request/event, in the form of
/// a `&'static [ArgumentType]`. If it returns `None`, meaning that
/// the couple object/opcode does not exist, the parsing will be
/// prematurely interrupted and this method will return a
/// `MessageParseError::Malformed` error.
/// - The second closure is charged to process the parsed message. If it
/// returns `false`, the iteration will be prematurely stopped.
///
/// In both cases of early stopping, the remaining unused data will be left
/// in the buffers, and will start to be processed at the next call of this
/// method.
///
/// There are 3 possibilities of return value:
///
/// - `Ok(Ok(n))`: no error occurred, `n` messages where processed
/// - `Ok(Err(MessageParseError::Malformed))`: a malformed message was encountered
/// (this is a protocol error and is supposed to be fatal to the connection).
/// - `Err(e)`: an I/O error occurred reading from the socked, details are in `e`
/// (this can be a "wouldblock" error, which just means that no message is available
/// to read)
pub fn read_messages<F1, F2>(
&mut self,
mut signature: F1,
mut callback: F2,
) -> NixResult<Result<usize, MessageParseError>>
where
F1: FnMut(u32, u16) -> Option<&'static [ArgumentType]>,
F2: FnMut(Message) -> bool,
{
// message parsing
let mut dispatched = 0;
loop {
let mut err = None;
// first parse any leftover messages
loop {
match self.read_one_message(&mut signature) {
Ok(msg) => {
let keep_going = callback(msg);
dispatched += 1;
if !keep_going {
break;
}
}
Err(e) => {
err = Some(e);
break;
}
}
}
// copy back any leftover content to the front of the buffer
self.in_data.move_to_front();
self.in_fds.move_to_front();
if let Some(MessageParseError::Malformed) = err {
// early stop here
return Ok(Err(MessageParseError::Malformed));
}
if err.is_none() && self.in_data.has_content() {
// we stopped reading without error while there is content? That means
// the user requested an early stopping
return Ok(Ok(dispatched));
}
// now, try to get more data
match self.fill_incoming_buffers() {
Ok(()) => (),
Err(e @ ::nix::Error::Sys(::nix::errno::Errno::EAGAIN)) => {
// stop looping, returning Ok() or EAGAIN depending on whether messages
// were dispatched
if dispatched == 0 {
return Err(e);
} else {
break;
}
}
Err(e) => return Err(e),
}
}
Ok(Ok(dispatched))
}
}
/*
* Buffer
*/
struct Buffer<T: Copy> {
storage: Vec<T>,
occupied: usize,
offset: usize,
}
impl<T: Copy + Default> Buffer<T> {
fn new(size: usize) -> Buffer<T> {
Buffer {
storage: vec![T::default(); size],
occupied: 0,
offset: 0,
}
}
/// Check if this buffer has content to read
fn has_content(&self) -> bool {
self.occupied > self.offset
}
/// Advance the internal counter of occupied space
fn advance(&mut self, bytes: usize) {
self.occupied += bytes;
}
/// Advance the read offset of current occupied space
fn offset(&mut self, bytes: usize) {
self.offset += bytes;
}
/// Clears the contents of the buffer
///
/// This only sets the counter of occupied space back to zero,
/// allowing previous content to be overwritten.
fn clear(&mut self) {
self.occupied = 0;
self.offset = 0;
}
/// Get the current contents of the occupied space of the buffer
fn get_contents(&self) -> &[T] {
&self.storage[(self.offset)..(self.occupied)]
}
/// Get mutable access to the unoccupied space of the buffer
fn get_writable_storage(&mut self) -> &mut [T] {
&mut self.storage[(self.occupied)..]
}
/// Move the unread contents of the buffer to the front, to ensure
/// maximal write space availability
fn move_to_front(&mut self) {
unsafe {
::std::ptr::copy(
&self.storage[self.offset] as *const T,
&mut self.storage[0] as *mut T,
self.occupied - self.offset,
);
}
self.occupied -= self.offset;
self.offset = 0;
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::wire::{Argument, ArgumentType, Message};
use std::ffi::CString;
use smallvec::smallvec;
fn same_file(a: RawFd, b: RawFd) -> bool {
let stat1 = ::nix::sys::stat::fstat(a).unwrap();
let stat2 = ::nix::sys::stat::fstat(b).unwrap();
stat1.st_dev == stat2.st_dev && stat1.st_ino == stat2.st_ino
}
// check if two messages are equal
//
// if arguments contain FDs, check that the fd point to
// the same file, rather than are the same number.
fn assert_eq_msgs(msg1: &Message, msg2: &Message) {
assert_eq!(msg1.sender_id, msg2.sender_id);
assert_eq!(msg1.opcode, msg2.opcode);
assert_eq!(msg1.args.len(), msg2.args.len());
for (arg1, arg2) in msg1.args.iter().zip(msg2.args.iter()) {
if let (&Argument::Fd(fd1), &Argument::Fd(fd2)) = (arg1, arg2) {
assert!(same_file(fd1, fd2));
} else {
assert_eq!(arg1, arg2);
}
}
}
#[test]
fn write_read_cycle() {
let msg = Message {
sender_id: 42,
opcode: 7,
args: smallvec![
Argument::Uint(3),
Argument::Fixed(-89),
Argument::Str(Box::new(CString::new(&b"I like trains!"[..]).unwrap())),
Argument::Array(vec![1, 2, 3, 4, 5, 6, 7, 8, 9].into()),
Argument::Object(88),
Argument::NewId(56),
Argument::Int(-25),
],
};
let (client, server) = ::std::os::unix::net::UnixStream::pair().unwrap();
let mut client = BufferedSocket::new(unsafe { Socket::from_raw_fd(client.into_raw_fd()) });
let mut server = BufferedSocket::new(unsafe { Socket::from_raw_fd(server.into_raw_fd()) });
client.write_message(&msg).unwrap();
client.flush().unwrap();
static SIGNATURE: &'static [ArgumentType] = &[
ArgumentType::Uint,
ArgumentType::Fixed,
ArgumentType::Str,
ArgumentType::Array,
ArgumentType::Object,
ArgumentType::NewId,
ArgumentType::Int,
];
let ret = server
.read_messages(
|sender_id, opcode| {
if sender_id == 42 && opcode == 7 {
Some(SIGNATURE)
} else {
None
}
},
|message| {
assert_eq_msgs(&message, &msg);
true
},
)
.unwrap()
.unwrap();
assert_eq!(ret, 1);
}
#[test]
fn write_read_cycle_fd() {
let msg = Message {
sender_id: 42,
opcode: 7,
args: smallvec![
Argument::Fd(1), // stdin
Argument::Fd(0), // stdout
],
};
let (client, server) = ::std::os::unix::net::UnixStream::pair().unwrap();
let mut client = BufferedSocket::new(unsafe { Socket::from_raw_fd(client.into_raw_fd()) });
let mut server = BufferedSocket::new(unsafe { Socket::from_raw_fd(server.into_raw_fd()) });
client.write_message(&msg).unwrap();
client.flush().unwrap();
static SIGNATURE: &'static [ArgumentType] = &[ArgumentType::Fd, ArgumentType::Fd];
let ret = server
.read_messages(
|sender_id, opcode| {
if sender_id == 42 && opcode == 7 {
Some(SIGNATURE)
} else {
None
}
},
|message| {
assert_eq_msgs(&message, &msg);
true
},
)
.unwrap()
.unwrap();
assert_eq!(ret, 1);
}
#[test]
fn write_read_cycle_multiple() {
let messages = [
Message {
sender_id: 42,
opcode: 0,
args: smallvec![
Argument::Int(42),
Argument::Str(Box::new(CString::new(&b"I like trains"[..]).unwrap())),
],
},
Message {
sender_id: 42,
opcode: 1,
args: smallvec![
Argument::Fd(1), // stdin
Argument::Fd(0), // stdout
],
},
Message {
sender_id: 42,
opcode: 2,
args: smallvec![
Argument::Uint(3),
Argument::Fd(2), // stderr
],
},
];
static SIGNATURES: &'static [&'static [ArgumentType]] = &[
&[ArgumentType::Int, ArgumentType::Str],
&[ArgumentType::Fd, ArgumentType::Fd],
&[ArgumentType::Uint, ArgumentType::Fd],
];
let (client, server) = ::std::os::unix::net::UnixStream::pair().unwrap();
let mut client = BufferedSocket::new(unsafe { Socket::from_raw_fd(client.into_raw_fd()) });
let mut server = BufferedSocket::new(unsafe { Socket::from_raw_fd(server.into_raw_fd()) });
for msg in &messages {
client.write_message(msg).unwrap();
}
client.flush().unwrap();
let mut recv_msgs = Vec::new();
let ret = server
.read_messages(
|sender_id, opcode| {
if sender_id == 42 {
Some(SIGNATURES[opcode as usize])
} else {
None
}
},
|message| {
recv_msgs.push(message);
true
},
)
.unwrap()
.unwrap();
assert_eq!(ret, 3);
assert_eq!(recv_msgs.len(), 3);
for (msg1, msg2) in messages.iter().zip(recv_msgs.iter()) {
assert_eq_msgs(msg1, msg2);
}
}
#[test]
fn parse_with_string_len_multiple_of_4() {
let msg = Message {
sender_id: 2,
opcode: 0,
args: smallvec![
Argument::Uint(18),
Argument::Str(Box::new(CString::new(&b"wl_shell"[..]).unwrap())),
Argument::Uint(1),
],
};
let (client, server) = ::std::os::unix::net::UnixStream::pair().unwrap();
let mut client = BufferedSocket::new(unsafe { Socket::from_raw_fd(client.into_raw_fd()) });
let mut server = BufferedSocket::new(unsafe { Socket::from_raw_fd(server.into_raw_fd()) });
client.write_message(&msg).unwrap();
client.flush().unwrap();
static SIGNATURE: &'static [ArgumentType] =
&[ArgumentType::Uint, ArgumentType::Str, ArgumentType::Uint];
let ret = server
.read_messages(
|sender_id, opcode| {
if sender_id == 2 && opcode == 0 {
Some(SIGNATURE)
} else {
None
}
},
|message| {
assert_eq_msgs(&message, &msg);
true
},
)
.unwrap()
.unwrap();
assert_eq!(ret, 1);
}
}