waybackend 0.10.1

//! # Waybackend
//!
//! Welcome. Waybackend is a low-level wayland backend that allows you to do wayland stuff without
//! wrapping everything in `Arc`s.
//!
//! Start by calling [`waybackend::connect`](crate::connect).
#![no_std]

extern crate alloc;

use rustix::fd::OwnedFd;

use types::ObjectId;

mod wayland;

pub use bitflags;
pub mod objman;
pub use rustix;
pub mod shm;
pub mod types;
pub mod wire;

pub use rustix::io::Errno;

/// This struct holds the message builder and the wayland file descriptor
///
/// To create this struct, use the [`connect()`] function.
pub struct Waybackend {
    /// the message builder incrementally builds up wire messages
    pub wire_msg_builder: wire::MessageBuilder,
    /// the wayland file descriptor. You can pass this to `poll` to poll events
    pub wayland_fd: OwnedFd,
}

impl Waybackend {
    #[inline]
    #[must_use]
    fn new(wayland_fd: OwnedFd) -> Self {
        Self {
            wire_msg_builder: wire::MessageBuilder::new(),
            wayland_fd,
        }
    }

    #[inline]
    pub fn flush(&mut self) -> rustix::io::Result<()> {
        self.wire_msg_builder.flush(&self.wayland_fd)
    }
}

use core::num::NonZeroU32;

/// The wayland display global object always has the same id: 1
pub const WL_DISPLAY: types::ObjectId = types::ObjectId::new(NonZeroU32::new(1).unwrap());

#[derive(Debug)]
pub enum ConnectionError {
    InvalidWaylandSocketEnvVar,
    InvalidSocketAddrFamily(rustix::net::AddressFamily),
    GetSocketNameFailed(rustix::io::Errno),
    SocketCreationFailed(rustix::io::Errno),
    ConnectionFailed(rustix::io::Errno),
}

impl core::fmt::Display for ConnectionError {
    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
        match self {
            ConnectionError::InvalidWaylandSocketEnvVar => f.write_str(
                "WAYLAND_SOCKET environment variable contains a value we failed to parse",
            ),
            ConnectionError::InvalidSocketAddrFamily(actual) => write!(
                f,
                "socket address in WAYLAND_SOCKET is not a unix socket. It's actual address family is: {actual:?}"
            ),
            ConnectionError::GetSocketNameFailed(errno) => {
                write!(f, "failed to get socket name: {errno}")
            }
            ConnectionError::SocketCreationFailed(errno) => {
                write!(f, "failed to create socket: {errno}")
            }
            ConnectionError::ConnectionFailed(errno) => {
                write!(f, "failed to connect to the unix stream: {errno}")
            }
        }
    }
}

impl core::error::Error for ConnectionError {}

#[cfg(feature = "libc")]
fn get_env(var: &core::ffi::CStr) -> Option<&'static core::ffi::CStr> {
    unsafe {
        let ptr = libc::getenv(var.as_ptr());
        if !ptr.is_null() {
            Some(core::ffi::CStr::from_ptr(ptr))
        } else {
            None
        }
    }
}

