/*
* DWN - Desktop Window Manager
* retoor <retoor@molodetz.nl>
* Screenshot capture implementation using X11 and libpng
*/
#include "screenshot.h"
#include "dwn.h"
#include "client.h"
#include "workspace.h"
#include "util.h"
#include <stdlib.h>
#include <string.h>
#include <stdatomic.h>
#include <png.h>
static bool initialized = false;
static pthread_mutex_t screenshot_mutex = PTHREAD_MUTEX_INITIALIZER;
static ScreenshotRequest *screenshot_queue = NULL;
static atomic_int screenshot_shutting_down = 0;
static const char base64_chars[] =
"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
bool screenshot_init(void)
{
atomic_store(&screenshot_shutting_down, 0);
initialized = true;
LOG_INFO("Screenshot module initialized");
return true;
}
void screenshot_cleanup(void)
{
atomic_store(&screenshot_shutting_down, 1);
pthread_mutex_lock(&screenshot_mutex);
while (screenshot_queue != NULL) {
ScreenshotRequest *req = screenshot_queue;
screenshot_queue = req->next;
pthread_mutex_unlock(&screenshot_mutex);
pthread_join(req->thread, NULL);
screenshot_result_free(&req->result);
dwn_free(req);
pthread_mutex_lock(&screenshot_mutex);
}
pthread_mutex_unlock(&screenshot_mutex);
initialized = false;
}
static size_t base64_encoded_size(size_t input_size)
{
return 4 * ((input_size + 2) / 3);
}
char *screenshot_to_base64(const ScreenshotResult *result, size_t *out_len)
{
if (result == NULL || result->data == NULL || result->size == 0) {
return NULL;
}
size_t encoded_size = base64_encoded_size(result->size);
char *encoded = dwn_malloc(encoded_size + 1);
if (encoded == NULL) {
return NULL;
}
const unsigned char *input = result->data;
size_t input_len = result->size;
char *output = encoded;
for (size_t i = 0; i < input_len; i += 3) {
unsigned int n = ((unsigned int)input[i]) << 16;
if (i + 1 < input_len) {
n |= ((unsigned int)input[i + 1]) << 8;
}
if (i + 2 < input_len) {
n |= (unsigned int)input[i + 2];
}
*output++ = base64_chars[(n >> 18) & 0x3F];
*output++ = base64_chars[(n >> 12) & 0x3F];
*output++ = (i + 1 < input_len) ? base64_chars[(n >> 6) & 0x3F] : '=';
*output++ = (i + 2 < input_len) ? base64_chars[n & 0x3F] : '=';
}
*output = '\0';
if (out_len != NULL) {
*out_len = output - encoded;
}
return encoded;
}
void screenshot_result_free(ScreenshotResult *result)
{
if (result == NULL) {
return;
}
if (result->data != NULL) {
dwn_free(result->data);
result->data = NULL;
}
result->size = 0;
result->width = 0;
result->height = 0;
}
const char *screenshot_status_string(ScreenshotStatus status)
{
switch (status) {
case SCREENSHOT_OK:
return "OK";
case SCREENSHOT_ERROR_INVALID_ARG:
return "Invalid argument";
case SCREENSHOT_ERROR_X11:
return "X11 error";
case SCREENSHOT_ERROR_PNG:
return "PNG encoding error";
case SCREENSHOT_ERROR_NO_MEMORY:
return "Out of memory";
case SCREENSHOT_ERROR_WINDOW_NOT_FOUND:
return "Window not found";
default:
return "Unknown error";
}
}
typedef struct {
unsigned char *data;
size_t size;
size_t capacity;
} PngBuffer;
static void png_write_callback(png_structp png, png_bytep data, png_size_t length)
{
PngBuffer *buf = (PngBuffer *)png_get_io_ptr(png);
if (buf == NULL) {
return;
}
size_t new_size = buf->size + length;
if (new_size > buf->capacity) {
size_t new_capacity = buf->capacity * 2;
if (new_capacity < new_size) {
new_capacity = new_size + 1024;
}
unsigned char *new_data = dwn_realloc(buf->data, new_capacity);
if (new_data == NULL) {
png_error(png, "Memory allocation failed");
return;
}
buf->data = new_data;
buf->capacity = new_capacity;
}
memcpy(buf->data + buf->size, data, length);
buf->size = new_size;
}
static void png_flush_callback(png_structp png)
{
(void)png;
}
static ScreenshotStatus ximage_to_png(XImage *img, ScreenshotResult *result)
{
if (img == NULL || result == NULL) {
return SCREENSHOT_ERROR_INVALID_ARG;
}
png_structp png = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
if (png == NULL) {
return SCREENSHOT_ERROR_PNG;
}
png_infop info = png_create_info_struct(png);
if (info == NULL) {
png_destroy_write_struct(&png, NULL);
return SCREENSHOT_ERROR_PNG;
}
if (setjmp(png_jmpbuf(png))) {
png_destroy_write_struct(&png, &info);
return SCREENSHOT_ERROR_PNG;
}
PngBuffer buf = {0};
buf.capacity = (size_t)(img->width * img->height * 3 / 2 + 1024);
buf.data = dwn_malloc(buf.capacity);
if (buf.data == NULL) {
png_destroy_write_struct(&png, &info);
return SCREENSHOT_ERROR_NO_MEMORY;
}
png_set_write_fn(png, &buf, png_write_callback, png_flush_callback);
png_set_IHDR(png, info, img->width, img->height, 8,
PNG_COLOR_TYPE_RGB, PNG_INTERLACE_NONE,
PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);
png_write_info(png, info);
unsigned char *row = dwn_malloc(img->width * 3);
if (row == NULL) {
dwn_free(buf.data);
png_destroy_write_struct(&png, &info);
return SCREENSHOT_ERROR_NO_MEMORY;
}
unsigned long red_mask = img->red_mask;
unsigned long green_mask = img->green_mask;
unsigned long blue_mask = img->blue_mask;
int red_shift = 0;
int green_shift = 0;
int blue_shift = 0;
while ((red_mask & 1) == 0 && red_shift < 32) {
red_shift++;
red_mask >>= 1;
}
while ((green_mask & 1) == 0 && green_shift < 32) {
green_shift++;
green_mask >>= 1;
}
while ((blue_mask & 1) == 0 && blue_shift < 32) {
blue_shift++;
blue_mask >>= 1;
}
for (int y = 0; y < img->height; y++) {
for (int x = 0; x < img->width; x++) {
unsigned long pixel = XGetPixel(img, x, y);
row[x * 3 + 0] = (pixel >> red_shift) & 0xFF;
row[x * 3 + 1] = (pixel >> green_shift) & 0xFF;
row[x * 3 + 2] = (pixel >> blue_shift) & 0xFF;
}
png_write_row(png, row);
}
png_write_end(png, info);
dwn_free(row);
png_destroy_write_struct(&png, &info);
result->data = buf.data;
result->size = buf.size;
result->width = img->width;
result->height = img->height;
return SCREENSHOT_OK;
}
ScreenshotStatus screenshot_capture_fullscreen(ScreenshotResult *result)
{
if (result == NULL) {
return SCREENSHOT_ERROR_INVALID_ARG;
}
if (dwn == NULL || dwn->display == NULL) {
return SCREENSHOT_ERROR_X11;
}
memset(result, 0, sizeof(ScreenshotResult));
XImage *img = XGetImage(dwn->display, dwn->root,
0, 0, dwn->screen_width, dwn->screen_height,
AllPlanes, ZPixmap);
if (img == NULL) {
LOG_ERROR("XGetImage failed for fullscreen capture");
return SCREENSHOT_ERROR_X11;
}
ScreenshotStatus status = ximage_to_png(img, result);
XDestroyImage(img);
return status;
}
ScreenshotStatus screenshot_capture_window(Window window, ScreenshotResult *result)
{
if (result == NULL) {
return SCREENSHOT_ERROR_INVALID_ARG;
}
if (dwn == NULL || dwn->display == NULL) {
return SCREENSHOT_ERROR_X11;
}
if (window == None) {
return SCREENSHOT_ERROR_WINDOW_NOT_FOUND;
}
memset(result, 0, sizeof(ScreenshotResult));
XWindowAttributes wa;
if (!XGetWindowAttributes(dwn->display, window, &wa)) {
return SCREENSHOT_ERROR_WINDOW_NOT_FOUND;
}
int x, y;
Window child;
XTranslateCoordinates(dwn->display, window, dwn->root, 0, 0, &x, &y, &child);
int width = wa.width;
int height = wa.height;
if (x < 0) {
width += x;
x = 0;
}
if (y < 0) {
height += y;
y = 0;
}
if (x + width > dwn->screen_width) {
width = dwn->screen_width - x;
}
if (y + height > dwn->screen_height) {
height = dwn->screen_height - y;
}
if (width <= 0 || height <= 0) {
return SCREENSHOT_ERROR_X11;
}
XImage *img = XGetImage(dwn->display, dwn->root, x, y, width, height,
AllPlanes, ZPixmap);
if (img == NULL) {
LOG_ERROR("XGetImage failed for window capture");
return SCREENSHOT_ERROR_X11;
}
ScreenshotStatus status = ximage_to_png(img, result);
XDestroyImage(img);
return status;
}
ScreenshotStatus screenshot_capture_active(ScreenshotResult *result)
{
if (result == NULL) {
return SCREENSHOT_ERROR_INVALID_ARG;
}
if (dwn == NULL) {
return SCREENSHOT_ERROR_X11;
}
Workspace *ws = workspace_get_current();
if (ws == NULL || ws->focused == NULL) {
return SCREENSHOT_ERROR_WINDOW_NOT_FOUND;
}
Client *c = ws->focused;
Window target = (c->frame != None) ? c->frame : c->window;
return screenshot_capture_window(target, result);
}
ScreenshotStatus screenshot_capture_area(int x, int y, int width, int height, ScreenshotResult *result)
{
if (result == NULL) {
return SCREENSHOT_ERROR_INVALID_ARG;
}
if (width <= 0 || height <= 0) {
return SCREENSHOT_ERROR_INVALID_ARG;
}
if (dwn == NULL || dwn->display == NULL) {
return SCREENSHOT_ERROR_X11;
}
memset(result, 0, sizeof(ScreenshotResult));
if (x < 0) {
width += x;
x = 0;
}
if (y < 0) {
height += y;
y = 0;
}
if (x + width > dwn->screen_width) {
width = dwn->screen_width - x;
}
if (y + height > dwn->screen_height) {
height = dwn->screen_height - y;
}
if (width <= 0 || height <= 0) {
return SCREENSHOT_ERROR_INVALID_ARG;
}
XImage *img = XGetImage(dwn->display, dwn->root, x, y, width, height,
AllPlanes, ZPixmap);
if (img == NULL) {
LOG_ERROR("XGetImage failed for area capture");
return SCREENSHOT_ERROR_X11;
}
ScreenshotStatus status = ximage_to_png(img, result);
XDestroyImage(img);
return status;
}
static void *screenshot_worker(void *arg)
{
ScreenshotRequest *req = (ScreenshotRequest *)arg;
if (atomic_load(&screenshot_shutting_down)) {
req->state = SCREENSHOT_STATE_ERROR;
req->status = SCREENSHOT_ERROR_X11;
return NULL;
}
switch (req->mode) {
case SCREENSHOT_MODE_FULLSCREEN:
req->status = screenshot_capture_fullscreen(&req->result);
break;
case SCREENSHOT_MODE_WINDOW:
req->status = screenshot_capture_window(req->window, &req->result);
break;
case SCREENSHOT_MODE_ACTIVE:
req->status = screenshot_capture_active(&req->result);
break;
case SCREENSHOT_MODE_AREA:
req->status = screenshot_capture_area(req->x, req->y,
req->width, req->height,
&req->result);
break;
}
req->state = (req->status == SCREENSHOT_OK)
? SCREENSHOT_STATE_COMPLETED
: SCREENSHOT_STATE_ERROR;
return NULL;
}
ScreenshotRequest *screenshot_capture_async(ScreenshotMode mode,
void (*callback)(ScreenshotRequest *))
{
if (!initialized || atomic_load(&screenshot_shutting_down)) {
return NULL;
}
ScreenshotRequest *req = dwn_calloc(1, sizeof(ScreenshotRequest));
if (req == NULL) {
return NULL;
}
req->mode = mode;
req->state = SCREENSHOT_STATE_PENDING;
req->callback = callback;
pthread_mutex_lock(&screenshot_mutex);
req->next = screenshot_queue;
screenshot_queue = req;
pthread_mutex_unlock(&screenshot_mutex);
if (pthread_create(&req->thread, NULL, screenshot_worker, req) != 0) {
pthread_mutex_lock(&screenshot_mutex);
ScreenshotRequest **pp = &screenshot_queue;
while (*pp != NULL) {
if (*pp == req) {
*pp = req->next;
break;
}
pp = &(*pp)->next;
}
pthread_mutex_unlock(&screenshot_mutex);
dwn_free(req);
return NULL;
}
LOG_DEBUG("Screenshot async request started (mode=%d)", mode);
return req;
}
ScreenshotRequest *screenshot_capture_window_async(Window window,
void (*callback)(ScreenshotRequest *))
{
if (!initialized || atomic_load(&screenshot_shutting_down)) {
return NULL;
}
ScreenshotRequest *req = dwn_calloc(1, sizeof(ScreenshotRequest));
if (req == NULL) {
return NULL;
}
req->mode = SCREENSHOT_MODE_WINDOW;
req->window = window;
req->state = SCREENSHOT_STATE_PENDING;
req->callback = callback;
pthread_mutex_lock(&screenshot_mutex);
req->next = screenshot_queue;
screenshot_queue = req;
pthread_mutex_unlock(&screenshot_mutex);
if (pthread_create(&req->thread, NULL, screenshot_worker, req) != 0) {
pthread_mutex_lock(&screenshot_mutex);
ScreenshotRequest **pp = &screenshot_queue;
while (*pp != NULL) {
if (*pp == req) {
*pp = req->next;
break;
}
pp = &(*pp)->next;
}
pthread_mutex_unlock(&screenshot_mutex);
dwn_free(req);
return NULL;
}
LOG_DEBUG("Screenshot async window request started");
return req;
}
ScreenshotRequest *screenshot_capture_area_async(int x, int y, int w, int h,
void (*callback)(ScreenshotRequest *))
{
if (!initialized || atomic_load(&screenshot_shutting_down)) {
return NULL;
}
ScreenshotRequest *req = dwn_calloc(1, sizeof(ScreenshotRequest));
if (req == NULL) {
return NULL;
}
req->mode = SCREENSHOT_MODE_AREA;
req->x = x;
req->y = y;
req->width = w;
req->height = h;
req->state = SCREENSHOT_STATE_PENDING;
req->callback = callback;
pthread_mutex_lock(&screenshot_mutex);
req->next = screenshot_queue;
screenshot_queue = req;
pthread_mutex_unlock(&screenshot_mutex);
if (pthread_create(&req->thread, NULL, screenshot_worker, req) != 0) {
pthread_mutex_lock(&screenshot_mutex);
ScreenshotRequest **pp = &screenshot_queue;
while (*pp != NULL) {
if (*pp == req) {
*pp = req->next;
break;
}
pp = &(*pp)->next;
}
pthread_mutex_unlock(&screenshot_mutex);
dwn_free(req);
return NULL;
}
LOG_DEBUG("Screenshot async area request started");
return req;
}
void screenshot_process_pending(void)
{
if (atomic_load(&screenshot_shutting_down)) {
return;
}
pthread_mutex_lock(&screenshot_mutex);
ScreenshotRequest **pp = &screenshot_queue;
while (*pp != NULL) {
ScreenshotRequest *req = *pp;
if (req->state == SCREENSHOT_STATE_COMPLETED ||
req->state == SCREENSHOT_STATE_ERROR) {
pthread_join(req->thread, NULL);
*pp = req->next;
pthread_mutex_unlock(&screenshot_mutex);
if (req->callback != NULL) {
req->callback(req);
}
pthread_mutex_lock(&screenshot_mutex);
} else {
pp = &(*pp)->next;
}
}
pthread_mutex_unlock(&screenshot_mutex);
}
void screenshot_cancel(ScreenshotRequest *req)
{
if (req == NULL) {
return;
}
pthread_mutex_lock(&screenshot_mutex);
ScreenshotRequest **pp = &screenshot_queue;
while (*pp != NULL) {
if (*pp == req) {
*pp = req->next;
break;
}
pp = &(*pp)->next;
}
pthread_mutex_unlock(&screenshot_mutex);
pthread_cancel(req->thread);
pthread_join(req->thread, NULL);
screenshot_result_free(&req->result);
dwn_free(req);
}
Reference in New Issue View Git Blame Copy Permalink