Skip to content

Commit

Permalink
refactor(ScreenCapturer): 支持electron memory cage
Browse files Browse the repository at this point in the history
  • Loading branch information
@nzh63
nzh63 committed Oct 22, 2023
1 parent 6611e8b commit b8503f9
Show file tree
Hide file tree
Showing 3 changed files with 65 additions and 31 deletions.
92 changes: 62 additions & 30 deletions native/addons/ScreenCapturer/capture.cpp
View file
Expand Up @@ -11,37 +11,70 @@ struct CaptureData {
bool error = false;
int width = 0;
int height = 0;
uint8_t *buffer = nullptr;
int bufferSize = 0;
void *data = nullptr;
napi_ref buffer = nullptr;
napi_deferred deferred = nullptr;
napi_value promise = nullptr;
napi_async_work work = nullptr;
};

void executeCapture(napi_env env, void *_data) {
void queryWindowSize(napi_env env, void *_data);
void mallocBuffer(napi_env env, napi_status status, void *_data);
void executeCapture(napi_env env, void *_data);
void completeCapture(napi_env env, napi_status status, void *_data);

void queryWindowSize(napi_env env, void *_data) {
auto &data = *(CaptureData *)_data;
RECT rect{0, 0, 0, 0};
DwmGetWindowAttribute(data.hwnd, DWMWA_EXTENDED_FRAME_BOUNDS, &rect, sizeof rect);
const int width = rect.right - rect.left;
const int height = rect.bottom - rect.top;
if (height == 0 || width == 0) {
if (height <= 0 || width <= 0) {
data.error = true;
return;
}
data.width = width;
data.height = height;
}

void mallocBuffer(napi_env env, napi_status status, void *_data) {
auto &data = *(CaptureData *)_data;
if (data.error) {
if (data.deferred) napi_reject_deferred(env, data.deferred, createError(env, nullptr, "Capture error"));
if (data.work) napi_delete_async_work(env, data.work);
if (data.buffer) napi_delete_reference(env, data.buffer);
delete (CaptureData *)_data;
return;
}
napi_value buffer = nullptr;
NAPI_CALL(napi_create_buffer(env, data.width * data.height * 4, &data.data, &buffer));
NAPI_CALL(napi_create_reference(env, buffer, 1, &data.buffer));

NAPI_CALL(napi_delete_async_work(env, data.work));
data.work = nullptr;
napi_value resource_name;
NAPI_CALL(napi_create_string_utf8(env, "capture", NAPI_AUTO_LENGTH, &resource_name));
NAPI_CALL(napi_create_async_work(env, nullptr, resource_name, executeCapture, completeCapture, &data, &data.work));
NAPI_CALL(napi_queue_async_work(env, data.work));
return;
err:
if (data.deferred) napi_reject_deferred(env, data.deferred, handelError(env));
if (data.work) napi_delete_async_work(env, data.work);
if (data.buffer) napi_delete_reference(env, data.buffer);
delete (CaptureData *)_data;
}

void executeCapture(napi_env env, void *_data) {
auto &data = *(CaptureData *)_data;
const auto hwndDC = GetWindowDC(data.hwnd);
const auto saveDC = CreateCompatibleDC(hwndDC);
const auto saveBitmap = CreateCompatibleBitmap(hwndDC, width, height);
const auto saveBitmap = CreateCompatibleBitmap(hwndDC, data.width, data.height);
SelectObject(saveDC, saveBitmap);
PrintWindow(data.hwnd, saveDC, PW_RENDERFULLCONTENT);
data.bufferSize = width * height * 4;
data.buffer = new uint8_t[data.bufferSize];
BITMAPFILEHEADER bmfHeader;
BITMAPINFOHEADER bi;
bi.biSize = sizeof(BITMAPINFOHEADER);
bi.biWidth = width;
bi.biHeight = height;
bi.biWidth = data.width;
bi.biHeight = data.height;
bi.biPlanes = 1;
bi.biBitCount = 32;
bi.biCompression = BI_RGB;
Expand All @@ -50,41 +83,37 @@ void executeCapture(napi_env env, void *_data) {
bi.biYPelsPerMeter = 0;
bi.biClrUsed = 0;
bi.biClrImportant = 0;
GetDIBits(hwndDC, saveBitmap, 0, height, data.buffer, (BITMAPINFO *)&bi, DIB_RGB_COLORS);
GetDIBits(hwndDC, saveBitmap, 0, data.height, data.data, (BITMAPINFO *)&bi, DIB_RGB_COLORS);
DeleteObject(saveBitmap);
DeleteDC(saveDC);
ReleaseDC(data.hwnd, hwndDC);
DeleteObject(saveBitmap);
}

void completeCapture(napi_env env, napi_status status, void *_data) {
const auto &data = *(CaptureData *)_data;
if (data.error) {
NAPI_CALL(napi_reject_deferred(env, data.deferred, createError(env, nullptr, "Capture error")));
return;
}
auto &data = *(CaptureData *)_data;
napi_value result, width, height, buffer;
NAPI_CALL(napi_create_uint32(env, data.width, &width));
NAPI_CALL(napi_create_uint32(env, data.height, &height));
NAPI_CALL(napi_create_external_buffer(
env, data.bufferSize, data.buffer,
[](napi_env env, void *finalize_data, void *finalize_hint) { delete[](uint8_t *) finalize_hint; },
(void *)data.buffer, &buffer));
NAPI_CALL(napi_get_reference_value(env, data.buffer, &buffer));
napi_property_descriptor desc[3] = {
{"width", nullptr, nullptr, nullptr, nullptr, width, napi_enumerable, nullptr},
{"height", nullptr, nullptr, nullptr, nullptr, height, napi_enumerable, nullptr},
{"buffer", nullptr, nullptr, nullptr, nullptr, buffer, napi_enumerable, nullptr}};
NAPI_CALL(napi_create_object(env, &result));
NAPI_CALL(napi_define_properties(env, result, 3, desc));
NAPI_CALL(napi_resolve_deferred(env, data.deferred, result));
data.deferred = nullptr;
NAPI_CALL(napi_delete_async_work(env, data.work));
data.work = nullptr;
NAPI_CALL(napi_delete_reference(env, data.buffer));
data.buffer = nullptr;

delete (CaptureData *)_data;
return;
err:
napi_delete_async_work(env, data.work);
napi_reject_deferred(env, data.deferred, handelError(env));
if (data.buffer)
delete[] data.buffer;
if (data.deferred) napi_reject_deferred(env, data.deferred, handelError(env));
if (data.work) napi_delete_async_work(env, data.work);
if (data.buffer) napi_delete_reference(env, data.buffer);
delete (CaptureData *)_data;
}

Expand All @@ -101,15 +130,18 @@ napi_value capture(napi_env env, napi_callback_info info) {

napi_value resource_name;
NAPI_CALL(napi_create_string_utf8(env, "capture", NAPI_AUTO_LENGTH, &resource_name));
NAPI_CALL(napi_create_async_work(env, nullptr, resource_name, executeCapture, completeCapture, data, &data->work));
NAPI_CALL(napi_create_async_work(env, nullptr, resource_name, queryWindowSize, mallocBuffer, data, &data->work));
NAPI_CALL(napi_queue_async_work(env, data->work));

return data->promise;
err:
if (data->deferred)
napi_reject_deferred(env, data->deferred, createError(env, nullptr, ""));
if (data->work)
napi_delete_async_work(env, data->work);
if (data->deferred) {
const napi_extended_error_info *error_info = nullptr;
napi_get_last_error_info(env, &error_info);
const char* msg = error_info->error_message ? error_info->error_message : "";
napi_reject_deferred(env, data->deferred, createError(env, nullptr, msg));
}
if (data->work) napi_delete_async_work(env, data->work);
delete data;
return throwError(env);
}
2 changes: 2 additions & 0 deletions src/main/extractor/OcrExtractor/ScreenCapturer.ts
View file
@@ -1,5 +1,6 @@
import sharp from 'sharp';
import { findWindow, capture, CaptureResult, HWND } from '@addons/ScreenCapturer';
import logger from '@logger/extractor/OcrExtractor/ScreenCapturer';

export class ScreenCapturer {
private static emptyImage = () => sharp(Buffer.alloc(1, 0), { raw: { width: 1, height: 1, channels: 1 } });
Expand All @@ -24,6 +25,7 @@ export class ScreenCapturer {
try {
result = await capture(this.hwnd);
} catch (e) {
logger('fail to capture screen: %O', e);
if (canRetry) return this.capture(false);
else return ScreenCapturer.emptyImage();
}
Expand Down
2 changes: 1 addition & 1 deletion src/main/providers/ocr/tesseract.ts
View file
Expand Up @@ -35,7 +35,7 @@ export default defineOcrProvider({
if (import.meta.env.DEV) {
worker.on('message', args => {
if (args.type === 'log') {
logger(args.value);
logger('%O', args.value);
}
});
}
Expand Down

0 comments on commit b8503f9

Please sign in to comment.