From 957615442bb7997e67ccf4587fc18b784bd6a4de Mon Sep 17 00:00:00 2001
From: Victoria Brekenfeld <victoria@system76.com>
Date: Mon, 6 Nov 2023 18:40:52 +0100
Subject: [PATCH] wayland: Add (currently optional) privileged filtering

---
 src/shell/mod.rs | 22 ++++++++--------------
 src/state.rs     | 34 ++++++++++++++++++++++++++++------
 2 files changed, 36 insertions(+), 20 deletions(-)

diff --git a/src/shell/mod.rs b/src/shell/mod.rs
index 87b0f287..5fdd1119 100644
--- a/src/shell/mod.rs
+++ b/src/shell/mod.rs
@@ -36,7 +36,7 @@ use smithay::{
 
 use crate::{
     config::{Config, KeyModifiers, KeyPattern},
-    state::client_has_security_context,
+    state::client_should_see_privileged_protocols,
     utils::prelude::*,
     wayland::protocols::{
         toplevel_info::ToplevelInfoState,
@@ -849,28 +849,22 @@ pub struct InvalidWorkspaceIndex;
 
 impl Shell {
     pub fn new(config: &Config, dh: &DisplayHandle) -> Self {
-        // TODO: Privileged protocols
-        let layer_shell_state = WlrLayerShellState::new::<State>(dh);
-        let xdg_shell_state = XdgShellState::new::<State>(dh);
-        let toplevel_info_state = ToplevelInfoState::new(
+        let layer_shell_state = WlrLayerShellState::new_with_filter::<State, _>(
             dh,
-            //|client| client.get_data::<ClientState>().map_or(false, |s| s.privileged),
-            client_has_security_context,
+            client_should_see_privileged_protocols,
         );
+        let xdg_shell_state = XdgShellState::new::<State>(dh);
+        let toplevel_info_state =
+            ToplevelInfoState::new(dh, client_should_see_privileged_protocols);
         let toplevel_management_state = ToplevelManagementState::new::<State, _>(
             dh,
             vec![
                 ManagementCapabilities::Close,
                 ManagementCapabilities::Activate,
             ],
-            //|client| client.get_data::<ClientState>().map_or(false, |s| s.privileged),
-            client_has_security_context,
-        );
-        let workspace_state = WorkspaceState::new(
-            dh,
-            //|client| client.get_data::<ClientState>().map_or(false, |s| s.privileged),
-            client_has_security_context,
+            client_should_see_privileged_protocols,
         );
+        let workspace_state = WorkspaceState::new(dh, client_should_see_privileged_protocols);
         let theme = cosmic::theme::system_preference();
 
         Shell {
diff --git a/src/state.rs b/src/state.rs
index e2cd81ff..cc763e67 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -279,12 +279,34 @@ impl BackendData {
     }
 }
 
-pub fn client_has_security_context(client: &Client) -> bool {
+pub fn client_has_no_security_context(client: &Client) -> bool {
     client
         .get_data::<ClientState>()
         .map_or(true, |client_state| client_state.security_context.is_none())
 }
 
+pub fn client_is_privileged(client: &Client) -> bool {
+    client
+        .get_data::<ClientState>()
+        .map_or(false, |client_state| client_state.privileged)
+}
+
+pub fn client_should_see_privileged_protocols(client: &Client) -> bool {
+    if std::env::var("COSMIC_ENABLE_WAYLAND_SECURITY")
+        .map(|x| {
+            x == "1"
+                || x.to_lowercase() == "true"
+                || x.to_lowercase() == "yes"
+                || x.to_lowercase() == "y"
+        })
+        .unwrap_or(false)
+    {
+        client_is_privileged(client)
+    } else {
+        client_has_no_security_context(client)
+    }
+}
+
 impl State {
     pub fn new(
         dh: &DisplayHandle,
@@ -306,14 +328,14 @@ impl State {
         let keyboard_shortcuts_inhibit_state = KeyboardShortcutsInhibitState::new::<Self>(dh);
         let output_state = OutputManagerState::new_with_xdg_output::<Self>(dh);
         let output_configuration_state =
-            OutputConfigurationState::new(dh, client_has_security_context);
+            OutputConfigurationState::new(dh, client_should_see_privileged_protocols);
         let presentation_state = PresentationState::new::<Self>(dh, clock.id() as u32);
         let primary_selection_state = PrimarySelectionState::new::<Self>(dh);
         let screencopy_state = ScreencopyState::new::<Self, _, _>(
             dh,
             vec![CursorMode::Embedded, CursorMode::Hidden],
-            client_has_security_context,
-        ); // TODO: privileged
+            client_should_see_privileged_protocols,
+        );
         let shm_state =
             ShmState::new::<Self>(dh, vec![wl_shm::Format::Xbgr8888, wl_shm::Format::Abgr8888]);
         let seat_state = SeatState::<Self>::new();
@@ -322,11 +344,11 @@ impl State {
         let kde_decoration_state = KdeDecorationState::new::<Self>(&dh, Mode::Client);
         let xdg_decoration_state = XdgDecorationState::new::<Self>(&dh);
         let session_lock_manager_state =
-            SessionLockManagerState::new::<Self, _>(&dh, client_has_security_context);
+            SessionLockManagerState::new::<Self, _>(&dh, client_should_see_privileged_protocols);
         XWaylandKeyboardGrabState::new::<Self>(&dh);
         PointerConstraintsState::new::<Self>(&dh);
         PointerGesturesState::new::<Self>(&dh);
-        SecurityContextState::new::<Self, _>(&dh, client_has_security_context);
+        SecurityContextState::new::<Self, _>(&dh, client_has_no_security_context);
 
         let shell = Shell::new(&config, dh);