src/objects.rs

@@ -75,8 +75,4 @@ impl WlObjects {
     pub fn lookup_global(&self, name: u32) -> Option<&str> {
         self.global_names.get(&name).map(|s| s.as_str())
     }
-
-    pub fn remove_global(&mut self, name: u32) {
-        self.global_names.remove(&name);
-    }
 }

src/proto.rs

@@ -25,6 +25,18 @@ macro_rules! bubble_malformed {
     }};
 }
 
+macro_rules! match_decoded {
+    (match $decoded:ident {$($t:ty => $act:block$(,)?)+}) => {
+        if let crate::proto::WaylandProtocolParsingOutcome::Ok($decoded) = $decoded {
+            $(
+                if let Some($decoded) = $decoded.downcast_ref::<$t>() {
+                    $act
+                }
+            )+
+        }
+    };
+}
+
 #[derive(PartialEq, Eq)]
 pub enum WlMsgType {
     Request,

src/state.rs

@@ -7,9 +7,8 @@ use crate::{
     config::Config,
     objects::WlObjects,
     proto::{
-        WL_REGISTRY, WaylandProtocolParsingOutcome, WlDisplayDeleteIdEvent,
-        WlDisplayGetRegistryRequest, WlRegistryBindRequest, WlRegistryGlobalEvent,
-        WlRegistryGlobalRemoveEvent,
+        WL_REGISTRY, WlDisplayDeleteIdEvent, WlDisplayGetRegistryRequest, WlRegistryBindRequest,
+        WlRegistryGlobalEvent,
     },
 };
 
@@ -28,11 +27,8 @@ impl WlMitmState {
 
     #[tracing::instrument(skip_all)]
     pub fn on_c2s_request(&mut self, raw_msg: &WlRawMsg) -> bool {
-        let msg = match crate::proto::decode_request(&self.objects, raw_msg) {
-            WaylandProtocolParsingOutcome::Ok(msg) => msg,
-            WaylandProtocolParsingOutcome::MalformedMessage => {
-                // Kill all malformed messages
-                // Note that they are different from messages whose object / message types are unknown
+        let msg = crate::proto::decode_request(&self.objects, raw_msg);
+        if let crate::proto::WaylandProtocolParsingOutcome::MalformedMessage = msg {
             error!(
                 obj_id = raw_msg.obj_id,
                 opcode = raw_msg.opcode,
@@ -41,36 +37,19 @@ impl WlMitmState {
             );
             return false;
         }
-            _ => {
-                // Pass through all unknown messages -- they could be from a Wayland protocol we haven't
-                // been built against!
-                // Note that this won't pass through messages for globals we haven't allowed:
-                // to use a global, a client must first _bind_ that global, and _that_ message is intercepted
-                // below. There, we match based on the textual representation of the interface, so it works
-                // even for globals from protocols we don't know.
-                // It does mean we can't filter against methods that create more objects _from_ that
-                // global, though.
-                return true;
-            }
-        };
 
-        if let Some(msg) = msg.downcast_ref::<WlDisplayGetRegistryRequest>() {
+        match_decoded! {
+            match msg {
+                WlDisplayGetRegistryRequest => {
                     self.objects.record_object(WL_REGISTRY, msg.registry);
-        } else if let Some(msg) = msg.downcast_ref::<WlRegistryBindRequest>() {
-            // If we have blocked this global, this lookup should return None, thus blocking client attempts
-            // to bind to a blocked global.
-            // Note that because we've removed said global from the registry, a client _SHOULD NOT_ be attempting
-            // to bind to it; if it does, it's likely a malicious client!
-            // So, we simply remove these messages from the stream, which will cause the Wayland server to error out.
+                }
+                WlRegistryBindRequest => {
                     let Some(interface) = self.objects.lookup_global(msg.name) else {
                         return false;
                     };
 
                     if interface != msg.id_interface_name {
-                error!(
-                    "Client binding to interface {}, but the interface name {} should correspond to {}",
-                    msg.id_interface_name, msg.name, interface
-                );
+                        error!("Client binding to interface {}, but the interface name {} should correspond to {}", msg.id_interface_name, msg.name, interface);
                         return false;
                     }
 
@@ -84,32 +63,27 @@ impl WlMitmState {
                         self.objects.record_object(t, msg.id);
                     }
                 }
+            }
+        }
 
         true
     }
 
     #[tracing::instrument(skip_all)]
     pub fn on_s2c_event(&mut self, raw_msg: &WlRawMsg) -> bool {
-        let msg = match crate::proto::decode_event(&self.objects, raw_msg) {
-            WaylandProtocolParsingOutcome::Ok(msg) => msg,
-            WaylandProtocolParsingOutcome::MalformedMessage => {
+        let msg = crate::proto::decode_event(&self.objects, raw_msg);
+        if let crate::proto::WaylandProtocolParsingOutcome::MalformedMessage = msg {
             error!(
                 obj_id = raw_msg.obj_id,
                 opcode = raw_msg.opcode,
-                    num_fds = raw_msg.fds.len(),
                 "Malformed event"
             );
             return false;
         }
-            _ => {
-                return true;
-            }
-        };
 
-        if let Some(msg) = msg.downcast_ref::<WlRegistryGlobalEvent>() {
-            // This event is how Wayland servers announce globals -- and they are the entrypoint to
-            // most extensions! You need at least one global registered for clients to be able to
-            // access methods from that extension; but those methods _could_ create more objects.
+        match_decoded! {
+            match msg {
+                WlRegistryGlobalEvent => {
                     debug!(
                         interface = msg.interface,
                         name = msg.name,
@@ -117,7 +91,8 @@ impl WlMitmState {
                         "got global"
                     );
 
-            // To block entire extensions, we just need to filter out their announced global objects.
+                    self.objects.record_global(msg.name, msg.interface);
+
                     if !self.config.filter.allowed_globals.contains(msg.interface) {
                         info!(
                             interface = msg.interface,
@@ -125,17 +100,14 @@ impl WlMitmState {
                         );
                         return false;
                     }
-
-            // Else, record the global object. These are the only ones we're ever going to allow through.
-            // We block bind requests on any interface that's not recorded here.
-            self.objects.record_global(msg.name, msg.interface);
-        } else if let Some(msg) = msg.downcast_ref::<WlRegistryGlobalRemoveEvent>() {
-            // Remove globals that the server has removed
-            self.objects.remove_global(msg.name);
-        } else if let Some(msg) = msg.downcast_ref::<WlDisplayDeleteIdEvent>() {
-            // Server has acknowledged deletion of an object
+                }
+                WlDisplayDeleteIdEvent => {
+                    // When an object is acknowledged to be deleted, remove it from our
+                    // internal cache of all registered objects
                     self.objects.remove_object(msg.id);
                 }
+            }
+        }
 
         true
     }