/// Mostly copy-pasted from `wayland-client.rs`
///
/// This will connect to the wayland server using several fallback heuristics.
///
/// It first tries the `WAYLAND_SOCKET` environment variable. Failing that, it tries to read the
/// `WAYLAND_DISPLAY` variable. If it isn't set, we default to `wayland-0`.
///
/// This will also create the auxiliary structs, [`ObjectManager`](objman::ObjectManager) and
/// [`Receiver`](wire::Receiver) that you will use to create new objects and receive messages
/// from the wire, respectively.
///
/// Returns a tuple with [`Waybackend`], that has  the wayland file descriptor and wire message
/// builder, and the two above mentioned structures.
///
/// The [`ObjectManager`](objman::ObjectManager) must be instantiated with an enum that represents
/// the wayland protocol you will be using (see the [`ObjectManager`](objman::ObjectManager)
/// documentation for details). Further, to instantiate it, you must pass the variant that
/// corresponds to the display. This is why this function also demands you pass in that variant as
/// a parameter. For example:
/// ```no_run
/// #[derive(Clone, Copy, PartialEq)]
/// enum WaylandProtocol {
///     Display,
///     Registry,
///     LayerShell,
///     //...
/// }
/// let (mut backend, mut objman, mut receiver) =
///     waybackend::connect(WaylandProtocol::Display).expect("failed to connect to wayland server");
/// ```
///
/// **NOTE**: this is only available when the `libc` feature is active, which is the default. This
/// is because we use `libc` to read the environment variables.
#[inline]
#[cfg(feature = "libc")]
pub fn connect<T: Copy + PartialEq>(
    display: T,
) -> Result<(Waybackend, objman::ObjectManager<T>, wire::Receiver), ConnectionError> {
    use ::alloc::string::ToString;
    use ::alloc::vec::Vec;
    use rustix::fd::{FromRawFd, OwnedFd};
    use rustix::net::AddressFamily;

    if let Some(txt) = get_env(c"WAYLAND_SOCKET") {
        // We should connect to the provided WAYLAND_SOCKET
        let fd = txt
            .to_str()
            .map(str::parse::<i32>)
            .map_err(|_| ConnectionError::InvalidWaylandSocketEnvVar)?
            .map_err(|_| ConnectionError::InvalidWaylandSocketEnvVar)?;

        let fd = unsafe { OwnedFd::from_raw_fd(fd) };
        match rustix::net::getsockname(&fd) {
            Ok(socket_addr) => {
                if socket_addr.address_family() == AddressFamily::UNIX {
                    Ok(unsafe { connect_from_fd(display, fd) })
                } else {
                    Err(ConnectionError::InvalidSocketAddrFamily(
                        socket_addr.address_family(),
                    ))
                }
            }
            Err(e) => Err(ConnectionError::GetSocketNameFailed(e)),
        }
    } else {
        let socket_name = get_env(c"WAYLAND_DISPLAY").unwrap_or_else(|| {
            log::warn!("WAYLAND_DISPLAY is not set! Defaulting to wayland-0");
            c"wayland-0"
        });

        let unix_addr = if socket_name.to_bytes()[0] == b'/' {
            rustix::net::SocketAddrUnix::new(socket_name).unwrap()
        } else {
            let mut socket_fullpath = Vec::new();
            match get_env(c"XDG_RUNTIME_DIR") {
                Some(socket_path) => {
                    socket_fullpath.extend_from_slice(socket_path.to_bytes());
                    socket_fullpath.push(b'/');
                }
                None => {
                    log::warn!("XDG_RUNTIME_DIR is not set! Defaulting to /run/user/UID");
                    let uid = rustix::process::getuid();
                    socket_fullpath.extend_from_slice(b"/run/user/");
                    socket_fullpath.extend_from_slice(uid.as_raw().to_string().as_bytes());
                    socket_fullpath.push(b'/');
                }
            }
            socket_fullpath.extend_from_slice(socket_name.to_bytes());
            rustix::net::SocketAddrUnix::new(socket_fullpath.as_slice()).unwrap()
        };

        let socket = rustix::net::socket_with(
            rustix::net::AddressFamily::UNIX,
            rustix::net::SocketType::STREAM,
            rustix::net::SocketFlags::CLOEXEC,
            None,
        )
        .map_err(ConnectionError::SocketCreationFailed)?;

        connect_to(display, socket, &unix_addr)
    }
}

/// This behave the same as the [`connect`](connect) function, but you must provide the wayland
/// socket file descriptor and socket address yourself manually.
#[inline]
pub fn connect_to<T, Addr>(
    display: T,
    wayland_fd: OwnedFd,
    addr: &Addr,
) -> Result<(Waybackend, objman::ObjectManager<T>, wire::Receiver), ConnectionError>
where
    T: Copy + PartialEq,
    Addr: rustix::net::addr::SocketAddrArg,
{
    rustix::net::connect(&wayland_fd, addr).map_err(ConnectionError::ConnectionFailed)?;
    let objman = objman::ObjectManager::new(display);
    let receiver = wire::Receiver::new();
    Ok((Waybackend::new(wayland_fd), objman, receiver))
}

/// This behave the same as the [`connect`](connect) function, but you must provide the wayland
/// socket file descriptor. This is a file descriptor that is already connected, typically found
/// through the WAYLAND_SOCKET environment variable.
///
/// # Safety
///
/// The file descriptor must the correct Wayland socket file descriptor:
///   * it is a Unix socket
///   * it is already connected to the server
///
/// We cannot really verify these conditions, so the function is unsafe. On the other hand, it need
/// not return a Result, since it will always succeed.
#[inline]
pub unsafe fn connect_from_fd<T>(
    display: T,
    wayland_fd: OwnedFd,
) -> (Waybackend, objman::ObjectManager<T>, wire::Receiver)
where
    T: Copy + PartialEq,
{
    let objman = objman::ObjectManager::new(display);
    let receiver = wire::Receiver::new();
    (Waybackend::new(wayland_fd), objman, receiver)
}

