canary-rs/apps/magpie/src/service/ipc.rs

// Copyright (c) 2022 Marceline Cramer
// SPDX-License-Identifier: AGPL-3.0-or-later

use std::collections::HashMap;
use std::ops::{Deref, DerefMut};
use std::path::{Path, PathBuf};
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::Arc;
use std::time::Duration;

use mio::net::{UnixListener, UnixStream};
use mio::{Events, Interest, Poll, Token, Waker};
use mio_signals::{Signal, Signals};
use parking_lot::RwLock;
use slab::Slab;

use crate::protocol::*;
use crate::service::window::{WindowMessage, WindowMessageSender};

const SOCK_NAME: &str = "magpie.sock";

#[derive(Debug)]
pub enum IpcMessage {
    PanelMessage { window: usize, message: Vec<u8> },
}

#[derive(Clone)]
pub struct IpcMessageSender {
    waker: Arc<Waker>,
    sender: Sender<IpcMessage>,
}

impl IpcMessageSender {
    pub fn send(&self, msg: IpcMessage) {
        let _ = self.sender.send(msg);
        let _ = self.waker.wake();
    }
}

/// Wraps [mio::net::UnixListener] with automatic file deletion on drop.
pub struct Listener {
    pub uds: UnixListener,
    pub path: PathBuf,
}

impl Drop for Listener {
    fn drop(&mut self) {
        match std::fs::remove_file(&self.path) {
            Ok(_) => {}
            Err(e) => log::error!("Could not delete UnixListener {:?}", e),
        }
    }
}

impl Deref for Listener {
    type Target = UnixListener;

    fn deref(&self) -> &UnixListener {
        &self.uds
    }
}

impl DerefMut for Listener {
    fn deref_mut(&mut self) -> &mut UnixListener {
        &mut self.uds
    }
}

impl Listener {
    fn new() -> std::io::Result<Self> {
        let sock_dir = std::env::var("XDG_RUNTIME_DIR").expect("XDG_RUNTIME_DIR not set");
        let sock_dir = Path::new(&sock_dir);
        let sock_path = sock_dir.join(SOCK_NAME);

        use std::io::{Error, ErrorKind};
        match UnixStream::connect(&sock_path) {
            Ok(_) => {
                log::warn!("Socket is already in use. Another instance of Magpie may be running.");
                let kind = ErrorKind::AddrInUse;
                let error = Error::new(kind, "Socket is already in use.");
                return Err(error);
            }
            Err(ref err) if err.kind() == ErrorKind::ConnectionRefused => {
                log::warn!("Found leftover socket; removing.");
                std::fs::remove_file(&sock_path)?;
            }
            Err(ref err) if err.kind() == ErrorKind::NotFound => {}
            Err(err) => return Err(err),
        }

        log::info!("Making socket at: {:?}", sock_path);
        let uds = UnixListener::bind(&sock_path)?;
        let path = sock_path.to_path_buf();
        Ok(Self { uds, path })
    }
}

pub struct IpcData {
    poll: Poll,
    window_to_client_panel: HashMap<usize, (usize, PanelId)>,
    next_window_id: usize,
}

impl IpcData {
    pub fn new_window_id(&mut self) -> usize {
        let id = self.next_window_id;
        self.next_window_id = self.next_window_id.wrapping_add(1);
        id
    }
}

pub struct Client {
    data: Arc<RwLock<IpcData>>,
    window_sender: WindowMessageSender,
    messenger: ServerMessenger<UnixStream>,
    token: Token,
    id_to_window: HashMap<u32, usize>,
}

impl Client {
    pub fn on_readable(&mut self) -> std::io::Result<bool> {
        if let Err(err) = self.messenger.flush_recv() {
            log::error!("flush_recv() error: {:?}", err);
        }

        while let Some(msg) = self.messenger.try_recv() {
            log::debug!("Client #{}: {:?}", self.token.0, msg);
            match msg {
                MagpieServerMsg::CreatePanel(CreatePanel {
                    id,
                    protocol,
                    script,
                    init_msg,
                }) => {
                    let mut data = self.data.write();

                    let window = data.new_window_id();
                    data.window_to_client_panel
                        .insert(window, (self.token.0, id));

                    if let Some(old_id) = self.id_to_window.insert(id, window) {
                        let msg = WindowMessage::CloseWindow { id: old_id };
                        let _ = self.window_sender.send_event(msg);
                    }

                    let msg = WindowMessage::OpenWindow {
                        id: window,
                        protocol,
                        script,
                        init_msg,
                    };
                    let _ = self.window_sender.send_event(msg);
                }
                MagpieServerMsg::ClosePanel(ClosePanel { id }) => {
                    if let Some(id) = self.id_to_window.get(&id).copied() {
                        let msg = WindowMessage::CloseWindow { id };
                        let _ = self.window_sender.send_event(msg);
                    }
                }
                MagpieServerMsg::SendMessage(SendMessage { id, msg }) => {
                    if let Some(id) = self.id_to_window.get(&id).cloned() {
                        let msg = WindowMessage::SendMessage { id, msg };
                        let _ = self.window_sender.send_event(msg);
                    }
                }
            }
        }

        Ok(self.messenger.is_closed())
    }

