windows_capture/
graphics_capture_api.rs

1use std::sync::{
2    atomic::{self, AtomicBool},
3    Arc,
4};
5
6use parking_lot::Mutex;
7use windows::{
8    core::{IInspectable, Interface, HSTRING},
9    Foundation::{EventRegistrationToken, Metadata::ApiInformation, TypedEventHandler},
10    Graphics::{
11        Capture::{Direct3D11CaptureFramePool, GraphicsCaptureItem, GraphicsCaptureSession},
12        DirectX::{Direct3D11::IDirect3DDevice, DirectXPixelFormat},
13    },
14    Win32::{
15        Foundation::{LPARAM, WPARAM},
16        Graphics::Direct3D11::{
17            ID3D11Device, ID3D11DeviceContext, ID3D11Texture2D, D3D11_TEXTURE2D_DESC,
18        },
19        System::WinRT::Direct3D11::IDirect3DDxgiInterfaceAccess,
20        UI::WindowsAndMessaging::{PostThreadMessageW, WM_QUIT},
21    },
22};
23
24use crate::{
25    capture::GraphicsCaptureApiHandler,
26    d3d11::{self, create_direct3d_device, SendDirectX},
27    frame::Frame,
28    settings::{ColorFormat, CursorCaptureSettings, DrawBorderSettings},
29};
30
31#[derive(thiserror::Error, Eq, PartialEq, Clone, Debug)]
32pub enum Error {
33    #[error("Graphics capture API is not supported")]
34    Unsupported,
35    #[error("Graphics capture API toggling cursor capture is not supported")]
36    CursorConfigUnsupported,
37    #[error("Graphics capture API toggling border capture is not supported")]
38    BorderConfigUnsupported,
39    #[error("Already started")]
40    AlreadyStarted,
41    #[error("DirectX error: {0}")]
42    DirectXError(#[from] d3d11::Error),
43    #[error("Windows API error: {0}")]
44    WindowsError(#[from] windows::core::Error),
45}
46
47/// Used to control the capture session
48pub struct InternalCaptureControl {
49    stop: Arc<AtomicBool>,
50}
51
52impl InternalCaptureControl {
53    /// Create a new `InternalCaptureControl` struct.
54    ///
55    /// # Arguments
56    ///
57    /// * `stop` - An `Arc<AtomicBool>` indicating whether the capture should stop.
58    ///
59    /// # Returns
60    ///
61    /// A new instance of `InternalCaptureControl`.
62    #[must_use]
63    #[inline]
64    pub const fn new(stop: Arc<AtomicBool>) -> Self {
65        Self { stop }
66    }
67
68    /// Gracefully stop the capture thread.
69    #[inline]
70    pub fn stop(self) {
71        self.stop.store(true, atomic::Ordering::Relaxed);
72    }
73}
74
75/// Represents the GraphicsCaptureApi struct.
76pub struct GraphicsCaptureApi {
77    /// The GraphicsCaptureItem associated with the GraphicsCaptureApi.
78    item: GraphicsCaptureItem,
79    /// The ID3D11Device associated with the GraphicsCaptureApi.
80    _d3d_device: ID3D11Device,
81    /// The IDirect3DDevice associated with the GraphicsCaptureApi.
82    _direct3d_device: IDirect3DDevice,
83    /// The ID3D11DeviceContext associated with the GraphicsCaptureApi.
84    _d3d_device_context: ID3D11DeviceContext,
85    /// The optional Arc<Direct3D11CaptureFramePool> associated with the GraphicsCaptureApi.
86    frame_pool: Option<Arc<Direct3D11CaptureFramePool>>,
87    /// The optional GraphicsCaptureSession associated with the GraphicsCaptureApi.
88    session: Option<GraphicsCaptureSession>,
89    /// The Arc<AtomicBool> used to halt the GraphicsCaptureApi.
90    halt: Arc<AtomicBool>,
91    /// Indicates whether the GraphicsCaptureApi is active or not.
92    active: bool,
93    /// The EventRegistrationToken associated with the capture closed event.
94    capture_closed_event_token: EventRegistrationToken,
95    /// The EventRegistrationToken associated with the frame arrived event.
96    frame_arrived_event_token: EventRegistrationToken,
97}
98
99impl GraphicsCaptureApi {
100    /// Create a new Graphics Capture API struct.
101    ///
102    /// # Arguments
103    ///
104    /// * `d3d_device` - The ID3D11Device to use for the capture.
105    /// * `d3d_device_context` - The ID3D11DeviceContext to use for the capture.
106    /// * `item` - The graphics capture item to capture.
107    /// * `callback` - The callback handler for capturing frames.
108    /// * `capture_cursor` - Optional flag to capture the cursor.
109    /// * `draw_border` - Optional flag to draw a border around the captured region.
110    /// * `color_format` - The color format for the captured frames.
111    /// * `thread_id` - The ID of the thread where the capture is running.
112    /// * `result` - The result of the capture operation.
113    ///
114    /// # Returns
115    ///
116    /// Returns a `Result` containing the new `GraphicsCaptureApi` struct if successful, or an `Error` if an error occurred.
117    #[allow(clippy::too_many_arguments)]
118    #[inline]
119    pub fn new<
120        T: GraphicsCaptureApiHandler<Error = E> + Send + 'static,
121        E: Send + Sync + 'static,
122    >(
123        d3d_device: ID3D11Device,
124        d3d_device_context: ID3D11DeviceContext,
125        item: GraphicsCaptureItem,
126        callback: Arc<Mutex<T>>,
127        cursor_capture: CursorCaptureSettings,
128        draw_border: DrawBorderSettings,
129        color_format: ColorFormat,
130        thread_id: u32,
131        result: Arc<Mutex<Option<E>>>,
132    ) -> Result<Self, Error> {
133        // Check support
134        if !Self::is_supported()? {
135            return Err(Error::Unsupported);
136        }
137
138        if cursor_capture != CursorCaptureSettings::Default
139            && !Self::is_cursor_settings_supported()?
140        {
141            return Err(Error::CursorConfigUnsupported);
142        }
143
144        if draw_border != DrawBorderSettings::Default && !Self::is_border_settings_supported()? {
145            return Err(Error::BorderConfigUnsupported);
146        }
147
148        // Create DirectX devices
149        let direct3d_device = create_direct3d_device(&d3d_device)?;
150
151        let pixel_format = DirectXPixelFormat(color_format as i32);
152
153        // Create frame pool
154        let frame_pool =
155            Direct3D11CaptureFramePool::Create(&direct3d_device, pixel_format, 1, item.Size()?)?;
156        let frame_pool = Arc::new(frame_pool);
157
158        // Create capture session
159        let session = frame_pool.CreateCaptureSession(&item)?;
160
161        // Preallocate memory
162        let mut buffer = vec![0u8; 3840 * 2160 * 4];
163
164        // Indicates if the capture is closed
165        let halt = Arc::new(AtomicBool::new(false));
166
167        // Set capture session closed event
168        let capture_closed_event_token = item.Closed(&TypedEventHandler::<
169            GraphicsCaptureItem,
170            IInspectable,
171        >::new({
172            // Init
173            let callback_closed = callback.clone();
174            let halt_closed = halt.clone();
175            let result_closed = result.clone();
176
177            move |_, _| {
178                halt_closed.store(true, atomic::Ordering::Relaxed);
179
180                // Notify the struct that the capture session is closed
181                let callback_closed = callback_closed.lock().on_closed();
182                if let Err(e) = callback_closed {
183                    *result_closed.lock() = Some(e);
184                }
185
186                // To stop message loop
187                unsafe {
188                    PostThreadMessageW(thread_id, WM_QUIT, WPARAM::default(), LPARAM::default())?;
189                };
190
191                Result::Ok(())
192            }
193        }))?;
194
195        // Set frame pool frame arrived event
196        let frame_arrived_event_token = frame_pool.FrameArrived(&TypedEventHandler::<
197            Direct3D11CaptureFramePool,
198            IInspectable,
199        >::new({
200            // Init
201            let frame_pool_recreate = frame_pool.clone();
202            let halt_frame_pool = halt.clone();
203            let d3d_device_frame_pool = d3d_device.clone();
204            let context = d3d_device_context.clone();
205            let result_frame_pool = result;
206
207            let mut last_size = item.Size()?;
208            let callback_frame_pool = callback;
209            let direct3d_device_recreate = SendDirectX::new(direct3d_device.clone());
210
211            move |frame, _| {
212                // Return early if the capture is closed
213                if halt_frame_pool.load(atomic::Ordering::Relaxed) {
214                    return Ok(());
215                }
216
217                // Get frame
218                let frame = frame
219                    .as_ref()
220                    .expect("FrameArrived parameter was None this should never happen.")
221                    .TryGetNextFrame()?;
222                let timespan = frame.SystemRelativeTime()?;
223
224                // Get frame content size
225                let frame_content_size = frame.ContentSize()?;
226
227                // Get frame surface
228                let frame_surface = frame.Surface()?;
229
230                // Convert surface to texture
231                let frame_dxgi_interface = frame_surface.cast::<IDirect3DDxgiInterfaceAccess>()?;
232                let frame_texture =
233                    unsafe { frame_dxgi_interface.GetInterface::<ID3D11Texture2D>()? };
234
235                // Get texture settings
236                let mut desc = D3D11_TEXTURE2D_DESC::default();
237                unsafe { frame_texture.GetDesc(&mut desc) }
238
239                // Check if the size has been changed
240                if frame_content_size.Width != last_size.Width
241                    || frame_content_size.Height != last_size.Height
242                {
243                    let direct3d_device_recreate = &direct3d_device_recreate;
244                    frame_pool_recreate.Recreate(
245                        &direct3d_device_recreate.0,
246                        pixel_format,
247                        1,
248                        frame_content_size,
249                    )?;
250
251                    last_size = frame_content_size;
252
253                    return Ok(());
254                }
255
256                // Set width & height
257                let texture_width = desc.Width;
258                let texture_height = desc.Height;
259
260                // Create a frame
261                let mut frame = Frame::new(
262                    &d3d_device_frame_pool,
263                    frame_surface,
264                    frame_texture,
265                    timespan,
266                    &context,
267                    &mut buffer,
268                    texture_width,
269                    texture_height,
270                    color_format,
271                );
272
273                // Init internal capture control
274                let stop = Arc::new(AtomicBool::new(false));
275                let internal_capture_control = InternalCaptureControl::new(stop.clone());
276
277                // Send the frame to the callback struct
278                let result = callback_frame_pool
279                    .lock()
280                    .on_frame_arrived(&mut frame, internal_capture_control);
281
282                if stop.load(atomic::Ordering::Relaxed) || result.is_err() {
283                    if let Err(e) = result {
284                        *result_frame_pool.lock() = Some(e);
285                    }
286
287                    halt_frame_pool.store(true, atomic::Ordering::Relaxed);
288
289                    // To stop the message loop
290                    unsafe {
291                        PostThreadMessageW(
292                            thread_id,
293                            WM_QUIT,
294                            WPARAM::default(),
295                            LPARAM::default(),
296                        )?;
297                    };
298                }
299
300                Result::Ok(())
301            }
302        }))?;
303
304        if cursor_capture != CursorCaptureSettings::Default {
305            if Self::is_cursor_settings_supported()? {
306                match cursor_capture {
307                    CursorCaptureSettings::Default => (),
308                    CursorCaptureSettings::WithCursor => session.SetIsCursorCaptureEnabled(true)?,
309                    CursorCaptureSettings::WithoutCursor => {
310                        session.SetIsCursorCaptureEnabled(false)?
311                    }
312                };
313            } else {
314                return Err(Error::CursorConfigUnsupported);
315            }
316        }
317
318        if draw_border != DrawBorderSettings::Default {
319            if Self::is_border_settings_supported()? {
320                match draw_border {
321                    DrawBorderSettings::Default => (),
322                    DrawBorderSettings::WithBorder => {
323                        session.SetIsBorderRequired(true)?;
324                    }
325                    DrawBorderSettings::WithoutBorder => session.SetIsBorderRequired(false)?,
326                }
327            } else {
328                return Err(Error::BorderConfigUnsupported);
329            }
330        }
331
332        Ok(Self {
333            item,
334            _d3d_device: d3d_device,
335            _direct3d_device: direct3d_device,
336            _d3d_device_context: d3d_device_context,
337            frame_pool: Some(frame_pool),
338            session: Some(session),
339            halt,
340            active: false,
341            frame_arrived_event_token,
342            capture_closed_event_token,
343        })
344    }
345
346    /// Start the capture.
347    ///
348    /// # Returns
349    ///
350    /// Returns `Ok(())` if the capture started successfully, or an `Error` if an error occurred.
351    #[inline]
352    pub fn start_capture(&mut self) -> Result<(), Error> {
353        if self.active {
354            return Err(Error::AlreadyStarted);
355        }
356        self.active = true;
357
358        self.session.as_ref().unwrap().StartCapture()?;
359
360        Ok(())
361    }
362
363    /// Stop the capture.
364    #[inline]
365    pub fn stop_capture(mut self) {
366        if let Some(frame_pool) = self.frame_pool.take() {
367            frame_pool
368                .RemoveFrameArrived(self.frame_arrived_event_token)
369                .expect("Failed to remove Frame Arrived event handler");
370
371            frame_pool.Close().expect("Failed to Close Frame Pool");
372        }
373
374        if let Some(session) = self.session.take() {
375            session.Close().expect("Failed to Close Capture Session");
376        }
377
378        self.item
379            .RemoveClosed(self.capture_closed_event_token)
380            .expect("Failed to remove Capture Session Closed event handler");
381    }
382
383    /// Get the halt handle.
384    ///
385    /// # Returns
386    ///
387    /// Returns an `Arc<AtomicBool>` representing the halt handle.
388    #[must_use]
389    #[inline]
390    pub fn halt_handle(&self) -> Arc<AtomicBool> {
391        self.halt.clone()
392    }
393
394    /// Check if the Windows Graphics Capture API is supported.
395    ///
396    /// # Returns
397    ///
398    /// Returns `Ok(true)` if the API is supported, `Ok(false)` if the API is not supported, or an `Error` if an error occurred.
399    #[inline]
400    pub fn is_supported() -> Result<bool, Error> {
401        Ok(ApiInformation::IsApiContractPresentByMajor(
402            &HSTRING::from("Windows.Foundation.UniversalApiContract"),
403            8,
404        )? && GraphicsCaptureSession::IsSupported()?)
405    }
406
407    /// Check if you can change the cursor capture setting.
408    ///
409    /// # Returns
410    ///
411    /// Returns `true` if toggling the cursor capture is supported, `false` otherwise.
412    #[inline]
413    pub fn is_cursor_settings_supported() -> Result<bool, Error> {
414        Ok(ApiInformation::IsPropertyPresent(
415            &HSTRING::from("Windows.Graphics.Capture.GraphicsCaptureSession"),
416            &HSTRING::from("IsCursorCaptureEnabled"),
417        )? && Self::is_supported()?)
418    }
419
420    /// Check if you can change the border capture setting.
421    ///
422    /// # Returns
423    ///
424    /// Returns `true` if toggling the border capture is supported, `false` otherwise.
425    #[inline]
426    pub fn is_border_settings_supported() -> Result<bool, Error> {
427        Ok(ApiInformation::IsPropertyPresent(
428            &HSTRING::from("Windows.Graphics.Capture.GraphicsCaptureSession"),
429            &HSTRING::from("IsBorderRequired"),
430        )? && Self::is_supported()?)
431    }
432}
433
434impl Drop for GraphicsCaptureApi {
435    fn drop(&mut self) {
436        if let Some(frame_pool) = self.frame_pool.take() {
437            frame_pool
438                .RemoveFrameArrived(self.frame_arrived_event_token)
439                .expect("Failed to remove Frame Arrived event handler");
440
441            frame_pool.Close().expect("Failed to Close Frame Pool");
442        }
443
444        if let Some(session) = self.session.take() {
445            session.Close().expect("Failed to Close Capture Session");
446        }
447
448        self.item
449            .RemoveClosed(self.capture_closed_event_token)
450            .expect("Failed to remove Capture Session Closed event handler");
451    }
452}