Save/load Pipewire monitor tokens (#14)

Co-authored-by: Aleksander <aleksander@olekolek1000.com>
galister · Jan 26, 2024 · a7107b1 · a7107b1
1 parent 542d2d7
commit a7107b1
Show file tree

Hide file tree

Showing 5 changed files with 103 additions and 16 deletions.
diff --git a/src/config.rs b/src/config.rs
@@ -1,5 +1,6 @@
 use crate::config_io;
 use crate::config_io::get_conf_d_path;
+use crate::desktop::def_pw_tokens;
 use crate::keyboard;
 use crate::load_with_fallback;
 use log::error;
@@ -52,6 +53,9 @@ pub struct GeneralConfig {
 
     #[serde(default = "def_one")]
     pub watch_scale: f32,
+
+    #[serde(default = "def_pw_tokens")]
+    pub pw_tokens: Vec<(String, String)>,
 }
 
 impl GeneralConfig {

diff --git a/src/desktop/capture/pw_capture.rs b/src/desktop/capture/pw_capture.rs
@@ -1,3 +1,4 @@
+use std::collections::BTreeMap;
 use std::io::Cursor;
 use std::mem::MaybeUninit;
 use std::ptr::null_mut;
@@ -41,10 +42,23 @@ use stereokit::StereoKitMultiThread;
 
 static FORMATS: Lazy<Arc<Vec<DrmFormat>>> = Lazy::new(|| Arc::new(load_dmabuf_formats()));
 
-pub async fn pipewire_select_screen(token: Option<&str>) -> Result<u32, ashpd::Error> {
+pub async fn pipewire_select_screen(
+    display_name: &str,
+    token_store: &mut BTreeMap<String, String>,
+) -> Result<u32, ashpd::Error> {
     let proxy = Screencast::new().await?;
     let session = proxy.create_session().await?;
 
+    // Find existing token by display
+    let token = token_store.get(display_name).map(|s| s.as_str());
+
+    if let Some(t) = token {
+        println!(
+            "Found existing Pipewire token for display {}: {}",
+            display_name, t
+        );
+    }
+
     proxy
         .select_sources(
             &session,
@@ -61,6 +75,15 @@ pub async fn pipewire_select_screen(token: Option<&str>) -> Result<u32, ashpd::E
         .await?
         .response()?;
 
+    if let Some(restore_token) = response.restore_token() {
+        if token_store
+            .insert(String::from(display_name), String::from(restore_token))
+            .is_none()
+        {
+            println!("Adding token {}", restore_token);
+        }
+    }
+
     if let Some(stream) = response.streams().first() {
         return Ok(stream.pipe_wire_node_id());
     }

diff --git a/src/desktop/mod.rs b/src/desktop/mod.rs
@@ -1,15 +1,18 @@
 use std::{
+    collections::BTreeMap,
+    error::Error,
     f32::consts::PI,
-    fs::read_to_string,
-    path::Path,
+    path::PathBuf,
     time::{Duration, Instant},
 };
 
 use glam::{vec2, Affine2, Quat, Vec2, Vec3};
 use log::{info, warn};
+use serde::{Deserialize, Serialize};
 use wayland_client::protocol::wl_output::Transform;
 
 use crate::{
+    config_io,
     desktop::capture::{
         pw_capture::{pipewire_select_screen, PipewireCapture},
         wlr_dmabuf_capture::WlrDmabufCapture,
@@ -112,8 +115,50 @@ impl InteractionHandler for ScreenInteractionHandler {
     fn on_left(&mut self, _hand: usize) {}
 }
 
+pub fn def_pw_tokens() -> Vec<(String, String)> {
+    Vec::new()
+}
+
+#[derive(Deserialize, Serialize, Default)]
+pub struct TokenConf {
+    #[serde(default = "def_pw_tokens")]
+    pub pw_tokens: Vec<(String, String)>,
+}
+
+fn get_pw_token_path() -> PathBuf {
+    let mut path = config_io::get_conf_d_path();
+    path.push("pw_tokens.yaml");
+    path
+}
+
+pub fn save_pw_token_config(tokens: &BTreeMap<String, String>) -> Result<(), Box<dyn Error>> {
+    let mut conf = TokenConf::default();
+
+    for (name, token) in tokens {
+        conf.pw_tokens.push((name.clone(), token.clone()));
+    }
+
+    let yaml = serde_yaml::to_string(&conf)?;
+    std::fs::write(get_pw_token_path(), yaml)?;
+
+    Ok(())
+}
+
+pub fn load_pw_token_config() -> Result<BTreeMap<String, String>, Box<dyn Error>> {
+    let mut map: BTreeMap<String, String> = BTreeMap::new();
+
+    let yaml = std::fs::read_to_string(get_pw_token_path())?;
+    let conf: TokenConf = serde_yaml::from_str(yaml.as_str())?;
+
+    for (name, token) in conf.pw_tokens {
+        map.insert(name, token);
+    }
+
+    Ok(map)
+}
+
 pub async fn try_create_screen(
-    wl: &WlClientState,
+    wl: &mut WlClientState,
     idx: usize,
     session: &AppSession,
 ) -> Option<OverlayData> {
@@ -132,11 +177,8 @@ pub async fn try_create_screen(
         capture = WlrDmabufCapture::try_new(wl, output);
     } else {
         info!("{}: Using Pipewire capture", &output.name);
-        let file_name = format!("{}.token", &output.name);
-        let full_path = Path::new(&session.config_root_path).join(file_name);
-        let token = read_to_string(full_path).ok();
 
-        if let Ok(node_id) = pipewire_select_screen(token.as_deref()).await {
+        if let Ok(node_id) = pipewire_select_screen(output.name.as_ref(), &mut wl.pw_tokens).await {
             info!("Node id: {}", node_id);
             capture = Some(Box::new(PipewireCapture::new(
                 output.name.clone(),

diff --git a/src/desktop/wl_client.rs b/src/desktop/wl_client.rs
@@ -1,5 +1,6 @@
-use glam::{Vec2, vec2};
+use glam::{vec2, Vec2};
 use log::warn;
+use std::collections::BTreeMap;
 use std::os::fd::IntoRawFd;
 use std::sync::{Arc, Mutex};
 
@@ -19,7 +20,7 @@ use wayland_client::{
         wl_output::{self, Transform, WlOutput},
         wl_registry::WlRegistry,
     },
-    Connection, Dispatch, EventQueue, Proxy, QueueHandle, 
+    Connection, Dispatch, EventQueue, Proxy, QueueHandle,
 };
 
 use crate::desktop::frame::{FramePlane, FRAME_FAILED};
@@ -46,6 +47,7 @@ pub struct WlClientState {
     pub desktop_rect: (i32, i32),
     pub queue: Arc<Mutex<EventQueue<Self>>>,
     pub queue_handle: QueueHandle<Self>,
+    pub pw_tokens: BTreeMap<String /* display name */, String /* token */>,
 }
 
 impl WlClientState {
@@ -64,15 +66,16 @@ impl WlClientState {
             desktop_rect: (0, 0),
             queue: Arc::new(Mutex::new(queue)),
             queue_handle: qh.clone(),
+            pw_tokens: BTreeMap::new(),
         };
 
         for o in globals.contents().clone_list().iter() {
             if o.interface == WlOutput::interface().name {
                 let wl_output: WlOutput = globals.registry().bind(o.name, o.version, &qh, o.name);
 
                 state.xdg_output_mgr.get_xdg_output(&wl_output, &qh, o.name);
-                
-                let unknown : Arc<str> = "Unknown".into();
+
+                let unknown: Arc<str> = "Unknown".into();
 
                 let output = OutputState {
                     wl_output,
@@ -151,7 +154,6 @@ impl Dispatch<ZxdgOutputV1, u32> for WlClientState {
             }
             zxdg_output_v1::Event::Done => {
                 if let Some(output) = state.outputs.iter_mut().find(|o| o.id == *data) {
-
                     if output.logical_size.x < 0. {
                         output.logical_pos.x += output.logical_size.x;
                         output.logical_size.x *= -1.;

diff --git a/src/main.rs b/src/main.rs
@@ -6,13 +6,16 @@ use std::{
 };
 
 use config::GeneralConfig;
-use desktop::{try_create_screen, wl_client::WlClientState};
+use desktop::{
+    load_pw_token_config, save_pw_token_config, try_create_screen, wl_client::WlClientState,
+};
 use gl::{egl::gl_init, GlRenderer, PANEL_SHADER_BYTES};
 use glam::{Quat, Vec3};
 use gui::font::FontCache;
 use input::INPUT;
 use interactions::InputState;
 use keyboard::create_keyboard;
+use log::error;
 use once_cell::sync::Lazy;
 use overlay::OverlayData;
 use stereokit::*;
@@ -149,7 +152,7 @@ fn main() {
     let mut overlays: Vec<OverlayData> = vec![];
     let mut screens: Vec<(usize, Arc<str>)> = vec![];
 
-    let wl = WlClientState::new();
+    let mut wl = WlClientState::new();
 
     if let Ok(mut uinput) = INPUT.lock() {
         uinput.set_desktop_extent(wl.get_desktop_extent());
@@ -161,8 +164,14 @@ fn main() {
     keyboard.want_visible = true;
     overlays.push(keyboard);
 
+    if let Ok(pw_tokens) = load_pw_token_config() {
+        wl.pw_tokens = pw_tokens;
+    }
+
+    let pw_tokens_copy = wl.pw_tokens.clone();
+
     for i in 0..wl.outputs.len() {
-        let maybe_screen = rt.block_on(try_create_screen(&wl, i, &session));
+        let maybe_screen = rt.block_on(try_create_screen(&mut wl, i, &session));
         if let Some(mut screen) = maybe_screen {
             screen.want_visible = session.show_screens.iter().any(|s| s == &*screen.name);
 
@@ -171,6 +180,13 @@ fn main() {
         }
     }
 
+    if pw_tokens_copy != wl.pw_tokens {
+        // Token list changed, re-create token config file
+        if let Err(err) = save_pw_token_config(&wl.pw_tokens) {
+            error!("Failed to save Pipewire token config: {}", err);
+        }
+    }
+
     overlays[0] = create_watch(&session, screens);
 
     let panel_shader = sk.shader_create_mem(PANEL_SHADER_BYTES).unwrap();