#[derive(Debug)]
pub enum RoundtripError<'a> {
    MessageFromUnknownObject(ObjectId),
    RustixIoErrno(Errno),
    UnexpectedDeleteId(u32),
    UnexpectedGlobalRemove(u32),
    WaylandError(ObjectId, u32, &'a str),
    WireError(wire::Error),
}

impl<'a> core::fmt::Display for RoundtripError<'a> {
    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
        use RoundtripError::*;
        match self {
            MessageFromUnknownObject(id) => {
                write!(f, "received message from unknown object of id: {id}")
            }
            RustixIoErrno(errno) => write!(f, "io error: {errno}"),
            UnexpectedDeleteId(id) => {
                write!(
                    f,
                    "Received a delete_id message from the display for id: {id}.\n\
                    This should never happen during the roundtrip initialization.\n\
                    This wayland implementation is probably fucked."
                )
            }
            UnexpectedGlobalRemove(name) => write!(
                f,
                "Received a request to remove the {name} global. Per the protocol,\
                we should not receive global_remove events from the registry\
                before we getting the done callback event from the sync request",
            ),
            WaylandError(id, code, msg) => write!(
                f,
                "Wayland protocol error. Object: {id}. Code: {code}. Message: {msg}"
            ),
            WireError(error) => write!(f, "roundtrip failed due to wayland wire error: {error}"),
        }
    }
}

impl<'a> core::error::Error for RoundtripError<'a> {}

/// A wayland global. Use this to bind the interfaces you want
#[derive(Debug)]
pub struct Global<'a> {
    name: u32,
    interface: &'a str,
    version: u32,
}

impl<'a> Global<'a> {
    #[inline]
    pub fn name(&self) -> u32 {
        self.name
    }

    #[inline]
    pub fn interface(&self) -> &str {
        self.interface
    }

    #[inline]
    pub fn version(&self) -> u32 {
        self.version
    }

    #[inline]
    pub fn bind<T: Copy + PartialEq>(
        &self,
        backend: &mut crate::Waybackend,
        registry: ObjectId,
        objman: &mut objman::ObjectManager<T>,
        object: T,
    ) -> Result<(), Errno> {
        let id = objman.create(object);
        wayland::wl_registry::req::bind(
            backend,
            registry,
            self.name,
            id,
            self.interface,
            self.version,
        )
    }
}

/// Does a roundtrip to gather all the globals during program initialization
///
/// `global_handler` is a function that will be called for every global the compositor advertises
#[inline]
pub fn roundtrip<'a>(
    backend: &mut Waybackend,
    receiver: &'a mut wire::Receiver,
    registry_id: ObjectId,
    callback_id: ObjectId,
    mut global_handler: impl FnMut(&mut Waybackend, Global<'a>),
) -> Result<(), RoundtripError<'a>> {
    use RoundtripError::*;
    wayland::wl_display::req::get_registry(backend, WL_DISPLAY, registry_id)
        .map_err(RustixIoErrno)?;
    wayland::wl_display::req::sync(backend, WL_DISPLAY, callback_id).map_err(RustixIoErrno)?;
    backend.flush().map_err(RustixIoErrno)?;

    // Note: we do this manually (instead of using `waybackend`'s callbacks), to comply with Rust's
    // borrow checking without having to allocate anything. Because each message is borrowing from
    // the receiver, we would need a self-referential struct to implement this with an iterator or
    // a separate struct (we would need a mutable reference to the receiver and the messages, which
    // wouldn't work).
    loop {
        let mut msgs = receiver.recv(&backend.wayland_fd).map_err(WireError)?;
        while let Some(sender_id) = msgs.next() {
            let sender_id = sender_id.map_err(WireError)?;
            match sender_id {
                WL_DISPLAY => match msgs.op() {
                    // error
                    0 => {
                        msgs.validate_len([false, false, true]).map_err(WireError)?;
                        let object_id = msgs
                            .next_object()
                            .ok_or(WireError(wire::Error::NullObjectId))?;
                        let code = msgs.next_u32();
                        let message = msgs.next_string().map_err(WireError)?;
                        return Err(WaylandError(object_id, code, message));
                    }
                    // delete_id
                    1 => {
                        let id = msgs.next_u32();
                        return Err(UnexpectedDeleteId(id));
                    }
                    otherwise => {
                        return Err(WireError(wire::Error::UnrecognizedEventOpCode(
                            "wl_display",
                            otherwise,
                        )));
                    }
                },
                id if id == registry_id => match msgs.op() {
                    // global
                    0 => {
                        msgs.validate_len([false, true, false]).map_err(WireError)?;
                        let name = msgs.next_u32();
                        let interface = msgs.next_string().map_err(WireError)?;
                        let version = msgs.next_u32();
                        let global = Global {
                            name,
                            interface,
                            version,
                        };
                        global_handler(backend, global);
                    }
                    // global_remove
                    1 => {
                        let name = msgs.next_u32();
                        return Err(UnexpectedGlobalRemove(name));
                    }
                    otherwise => {
                        return Err(WireError(wire::Error::UnrecognizedEventOpCode(
                            "wl_registry",
                            otherwise,
                        )));
                    }
                },
                id if id == callback_id => match msgs.op() {
                    // done
                    0 => {
                        // consume the callback data from the msg
                        let _callback_data = msgs.next_u32();
                        return Ok(());
                    }
                    otherwise => {
                        return Err(WireError(wire::Error::UnrecognizedEventOpCode(
                            "wl_callback",
                            otherwise,
                        )));
                    }
                },
                otherwise => return Err(MessageFromUnknownObject(otherwise)),
            }
        }
    }
}

