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

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::{CreatePanel, MagpieServerMsg, SendMessage, ServerMessenger};
use crate::service::window::{WindowMessage, WindowMessageSender};

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

pub enum IpcMessage {}

pub struct IpcMessageSender {
    waker: Waker,
    sender: Sender<IpcMessage>,
}

/// 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) => eprintln!("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
    }
}

pub struct IpcData {
    poll: Poll,
    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 Drop for Client {
    fn drop(&mut self) {
        println!("Client #{} disconnected", self.token.0);
        let data = self.data.write();
        let _ = data
            .poll
            .registry()
            .deregister(&mut self.messenger.transport);

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

impl Client {
    pub fn on_readable(&mut self) -> std::io::Result<bool> {
        self.messenger.flush_recv()?;

        while let Some(msg) = self.messenger.recv() {
            println!("Client #{}: {:?}", self.token.0, msg);
            match msg {
                MagpieServerMsg::CreatePanel(CreatePanel { id, script }) => {
                    let window = self.data.write().new_window_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, script };
                    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 struct Ipc {
    pub data: Arc<RwLock<IpcData>>,
    pub window_sender: WindowMessageSender,
    pub message_recv: Receiver<IpcMessage>,
    pub events: Events,
    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 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);
        eprintln!("Making socket at: {:?}", sock_path);

        let mut listener = Listener {
            uds: UnixListener::bind(&sock_path)?,
            path: sock_path.to_path_buf(),
        };

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

        let events = Events::with_capacity(128);
        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: Waker::new(registry, message_recv_token)?,
            sender,
        };

        let data = IpcData {
            poll,
            next_window_id: 0,
        };

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

        Ok((ipc, sender))
    }

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

        for event in self.events.iter() {
            if event.token() == self.listener_token {
                loop {
                    match self.listener.accept() {
                        Ok((mut connection, address)) => {
                            let token = Token(self.clients.vacant_key());
                            println!(
                                "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()? {
                    eprintln!("Received {:?} signal; exiting...", received);
                    let _ = self.window_sender.send_event(WindowMessage::Quit);
                    self.quit = true;
                }
            } 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);
                }
            } else {
                panic!("Unrecognized event token: {:?}", event);
            }
        }

        Ok(())
    }

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