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

use std::io::Read;
use std::ops::{Deref, DerefMut};
use std::path::{Path, PathBuf};
use std::str::from_utf8;
use std::sync::mpsc::{channel, Receiver, Sender};
use std::time::Duration;

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

use crate::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 Client {
    connection: UnixStream,
}

impl Client {
    pub fn new(connection: UnixStream) -> Self {
        Self { connection }
    }

    pub fn on_readable(&mut self) -> std::io::Result<bool> {
        let mut connection_closed = false;
        let mut received_data = vec![0; 4096];
        let mut bytes_read = 0;

        loop {
            match self.connection.read(&mut received_data[bytes_read..]) {
                Ok(0) => {
                    connection_closed = true;
                    break;
                }
                Ok(n) => {
                    bytes_read += n;
                    if bytes_read >= received_data.len() {
                        received_data.resize(received_data.len() + 1024, 0);
                    }
                }
                Err(ref err) if err.kind() == std::io::ErrorKind::WouldBlock => break,
                Err(ref err) if err.kind() == std::io::ErrorKind::Interrupted => continue,
                Err(err) => return Err(err),
            }
        }

        if bytes_read > 0 {
            let received_data = &received_data[..bytes_read];
            if let Ok(str_buf) = from_utf8(received_data) {
                println!("Received data: {}", str_buf.trim_end());
            } else {
                println!("Received (non-UTF-8) data: {:?}", received_data);
            }
        }

        Ok(connection_closed)
    }
}

pub struct Ipc {
    pub message_recv: Receiver<IpcMessage>,
    pub window_sender: WindowMessageSender,
    pub poll: Poll,
    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 ipc = Self {
            message_recv,
            window_sender,
            poll,
            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.poll.poll(&mut self.events, timeout)?;

        for event in self.events.iter() {
            if event.token() == self.listener_token {
                loop {
                    match self.listener.accept() {
                        Ok((connection, address)) => {
                            let token = Token(self.clients.vacant_key());
                            println!(
                                "Accepting connection (Client #{}) from {:?}",
                                token.0, address
                            );
                            let mut client = Client::new(connection);
                            let interest = Interest::READABLE;
                            self.poll.registry().register(
                                &mut client.connection,
                                token,
                                interest,
                            )?;
                            self.clients.insert(client);
                        }
                        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 {
                    println!("Client #{} disconnected", event.token().0);
                    let mut client = self.clients.remove(event.token().0);
                    self.poll.registry().deregister(&mut client.connection)?;
                }
            } 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);
                }
            }
        }
    }
}