1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

use std::slice;

use windows::core::PCSTR;
use windows::Win32::Foundation::{CloseHandle, HANDLE};
use windows::Win32::System::Memory::{
  CreateFileMappingA, MapViewOfFile, UnmapViewOfFile, FILE_MAP_ALL_ACCESS, MEMORYMAPPEDVIEW_HANDLE,
};
use windows::Win32::{
  Foundation::INVALID_HANDLE_VALUE,
  Graphics::{
    Direct3D11::ID3D11Texture2D,
    Dxgi::{DXGI_OUTDUPL_FRAME_INFO, DXGI_OUTPUT_DESC},
  },
  System::Memory::PAGE_READWRITE,
};

use crate::duplicate_context::DuplicateContext;
use crate::utils::{calc_buffer_size, Result};

use super::model::Capturer;

/// Capture screen to a chunk of shared memory.
pub struct SharedCapturer<'a> {
  buffer: *mut u8,
  buffer_size: usize,
  file: HANDLE,
  ctx: &'a DuplicateContext,
  texture: ID3D11Texture2D,
}

impl<'a> SharedCapturer<'a> {
  pub fn new(ctx: &'a DuplicateContext, name: &str) -> Result<Self> {
    let (buffer, buffer_size, file, texture) = Self::allocate(ctx, name)?;
    Ok(Self {
      buffer,
      buffer_size,
      file,
      texture,
      ctx,
    })
  }

  fn allocate(
    ctx: &'a DuplicateContext,
    name: &str,
  ) -> Result<(*mut u8, usize, HANDLE, ID3D11Texture2D)> {
    let (texture, desc) = ctx.create_readable_texture()?;
    let buffer_size = calc_buffer_size(desc);

    unsafe {
      let file = CreateFileMappingA(
        INVALID_HANDLE_VALUE,
        None,
        PAGE_READWRITE,
        0,
        buffer_size as u32,
        PCSTR(name.as_ptr()),
      )
      .map_err(|_| "CreateFileMappingA failed")?;

      let buffer = MapViewOfFile(
        file,                // handle to map object
        FILE_MAP_ALL_ACCESS, // read/write permission
        0,
        0,
        buffer_size,
      )
      .map_err(|_| "MapViewOfFile failed")?
      .0 as *mut u8;
      Ok((buffer, buffer_size, file, texture))
    }
  }

  fn free(&self) {
    unsafe {
      UnmapViewOfFile(MEMORYMAPPEDVIEW_HANDLE(self.buffer as isize));
      CloseHandle(self.file);
    }
  }
}

impl<'a> Capturer for SharedCapturer<'a> {
  fn get_buffer(&self) -> &[u8] {
    unsafe { slice::from_raw_parts(self.buffer, self.buffer_size) }
  }

  fn get_desc(&self) -> Result<DXGI_OUTPUT_DESC> {
    self.ctx.get_desc()
  }

  fn check_buffer(&self) -> Result<()> {
    if self.buffer_size < calc_buffer_size(self.get_desc()?) {
      return Err("Invalid buffer length");
    } else {
      Ok(())
    }
  }

  fn capture(&mut self) -> Result<DXGI_OUTDUPL_FRAME_INFO> {
    self
      .ctx
      .capture_frame(self.buffer, self.buffer_size, &self.texture)
  }

  fn safe_capture(&mut self) -> Result<DXGI_OUTDUPL_FRAME_INFO> {
    self.check_buffer()?;
    self.capture()
  }
}

impl DuplicateContext {
  pub fn shared_capturer(&self, name: &str) -> Result<SharedCapturer> {
    SharedCapturer::new(self, name)
  }
}

impl<'a> Drop for SharedCapturer<'a> {
  fn drop(&mut self) {
    self.free()
  }
}