1
0
Fork 0
forked from wry/wry
wry/src/utils/buffd/wl_buf_in.rs
2026-03-29 13:38:01 +02:00

125 lines
3.4 KiB
Rust

use {
crate::{
io_uring::IoUring,
object::ObjectId,
utils::{
buf::Buf,
buffd::{BufFdError, MAX_IN_FD},
},
},
std::{collections::VecDeque, ptr, rc::Rc, slice},
uapi::OwnedFd,
};
const WORD_SIZE: usize = 4;
const WORD_ALIGN: usize = 4;
const HEADER_WORDS: usize = 2;
const HEADER_SIZE: usize = HEADER_WORDS * WORD_SIZE;
const MAX_MESSAGE_SIZE: usize = 4096;
const BUF_SIZE: usize = 2 * MAX_MESSAGE_SIZE;
pub struct WlBufFdIn {
fd: Rc<OwnedFd>,
ring: Rc<IoUring>,
fds: VecDeque<Rc<OwnedFd>>,
buf: Buf,
lo: usize,
len: usize,
}
pub struct WlMessage<'a> {
pub obj_id: ObjectId,
pub message: u32,
pub body: &'a [u32],
pub fds: &'a mut VecDeque<Rc<OwnedFd>>,
}
impl WlBufFdIn {
pub fn new(fd: &Rc<OwnedFd>, ring: &Rc<IoUring>) -> Self {
let buf = Buf::new(BUF_SIZE);
assert_eq!(buf.as_ptr() as usize % WORD_ALIGN, 0);
Self {
fd: fd.clone(),
ring: ring.clone(),
fds: Default::default(),
buf,
lo: Default::default(),
len: Default::default(),
}
}
pub async fn read_message(&mut self) -> Result<WlMessage<'_>, BufFdError> {
if self.len == 0 {
self.lo = 0;
}
if self.len < HEADER_SIZE {
if self.lo > 0 {
self.compact();
}
while self.len < HEADER_SIZE {
self.recvmsg().await?;
}
}
let hdr: &[u32] =
unsafe { slice::from_raw_parts(self.buf[self.lo..].as_ptr().cast(), HEADER_WORDS) };
let obj_id = ObjectId::from_raw(hdr[0]);
let len = (hdr[1] >> 16) as usize;
let message = hdr[1] & 0xffff;
if len & 3 != 0 {
return Err(BufFdError::UnalignedMessageSize);
}
if len > MAX_MESSAGE_SIZE {
return Err(BufFdError::MessageTooLarge);
}
if len < HEADER_SIZE {
return Err(BufFdError::MessageTooSmall);
}
if len > self.len {
if self.lo + self.len >= MAX_MESSAGE_SIZE {
self.compact();
}
while len > self.len {
self.recvmsg().await?;
}
}
let body: &[u32] = unsafe {
let words = (len - HEADER_SIZE) >> 2;
slice::from_raw_parts(self.buf[self.lo + HEADER_SIZE..].as_ptr().cast(), words)
};
self.lo += len;
self.len -= len;
Ok(WlMessage {
obj_id,
message,
body,
fds: &mut self.fds,
})
}
#[inline(always)]
fn compact(&mut self) {
unsafe {
let dst = self.buf.as_mut_ptr();
let src = dst.add(self.lo);
ptr::copy(src, dst, self.len);
self.lo = 0;
}
}
async fn recvmsg(&mut self) -> Result<(), BufFdError> {
let mut buf = self.buf.slice(self.lo + self.len..);
match self
.ring
.recvmsg(&self.fd, slice::from_mut(&mut buf), &mut self.fds)
.await
{
Ok(0) => return Err(BufFdError::Closed),
Ok(n) => self.len += n,
Err(e) => return Err(BufFdError::Ring(e)),
}
if self.fds.len() > MAX_IN_FD {
return Err(BufFdError::TooManyFds);
}
Ok(())
}
}