feat #40: support to create snapshot with specific source id

include/traa/error.h

@@ -165,6 +165,13 @@ typedef enum traa_error {
    */
   TRAA_ERROR_ENUM_SCREEN_SOURCE_INFO_FAILED = 21,
 
+  /**
+   * @brief Invalid source id error.
+   *
+   * This is an invalid source id error.
+   */
+  TRAA_ERROR_INVALID_SOURCE_ID = 22,
+
   /**
    * @brief Error count.
    *

include/traa/traa.h

@@ -126,6 +126,32 @@ TRAA_API int TRAA_CALL traa_enum_screen_source_info(const traa_size icon_size,
  * information.
  */
 TRAA_API int TRAA_CALL traa_free_screen_source_info(traa_screen_source_info infos[], int count);
+
+/**
+ * @brief Creates a snapshot of the specified source.
+ *
+ * This function creates a snapshot of the specified source and returns the snapshot data.
+ *
+ * @param source_id The ID of the source to create a snapshot of.
+ * @param snapshot_size The size of the snapshot.
+ * @param data A pointer to a pointer to the snapshot data.
+ * @param data_size The size of the snapshot data.
+ * @param actual_size The actual size of the snapshot data.
+ * @return An integer value indicating the success or failure of the operation.
+ *         A return value of 0 indicates success, while a non-zero value
+ *         indicates failure.
+ */
+TRAA_API int TRAA_CALL traa_create_snapshot(const int64_t source_id, const traa_size snapshot_size,
+                                            uint8_t **data, int *data_size, traa_size *actual_size);
+
+/**
+ * @brief Frees the memory allocated for the snapshot data.
+ *
+ * This function frees the memory allocated for the snapshot data.
+ *
+ * @param data A pointer to the snapshot data to free.
+ */
+TRAA_API void TRAA_CALL traa_free_snapshot(uint8_t *data);
 #endif // (defined(_WIN32) || defined(__APPLE__) || defined(__linux__)) && !defined(__ANDROID__) &&
        // (!defined(TARGET_OS_IPHONE) || !TARGET_OS_IPHONE) &&
        // (!defined(TARGET_OS_VISION) || !TARGET_OS_VISION)

src/base/devices/screen/darwin/enumerator_darwin.mm

@@ -80,6 +80,54 @@
   return nil;
 }
 
