libdrmtap-sys 0.4.2

/*
 * libdrmtap — DRM/KMS screen capture library for Linux
 * https://github.com/fxd0h/libdrmtap
 *
 * Copyright (c) 2026 Mariano Abad <weimaraner@gmail.com>
 * SPDX-License-Identifier: MIT
 */

/**
 * @file privilege_helper.c
 * @brief Auto-spawn privileged helper and SCM_RIGHTS fd passing
 *
 * When the library detects that handles[0] == 0 (missing CAP_SYS_ADMIN),
 * it automatically spawns drmtap-helper via socketpair + fork/exec.
 * The helper opens the DRM device with CAP_SYS_ADMIN and passes back
 * DMA-BUF file descriptors via SCM_RIGHTS.
 *
 * Protocol:
 *   - Library creates socketpair(AF_UNIX, SOCK_STREAM)
 *   - fork/exec the helper with the child socket on fd 3
 *   - Send CMD_GRAB (0x01) to request a DMA-BUF fd
 *   - Receive status byte + optional fd via SCM_RIGHTS recvmsg
 */

#define _GNU_SOURCE
#define _POSIX_C_SOURCE 200809L

#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <errno.h>
#include <fcntl.h>
#include <unistd.h>
#include <signal.h>
#include <sys/socket.h>
#include <sys/wait.h>

#include "drmtap_internal.h"

/* Must match values in drmtap-helper.c */
#define HELPER_SOCKET_FD 3
#define CMD_GRAB 0x01
#define CMD_GET_CURSOR 0x02
#define CMD_QUIT 0xFF
#define RESP_OK    0x00

/* ========================================================================= */
/* Helper search paths                                                       */
/* ========================================================================= */

static const char *helper_search_paths[] = {
    NULL,  /* slot 0: ctx->helper_path (if set) */
    "/usr/lib/rustdesk/drmtap-helper",
    "/usr/libexec/drmtap-helper",
    "/usr/local/libexec/drmtap-helper",
    "/usr/local/bin/drmtap-helper",
    "/usr/bin/drmtap-helper",
    "/usr/lib/drmtap/drmtap-helper",
    NULL
};

// Find the helper binary
static const char *find_helper(drmtap_ctx *ctx) {
    /* Check configured path first */
    if (ctx->helper_path[0]) {
        if (access(ctx->helper_path, X_OK) == 0) {
            return ctx->helper_path;
        }
        drmtap_debug_log(ctx, "configured helper not found: %s",
                         ctx->helper_path);
    }

    /* Search standard paths */
    for (int i = 1; helper_search_paths[i]; i++) {
        if (access(helper_search_paths[i], X_OK) == 0) {
            drmtap_debug_log(ctx, "found helper: %s", helper_search_paths[i]);
            return helper_search_paths[i];
        }
    }

    return NULL;
}

/* ========================================================================= */
/* Helper lifecycle                                                          */
/* ========================================================================= */

// Spawn the helper binary via socketpair + fork/exec
// Sets ctx->helper_fd and ctx->helper_pid on success
int drmtap_helper_spawn(drmtap_ctx *ctx) {
    if (ctx->helper_fd >= 0) {
        /* Already running */
        return 0;
    }

    const char *helper_path = find_helper(ctx);
    if (!helper_path) {
        drmtap_set_error(ctx,
            "drmtap-helper not found. Install it with:\n"
            "  sudo cp drmtap-helper /usr/libexec/drmtap-helper\n"
            "  sudo setcap cap_sys_admin+ep /usr/libexec/drmtap-helper");
        return -EACCES;
    }

    /* Create socket pair for IPC */
    int socks[2];
    if (socketpair(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0, socks) < 0) {
        drmtap_set_error(ctx, "socketpair failed: %s", strerror(errno));
        return -errno;
    }

    pid_t pid = fork();
    if (pid < 0) {
        drmtap_set_error(ctx, "fork failed: %s", strerror(errno));
        close(socks[0]);
        close(socks[1]);
        return -errno;
    }

    if (pid == 0) {
        /* ---- Child process ---- */
        close(socks[0]);  /* Close parent's end */

        /* Move child socket to fd HELPER_SOCKET_FD */
        if (socks[1] != HELPER_SOCKET_FD) {
            if (dup2(socks[1], HELPER_SOCKET_FD) < 0) {
                _exit(127);
            }
            close(socks[1]);
        }
        /* Clear CLOEXEC on the helper socket so it survives exec */
        fcntl(HELPER_SOCKET_FD, F_SETFD, 0);

        /* Pass device path as argv[1] */
        execl(helper_path, "drmtap-helper", ctx->device_path, NULL);

        /* exec failed */
        _exit(127);
    }

    /* ---- Parent process ---- */
    close(socks[1]);  /* Close child's end */

    ctx->helper_fd = socks[0];
    ctx->helper_pid = pid;

    drmtap_debug_log(ctx, "spawned helper pid=%d from %s", pid, helper_path);
    return 0;
}

