Administrace komentářů

#include <assert.h>
#include <errno.h>
#include <limits.h>
#include <stdio.h>
#include <stdlib.h>
#include <xcb/xcb.h>
#include <xcb/shm.h>

#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>

#include <sys/ipc.h>
#include <sys/shm.h>


enum {
    SCR_ERR_OK,
    SCR_ERR_CONNECTION_FAIL,
    SCR_ERR_CONNECTION_ERR,
    SCR_ERR_INVALID_SCREEN,
    SCR_ERR_SHM_NOT_SUPPORTED,
    SCR_ERR_INVALID_VISUAL,
    SCR_ERR_INVALID_PIXMAP_FORMAT,
    SCR_ERR_UNSUPPORTED_VISUAL,
    SCR_ERR_IMAGE_EMPTY,
    SCR_ERR_IMAGE_TOO_LARGE,
    
    SCR_ERR_SHMGET_FAILED,
    SCR_ERR_SHM_ATTACH_FAILED,
    SCR_ERR_GET_IMAGE_FAILED,
    SCR_ERR_GET_IMAGE_RACE_CONDITION
};

struct screenshoter
{
    xcb_connection_t *con;
    xcb_setup_t const *setup;
    xcb_screen_t *screen;
    int screen_num;

    xcb_depth_t const* depth;
    xcb_visualtype_t const* visual;
    xcb_format_t const* pixformat;

    xcb_query_extension_reply_t const* shm_ext;

    xcb_shm_seg_t shm_seg;

    struct {
        unsigned char *data;
        size_t stride, size;
    } image;
};

int receive_screenshot(struct screenshoter *scr)
{
    xcb_shm_get_image_cookie_t img_ck = xcb_shm_get_image(scr->con, scr->screen->root,
            0, 0, scr->screen->width_in_pixels, scr->screen->height_in_pixels,
            UINT32_MAX, XCB_IMAGE_FORMAT_XY_PIXMAP, scr->shm_seg, 0);
    xcb_generic_error_t *err = NULL;
    xcb_shm_get_image_reply_t *img_rep = xcb_shm_get_image_reply(scr->con, img_ck, &err);
    if (!img_rep) {
        if (err) free(err);
        return SCR_ERR_GET_IMAGE_FAILED;
    }

    if (img_rep->size != scr->image.size ||
            img_rep->visual != scr->screen->root_visual ||
            img_rep->depth != scr->depth->depth)
    {
        free(img_rep);
        return SCR_ERR_GET_IMAGE_RACE_CONDITION;
    }

    free(img_rep);
    return 0;
}

int create_shm_seg_and_image(struct screenshoter *scr)
{
    // check if sreen has non-zero dimensions
    if (scr->screen->width_in_pixels == 0 ||
            scr->screen->height_in_pixels == 0)
        return SCR_ERR_IMAGE_EMPTY;

    // XCB doesn't run on platforms with non-8bit char, so char is 8bits
    // We only support depth, bits_per_pixel and scanline_pad that is multiple
    // of whole bytes
    //
    // NOTE: shm extension seems to return packed image data, bits_per_pixel and scanline_pad
    //       are not used
    if (scr->depth->depth % 8 != 0)
        return SCR_ERR_UNSUPPORTED_VISUAL;

    size_t bytes_per_pixel = scr->depth->depth / 8;

    // Calculate stride, first multiply width by bytes_per_pixel
    size_t stride = scr->screen->width_in_pixels;
    if (SIZE_MAX / stride < bytes_per_pixel)
        return SCR_ERR_IMAGE_TOO_LARGE;
    stride *= bytes_per_pixel;

    // now multiply by height, to get image size in bytes
    scr->image.stride = stride;
    if (SIZE_MAX / stride < scr->screen->height_in_pixels)
        return SCR_ERR_IMAGE_TOO_LARGE;
    scr->image.size = stride * scr->screen->height_in_pixels;

    int shmid = shmget(IPC_PRIVATE, scr->image.size, IPC_CREAT | 0777);
    if (shmid == -1)
        return SCR_ERR_SHMGET_FAILED;

    scr->image.data = shmat(shmid, NULL, 0);
    if (scr->image.data == (void*)-1) {
        shmctl(shmid, IPC_RMID, NULL);
        return SCR_ERR_SHM_ATTACH_FAILED;
    }

    // Generate shm segment id
    scr->shm_seg = xcb_generate_id(scr->con);
    // attach our file to it, use checked call so we can block for reply
    xcb_void_cookie_t attach_ck = xcb_shm_attach_checked(scr->con,
            scr->shm_seg, shmid, 0);
    // wait for reply
    xcb_generic_error_t *attach_err = xcb_request_check(scr->con, attach_ck);
    // we don't need shmid anymore after the request was processed by X
    shmctl(shmid, IPC_RMID, NULL);
    // check attach error
    if (attach_err) {
        free(attach_err);
        shmdt(scr->image.data);
        return SCR_ERR_SHM_ATTACH_FAILED;
    }

    return 0;
}