    pub fn disconnect(mut self) {
        log::info!("Client #{} disconnected", self.token.0);

        let mut transport = self.messenger.into_transport();
        let mut data = self.data.write();
        let _ = data.poll.registry().deregister(&mut transport);

        for (_id, window) in self.id_to_window.drain() {
            let msg = WindowMessage::CloseWindow { id: window };
            let _ = self.window_sender.send_event(msg);
            data.window_to_client_panel.remove(&window);
        }
    }
}

pub struct Ipc {
    pub data: Arc<RwLock<IpcData>>,
    pub window_sender: WindowMessageSender,
    pub message_recv: Receiver<IpcMessage>,
    pub quit: bool,
    pub listener: Listener,
    pub signals: Signals,
    pub listener_token: Token,
    pub signals_token: Token,
    pub message_recv_token: Token,
    pub clients: Slab<Client>,
}

impl Ipc {
    pub fn new(window_sender: WindowMessageSender) -> std::io::Result<(Self, IpcMessageSender)> {
        let mut listener = Listener::new()?;

        let mut signals = Signals::new(Signal::Interrupt | Signal::Quit)?;

        let poll = Poll::new()?;
        let listener_token = Token(usize::MAX);
        let signals_token = Token(listener_token.0 - 1);
        let message_recv_token = Token(signals_token.0 - 1);

        let registry = poll.registry();
        let interest = Interest::READABLE;
        registry.register(&mut listener.uds, listener_token, interest)?;
        registry.register(&mut signals, signals_token, interest)?;

        let (sender, message_recv) = channel();

        let sender = IpcMessageSender {
            waker: Arc::new(Waker::new(registry, message_recv_token)?),
            sender,
        };

        let data = IpcData {
            poll,
            window_to_client_panel: HashMap::new(),
            next_window_id: 0,
        };

        let ipc = Self {
            data: Arc::new(RwLock::new(data)),
            window_sender,
            message_recv,
            quit: false,
            listener,
            signals,
            listener_token,
            signals_token,
            message_recv_token,
            clients: Default::default(),
        };

        Ok((ipc, sender))
    }

    pub fn on_message(&mut self, msg: IpcMessage) -> std::io::Result<()> {
        match msg {
            IpcMessage::PanelMessage { window, message } => {
                let data = self.data.read();
                let (client, panel) = *data.window_to_client_panel.get(&window).unwrap();
                let client = self.clients.get_mut(client).unwrap();
                let reply = RecvMessage {
                    id: panel,
                    msg: message,
                };
                client
                    .messenger
                    .send(&MagpieClientMsg::RecvMessage(reply))?;
            }
        }

        Ok(())
    }

    pub fn poll(&mut self, events: &mut Events, timeout: Option<Duration>) -> std::io::Result<()> {
        self.data.write().poll.poll(events, timeout)?;

        for event in events.iter() {
            if event.token() == self.listener_token {
                loop {
                    match self.listener.accept() {
                        Ok((mut connection, address)) => {
                            let token = Token(self.clients.vacant_key());
                            log::info!(
                                "Accepting connection (Client #{}) from {:?}",
                                token.0,
                                address
                            );

                            let interest = Interest::READABLE;
                            self.data.write().poll.registry().register(
                                &mut connection,
                                token,
                                interest,
                            )?;

                            self.clients.insert(Client {
                                messenger: ServerMessenger::new(connection),
                                token,
                                data: self.data.clone(),
                                window_sender: self.window_sender.clone(),
                                id_to_window: Default::default(),
                            });
                        }
                        Err(ref err) if err.kind() == std::io::ErrorKind::WouldBlock => break,
                        Err(err) => return Err(err),
                    }
                }
            } else if event.token() == self.signals_token {
                while let Some(received) = self.signals.receive()? {
                    log::info!("Received {:?} signal; exiting...", received);
                    let _ = self.window_sender.send_event(WindowMessage::Quit);
                    self.quit = true;
                }
            } else if event.token() == self.message_recv_token {
                while let Ok(received) = self.message_recv.try_recv() {
                    self.on_message(received)?;
                }
            } else if let Some(client) = self.clients.get_mut(event.token().0) {
                let disconnected = client.on_readable()?;
                if disconnected {
                    self.clients.remove(event.token().0).disconnect();
                }
            } else {
                log::error!("Unrecognized event token: {:?}", event);
            }
        }

        Ok(())
    }

    pub fn run(mut self) {
        let mut events = Events::with_capacity(128);
        while !self.quit {
            let wait = Duration::from_millis(100);
            match self.poll(&mut events, Some(wait)) {
                Ok(_) => {}
                Err(e) => {
                    log::error!("IPC poll error: {:?}", e);
                }
            }
        }
    }
}