diff --git a/agent/src/capture.rs b/agent/src/capture.rs new file mode 100644 index 0000000..381546a --- /dev/null +++ b/agent/src/capture.rs @@ -0,0 +1,234 @@ +//! Screen capture module. +//! +//! Uses the `scrap` crate to capture display frames efficiently via platform APIs +//! (DXGI on Windows, X11 on Linux, CoreGraphics on macOS). Frames are converted +//! from BGRA to RGB and encoded as JPEG with configurable quality. + +use anyhow::{Context, Result}; +use base64::{Engine, engine::general_purpose::STANDARD}; +use image::{ImageBuffer, Rgb}; +use log::{info, warn}; +use scrap::{Capturer, Display}; +use std::time::Instant; + +/// Manages screen capture for a specific display. +pub struct ScreenCapture { + capturer: Capturer, + width: usize, + height: usize, + quality: u8, + frame_count: u64, + bytes_sent: u64, +} + +impl ScreenCapture { + /// Initialize capture for the given display index. + /// + /// `display_idx` 0 = primary monitor, 1 = second, etc. + /// `quality` is the JPEG encoding quality (1–100). + pub fn new(display_idx: usize, quality: u8) -> Result { + // Get all available displays. + let displays = Display::all() + .map_err(|e| anyhow::anyhow!("failed to enumerate displays: {}", e))?; + + if displays.is_empty() { + anyhow::bail!("no displays detected — is a display / X server available?"); + } + if display_idx >= displays.len() { + anyhow::bail!( + "display index {} out of range ({} displays available)", + display_idx, + displays.len() + ); + } + + let display = displays.into_iter().nth(display_idx).unwrap(); + let width = display.width(); + let height = display.height(); + + info!( + "initializing capture: display {} ({}x{})", + display_idx, width, height + ); + + let capturer = Capturer::new(display) + .map_err(|e| anyhow::anyhow!("failed to create capturer: {}", e))?; + + // The first call to `frame()` may fail on some platforms because the + // internal buffer isn't ready yet. Spin until we get a frame. + let capturer = wait_for_first_frame(capturer); + + Ok(Self { + capturer, + width, + height, + quality: quality.clamp(1, 100), + frame_count: 0, + bytes_sent: 0, + }) + } + + /// Capture a single frame and return it as a base64-encoded JPEG string. + /// + /// This method: + /// 1. Waits for the next frame from the OS (blocks briefly). + /// 2. Converts BGRA pixel data to RGB. + /// 3. Encodes as JPEG with the configured quality. + /// 4. Base64-encodes the result. + pub fn capture_frame(&mut self) -> Result { + let start = Instant::now(); + + // Capture the raw frame (BGRA format on all platforms). + let frame = self + .capturer + .frame() + .map_err(|e| anyhow::anyhow!("frame capture failed: {}", e))?; + + let frame_len = frame.len(); + let expected_len = self.width * self.height * 4; // BGRA = 4 bytes/pixel + if frame.len() < expected_len { + anyhow::bail!( + "frame size mismatch: got {} bytes, expected at least {}", + frame_len, + expected_len + ); + } + + // Convert BGRA → RGB (drop the alpha channel). + let rgb_data = bgra_to_rgb(frame, self.width, self.height); + + // Encode as JPEG. + let jpeg_bytes = encode_jpeg( + &rgb_data, + self.width as u32, + self.height as u32, + self.quality, + ) + .context("JPEG encoding failed")?; + + // Base64 encode. + let b64 = STANDARD.encode(&jpeg_bytes); + + // Track stats. + self.frame_count += 1; + self.bytes_sent += jpeg_bytes.len() as u64; + + let elapsed = start.elapsed(); + log::trace!( + "captured frame: {}x{} → {} bytes JPEG ({} bytes base64) in {:.1}ms", + self.width, + self.height, + jpeg_bytes.len(), + b64.len(), + elapsed.as_secs_f64() * 1000.0 + ); + + Ok(b64) + } + + /// Get the capture resolution as a string (e.g., "1920x1080"). + pub fn resolution(&self) -> String { + format!("{}x{}", self.width, self.height) + } + + /// Get the display width. + pub fn width(&self) -> usize { + self.width + } + + /// Get the display height. + pub fn height(&self) -> usize { + self.height + } + + /// Get total frames captured since start. + pub fn frame_count(&self) -> u64 { + self.frame_count + } + + /// Get total bytes sent (JPEG, before base64). + pub fn bytes_sent(&self) -> u64 { + self.bytes_sent + } +} + +/// Spin until the capturer produces its first valid frame. +/// On some platforms the first few `frame()` calls fail. +fn wait_for_first_frame(mut capturer: Capturer) -> Capturer { + let max_attempts = 30; + for i in 0..max_attempts { + match capturer.frame() { + Ok(_) => { + info!("capture ready (took {} attempts)", i + 1); + return capturer; + } + Err(e) => { + if i < max_attempts - 1 { + log::debug!("waiting for first frame (attempt {}): {}", i + 1, e); + std::thread::sleep(std::time::Duration::from_millis(100)); + } else { + warn!("first frame never ready after {} attempts: {}", max_attempts, e); + return capturer; + } + } + } + } + capturer +} + +/// Convert BGRA pixel data to RGB by dropping the alpha channel. +fn bgra_to_rgb(bgra: &[u8], width: usize, height: usize) -> Vec { + let mut rgb = Vec::with_capacity(width * height * 3); + for chunk in bgra.chunks_exact(4) { + // BGRA → RGB: reverse B,G,R order to R,G,B, skip A + rgb.push(chunk[2]); // R + rgb.push(chunk[1]); // G + rgb.push(chunk[0]); // B + } + rgb +} + +/// Encode raw RGB data as JPEG with the given quality. +fn encode_jpeg( + rgb_data: &[u8], + width: u32, + height: u32, + quality: u8, +) -> Result> { + let img_buffer = ImageBuffer::>::from_raw(width, height, rgb_data) + .context("failed to create image buffer from RGB data")?; + + let mut jpeg_bytes = Vec::new(); + let encoder = image::codecs::jpeg::JpegEncoder::new_with_quality(&mut jpeg_bytes, quality); + encoder + .encode_image(&img_buffer) + .context("JPEG encode failed")?; + + Ok(jpeg_bytes) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_bgra_to_rgb() { + // BGRA pixel: Blue=10, Green=20, Red=30, Alpha=255 + let bgra = vec![10, 20, 30, 255, 40, 50, 60, 128]; + let rgb = bgra_to_rgb(&bgra, 2, 1); + assert_eq!(rgb, vec![30, 20, 10, 60, 50, 40]); // RGB reversed, alpha dropped + } + + #[test] + fn test_encode_jpeg() { + // Create a 4x4 red image + let rgb_data = vec![255u8; 4 * 4 * 3]; + let result = encode_jpeg(&rgb_data, 4, 4, 80); + assert!(result.is_ok()); + let jpeg = result.unwrap(); + assert!(!jpeg.is_empty()); + // JPEG files start with FF D8 + assert_eq!(jpeg[0], 0xFF); + assert_eq!(jpeg[1], 0xD8); + } +}