void free_shm_seg_and_image(struct screenshoter *scr)
{
    xcb_shm_detach(scr->con, scr->shm_seg);
    shmdt(scr->image.data);
}

int find_default_screen(struct screenshoter *scr)
{
    if (scr->screen_num < 0)
        return SCR_ERR_INVALID_SCREEN;
    
    xcb_screen_iterator_t iter = xcb_setup_roots_iterator(scr->setup);
    for (int i = 0; i < scr->screen_num; ++i) {
        if (iter.rem == 0)
            return SCR_ERR_INVALID_SCREEN;
        xcb_screen_next(&iter);
    }
    scr->screen = iter.data;
    return 0;
}

int detect_root_visual(struct screenshoter *scr)
{
    xcb_depth_iterator_t depth_iter = xcb_screen_allowed_depths_iterator(scr->screen);
    for (;;) {
        xcb_depth_t const *depth = depth_iter.data;
        int visuals_len = xcb_depth_visuals_length(depth);
        if (visuals_len <= 0)
            break;
        xcb_visualtype_t const *visuals = xcb_depth_visuals(depth);

        for (int i = 0; i < visuals_len; ++i) {
            if (visuals[i].visual_id == scr->screen->root_visual) {
                scr->depth = depth;
                scr->visual = &visuals[i];
                return 0;
            }
        }

        if (depth_iter.rem == 0)
            break;
        xcb_depth_next(&depth_iter);
    }

    return SCR_ERR_INVALID_VISUAL;
}

int detect_pixmap_format(struct screenshoter *scr)
{
    int formats_len = xcb_setup_pixmap_formats_length(scr->setup);
    if (formats_len > 0) {
        xcb_format_t const *formats = xcb_setup_pixmap_formats(scr->setup);
        for (int i = 0; i < formats_len; ++i) {
            if (formats[i].depth == scr->depth->depth) {
                scr->pixformat = &formats[i];
                return 0;
            }
        }
    }
    return SCR_ERR_INVALID_PIXMAP_FORMAT;
}

int get_screenshot(void)
{
    struct screenshoter scr;
    int err = 0;

    // connect
    scr.con = xcb_connect(NULL, &scr.screen_num);
    if (!scr.con)
        return SCR_ERR_CONNECTION_FAIL;

    // check errors
    if (xcb_connection_has_error(scr.con)) {
        err = SCR_ERR_CONNECTION_ERR;
        goto retDisconnect;
    }

    // get setup and prefetch shm extension
    scr.setup = xcb_get_setup(scr.con);
    xcb_prefetch_extension_data(scr.con, &xcb_shm_id);
    xcb_flush(scr.con);

    // find screen
    err = find_default_screen(&scr);
    if (err != 0)
        goto retDisconnect;

    // detect visual
    err = detect_root_visual(&scr);
    if (err != 0)
        goto retDisconnect;

    // and pixmap format for root depth
    err = detect_pixmap_format(&scr);
    if (err != 0)
        goto retDisconnect;

    // get shm extension
    scr.shm_ext = xcb_get_extension_data(scr.con, &xcb_shm_id);
    if (!scr.shm_ext->present) {
        err = SCR_ERR_SHM_NOT_SUPPORTED;
        goto retDisconnect;
    }

    // create shm segment
    err = create_shm_seg_and_image(&scr);
    if (err != 0)
        goto retDisconnect;

    err = receive_screenshot(&scr);
    if (err == 0) {
        // TODO: process image data here

        // Image data is in scr.image.data, its bytesize is in scr.image.size,
        // its stride bytesize is in scr.image.stride, its width is in
        // scr.screen->width_in_pixels, height in scr.screen->height_in_pixels,
        // depth (in bits) is in scr.depth->depth. Visual format is in
        // scr.visual (red_mask, green_mask, blue_mask). Image data byteorder is
        // in scr.setup->image_byte_order
    }

    free_shm_seg_and_image(&scr);
    xcb_flush(scr.con);
retDisconnect:
    xcb_disconnect(scr.con);
    return err;
}


int main()
{
    int err = get_screenshot();
    fprintf(stderr, "err = %d\n", err);

    return 0;
}