// Kill the helper process
void drmtap_helper_stop(drmtap_ctx *ctx) {
    if (ctx->helper_fd >= 0) {
        /* Send quit command (best effort) */
        uint8_t cmd = CMD_QUIT;
        ssize_t n = send(ctx->helper_fd, &cmd, 1, MSG_NOSIGNAL);
        (void)n;

        close(ctx->helper_fd);
        ctx->helper_fd = -1;
    }

    if (ctx->helper_pid > 0) {
        /* Give helper 100ms to exit, then SIGKILL */
        int status;
        usleep(100000);
        if (waitpid(ctx->helper_pid, &status, WNOHANG) == 0) {
            kill(ctx->helper_pid, SIGKILL);
            waitpid(ctx->helper_pid, &status, 0);
        }
        ctx->helper_pid = -1;
    }
}

/* ========================================================================= */
/* SCM_RIGHTS fd receiving                                                   */
/* ========================================================================= */

// Receive exactly len bytes from socket, handling partial reads
static int recv_all(int sock, void *buf, size_t len) {
    uint8_t *p = (uint8_t *)buf;
    size_t received = 0;
    while (received < len) {
        ssize_t n = recv(sock, p + received, len - received, 0);
        if (n <= 0) {
            return -1;
        }
        received += (size_t)n;
    }
    return 0;
}

/* Receive metadata + optional DMA-BUF fd via SCM_RIGHTS (V3 protocol).
 * When helper sends FLAG_HAS_DMABUF, the metadata arrives via sendmsg
 * with the DMA-BUF fd attached as ancillary data. */
static int recv_fd_and_meta(int sock, void *meta_buf, size_t meta_len, int *out_fd) {
    struct msghdr msg = {0};
    struct iovec iov;
    union {
        struct cmsghdr align;
        char buf[CMSG_SPACE(sizeof(int))];
    } cmsg_buf;

    *out_fd = -1;

    iov.iov_base = meta_buf;
    iov.iov_len = meta_len;
    msg.msg_iov = &iov;
    msg.msg_iovlen = 1;
    msg.msg_control = cmsg_buf.buf;
    msg.msg_controllen = sizeof(cmsg_buf.buf);

    ssize_t n = recvmsg(sock, &msg, 0);
    if (n <= 0) {
        return -1;
    }

    /* Extract fd from ancillary data if present */
    struct cmsghdr *cmsg = CMSG_FIRSTHDR(&msg);
    if (cmsg && cmsg->cmsg_level == SOL_SOCKET &&
        cmsg->cmsg_type == SCM_RIGHTS &&
        cmsg->cmsg_len == CMSG_LEN(sizeof(int))) {
        memcpy(out_fd, CMSG_DATA(cmsg), sizeof(int));
    }

    /* We may have received less than meta_len in the first recvmsg.
     * Read the rest with plain recv. */
    if ((size_t)n < meta_len) {
        if (recv_all(sock, (uint8_t *)meta_buf + n, meta_len - n) < 0) {
            if (*out_fd >= 0) close(*out_fd);
            *out_fd = -1;
            return -1;
        }
    }

    return 0;
}

/* ========================================================================= */
/* Public helper API (called from drm_grab.c)                                */
/* ========================================================================= */