/// This macro can help you bind all the globals you want to their equivalent `enum`. It can also
/// do any arbitrary processing of globals through a fallback function.
///
/// # Panics
///
/// This macro will panic if binding the global fails.
///
/// Example usage:
/// ```no_run
/// # mod wayland {include!("../doc/wayland_protocols.rs");}
/// #[derive(Clone, Copy, PartialEq)]
/// enum WaylandObject {
///      Display,
///      Registry,
///      Callback,
///      Compositor,
///      Shm,
/// }
///
/// let (mut backend, mut objman, mut receiver) =
///      waybackend::connect(WaylandObject::Display).expect("failed to connect to wayland server");
/// let registry = objman.create(WaylandObject::Registry);
/// let callback = objman.create(WaylandObject::Callback);
///
/// // Here, we assume the wayland protocol code was generated in a `wayland` module.
/// // This is necessary because otherwise the macro cannot find the `NAME` constant in the
/// // interfaces
/// use wayland::*;
/// waybackend::roundtrip(
///     &mut backend,
///     &mut receiver,
///     registry,
///     callback,
///     |backend, global| {
///         waybackend::bind_globals!(
///             backend,
///             objman,
///             registry,
///             global,
///             // this is the fallback function for globals that will not be merely binded
///             // you can use this to, for example, gather all wl_output globals to do some
///             // preprocessing, if needed.
///             // In this example, we show how manually binding a global looks like. But it is
///             // more convenient to use the `bind_globals` macro automatic feature for this (look
///             // below)
///             // preprocessing, if needed.
///             |backend, objman, global: waybackend::Global| {
///                 match global.interface() {
///                     wl_compositor::NAME => {
///                         global.bind(
///                             backend,
///                             registry,
///                             objman,
///                             WaylandObject::Compositor
///                         ).expect("failed to bind global");
///                     }
///                     otherwise => println!("{global:?}"),
///                 }
///             },
///             // This is how you let the macro automatically bind a global for you. It
///             // essentially generates the same code we have manually written above for the
///             // `wl_compositor` interface
///             (wl_compositor, WaylandObject::Compositor),
///            );
///        },
///    )
///    .expect("failed to do initial roundtrip");
/// ```
#[macro_export]
macro_rules! bind_globals {
    (
        $backend:ident,
        $objman:ident,
        $registry:ident,
        $global:ident,
        $fallback: expr,
        $(($interface:ident, $object:path)),*$(,)?
    ) => {
        match $global.interface() {
            $(
                $interface::NAME => {
                    $global.bind($backend, $registry, &mut $objman, $object)
                        .expect("failed to bind global");
                }
            )*
            _ => $fallback($backend, &mut $objman, $global),
        }
    }
}

