From d6d075ea34d6ef6a5398b733556ea1c29d4372fb Mon Sep 17 00:00:00 2001
From: Val Packett <val@invisiblethingslab.com>
Date: Wed, 16 Jul 2025 02:33:12 -0300
Subject: [PATCH] Add message rewriting module for editing messages

Useful for proxying individual destinations
---
 sidebus-common/src/lib.rs     |   1 +
 sidebus-common/src/rewrite.rs | 147 ++++++++++++++++++++++++++++++++++
 2 files changed, 148 insertions(+)
 create mode 100644 sidebus-common/src/rewrite.rs

diff --git a/sidebus-common/src/lib.rs b/sidebus-common/src/lib.rs
index 1a3f766..df30f58 100644
--- a/sidebus-common/src/lib.rs
+++ b/sidebus-common/src/lib.rs
@@ -1 +1,2 @@
 pub mod raw;
+pub mod rewrite;
diff --git a/sidebus-common/src/rewrite.rs b/sidebus-common/src/rewrite.rs
new file mode 100644
index 0000000..a21e255
--- /dev/null
+++ b/sidebus-common/src/rewrite.rs
@@ -0,0 +1,147 @@
+//! Utility for rewriting zbus::Messages (since they're immutable, we need to do this to "edit" them)
+
+pub struct RewriteMessage<'a> {
+    builder: zbus::message::Builder<'a>,
+    header: zbus::message::Header<'a>,
+    body: zbus::message::Body,
+}
+
+impl<'a> RewriteMessage<'a> {
+    #[inline]
+    pub fn new(original: &'a zbus::Message) -> zbus::Result<Self> {
+        let header = original.header();
+        let mut builder = match original.message_type() {
+            zbus::message::Type::MethodCall => {
+                zbus::Message::method_call(header.path().unwrap(), header.member().unwrap())?
+            }
+            zbus::message::Type::MethodReturn => {
+                zbus::Message::method_return(&header)?.reply_serial(header.reply_serial())
+            }
+            zbus::message::Type::Error => {
+                zbus::Message::error(&header, header.error_name().unwrap())?
+                    .reply_serial(header.reply_serial())
+            }
+            zbus::message::Type::Signal => zbus::Message::signal(
+                header.path().unwrap(),
+                header.interface().unwrap(),
+                header.member().unwrap(),
+            )?,
+        }
+        .serial(header.primary().serial_num())
+        .endian(header.primary().endian_sig().into());
+        for flag in header.primary().flags().iter() {
+            builder = builder.with_flags(flag)?;
+        }
+        if let Some(sender) = header.sender() {
+            builder = builder.sender(sender)?;
+        }
+        if let Some(destination) = header.destination() {
+            builder = builder.destination(destination)?;
+        }
+        if let Some(iface) = header.interface() {
+            builder = builder.interface(iface)?;
+        }
+        Ok(RewriteMessage {
+            builder,
+            header,
+            body: original.body(),
+        })
+    }
+
+    #[inline]
+    pub fn sender_unless_well_known(
+        mut self,
+        sender: Option<&zbus::names::UniqueName<'a>>,
+    ) -> zbus::Result<Self> {
+        if let Some(new_sender) = self
+            .header
+            .sender()
+            .filter(|s| !s.starts_with(":"))
+            .or(sender)
+        {
+            self.builder = self.builder.sender(new_sender)?;
+        }
+        Ok(self)
+    }
+
+    #[inline]
+    pub fn destination_unless_well_known(
+        mut self,
+        destination: Option<&zbus::names::BusName<'a>>,
+    ) -> zbus::Result<Self> {
+        if let Some(new_destination) = self
+            .header
+            .destination()
+            .filter(|s| !s.starts_with(":"))
+            .or(destination)
+        {
+            self.builder = self.builder.destination(new_destination)?;
+        }
+        Ok(self)
+    }
+
+    #[inline]
+    pub fn build(self) -> zbus::Result<zbus::Message> {
+        unsafe {
+            self.builder
+                .build_raw_body(self.body.data(), self.body.signature(), vec![])
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use zbus::names::{
+        BusName, InterfaceName, MemberName, OwnedUniqueName, UniqueName, WellKnownName,
+    };
+
+    use super::*;
+
+    #[test]
+    fn test_rewrite() -> zbus::Result<()> {
+        let msg = zbus::Message::method_call("/org/freedesktop/portal/desktop", "OpenURI")?
+            .interface("org.freedesktop.portal.OpenURI")?
+            .destination("org.freedesktop.portal.Desktop")?
+            .serial(1337u32.try_into().unwrap())
+            .sender(UniqueName::from_static_str_unchecked(":12.34"))?
+            .build(&("what", "ever"))?;
+        let new_sender = OwnedUniqueName::try_from(":rw.1234")?;
+        let new_msg = RewriteMessage::new(&msg.clone())?
+            .sender_unless_well_known(Some(&new_sender))?
+            .destination_unless_well_known(Some(&BusName::WellKnown(
+                WellKnownName::from_static_str("lol.clan.Test")?,
+            )))?
+            .build()?;
+        let new_hdr = new_msg.header();
+        assert_eq!(new_hdr.message_type(), zbus::message::Type::MethodCall);
+        assert_eq!(new_hdr.reply_serial(), None);
+        assert_eq!(
+            new_hdr.interface(),
+            Some(InterfaceName::from_static_str_unchecked(
+                "org.freedesktop.portal.OpenURI"
+            ))
+            .as_ref()
+        );
+        assert_eq!(
+            new_hdr.member(),
+            Some(MemberName::from_static_str_unchecked("OpenURI")).as_ref()
+        );
+        assert_eq!(
+            new_hdr.sender(),
+            Some(UniqueName::from_static_str_unchecked(":rw.1234")).as_ref()
+        );
+        assert_eq!(
+            new_hdr.destination(),
+            // Not rewritten, as the original was well-known (not a :12.34 type thing)
+            Some(BusName::WellKnown(
+                WellKnownName::from_static_str_unchecked("org.freedesktop.portal.Desktop")
+            ))
+            .as_ref()
+        );
+        assert_eq!(
+            new_msg.body().deserialize::<(&str, &str)>()?,
+            ("what", "ever")
+        );
+        Ok(())
+    }
+}