feat(tcp): add listener + export socket schema

This commit adds a TCP listener that can be optionally exposed on a port
provided by the user using the --tcp-port flag. If the flag is not
provided, the TCP listener will not be started.

Client state is tracked using the connecting address, and clients are
purged if they send unrecognised messages.

A JSON Schema of the SocketMessage enum can be exported via komorebic
and be used to generate type definitions in various programming
languages.

This commit also makes some improvements to the handling of 'komorebic
start'.

The previous backoff approach was not working as once the Windows API
denies access to the process for any call, no amount of retries with the
same process id will result in approval.

Therefore, 'komorebic start' now checks if the process has been started,
and if it hasn't (ie. it has errored out because of an access denied
error), it will continue to restart the process until all the komorebi
startup calls to the Windows API succeed.

resolve #176

komorebi/Cargo.toml

@@ -25,6 +25,7 @@ hotwatch = "0.4"
 lazy_static = "1"
 miow = "0.4"
 nanoid = "0.4"
+net2 = "0.2"
 os_info = "3.4"
 parking_lot = { version = "0.12", features = ["deadlock_detection"] }
 paste = "1"

komorebi/src/main.rs

@@ -4,6 +4,7 @@
 use std::collections::HashMap;
 use std::fs::File;
 use std::io::Write;
+use std::net::TcpStream;
 use std::path::PathBuf;
 use std::process::Command;
 use std::sync::atomic::AtomicBool;
@@ -42,6 +43,7 @@ use komorebi_core::Rect;
 use komorebi_core::SocketMessage;
 
 use crate::process_command::listen_for_commands;
+use crate::process_command::listen_for_commands_tcp;
 use crate::process_event::listen_for_events;
 use crate::process_movement::listen_for_movements;
 use crate::window_manager::State;
@@ -103,6 +105,8 @@ lazy_static! {
     ]));
     static ref SUBSCRIPTION_PIPES: Arc<Mutex<HashMap<String, File>>> =
         Arc::new(Mutex::new(HashMap::new()));