/// This macro can help you dispatch all the received wayland matches through the appropriate
/// protocol event handler.
///
/// # Panics
///
/// This macro will panic if dispatching the event failed. This can happen on malformed messages
/// from the wayland server.
///
/// Example usage:
/// ```no_run
/// use waybackend::{
///     Waybackend,
///     objman::ObjectManager,
///     types::ObjectId
/// };
///
/// #[derive(Clone, Copy, PartialEq)]
/// enum WaylandObject {
///      Display,
///      Registry,
///      Callback,
///      Compositor,
/// }
///
/// struct App {
///     backend: Waybackend,
///     objman: ObjectManager<WaylandObject>,
///     _registry: ObjectId,
///     callback: ObjectId,
///     should_exit: bool,
/// }
///
/// // implement all relevant handlers for App...
///
/// # impl App {
/// #     fn new(
/// #         mut backend: Waybackend,
/// #         mut objman: ObjectManager<WaylandObject>,
/// #     ) -> Self {
/// #         let registry = objman.get_first(WaylandObject::Registry).unwrap();
/// #         let callback = objman.create(WaylandObject::Callback);
/// #         wl_display::req::sync(&mut backend, waybackend::WL_DISPLAY, callback).unwrap();
/// #
/// #         Self {
/// #             backend,
/// #             objman,
/// #             _registry: registry,
/// #             callback,
/// #             should_exit: false,
/// #         }
/// #     }
/// # }
/// #
/// # impl wl_display::EvHandler for App {
/// #     fn error(&mut self, _sender_id: ObjectId, object_id: ObjectId, code: u32, message: &str) {
/// #         panic!("wayland error for object {object_id}, code {code}: {message}");
/// #     }
/// #
/// #     fn delete_id(&mut self, _sender_id: ObjectId, id: u32) {
/// #         self.objman.remove(id);
/// #     }
/// # }
/// #
/// # impl wl_registry::EvHandler for App {
/// #     fn global(&mut self, _sender_id: ObjectId, name: u32, interface: &str, version: u32) {
/// #         println!("GLOBAL: {interface} ({version}). Name: {name}")
/// #     }
/// #
/// #     fn global_remove(&mut self, _sender_id: ObjectId, name: u32) {
/// #         println!("GLOBAL_REMOVE: {name}")
/// #     }
/// # }
/// #
/// # impl wl_callback::EvHandler for App {
/// #     fn done(&mut self, sender_id: ObjectId, _callback_data: u32) {
/// #         println!("CALLBACK DONE: {sender_id}");
/// #         if sender_id == self.callback {
/// #             self.should_exit = true;
/// #         }
/// #     }
/// # }
/// #
/// # impl wl_compositor::EvHandler for App {}
/// #
/// # mod wayland {include!("../doc/wayland_protocols.rs");}
/// #
/// # let (mut backend, mut objman, mut receiver) =
/// #      waybackend::connect(WaylandObject::Display).expect("failed to connect to wayland server");
/// # let registry = objman.create(WaylandObject::Registry);
/// # let callback = objman.create(WaylandObject::Callback);
///
/// # use wayland::*;
/// # waybackend::roundtrip(
/// #     &mut backend,
/// #     &mut receiver,
/// #     registry,
/// #     callback,
/// #     |backend, global| {
/// #         waybackend::bind_globals!(
/// #             backend,
/// #             objman,
/// #             registry,
/// #             global,
/// #             |_backend, _objman, _global| {},
/// #             (wl_compositor, WaylandObject::Compositor),
/// #            );
/// #        },
/// #    )
/// #    .expect("failed to do initial roundtrip");
///
/// // then, your event loop can look kind of like this:
///
/// let mut app = App::new(backend, objman);
/// app.backend.flush().unwrap();
/// while !app.should_exit {
///     let mut msgs = receiver.recv(&app.backend.wayland_fd).unwrap();
///     while let Some(sender_id) = msgs.next() {
///         let sender_id = sender_id.unwrap();
///         // Otherwise, we get the object type from the ObjectManager, and
///         // dispatch the handler accordingly
///         let sender = app.objman.get(sender_id).unwrap();
///         waybackend::match_enum_with_interface!(
///             app,
///             sender,
///             msgs,
///             (WaylandObject::Display, wl_display),
///             (WaylandObject::Registry, wl_registry),
///             (WaylandObject::Callback, wl_callback),
///             (WaylandObject::Compositor, wl_compositor),
///         );
///     }
/// }
/// ```
#[macro_export]
macro_rules! match_enum_with_interface {
    ($handler:ident, $object:ident, $msgs:ident, $(($variant:path, $interface:ident)),*$(,)?) => {
        match $object {
            $(
                $variant => {
                    $interface::event(&mut $handler, &mut $msgs)
                        .expect("failed to dispatch event handler");
                }
            )*
        }
    }
}