+bool is_source_id_valid(const int64_t source_id, bool &is_window) {
+  if (source_id <= 0) {
+    return false;
+  }
+
+  // try to find source_id in the screen list
+  NSArray *screen_array = NSScreen.screens;
+  for (size_t i = 0; i < screen_array.count; i++) {
+    NSScreen *screen = screen_array[i];
+    CGDirectDisplayID screen_id = static_cast<CGDirectDisplayID>(
+        [[screen.deviceDescription objectForKey:@"NSScreenNumber"] unsignedIntValue]);
+    if (source_id == static_cast<int64_t>(screen_id)) {
+      is_window = false;
+      return true;
+    }
+  }
+
+  // try to find source_id in the window list
+  CFArrayRef window_array = CGWindowListCopyWindowInfo(
+      kCGWindowListOptionAll | kCGWindowListExcludeDesktopElements, kCGNullWindowID);
+  for (CFIndex i = 0; i < CFArrayGetCount(window_array); i++) {
+    CFDictionaryRef window =
+        reinterpret_cast<CFDictionaryRef>(CFArrayGetValueAtIndex(window_array, i));
+    if (!window) {
+      continue;
+    }
+
+    CFNumberRef window_id =
+        reinterpret_cast<CFNumberRef>(CFDictionaryGetValue(window, kCGWindowNumber));
+    if (!window_id) {
+      continue;
+    }
+
+    int64_t window_id_value;
+    if (!CFNumberGetValue(window_id, kCFNumberSInt64Type, &window_id_value) ||
+        window_id_value <= 0) {
+      continue;
+    }
+
+    if (source_id == window_id_value) {
+      is_window = true;
+      return true;
+    }
+  }
+
+  return false;
+}
+
 void dump_image_to_file(NSBitmapImageRep *image_rep, const char *file_name) {
   if (image_rep == nil || file_name == nullptr) {
     return;
@@ -414,14 +462,14 @@ int enum_screens(const traa_size thumbnail_size, const unsigned int external_fla
   }
 
   if (sources.size() == 0) {
-    return traa_error::TRAA_ERROR_NONE;
+    return TRAA_ERROR_NONE;
   }
 
   *count = static_cast<int>(sources.size());
   *infos = reinterpret_cast<traa_screen_source_info *>(new traa_screen_source_info[sources.size()]);
   if (*infos == nullptr) {
     LOG_ERROR("alloca memroy for infos failed");
-    return traa_error::TRAA_ERROR_OUT_OF_MEMORY;
+    return TRAA_ERROR_OUT_OF_MEMORY;
   }
 
   for (size_t i = 0; i < sources.size(); ++i) {
@@ -442,5 +490,56 @@ int enum_screens(const traa_size thumbnail_size, const unsigned int external_fla
   return TRAA_ERROR_NONE;
 }
 
+int screen_source_info_enumerator::create_snapshot(const int64_t source_id,
+                                                   const traa_size snapshot_size, uint8_t **data,
+                                                   int *data_size, traa_size *actual_size) {
+  if (source_id < 0 || data == nullptr || data_size == nullptr || actual_size == nullptr) {
+    return TRAA_ERROR_INVALID_ARGUMENT;
+  }
+
+  if (__builtin_available(macOS 11, *)) {
+    if (!CGPreflightScreenCaptureAccess()) {
+      CGRequestScreenCaptureAccess();
+      return TRAA_ERROR_PERMISSION_DENIED;
+    }
+  }
+
+  bool is_window = false;
+  if (!is_source_id_valid(source_id, is_window)) {
+    return TRAA_ERROR_INVALID_SOURCE_ID;
+  }
+
+  NSBitmapImageRep *image_rep = get_image_rep(is_window, static_cast<CGWindowID>(source_id));
+  if (!image_rep) {
+    return TRAA_ERROR_INVALID_SOURCE_ID;
+  }
+
+  NSImage *image = [NSImage new];
+  [image addRepresentation:image_rep];
+
+  desktop_size scaled_size =
+      calc_scaled_size(desktop_size(image.size.width, image.size.height), snapshot_size);
+  CGSize scaled_ns_size = CGSizeMake(scaled_size.width(), scaled_size.height());
+  NSBitmapImageRep *scaled_image_rep = scale_image(image_rep, scaled_ns_size);
+  if (!scaled_image_rep) {
+    return TRAA_ERROR_INVALID_SOURCE_ID;
+  }
+
+  size_t thumbnail_data_size = [scaled_image_rep pixelsWide] * [scaled_image_rep pixelsHigh] *
+                               [scaled_image_rep samplesPerPixel] * sizeof(unsigned char);
+
+  *data = new uint8_t[thumbnail_data_size];
+  *data_size = static_cast<int>(thumbnail_data_size);
+  if (*data == nullptr) {
+    LOG_ERROR("alloc memory for data failed");
+    return TRAA_ERROR_OUT_OF_MEMORY;
+  }
+
+  memcpy(*data, [scaled_image_rep bitmapData], thumbnail_data_size);
+  *actual_size = scaled_size.to_traa_size();
+
+  return TRAA_ERROR_NONE;
+}
+
 } // namespace base
 } // namespace traa

src/base/devices/screen/enumerator.cc

@@ -24,6 +24,10 @@ int screen_source_info_enumerator::free_screen_source_info(traa_screen_source_in
 
   return traa_error::TRAA_ERROR_NONE;
 }
+
+void screen_source_info_enumerator::free_snapshot(uint8_t *data) {
+  delete[] data;
+}
 #endif // (defined(_WIN32) || defined(__APPLE__) || defined(__linux__)) && !defined(__ANDROID__) &&
        // (!defined(TARGET_OS_IPHONE) || !TARGET_OS_IPHONE) &&
        // (!defined(TARGET_OS_VISION) || !TARGET_OS_VISION)

src/base/devices/screen/enumerator.h

@@ -50,6 +50,27 @@ class screen_source_info_enumerator {
    * @return traa_error::TRAA_ERROR_NONE if successful, otherwise an error code.
    */
   static int free_screen_source_info(traa_screen_source_info infos[], int count);
+
+  /**
+   * @brief Creates a snapshot of the specified source.
+   *
+   * @param source_id The ID of the source to create a snapshot of.
+   * @param snapshot_size The size of the snapshot.
+   * @param data A pointer to a pointer to the snapshot data.
+   * @param data_size The size of the snapshot data.
+   * @param actual_size The actual size of the snapshot data.
+   *
+   * @return traa_error::TRAA_ERROR_NONE if successful, otherwise an error code.
+   */
+  static int create_snapshot(const int64_t source_id, const traa_size snapshot_size, uint8_t **data,
+                             int *data_size, traa_size *actual_size);
+
+  /**
+   * @brief Frees the snapshot data.
+   *
+   * @param data The snapshot data to free.
+   */
+  static void free_snapshot(uint8_t *data);
 #endif // (defined(_WIN32) || defined(__APPLE__) || defined(__linux__)) && !defined(__ANDROID__) &&
        // (!defined(TARGET_OS_IPHONE) || !TARGET_OS_IPHONE) &&
        // (!defined(TARGET_OS_VISION) || !TARGET_OS_VISION)

src/base/devices/screen/linux/enumerator_linux.cc

@@ -31,5 +31,11 @@ int screen_source_info_enumerator::enum_screen_source_info(const traa_size icon_
   return traa_error::TRAA_ERROR_ENUM_SCREEN_SOURCE_INFO_FAILED;
 }
 
+int screen_source_info_enumerator::create_snapshot(const int64_t source_id,
+                                                   const traa_size snapshot_size, uint8_t **data,
+                                                   int *data_size, traa_size *actual_size) {
+  return traa_error::TRAA_ERROR_NOT_IMPLEMENTED;
+}
+
 } // namespace base
 } // namespace traa