+    static ref TCP_CONNECTIONS: Arc<Mutex<HashMap<String, TcpStream>>> =
+        Arc::new(Mutex::new(HashMap::new()));
     static ref HIDING_BEHAVIOUR: Arc<Mutex<HidingBehaviour>> =
         Arc::new(Mutex::new(HidingBehaviour::Minimize));
     static ref HOME_DIR: PathBuf = {
@@ -391,19 +395,33 @@ struct Opts {
     /// Wait for 'komorebic complete-configuration' to be sent before processing events
     #[clap(action, short, long)]
     await_configuration: bool,
+    /// Start a TCP server on the given port to allow the direct sending of SocketMessages
+    #[clap(action, short, long)]
+    tcp_port: Option<usize>,
 }
 
 #[tracing::instrument]
+#[allow(clippy::nonminimal_bool)]
 fn main() -> Result<()> {
     let opts: Opts = Opts::parse();
     CUSTOM_FFM.store(opts.focus_follows_mouse, Ordering::SeqCst);
 
     let arg_count = std::env::args().count();
+
     let has_valid_args = arg_count == 1
-        || (arg_count == 2 && (opts.await_configuration || opts.focus_follows_mouse))
-        || (arg_count == 3 && opts.await_configuration && opts.focus_follows_mouse);
+        || (arg_count == 2
+            && (opts.await_configuration || opts.focus_follows_mouse || opts.tcp_port.is_some()))
+        || (arg_count == 3 && opts.await_configuration && opts.focus_follows_mouse)
+        || (arg_count == 3 && opts.tcp_port.is_some() && opts.focus_follows_mouse)
+        || (arg_count == 3 && opts.tcp_port.is_some() && opts.await_configuration)
+        || (arg_count == 4
+            && (opts.focus_follows_mouse && opts.await_configuration && opts.tcp_port.is_some()));
 
     if has_valid_args {
+        let process_id = WindowsApi::current_process_id();
+        WindowsApi::allow_set_foreground_window(process_id)?;
+        WindowsApi::set_process_dpi_awareness_context()?;
+
         let session_id = WindowsApi::process_id_to_session_id()?;
         SESSION_ID.store(session_id, Ordering::SeqCst);
 
@@ -432,42 +450,6 @@ fn main() -> Result<()> {
         #[cfg(feature = "deadlock_detection")]
         detect_deadlocks();
 
-        let process_id = WindowsApi::current_process_id();
-
-        {
-            let mut proceed = false;
-            let backoff = Backoff::new();
-
-            while !proceed {
-                if WindowsApi::allow_set_foreground_window(process_id).is_ok() {
-                    proceed = true;
-                } else {
-                    tracing::warn!(
-                        "could not allow komorebi to set foreground windows, retrying..."
-                    );
-
-                    backoff.snooze();
-                }
-            }
-        }
-
-        {
-            let mut proceed = false;
-            let backoff = Backoff::new();
-
-            while !proceed {
-                if WindowsApi::set_process_dpi_awareness_context().is_ok() {
-                    proceed = true;
-                } else {
-                    tracing::warn!(
-                        "could not allow komorebi to set itself as dpi-aware, retrying..."
-                    );
-
-                    backoff.snooze();
-                }
-            }
-        }
-
         let (outgoing, incoming): (Sender<WindowManagerEvent>, Receiver<WindowManagerEvent>) =
             crossbeam_channel::unbounded();
 
@@ -485,6 +467,10 @@ fn main() -> Result<()> {
             INITIAL_CONFIGURATION_LOADED.store(true, Ordering::SeqCst);
         };
 
+        if let Some(port) = opts.tcp_port {
+            listen_for_commands_tcp(wm.clone(), port);
+        }
+
         std::thread::spawn(|| {
             load_configuration().expect("could not load configuration");
         });

komorebi/src/process_command.rs

@@ -2,15 +2,20 @@ use std::fs::File;
 use std::fs::OpenOptions;
 use std::io::BufRead;
 use std::io::BufReader;
+use std::io::Read;
 use std::io::Write;
+use std::net::TcpListener;
+use std::net::TcpStream;
 use std::num::NonZeroUsize;
 use std::str::FromStr;
 use std::sync::atomic::Ordering;
 use std::sync::Arc;
+use std::time::Duration;
 
 use color_eyre::eyre::anyhow;
 use color_eyre::Result;
 use miow::pipe::connect;
+use net2::TcpStreamExt;
 use parking_lot::Mutex;
 use schemars::schema_for;
 use uds_windows::UnixStream;
@@ -52,6 +57,7 @@ use crate::LAYERED_WHITELIST;
 use crate::MANAGE_IDENTIFIERS;
 use crate::OBJECT_NAME_CHANGE_ON_LAUNCH;
 use crate::SUBSCRIPTION_PIPES;
+use crate::TCP_CONNECTIONS;
 use crate::TRAY_AND_MULTI_WINDOW_IDENTIFIERS;
 use crate::WORKSPACE_RULES;
 
@@ -64,10 +70,10 @@ pub fn listen_for_commands(wm: Arc<Mutex<WindowManager>>) {
         .expect("could not clone unix listener");
 
     std::thread::spawn(move || {
-        tracing::info!("listening");
+        tracing::info!("listening on komorebi.sock");
         for client in listener.incoming() {
             match client {
-                Ok(stream) => match wm.lock().read_commands(stream) {
+                Ok(stream) => match read_commands_uds(&wm, stream) {
                     Ok(()) => {}
                     Err(error) => tracing::error!("{}", error),
                 },
@@ -80,6 +86,48 @@ pub fn listen_for_commands(wm: Arc<Mutex<WindowManager>>) {
     });
 }
 
+#[tracing::instrument]
+pub fn listen_for_commands_tcp(wm: Arc<Mutex<WindowManager>>, port: usize) {
+    let listener =
+        TcpListener::bind(format!("0.0.0.0:{}", port)).expect("could not start tcp server");
+
+    std::thread::spawn(move || {
+        tracing::info!("listening on 0.0.0.0:43663");
+        for client in listener.incoming() {
+            match client {
+                Ok(mut stream) => {
+                    stream
+                        .set_keepalive(Some(Duration::from_secs(30)))
+                        .expect("TCP keepalive should be set");
+
+                    let addr = stream
+                        .peer_addr()
+                        .expect("incoming connection should have an address")
+                        .to_string();
+
+                    let mut connections = TCP_CONNECTIONS.lock();
+
+                    connections.insert(
+                        addr.clone(),
+                        stream.try_clone().expect("stream should be cloneable"),
+                    );
+
+                    tracing::info!("listening for incoming tcp messages from {}", &addr);
+
+                    match read_commands_tcp(&wm, &mut stream, &addr) {
+                        Ok(()) => {}
+                        Err(error) => tracing::error!("{}", error),
+                    }
+                }
+                Err(error) => {
+                    tracing::error!("{}", error);
+                    break;
+                }
+            }
+        }
+    });
+}
+
 impl WindowManager {
     #[tracing::instrument(skip(self))]
     pub fn process_command(&mut self, message: SocketMessage) -> Result<()> {
@@ -787,6 +835,16 @@ impl WindowManager {
                 socket.push("komorebic.sock");
                 let socket = socket.as_path();
 
+                let mut stream = UnixStream::connect(&socket)?;
+                stream.write_all(schema.as_bytes())?;
+            }
+            SocketMessage::SocketSchema => {
+                let socket_message = schema_for!(SocketMessage);
+                let schema = serde_json::to_string_pretty(&socket_message)?;
+                let mut socket = HOME_DIR.clone();
+                socket.push("komorebic.sock");
+                let socket = socket.as_path();
+
                 let mut stream = UnixStream::connect(&socket)?;
                 stream.write_all(schema.as_bytes())?;
             }
@@ -850,32 +908,87 @@ impl WindowManager {
         tracing::info!("processed");
         Ok(())
     }
+}
 
-    #[tracing::instrument(skip(self, stream))]
-    pub fn read_commands(&mut self, stream: UnixStream) -> Result<()> {
-        let stream = BufReader::new(stream);
-        for line in stream.lines() {
-            let message = SocketMessage::from_str(&line?)?;
+pub fn read_commands_uds(wm: &Arc<Mutex<WindowManager>>, stream: UnixStream) -> Result<()> {
+    let stream = BufReader::new(stream);
+    for line in stream.lines() {
+        let message = SocketMessage::from_str(&line?)?;
 
-            if self.is_paused {
-                return match message {
-                    SocketMessage::TogglePause | SocketMessage::State | SocketMessage::Stop => {
-                        Ok(self.process_command(message)?)
-                    }
-                    _ => {
-                        tracing::trace!("ignoring while paused");
-                        Ok(())
-                    }
-                };
-            }
+        let mut wm = wm.lock();
 
-            self.process_command(message.clone())?;
-            notify_subscribers(&serde_json::to_string(&Notification {
-                event: NotificationEvent::Socket(message.clone()),
-                state: self.as_ref().into(),
-            })?)?;
+        if wm.is_paused {
+            return match message {
+                SocketMessage::TogglePause | SocketMessage::State | SocketMessage::Stop => {
+                    Ok(wm.process_command(message)?)
+                }
+                _ => {
+                    tracing::trace!("ignoring while paused");
+                    Ok(())
+                }
+            };
         }
 
-        Ok(())
+        wm.process_command(message.clone())?;
+        notify_subscribers(&serde_json::to_string(&Notification {
+            event: NotificationEvent::Socket(message.clone()),
+            state: wm.as_ref().into(),
+        })?)?;
     }
+
+    Ok(())
+}
+
+pub fn read_commands_tcp(
+    wm: &Arc<Mutex<WindowManager>>,
+    stream: &mut TcpStream,
+    addr: &str,
+) -> Result<()> {
+    let mut stream = BufReader::new(stream);
+
+    loop {
+        let mut buf = vec![0; 1024];
+        match stream.read(&mut buf) {
+            Err(..) => {
+                tracing::warn!("removing disconnected tcp client: {addr}");
+                let mut connections = TCP_CONNECTIONS.lock();
+                connections.remove(addr);
+                break;
+            }
+            Ok(size) => {
+                let message = if let Ok(message) =
+                    SocketMessage::from_str(&String::from_utf8_lossy(&buf[..size]))
+                {
+                    message
+                } else {
+                    tracing::warn!("client sent an invalid message, disconnecting: {addr}");
+                    let mut connections = TCP_CONNECTIONS.lock();
+                    connections.remove(addr);
+                    break;
+                };
+
+                let mut wm = wm.lock();
+
+                if wm.is_paused {
+                    return match message {
+                        SocketMessage::TogglePause | SocketMessage::State | SocketMessage::Stop => {
+                            Ok(wm.process_command(message)?)
+                        }
+                        _ => {
+                            tracing::trace!("ignoring while paused");
+                            Ok(())
+                        }
+                    };
+                }
+
+                wm.process_command(message.clone())?;
+                notify_subscribers(&serde_json::to_string(&Notification {
+                    event: NotificationEvent::Socket(message.clone()),
+                    state: wm.as_ref().into(),
+                })?)?;
+            }
+        }
+    }
+
+    Ok(())
 }