// Request metadata + pixel data from the helper.
// The helper reads the framebuffer in its own process (fresh data)
// and sends metadata followed by raw pixel bytes through the socket.
// Returns 0 on success, negative errno on error.
int drmtap_helper_grab(drmtap_ctx *ctx, helper_grab_result_t *result,
                        void *pixel_buf, size_t buf_size) {
    if (ctx->helper_fd < 0) {
        int ret = drmtap_helper_spawn(ctx);
        if (ret < 0) {
            return ret;
        }
    }

    /* Send grab command */
    helper_cmd_grab_t hcmd = {
        .cmd = CMD_GRAB,
        .crtc_id = ctx->crtc_id
    };
    ssize_t n = send(ctx->helper_fd, &hcmd, sizeof(hcmd), MSG_NOSIGNAL);
    if (n != sizeof(hcmd)) {
        drmtap_debug_log(ctx, "helper send failed, trying respawn");
        drmtap_helper_stop(ctx);

        int ret = drmtap_helper_spawn(ctx);
        if (ret < 0) {
            return ret;
        }

        n = send(ctx->helper_fd, &hcmd, sizeof(hcmd), MSG_NOSIGNAL);
        if (n != sizeof(hcmd)) {
            drmtap_set_error(ctx, "helper communication failed after respawn");
            return -EIO;
        }
    }

    /* Receive metadata (possibly with DMA-BUF fd via SCM_RIGHTS) */
    memset(result, 0, sizeof(*result));
    result->dmabuf_fd = -1;
    int dmabuf_fd = -1;
    /* Recv into wire portion only — dmabuf_fd is local, not on the wire */
    if (recv_fd_and_meta(ctx->helper_fd, &result->wire, sizeof(result->wire), &dmabuf_fd) < 0) {
        drmtap_set_error(ctx, "helper metadata recv failed");
        return -EIO;
    }

    /* Check for DMA-BUF fd mode (V3 protocol) */
    if (result->wire.flags & HELPER_FLAG_HAS_DMABUF) {
        if (dmabuf_fd < 0) {
            drmtap_set_error(ctx, "helper signaled DMA-BUF but no fd received");
            return -EIO;
        }
        /* Store the fd in the result for the caller */
        result->dmabuf_fd = dmabuf_fd;
        drmtap_debug_log(ctx,
            "helper: %ux%u fb=%u DMA-BUF fd=%d mod=0x%lx (V3 SCM_RIGHTS)",
            result->wire.width, result->wire.height, result->wire.fb_id, dmabuf_fd,
            (unsigned long)result->wire.modifier);
        return 0;
    }

    /* Close any unexpected fd */
    if (dmabuf_fd >= 0) {
        close(dmabuf_fd);
    }

    /* V2 fallback: data_size == 0 means error */
    if (result->wire.data_size == 0) {
        drmtap_set_error(ctx, "helper returned error");
        return -EIO;
    }

    result->dmabuf_fd = -1;
    drmtap_debug_log(ctx, "helper: %ux%u fb=%u data_size=%u seq=%u (V2 pixels)",
                     result->wire.width, result->wire.height,
                     result->wire.fb_id, result->wire.data_size, result->wire.seq);

    /* Receive pixel data directly into caller's buffer */
    if (result->wire.data_size > buf_size) {
        drmtap_set_error(ctx, "helper data_size %u exceeds buffer %zu",
                         result->wire.data_size, buf_size);
        size_t remaining = result->wire.data_size;
        char drain[4096];
        while (remaining > 0) {
            size_t chunk = remaining < sizeof(drain) ? remaining : sizeof(drain);
            recv(ctx->helper_fd, drain, chunk, 0);
            remaining -= chunk;
        }
        return -ENOSPC;
    }

    if (recv_all(ctx->helper_fd, pixel_buf, result->wire.data_size) < 0) {
        drmtap_set_error(ctx, "helper pixel recv failed");
        return -EIO;
    }

    return 0;
}

int drmtap_helper_get_cursor(drmtap_ctx *ctx, drmtap_cursor_info *cursor) {
    memset(cursor, 0, sizeof(*cursor));
    cursor->pixels = NULL;

    if (ctx->helper_fd < 0) {
        int ret = drmtap_helper_spawn(ctx);
        if (ret < 0) {
            return ret;
        }
    }

    helper_cmd_grab_t hcmd = {
        .cmd = CMD_GET_CURSOR,
        .crtc_id = ctx->crtc_id,
    };
    ssize_t n = send(ctx->helper_fd, &hcmd, sizeof(hcmd), MSG_NOSIGNAL);
    if (n != sizeof(hcmd)) {
        drmtap_helper_stop(ctx);
        if (drmtap_helper_spawn(ctx) < 0) {
            return -EIO;
        }
        n = send(ctx->helper_fd, &hcmd, sizeof(hcmd), MSG_NOSIGNAL);
        if (n != sizeof(hcmd)) {
            return -EIO;
        }
    }

    helper_cursor_wire_t w;
    memset(&w, 0, sizeof(w));
    if (recv_all(ctx->helper_fd, &w, sizeof(w)) < 0) {
        drmtap_set_error(ctx, "helper cursor metadata recv failed");
        return -EIO;
    }

    cursor->x = w.x;
    cursor->y = w.y;
    cursor->hot_x = w.hot_x;
    cursor->hot_y = w.hot_y;
    cursor->width = w.width;
    cursor->height = w.height;
    cursor->visible = w.visible ? 1 : 0;

    if (w.visible && w.data_size > 0) {
        /* Drain oversized payloads defensively (cursors are tiny, ~64x64). */
        if (w.data_size > 256u * 256u * 4u) {
            size_t rem = w.data_size;
            char drn[4096];
            while (rem > 0) {
                size_t c = rem < sizeof(drn) ? rem : sizeof(drn);
                if (recv(ctx->helper_fd, drn, c, 0) <= 0) break;
                rem -= c;
            }
            cursor->visible = 0;
            return 0;
        }
        cursor->pixels = (uint32_t *)malloc(w.data_size);
        if (!cursor->pixels) {
            return -ENOMEM;
        }
        if (recv_all(ctx->helper_fd, cursor->pixels, w.data_size) < 0) {
            free(cursor->pixels);
            cursor->pixels = NULL;
            return -EIO;
        }
    }

    return 0;
}