diff --git a/src/buffer.rs b/src/buffer.rs index 991100d..48b8b7a 100644 --- a/src/buffer.rs +++ b/src/buffer.rs @@ -1,34 +1,27 @@ // SPDX-License-Identifier: GPL-3.0-or-later // Copyright (C) 2026, Nathan Gill -use std::{ - os::{fd::AsFd, raw::c_void}, - ptr::NonNull, - sync::{ - Arc, - atomic::{AtomicBool, Ordering}, - }, +use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, }; -use nix::{ - sys::mman::{MapFlags, ProtFlags, mmap, munmap}, - unistd::ftruncate, -}; -use tracing::warn; +use tracing::{trace, warn}; use wayland_client::{ Dispatch, QueueHandle, protocol::{wl_buffer, wl_shm, wl_surface}, }; -use crate::{state::NLockState, util::open_shm}; +use crate::{shm::NLockShm, state::NLockState, surface::NLockSurfaceTracking}; pub struct NLockBuffer { buffer: wl_buffer::WlBuffer, - data: NonNull, + + // required to keep shm mapping alive + _shm: NLockShm, pub width: i32, pub height: i32, - pub size: usize, pub state: Arc, pub surface: cairo::ImageSurface, pub context: cairo::Context, @@ -38,6 +31,13 @@ pub struct NLockBufferState { pub in_use: AtomicBool, } +pub struct NLockCommitArgs<'a> { + pub surface: &'a wl_surface::WlSurface, + pub scale: i32, + pub output: u32, + pub tracking: &'a mut NLockSurfaceTracking, +} + pub struct NLockBufferGuard<'a> { wl_buffer: &'a wl_buffer::WlBuffer, state: &'a Arc, @@ -47,13 +47,16 @@ pub struct NLockBufferGuard<'a> { impl<'a> NLockBufferGuard<'a> { /// Attaches, damages, and commits the current buffer onto the specified /// surface. - pub fn commit_to(&mut self, surface: &wl_surface::WlSurface, scale: i32) { - surface.attach(Some(self.wl_buffer), 0, 0); - surface.set_buffer_scale(scale); - surface.damage(0, 0, i32::MAX, i32::MAX); - surface.commit(); + pub fn commit_to(&mut self, args: NLockCommitArgs, qh: &QueueHandle) { + args.surface.set_buffer_scale(args.scale); + args.surface.attach(Some(self.wl_buffer), 0, 0); + args.surface.damage(0, 0, i32::MAX, i32::MAX); + args.surface.frame(qh, args.output); + args.surface.commit(); self.committed = true; + args.tracking.ready = false; + args.tracking.dirty = false; } } @@ -67,7 +70,7 @@ impl<'a> Drop for NLockBufferGuard<'a> { impl NLockBuffer { pub fn new( - shm: &wl_shm::WlShm, + wl_shm: &wl_shm::WlShm, width: i32, height: i32, format: wl_shm::Format, @@ -84,26 +87,14 @@ impl NLockBuffer { let stride = width * 4; let size = stride * height; - let fd = open_shm()?; - ftruncate(&fd, size as i64).ok()?; - - let data = unsafe { - mmap( - None, - std::num::NonZeroUsize::new(size as usize)?, - ProtFlags::PROT_READ | ProtFlags::PROT_WRITE, - MapFlags::MAP_SHARED, - &fd, - 0, - ) - .ok()? - }; + let mut shm = NLockShm::new(size as i64)?; + let data = shm.map().ok()?; let state = Arc::new(NLockBufferState { in_use: AtomicBool::new(false), }); - let pool = shm.create_pool(fd.as_fd(), size, qh, ()); + let pool = wl_shm.create_pool(shm.fd(), size, qh, ()); let buffer = pool.create_buffer(0, width, height, stride, format, qh, state.clone()); pool.destroy(); @@ -123,10 +114,9 @@ impl NLockBuffer { Some(Self { buffer, - data, + _shm: shm, width, height, - size: size as usize, state, surface, context, @@ -150,7 +140,6 @@ impl NLockBuffer { impl Drop for NLockBuffer { fn drop(&mut self) { self.buffer.destroy(); - let _ = unsafe { munmap(self.data, self.size) }; } } @@ -164,6 +153,7 @@ impl Dispatch> for NLockState { _: &QueueHandle, ) { if let wl_buffer::Event::Release = event { + trace!("release {:p}", Arc::as_ptr(data),); data.in_use.store(false, Ordering::Release); } } diff --git a/src/comm.rs b/src/comm.rs index 8be8cad..937d3c1 100644 --- a/src/comm.rs +++ b/src/comm.rs @@ -36,7 +36,7 @@ where let size_buf = msg_buf.len().to_ne_bytes(); write(&self.tx, &size_buf)?; - write(&self.tx, msg_buf)?; + write(&self.tx, &msg_buf)?; Ok(()) } @@ -93,37 +93,37 @@ where pub trait AsBytes { /// Convert to a bytes-like representation of the object - fn as_bytes(&self) -> &[u8]; + fn as_bytes(&self) -> Vec; } impl AsBytes for String { - fn as_bytes(&self) -> &[u8] { - self.as_bytes() + fn as_bytes(&self) -> Vec { + self.as_bytes().to_vec() } } impl AsBytes for bool { - fn as_bytes(&self) -> &[u8] { + fn as_bytes(&self) -> Vec { match self { - false => &[0u8], - true => &[1u8], + false => vec![0u8], + true => vec![1u8], } } } impl AsBytes for AuthState { - fn as_bytes(&self) -> &[u8] { + fn as_bytes(&self) -> Vec { match self { - Self::Idle => &[0u8], - Self::Success => &[1u8], - Self::Fail => &[2u8], + Self::Idle => vec![0u8], + Self::Success => vec![1u8], + Self::Fail => vec![2u8], } } } impl AsBytes for () { - fn as_bytes(&self) -> &[u8] { - &[0u8] + fn as_bytes(&self) -> Vec { + vec![0u8] } } diff --git a/src/event.rs b/src/event.rs index 65c8344..24b775b 100644 --- a/src/event.rs +++ b/src/event.rs @@ -25,6 +25,7 @@ pub enum EventType { KeyboardRepeat = 1, AuthStateChanged = 2, Debug = 3, + Interrupt = 4, } impl NLockState { @@ -90,6 +91,10 @@ impl NLockState { warn!("Received debug event, but not in debug mode, ignoring"); } } + EventType::Interrupt => { + // interrupt the event poll, something happened + self.interrupt.read()?; + } _ => {} }, Event::Timeout { tag } => { @@ -111,25 +116,32 @@ impl NLockState { } fn re_render(&mut self, qh: &QueueHandle) { - // Re-render only if state was updated - if self.state_changed.load(Ordering::Relaxed) - && let Some(shm) = &self.shm - { - let auth_state = self.auth_state.clone().load(Ordering::Relaxed); - - for i in 0..self.surfaces.len() { - self.surfaces[i].render( - &self.config, - auth_state, - self.password.chars().count(), - self.background_image.as_ref(), - shm, - qh, - ); + if self.state_changed.load(Ordering::Relaxed) { + // mark as dirty when state changes + for (_, surface) in self.surfaces.iter_mut() { + surface.tracking.dirty = true; } self.state_changed.store(false, Ordering::Relaxed); } + + if let Some(shm) = &self.shm { + let auth_state = self.auth_state.clone().load(Ordering::Relaxed); + + for (_, surface) in self.surfaces.iter_mut() { + // only render when surface advertised as available + if surface.tracking.dirty && surface.tracking.ready { + surface.render( + &self.config, + auth_state, + self.password.chars().count(), + self.background_image.as_ref(), + shm, + qh, + ); + } + } + } } pub fn event_loop_cycle(&mut self, event_queue: &mut EventQueue) -> Result<()> { diff --git a/src/main.rs b/src/main.rs index e89ab6d..6bd7a0b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -13,6 +13,7 @@ pub mod event_loop; pub mod font; pub mod render; pub mod seat; +pub mod shm; pub mod signal; pub mod state; pub mod surface; @@ -99,7 +100,7 @@ fn start(config: NLockConfig, debug: bool) -> Result<()> { } } - state.unlock(&qh); + state.unlock(); event_queue.roundtrip(&mut state)?; if let Err(e) = auth_comm.stop.write(true) { diff --git a/src/shm.rs b/src/shm.rs new file mode 100644 index 0000000..6a6a07e --- /dev/null +++ b/src/shm.rs @@ -0,0 +1,110 @@ +// SPDX-License-Identifier: GPL-3.0-or-later +// Copyright (c) 2026, Nathan Gill + +use std::{ + os::{ + fd::{AsFd, BorrowedFd, OwnedFd}, + raw::c_void, + }, + ptr::NonNull, + sync::atomic::{AtomicU64, Ordering}, +}; + +use anyhow::{Result, anyhow}; +use nix::{ + fcntl::OFlag, + sys::{ + mman::{MapFlags, ProtFlags, mmap, munmap, shm_open, shm_unlink}, + stat::Mode, + }, + unistd::{ftruncate, getpid}, +}; +use tracing::{debug, warn}; + +static SHM_NUM: AtomicU64 = AtomicU64::new(0); + +pub struct NLockShm { + fd: OwnedFd, + name: String, + size: usize, + data: Option>, +} + +impl NLockShm { + pub fn new(size: i64) -> Option { + if size <= 0 { + return None; + } + + let name = format!( + "/nlock-{}-{}", + getpid(), + SHM_NUM.fetch_add(1, Ordering::Relaxed), + ); + debug!("Trying shm name '{}'", name); + + let fd = match shm_open( + name.as_str(), + OFlag::O_RDWR | OFlag::O_CREAT | OFlag::O_EXCL, + Mode::S_IRUSR | Mode::S_IWUSR, + ) { + Ok(f) => f, + Err(e) => { + warn!("Failed to open shm '{}': {:?}", name, e); + return None; + } + }; + + ftruncate(&fd, size).ok()?; + + Some(Self { + fd, + name, + size: size as usize, + data: None, + }) + } + + pub fn map(&mut self) -> Result> { + if let Some(d) = self.data { + return Err(anyhow!("shm '{}' already mapped at {:p}", self.name, d)); + } + + let data = unsafe { + mmap( + None, + std::num::NonZeroUsize::new(self.size).ok_or(anyhow!("shm size was zero"))?, + ProtFlags::PROT_READ | ProtFlags::PROT_WRITE, + MapFlags::MAP_SHARED, + &self.fd, + 0, + ) + }?; + + self.data = Some(data); + + Ok(data) + } + + pub fn data(&self) -> Option> { + self.data + } + + pub fn fd(&self) -> BorrowedFd<'_> { + self.fd.as_fd() + } +} + +impl Drop for NLockShm { + fn drop(&mut self) { + if let Some(data) = self.data + && let Err(e) = unsafe { munmap(data, self.size) } + { + warn!("Failed to unmap shm '{}': {:?}", self.name, e); + } + + if let Err(e) = shm_unlink(self.name.as_str()) { + warn!("Failed to unlink shm '{}': {:?}", self.name, e); + } + } +} diff --git a/src/state.rs b/src/state.rs index cc735fe..ec7a6f8 100644 --- a/src/state.rs +++ b/src/state.rs @@ -2,6 +2,7 @@ // Copyright (C) 2026, Nathan Gill use std::{ + collections::HashMap, fs::File, io::Seek, os::fd::AsRawFd, @@ -10,7 +11,7 @@ use std::{ use anyhow::{Result, anyhow, bail}; use cairo::ImageSurface; -use tracing::{debug, warn}; +use tracing::{debug, trace, warn}; use wayland_client::protocol::{wl_region, wl_subcompositor, wl_subsurface}; use wayland_client::{ Connection, Dispatch, QueueHandle, delegate_noop, @@ -65,7 +66,7 @@ pub struct NLockState { pub r_seat: Option, pub session_lock_manager: Option, pub session_lock: Option, - pub surfaces: Vec, + pub surfaces: HashMap, pub seat: NLockSeat, pub xkb: NLockXkb, pub password: Zeroizing, @@ -74,6 +75,7 @@ pub struct NLockState { pub background_image: Option, pub event_loop: NLockEventLoop, pub debug_comm: Option>>, + pub interrupt: PipeCommChannel<()>, } impl NLockState { @@ -84,6 +86,8 @@ impl NLockState { None }; + let interrupt = PipeCommChannel::new()?; + let mut s = Self { config: args.config, running: Arc::new(AtomicBool::new(true)), @@ -98,7 +102,7 @@ impl NLockState { r_seat: None, session_lock_manager: None, session_lock: None, - surfaces: Vec::new(), + surfaces: HashMap::new(), seat: NLockSeat::default(), xkb: NLockXkb::default(), password: Zeroizing::new("".to_string()), @@ -107,6 +111,7 @@ impl NLockState { background_image: None, event_loop: NLockEventLoop::default(), debug_comm, + interrupt, }; if let Err(e) = s.try_load_background_image() { @@ -122,6 +127,11 @@ impl NLockState { EventType::AuthStateChanged.into(), )?; + s.event_loop.add( + EventSource::Fd(s.interrupt.rx().as_raw_fd()), + EventType::Interrupt.into(), + )?; + if args.debug && let Some(debug_comm) = &s.debug_comm { @@ -149,7 +159,7 @@ impl NLockState { } } - pub fn unlock(&mut self, qh: &QueueHandle) { + pub fn unlock(&mut self) { if let Some(session_lock) = &self.session_lock { if self.locked { session_lock.unlock_and_destroy(); @@ -158,9 +168,8 @@ impl NLockState { } // free any held surfaces - self.surfaces = Vec::new(); + self.surfaces.clear(); - self.display.sync(qh, ()); self.session_lock = None; self.locked = false; self.unlocked = true; @@ -233,13 +242,12 @@ impl Dispatch for NLockState { _: &Connection, qh: &QueueHandle, ) { - if let wl_registry::Event::Global { - name, - interface, - version, - } = event - { - match &interface[..] { + match event { + wl_registry::Event::Global { + name, + interface, + version, + } => match &interface[..] { "wl_compositor" => { let compositor = registry.bind::(name, version, qh, ()); @@ -263,12 +271,15 @@ impl Dispatch for NLockState { state.r_seat = Some(seat); } "wl_output" => { - let index = state.surfaces.len(); + trace!("Output ( name: {:?}, ... )", name); + let output = - registry.bind::(name, version, qh, index); + registry.bind::(name, version, qh, name); - let surface = NLockSurface::new(output, index); - state.surfaces.push(surface); + state + .surfaces + .entry(name) + .or_insert_with(|| NLockSurface::new(output, name)); } "ext_session_lock_manager_v1" => { let session_lock_manager = registry @@ -281,7 +292,14 @@ impl Dispatch for NLockState { state.session_lock_manager = Some(session_lock_manager); } _ => {} + }, + wl_registry::Event::GlobalRemove { name } => { + trace!("GlobalRemove ( name: {:?}, ... )", name); + + // old surfaces dropped here + let _ = state.surfaces.remove(&name); } + _ => {} } } } @@ -292,10 +310,29 @@ delegate_noop!(NLockState: ignore wl_shm::WlShm); delegate_noop!(NLockState: ignore wl_surface::WlSurface); delegate_noop!(NLockState: ignore wl_subsurface::WlSubsurface); delegate_noop!(NLockState: ignore ext_session_lock_manager_v1::ExtSessionLockManagerV1); -delegate_noop!(NLockState: ignore wl_callback::WlCallback); delegate_noop!(NLockState: ignore wl_shm_pool::WlShmPool); delegate_noop!(NLockState: ignore wl_region::WlRegion); +impl Dispatch for NLockState { + fn event( + state: &mut Self, + _: &wl_callback::WlCallback, + event: ::Event, + data: &u32, + _: &Connection, + _: &QueueHandle, + ) { + if let wl_callback::Event::Done { callback_data: _ } = event + && let Some(surface) = state.surfaces.get_mut(data) + { + surface.tracking.ready = true; + if let Err(e) = state.interrupt.write(()) { + warn!("Failed to write interrupt in frame callback: {:?}", e); + } + } + } +} + impl Dispatch for NLockState { fn event( state: &mut Self, @@ -303,7 +340,7 @@ impl Dispatch for NLockState { event: ::Event, _: &(), _: &Connection, - qh: &QueueHandle, + _: &QueueHandle, ) { match event { ext_session_lock_v1::Event::Locked => { @@ -312,22 +349,27 @@ impl Dispatch for NLockState { debug!("Session is locked"); } ext_session_lock_v1::Event::Finished => { - state.unlock(qh); + state.unlock(); } _ => {} } } } -impl Dispatch for NLockState { +impl Dispatch for NLockState { fn event( state: &mut Self, _: &wl_output::WlOutput, event: ::Event, - data: &usize, + data: &u32, _: &Connection, qh: &QueueHandle, ) { + let Some(surface) = state.surfaces.get_mut(data) else { + warn!("could not find surface for {}", data); + return; + }; + match event { wl_output::Event::Geometry { x: _, @@ -339,29 +381,23 @@ impl Dispatch for NLockState { model: _, transform: _, } => { - state.surfaces[*data] - .set_subpixel_order(cairo::SubpixelOrder::from_wl_subpixel(subpixel)); + surface.set_subpixel_order(cairo::SubpixelOrder::from_wl_subpixel(subpixel)); - if let Err(e) = - state.surfaces[*data].set_physical_dimensions(physical_width, physical_height) - { + if let Err(e) = surface.set_physical_dimensions(physical_width, physical_height) { warn!("Failed to set output physical dimensions: {e}"); } } wl_output::Event::Name { name } => { debug!("Found output '{name}'"); - state.surfaces[*data].output_name = Some(name); + surface.output_name = Some(name); } wl_output::Event::Scale { factor } => { - if let Err(e) = state.surfaces[*data].set_scale(factor) { + if let Err(e) = surface.set_scale(factor) { warn!("Failed to set output scale: {e}"); } else { debug!( "Set output scale for '{}' to {factor}", - state.surfaces[*data] - .output_name - .as_ref() - .unwrap_or(&"".to_string()) + surface.output_name.as_ref().unwrap_or(&"".to_string()) ); } } @@ -369,12 +405,7 @@ impl Dispatch for NLockState { if let (Some(compositor), Some(subcompositor), Some(session_lock)) = (&state.compositor, &state.subcompositor, &state.session_lock) { - state.surfaces[*data].create_surface( - compositor, - subcompositor, - session_lock, - qh, - ); + surface.create_surface(compositor, subcompositor, session_lock, qh); } } _ => {} diff --git a/src/surface.rs b/src/surface.rs index 8968da8..b4cb90e 100644 --- a/src/surface.rs +++ b/src/surface.rs @@ -15,18 +15,33 @@ use wayland_protocols::ext::session_lock::v1::client::{ use crate::{ auth::AuthState, - buffer::NLockBuffer, + buffer::{NLockBuffer, NLockCommitArgs}, config::NLockConfig, render::{DEFAULT_DPI, NLockRenderBackgroundArgs, NLockRenderOverlayArgs, NLockRenderer}, state::NLockState, }; +pub struct NLockSurfaceTracking { + pub dirty: bool, + pub ready: bool, +} + +impl Default for NLockSurfaceTracking { + fn default() -> Self { + Self { + dirty: false, + ready: true, + } + } +} + pub struct NLockSurface { + pub tracking: NLockSurfaceTracking, pub created: bool, // Background rendering is expensive, only do it once. pub bg_rendered: bool, - pub index: usize, pub output_name: Option, + pub output_id: u32, output_scale: i32, width: Option, @@ -50,12 +65,13 @@ pub struct NLockSurface { } impl NLockSurface { - pub fn new(output: wl_output::WlOutput, index: usize) -> Self { + pub fn new(output: wl_output::WlOutput, output_id: u32) -> Self { Self { + tracking: NLockSurfaceTracking::default(), created: false, bg_rendered: false, - index, output_name: None, + output_id, output_scale: 1, width: None, height: None, @@ -210,6 +226,22 @@ impl NLockSurface { ) -> Option { let (width, height) = self.get_dimensions::().ok()?; + // expensive, don't run this in rel + #[cfg(debug_assertions)] + { + use std::sync::Arc; + + for (i, buf) in self.buffers.iter().enumerate() { + trace!( + "surface {} buffer {} {:p} in_use={}", + self.output_id, + i, + Arc::as_ptr(&buf.state), + buf.state.in_use.load(Ordering::Acquire), + ); + } + } + // The surface size changed, new buffers needed if let Some(last_width) = self.last_width && let Some(last_height) = self.last_height @@ -279,7 +311,7 @@ impl NLockSurface { && self.subsurface.is_some() { let lock_surface = - session_lock.get_lock_surface(surface, &self.output, qh, self.index); + session_lock.get_lock_surface(surface, &self.output, qh, self.output_id); self.lock_surface = Some(lock_surface); } else { warn!("Failed to create background, overlay, or sub surface"); @@ -367,10 +399,17 @@ impl NLockSurface { )?; context.restore()?; + let commit_args = NLockCommitArgs { + surface, + scale: self.output_scale, + output: self.output_id, + tracking: &mut self.tracking, + }; + let mut buf_guard = buffer .lock_buffer() .ok_or(anyhow!("Failed to lock buffer {}", idx))?; - buf_guard.commit_to(surface, self.output_scale); + buf_guard.commit_to(commit_args, qh); // Avoid rendering the background again self.bg_rendered = true; @@ -431,10 +470,17 @@ impl NLockSurface { // Ensure subsurface position is always set to 0,0 subsurface.set_position(0, 0); + let commit_args = NLockCommitArgs { + surface, + scale: self.output_scale, + output: self.output_id, + tracking: &mut self.tracking, + }; + let mut buf_guard = buffer .lock_buffer() .ok_or(anyhow!("Failed to lock buffer {}", idx))?; - buf_guard.commit_to(surface, self.output_scale); + buf_guard.commit_to(commit_args, qh); Ok(()) } @@ -453,12 +499,12 @@ impl Drop for NLockSurface { } } -impl Dispatch for NLockState { +impl Dispatch for NLockState { fn event( state: &mut Self, lock_surface: &ext_session_lock_surface_v1::ExtSessionLockSurfaceV1, event: ::Event, - data: &usize, + data: &u32, _: &wayland_client::Connection, qh: &QueueHandle, ) { @@ -469,13 +515,18 @@ impl Dispatch for N } = event && let Some(shm) = &state.shm { - let surface = &mut state.surfaces[*data]; + let Some(surface) = state.surfaces.get_mut(data) else { + warn!("could not find surface {}", *data); + return; + }; if let Err(e) = surface.set_raw_dimensions(width, height) { warn!("Failed to set surface dimensions: {e}"); return; } + trace!("configure {}", serial); + lock_surface.ack_configure(serial); let auth_state = state.auth_state.clone().load(Ordering::Relaxed); diff --git a/src/util.rs b/src/util.rs index 0ac9574..edcfbb2 100644 --- a/src/util.rs +++ b/src/util.rs @@ -3,21 +3,11 @@ use std::{ io::{self, Read}, - os::fd::OwnedFd, str::FromStr, }; use clap::ValueEnum; -use nix::{ - fcntl::OFlag, - sys::{ - mman::{shm_open, shm_unlink}, - stat::Mode, - }, - unistd::getpid, -}; use serde::{Deserialize, de}; -use tracing::debug; #[derive(Debug, Deserialize, Copy, Clone, PartialEq, ValueEnum)] #[serde(rename_all = "lowercase")] @@ -128,37 +118,6 @@ impl From for tracing::level_filters::LevelFilter { } } -pub fn open_shm() -> Option { - let mut retries = 100; - - loop { - let time = chrono::Local::now(); - let name = format!( - "/nlock-{}-{}-{}", - getpid(), - time.timestamp_micros(), - time.timestamp_subsec_nanos() - ); - debug!("Trying shm file name '{}'", name); - - if let Ok(fd) = shm_open( - name.as_str(), - OFlag::O_RDWR | OFlag::O_CREAT | OFlag::O_EXCL, - Mode::S_IRUSR | Mode::S_IWUSR, - ) { - let _ = shm_unlink(name.as_str()); - return Some(fd); - } - - retries -= 1; - if retries <= 0 { - break; - } - } - - None -} - const PNG_SIG: [u8; 8] = [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]; // Detect if a source stream starts with a PNG signature.