src/base/devices/screen/win/enumerator_win.cc

@@ -195,7 +195,7 @@ bool get_process_icon_data(LPCWSTR process_path, desktop_size icon_size, uint8_t
     return false;
   }
 
-  auto data_size = scaled_size.width() * scaled_size.height() * 4;
+  auto data_size = scaled_size.width() * scaled_size.height() * desktop_frame::k_bytes_per_pixel;
 
   *icon_data = new uint8_t[data_size];
   if (!*icon_data) {
@@ -213,6 +213,29 @@ bool get_process_icon_data(LPCWSTR process_path, desktop_size icon_size, uint8_t
   return true;
 }
 
+bool is_source_id_valid(const int64_t source_id, bool &is_window) {
+  if (source_id < 0) {
+    return false;
+  }
+
+  // try to find source_id in the window list
+  HWND window = reinterpret_cast<HWND>(source_id);
+  if (::IsWindow(window)) {
+    is_window = true;
+    return true;
+  }
+
+  // try to find source_id in the screen list
+  DISPLAY_DEVICEW device;
+  device.cb = sizeof(device);
+  if (::EnumDisplayDevicesW(NULL, static_cast<DWORD>(source_id), &device, 0)) {
+    is_window = false;
+    return true;
+  }
+
+  return false;
+}
+
 BOOL WINAPI enum_windows_cb(HWND window, LPARAM lParam) {
   auto *param = reinterpret_cast<enumerator_param *>(lParam);
 
@@ -411,10 +434,10 @@ int enum_windows(enumerator_param &param) {
   BOOL ret = ::EnumWindows(enum_windows_cb, reinterpret_cast<LPARAM>(&param));
   if (!ret) {
     LOG_ERROR("call ::EnumWindows failed: {}", ::GetLastError());
-    return traa_error::TRAA_ERROR_ENUM_SCREEN_SOURCE_INFO_FAILED;
+    return TRAA_ERROR_ENUM_SCREEN_SOURCE_INFO_FAILED;
   }
 
-  return traa_error::TRAA_ERROR_NONE;
+  return TRAA_ERROR_NONE;
 }
 
 int enum_screens(enumerator_param &param) {
@@ -437,7 +460,8 @@ int enum_screens(enumerator_param &param) {
     DEVMODEW device_mode;
     device_mode.dmSize = sizeof(device_mode);
     device_mode.dmDriverExtra = 0;
-    BOOL result = ::EnumDisplaySettingsExW(device.DeviceName, ENUM_CURRENT_SETTINGS, &device_mode, 0);
+    BOOL result =
+        ::EnumDisplaySettingsExW(device.DeviceName, ENUM_CURRENT_SETTINGS, &device_mode, 0);
     if (!result) {
       break;
     }
@@ -470,7 +494,7 @@ int enum_screens(enumerator_param &param) {
     param.infos.push_back(screen_info);
   }
 
-  return traa_error::TRAA_ERROR_NONE;
+  return TRAA_ERROR_NONE;
 }
 
 } // namespace
@@ -500,15 +524,15 @@ int screen_source_info_enumerator::enum_screen_source_info(const traa_size icon_
   }
 
   if (param.infos.size() == 0) {
-    return traa_error::TRAA_ERROR_NONE;
+    return TRAA_ERROR_NONE;
   }
 
   *count = static_cast<int>(param.infos.size());
   *infos =
       reinterpret_cast<traa_screen_source_info *>(new traa_screen_source_info[param.infos.size()]);
   if (*infos == nullptr) {
     LOG_ERROR("alloca memroy for infos failed: {}", ::GetLastError());
-    return traa_error::TRAA_ERROR_OUT_OF_MEMORY;
+    return TRAA_ERROR_OUT_OF_MEMORY;
   }
 
   for (size_t i = 0; i < param.infos.size(); ++i) {
@@ -526,7 +550,59 @@ int screen_source_info_enumerator::enum_screen_source_info(const traa_size icon_
     }
   }
 
-  return traa_error::TRAA_ERROR_NONE;
+  return TRAA_ERROR_NONE;
+}
+
+int screen_source_info_enumerator::create_snapshot(const int64_t source_id,
+                                                   const traa_size snapshot_size, uint8_t **data,
+                                                   int *data_size, traa_size *actual_size) {
+  if (source_id < 0 || data == nullptr || data_size == nullptr || actual_size == nullptr) {
+    return TRAA_ERROR_INVALID_ARGUMENT;
+  }
+
+  bool is_window = false;
+  if (!is_source_id_valid(source_id, is_window)) {
+    return TRAA_ERROR_INVALID_SOURCE_ID;
+  }
+
+  if (is_window) {
+    thumbnail thumbnail_instance;
+    thumbnail_instance.get_thumbnail_data(reinterpret_cast<HWND>(source_id), snapshot_size, data,
+                                          *actual_size);
+  } else {
+    DISPLAY_DEVICEW device;
+    device.cb = sizeof(device);
+    if (!::EnumDisplayDevicesW(NULL, static_cast<DWORD>(source_id), &device, 0)) {
+      LOG_ERROR("enum display devices failed: {}", ::GetLastError());
+      return TRAA_ERROR_INVALID_SOURCE_ID;
+    }
+
+    DEVMODEW device_mode;
+    device_mode.dmSize = sizeof(device_mode);
+    device_mode.dmDriverExtra = 0;
+    if (!::EnumDisplaySettingsExW(device.DeviceName, ENUM_CURRENT_SETTINGS, &device_mode, 0)) {
+      LOG_ERROR("enum display settings failed: {}", ::GetLastError());
+      return TRAA_ERROR_INVALID_SOURCE_ID;
+    }
+
+    desktop_rect screen_rect =
+        desktop_rect::make_ltrb(device_mode.dmPosition.x, device_mode.dmPosition.y,
+                                device_mode.dmPelsWidth, device_mode.dmPelsHeight);
+    if (screen_rect.is_empty()) {
+      return TRAA_ERROR_INVALID_SOURCE_ID;
+    }
+
+    // get the screen snapshot with GDI
+    bool ret = capture_utils::get_screen_image_by_gdi(screen_rect.to_traa_rect(), snapshot_size,
+                                                      data, *actual_size);
+    if (!ret) {
+      return TRAA_ERROR_INVALID_SOURCE_ID;
+    }
+  }
+
+  *data_size = actual_size->width * actual_size->height * desktop_frame::k_bytes_per_pixel;
+
+  return TRAA_ERROR_NONE;
 }
 
 } // namespace base

src/base/devices/screen/win/wgc/wgc_capturer_win.cc

@@ -366,7 +366,8 @@ void wgc_capturer_win::capture_frame() {
   }
 
   int64_t capture_time_ms = (time_nanos() - capture_start_time_nanos) / k_num_nanosecs_per_millisec;
-  TRAA_HISTOGRAM_COUNTS_1000("WebRTC.DesktopCapture.Win.WgcCapturerFrameTime", capture_time_ms);
+  TRAA_HISTOGRAM_COUNTS_1000("WebRTC.DesktopCapture.Win.WgcCapturerFrameTime",
+                             static_cast<int>(capture_time_ms));
 
   frame->set_capture_time_ms(capture_time_ms);
   frame->set_capturer_id(desktop_capture_id::k_capture_wgc);