Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

chore: remove unsafe from rawfd #66

Merged
merged 1 commit into from
Oct 7, 2023
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
139 changes: 124 additions & 15 deletions libwayshot/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,7 @@ mod screencopy;
use std::{
cmp,
fs::File,
os::{
fd::{AsFd, AsRawFd},
unix::prelude::FromRawFd,
},
os::fd::AsFd,
process::exit,
sync::atomic::{AtomicBool, Ordering},
};
Expand Down Expand Up @@ -175,17 +172,132 @@ impl WayshotConnection {
fd: T,
capture_region: Option<CaptureRegion>,
) -> Result<FrameFormat> {
self.capture_output_frame_shm_fd_inner(cursor_overlay, output, fd, None, capture_region)
// Connecting to wayland environment.
let mut state = CaptureFrameState {
formats: Vec::new(),
state: None,
buffer_done: AtomicBool::new(false),
};
let mut event_queue = self.conn.new_event_queue::<CaptureFrameState>();
let qh = event_queue.handle();

// Instantiating screencopy manager.
let screencopy_manager = match self.globals.bind::<ZwlrScreencopyManagerV1, _, _>(
&qh,
3..=3,
(),
) {
Ok(x) => x,
Err(e) => {
log::error!("Failed to create screencopy manager. Does your compositor implement ZwlrScreencopy?");
log::error!("err: {e}");
return Err(Error::ProtocolNotFound(
"ZwlrScreencopy Manager not found".to_string(),
));
}
};

// Capture output.
let frame: ZwlrScreencopyFrameV1 = if let Some(region) = capture_region {
screencopy_manager.capture_output_region(
cursor_overlay,
output,
region.x_coordinate,
region.y_coordinate,
region.width,
region.height,
&qh,
(),
)
} else {
screencopy_manager.capture_output(cursor_overlay, output, &qh, ())
};

// Empty internal event buffer until buffer_done is set to true which is when the Buffer done
// event is fired, aka the capture from the compositor is succesful.
while !state.buffer_done.load(Ordering::SeqCst) {
event_queue.blocking_dispatch(&mut state)?;
}

log::debug!(
"Received compositor frame buffer formats: {:#?}",
state.formats
);
// Filter advertised wl_shm formats and select the first one that matches.
let frame_format = state
.formats
.iter()
.find(|frame| {
matches!(
frame.format,
wl_shm::Format::Xbgr2101010
| wl_shm::Format::Abgr2101010
| wl_shm::Format::Argb8888
| wl_shm::Format::Xrgb8888
| wl_shm::Format::Xbgr8888
)
})
.copied();
log::debug!("Selected frame buffer format: {:#?}", frame_format);

// Check if frame format exists.
let frame_format = match frame_format {
Some(format) => format,
None => {
log::error!("No suitable frame format found");
return Err(Error::NoSupportedBufferFormat);
}
};

// Bytes of data in the frame = stride * height.
let frame_bytes = frame_format.stride * frame_format.height;

// Create an in memory file and return it's file descriptor.

// Instantiate shm global.
let shm = self.globals.bind::<WlShm, _, _>(&qh, 1..=1, ()).unwrap();
let shm_pool = shm.create_pool(fd.as_fd(), frame_bytes as i32, &qh, ());
let buffer = shm_pool.create_buffer(
0,
frame_format.width as i32,
frame_format.height as i32,
frame_format.stride as i32,
frame_format.format,
&qh,
(),
);

// Copy the pixel data advertised by the compositor into the buffer we just created.
frame.copy(&buffer);
// On copy the Ready / Failed events are fired by the frame object, so here we check for them.
loop {
// Basically reads, if frame state is not None then...
if let Some(state) = state.state {
match state {
FrameState::Failed => {
log::error!("Frame copy failed");
return Err(Error::FramecopyFailed);
}
FrameState::Finished => {
buffer.destroy();
shm_pool.destroy();
return Ok(frame_format);
}
}
}

event_queue.blocking_dispatch(&mut state)?;
}
}

fn capture_output_frame_shm_fd_inner<T: AsFd>(
fn capture_output_frame_shm_from_file(
&self,
cursor_overlay: i32,
output: &WlOutput,
fd: T,
file: Option<&File>,
file: &File,
capture_region: Option<CaptureRegion>,
) -> Result<FrameFormat> {
let fd = file.as_fd();
// Connecting to wayland environment.
let mut state = CaptureFrameState {
formats: Vec::new(),
Expand Down Expand Up @@ -265,9 +377,7 @@ impl WayshotConnection {

// Bytes of data in the frame = stride * height.
let frame_bytes = frame_format.stride * frame_format.height;
if let Some(file) = file {
file.set_len(frame_bytes as u64)?;
}
file.set_len(frame_bytes as u64)?;
// Create an in memory file and return it's file descriptor.

// Instantiate shm global.
Expand Down Expand Up @@ -317,13 +427,12 @@ impl WayshotConnection {
// Create an in memory file and return it's file descriptor.
let fd = create_shm_fd()?;
// Create a writeable memory map backed by a mem_file.
let mem_file = unsafe { File::from_raw_fd(fd.as_raw_fd()) };
let mem_file = File::from(fd);

let frame_format = self.capture_output_frame_shm_fd_inner(
let frame_format = self.capture_output_frame_shm_from_file(
cursor_overlay,
output,
fd,
Some(&mem_file),
&mem_file,
capture_region,
)?;

Expand Down
Loading