Support `WAYLAND_SOCKET` in winit backend

I know of a single compositor that supports `WAYLAND_SOCKET` but not
`WAYLAND_DISPLAY`: https://gitlab.freedesktop.org/mstoeckl/windowtolayer

This should also make niri more robust against accidentally setting
`WAYLAND_SOCKET` when starting as a session, before programs could fail
if they preffered `WAYLAND_SOCKET` over `WAYLAND_DISPLAY`

2 files changed, 9 insertions, 4 deletions

diff --git a/src/main.rs b/src/main.rs
index 55fc032f..51ff0a5e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -69,8 +69,8 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     if cli.session {
         // If we're starting as a session, assume that the intention is to start on a TTY. Remove
-        // DISPLAY or WAYLAND_DISPLAY from our environment if they are set, since they will cause
-        // the winit backend to be selected instead.
+        // DISPLAY, WAYLAND_DISPLAY or WAYLAND_SOCKET from our environment if they are set, since
+        // they will cause the winit backend to be selected instead.
         if env::var_os("DISPLAY").is_some() {
             warn!("running as a session but DISPLAY is set, removing it");
             env::remove_var("DISPLAY");
@@ -79,6 +79,10 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
             warn!("running as a session but WAYLAND_DISPLAY is set, removing it");
             env::remove_var("WAYLAND_DISPLAY");
         }
+        if env::var_os("WAYLAND_SOCKET").is_some() {
+            warn!("running as a session but WAYLAND_SOCKET is set, removing it");
+            env::remove_var("WAYLAND_SOCKET");
+        }
 
         // Set the current desktop for xdg-desktop-portal.
         env::set_var("XDG_CURRENT_DESKTOP", "niri");
diff --git a/src/niri.rs b/src/niri.rs
index 88dbd0ee..a4801cc5 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -547,8 +547,9 @@ impl State {
 
         let config = Rc::new(RefCell::new(config));
 
-        let has_display =
-            env::var_os("WAYLAND_DISPLAY").is_some() || env::var_os("DISPLAY").is_some();
+        let has_display = env::var_os("WAYLAND_DISPLAY").is_some()
+            || env::var_os("WAYLAND_SOCKET").is_some()
+            || env::var_os("DISPLAY").is_some();
 
         let mut backend = if headless {
             let headless = Headless::new();