anti-prestige-tool/portal_capture_frame.cpp

#include <QCoreApplication>
#include <QDBusArgument>
#include <QDBusConnection>
#include <QDBusError>
#include <QDBusInterface>
#include <QDBusMetaType>
#include <QDBusObjectPath>
#include <QDBusReply>
#include <QDBusUnixFileDescriptor>
#include <QDBusVariant>
#include <QFile>
#include <QImage>
#include <QJsonArray>
#include <QJsonDocument>
#include <QJsonObject>
#include <QRandomGenerator>
#include <QTextStream>
#include <QTimer>
#include <QVariantMap>

#include <gst/app/gstappsink.h>
#include <gst/gst.h>
#include <gst/video/video.h>

#include <unistd.h>

static const char *PORTAL_SERVICE = "org.freedesktop.portal.Desktop";
static const char *PORTAL_PATH = "/org/freedesktop/portal/desktop";
static const char *SCREENCAST_IFACE = "org.freedesktop.portal.ScreenCast";
static const uint SOURCE_TYPE_MONITOR = 1;
static const uint SOURCE_TYPE_WINDOW = 2;
static const uint CURSOR_MODE_HIDDEN = 1;

struct PipeWireStream
{
    uint nodeId = 0;
    QVariantMap properties;
};

Q_DECLARE_METATYPE(PipeWireStream)
Q_DECLARE_METATYPE(QList<PipeWireStream>)

const QDBusArgument &operator>>(const QDBusArgument &argument, PipeWireStream &stream)
{
    argument.beginStructure();
    argument >> stream.nodeId >> stream.properties;
    argument.endStructure();
    return argument;
}

QDBusArgument &operator<<(QDBusArgument &argument, const PipeWireStream &stream)
{
    argument.beginStructure();
    argument << stream.nodeId << stream.properties;
    argument.endStructure();
    return argument;
}

class RequestWatcher : public QObject
{
    Q_OBJECT

public:
    uint response = 2;
    QVariantMap results;
    bool done = false;
    QString error;

public slots:
    void Response(uint responseCode, const QVariantMap &responseResults)
    {
        response = responseCode;
        results = responseResults;
        done = true;
        QCoreApplication::quit();
    }
};

QDBusConnection sessionBus()
{
    QDBusConnection bus = QDBusConnection::sessionBus();
    if (bus.isConnected()) {
        return bus;
    }

    QString runtimeDir = qEnvironmentVariable("XDG_RUNTIME_DIR");
    if (runtimeDir.isEmpty()) {
        runtimeDir = QStringLiteral("/run/user/%1").arg(getuid());
    }

    return QDBusConnection::connectToBus(
        QStringLiteral("unix:path=%1/bus").arg(runtimeDir),
        QStringLiteral("reforger-portal-session"));
}

QString token()
{
    return QStringLiteral("reforger_%1_%2")
        .arg(QCoreApplication::applicationPid())
        .arg(QString::number(QRandomGenerator::global()->generate(), 16));
}

bool waitForRequest(QCoreApplication &app, QDBusConnection &bus, const QDBusObjectPath &handle, int timeoutMs, RequestWatcher &watcher, QTextStream &err)
{
    const bool connected = bus.connect(
        QString::fromLatin1(PORTAL_SERVICE),
        handle.path(),
        QStringLiteral("org.freedesktop.portal.Request"),
        QStringLiteral("Response"),
        &watcher,
        SLOT(Response(uint,QVariantMap)));
    if (!connected) {
        err << "failed to listen for portal request response on " << handle.path() << ": "
            << bus.lastError().message() << "\n";
        return false;
    }

    QTimer::singleShot(timeoutMs, &app, [&watcher]() {
        if (!watcher.done) {
            watcher.error = QStringLiteral("timed out waiting for portal response");
            watcher.done = true;
            QCoreApplication::quit();
        }
    });
    app.exec();
    bus.disconnect(
        QString::fromLatin1(PORTAL_SERVICE),
        handle.path(),
        QStringLiteral("org.freedesktop.portal.Request"),
        QStringLiteral("Response"),
        &watcher,
        SLOT(Response(uint,QVariantMap)));

    if (!watcher.error.isEmpty()) {
        err << watcher.error << "\n";
        return false;
    }
    if (watcher.response == 1) {
        err << "portal request was cancelled\n";
        return false;
    }
    if (watcher.response != 0) {
        err << "portal request failed with response code " << watcher.response << "\n";
        return false;
    }
    return true;
}

QVariantMap callRequest(
    QCoreApplication &app,
    QDBusConnection &bus,
    QDBusInterface &portal,
    const QString &method,
    const QList<QVariant> &args,
    int timeoutMs,
    QTextStream &err,
    bool *ok)
{
    *ok = false;
    QDBusMessage reply = portal.callWithArgumentList(QDBus::Block, method, args);
    if (reply.type() == QDBusMessage::ErrorMessage) {
        err << method << " failed: " << reply.errorName() << ": " << reply.errorMessage() << "\n";
        return {};
    }
    if (reply.arguments().isEmpty()) {
        err << method << " returned no request handle\n";
        return {};
    }

    QDBusObjectPath handle = qdbus_cast<QDBusObjectPath>(reply.arguments().at(0));
    if (handle.path().isEmpty()) {
        err << method << " returned an invalid request handle\n";
        return {};
    }

    RequestWatcher watcher;
    if (!waitForRequest(app, bus, handle, timeoutMs, watcher, err)) {
        return {};
    }
    *ok = true;
    return watcher.results;
}

uint portalUIntProperty(QDBusConnection &bus, const QString &name)
{
    QDBusInterface props(
        QString::fromLatin1(PORTAL_SERVICE),
        QString::fromLatin1(PORTAL_PATH),
        QStringLiteral("org.freedesktop.DBus.Properties"),
        bus);
    QDBusReply<QDBusVariant> reply = props.call(QStringLiteral("Get"), QString::fromLatin1(SCREENCAST_IFACE), name);
    if (!reply.isValid()) {
        return 0;
    }
    return reply.value().variant().toUInt();
}

QList<PipeWireStream> parseStreams(const QVariant &value)
{
    QList<PipeWireStream> streams = qdbus_cast<QList<PipeWireStream>>(value);
    if (!streams.isEmpty()) {
        return streams;
    }

    if (value.canConvert<QDBusArgument>()) {
        return qdbus_cast<QList<PipeWireStream>>(value.value<QDBusArgument>());
    }

    return {};
}

QString variantString(const QVariantMap &map, const QString &key)
{
    const QVariant value = map.value(key);
    if (value.canConvert<QDBusObjectPath>()) {
        return value.value<QDBusObjectPath>().path();
    }
    return value.toString();
}

void closePortalSession(QDBusConnection &bus, const QString &sessionHandle)
{
    if (sessionHandle.isEmpty()) {
        return;
    }

    QDBusInterface session(
        QString::fromLatin1(PORTAL_SERVICE),
        sessionHandle,
        QStringLiteral("org.freedesktop.portal.Session"),
        bus);
    if (session.isValid()) {
        session.call(QStringLiteral("Close"));
    }
}

int writePipeWireFrame(int fd, uint nodeId, const QString &outputPath, int timeoutMs, QTextStream &err)
{
    GstElement *pipeline = gst_pipeline_new("reforger-portal-capture");
    GstElement *source = gst_element_factory_make("pipewiresrc", "source");
    GstElement *convert = gst_element_factory_make("videoconvert", "convert");
    GstElement *filter = gst_element_factory_make("capsfilter", "filter");
    GstElement *sink = gst_element_factory_make("appsink", "sink");

    if (!pipeline || !source || !convert || !filter || !sink) {
        err << "failed to create GStreamer PipeWire capture pipeline\n";
        if (pipeline) {
            gst_object_unref(pipeline);
        }
        return 1;
    }

    const QByteArray node = QByteArray::number(nodeId);
    g_object_set(source, "fd", fd, "path", node.constData(), "num-buffers", 1, nullptr);

    GstCaps *caps = gst_caps_new_simple("video/x-raw", "format", G_TYPE_STRING, "BGRA", nullptr);
    g_object_set(filter, "caps", caps, nullptr);
    gst_caps_unref(caps);

    g_object_set(sink, "emit-signals", FALSE, "sync", FALSE, "max-buffers", 1, "drop", TRUE, nullptr);

    gst_bin_add_many(GST_BIN(pipeline), source, convert, filter, sink, nullptr);
    if (!gst_element_link_many(source, convert, filter, sink, nullptr)) {
        err << "failed to link GStreamer PipeWire capture pipeline\n";
        gst_object_unref(pipeline);
        return 1;
    }

    if (gst_element_set_state(pipeline, GST_STATE_PLAYING) == GST_STATE_CHANGE_FAILURE) {
        err << "failed to start GStreamer PipeWire capture pipeline\n";
        gst_object_unref(pipeline);
        return 1;
    }

    GstSample *sample = gst_app_sink_try_pull_sample(GST_APP_SINK(sink), timeoutMs * GST_MSECOND);
    if (!sample) {
        GstBus *bus = gst_element_get_bus(pipeline);
        GstMessage *message = gst_bus_pop_filtered(bus, static_cast<GstMessageType>(GST_MESSAGE_ERROR | GST_MESSAGE_EOS));
        if (message && GST_MESSAGE_TYPE(message) == GST_MESSAGE_ERROR) {
            GError *error = nullptr;
            gchar *debug = nullptr;
            gst_message_parse_error(message, &error, &debug);
            err << "GStreamer PipeWire capture failed: " << (error ? error->message : "unknown error") << "\n";
            if (debug) {
                err << debug << "\n";
            }
            if (error) {
                g_error_free(error);
            }
            g_free(debug);
        } else {
            err << "timed out waiting for a PipeWire video frame\n";
        }
        if (message) {
            gst_message_unref(message);
        }
        gst_object_unref(bus);
        gst_element_set_state(pipeline, GST_STATE_NULL);
        gst_object_unref(pipeline);
        return 1;
    }

    GstCaps *sampleCaps = gst_sample_get_caps(sample);
    GstVideoInfo info;
    if (!sampleCaps || !gst_video_info_from_caps(&info, sampleCaps)) {
        err << "failed to read video frame format from PipeWire sample\n";
        gst_sample_unref(sample);
        gst_element_set_state(pipeline, GST_STATE_NULL);
        gst_object_unref(pipeline);
        return 1;
    }

    GstBuffer *buffer = gst_sample_get_buffer(sample);
    GstMapInfo map;
    if (!buffer || !gst_buffer_map(buffer, &map, GST_MAP_READ)) {
        err << "failed to map PipeWire video frame\n";
        gst_sample_unref(sample);
        gst_element_set_state(pipeline, GST_STATE_NULL);
        gst_object_unref(pipeline);
        return 1;
    }

    const int width = GST_VIDEO_INFO_WIDTH(&info);
    const int height = GST_VIDEO_INFO_HEIGHT(&info);
    const int stride = GST_VIDEO_INFO_PLANE_STRIDE(&info, 0);
    QImage image(map.data, width, height, stride, QImage::Format_ARGB32);
    const bool saved = image.copy().save(outputPath, "PNG");

    gst_buffer_unmap(buffer, &map);
    gst_sample_unref(sample);
    gst_element_set_state(pipeline, GST_STATE_NULL);
    gst_object_unref(pipeline);

    if (!saved) {
        err << "failed to write PNG output: " << outputPath << "\n";
        return 1;
    }

    return 0;
}

int captureWindow(QCoreApplication &app, const QString &outputPath, const QString &restoreToken, bool persist, int timeoutMs, QTextStream &out, QTextStream &err)
{
    QDBusConnection bus = sessionBus();
    if (!bus.isConnected()) {
        err << "failed to connect to the session D-Bus: " << bus.lastError().message() << "\n";
        return 1;
    }

    const uint sourceTypes = portalUIntProperty(bus, QStringLiteral("AvailableSourceTypes"));
    if (!(sourceTypes & SOURCE_TYPE_WINDOW)) {
        err << "desktop portal does not advertise window capture support";
        if (sourceTypes) {
            err << " (AvailableSourceTypes=" << sourceTypes << ")";
        }
        err << "\n";
        return 1;
    }

    QDBusInterface portal(
        QString::fromLatin1(PORTAL_SERVICE),
        QString::fromLatin1(PORTAL_PATH),
        QString::fromLatin1(SCREENCAST_IFACE),
        bus);
    if (!portal.isValid()) {
        err << "failed to connect to desktop ScreenCast portal: " << portal.lastError().message() << "\n";
        return 1;
    }

    QVariantMap createOptions;
    createOptions.insert(QStringLiteral("handle_token"), token());
    createOptions.insert(QStringLiteral("session_handle_token"), token());
    bool ok = false;
    QVariantMap createResults = callRequest(app, bus, portal, QStringLiteral("CreateSession"), {createOptions}, timeoutMs, err, &ok);
    if (!ok) {
        return 1;
    }
    const QString sessionHandle = variantString(createResults, QStringLiteral("session_handle"));
    if (sessionHandle.isEmpty()) {
        err << "portal CreateSession did not return a session handle\n";
        return 1;
    }

    QVariantMap selectOptions;
    selectOptions.insert(QStringLiteral("handle_token"), token());
    selectOptions.insert(QStringLiteral("types"), SOURCE_TYPE_WINDOW);
    selectOptions.insert(QStringLiteral("multiple"), false);
    selectOptions.insert(QStringLiteral("cursor_mode"), CURSOR_MODE_HIDDEN);
    if (!restoreToken.isEmpty()) {
        selectOptions.insert(QStringLiteral("restore_token"), restoreToken);
    }
    if (persist) {
        selectOptions.insert(QStringLiteral("persist_mode"), 2u);
    }
    callRequest(app, bus, portal, QStringLiteral("SelectSources"), {QVariant::fromValue(QDBusObjectPath(sessionHandle)), selectOptions}, timeoutMs, err, &ok);
    if (!ok) {
        closePortalSession(bus, sessionHandle);
        return 1;
    }

    QVariantMap startOptions;
    startOptions.insert(QStringLiteral("handle_token"), token());
    QVariantMap startResults = callRequest(app, bus, portal, QStringLiteral("Start"), {QVariant::fromValue(QDBusObjectPath(sessionHandle)), QString(), startOptions}, timeoutMs, err, &ok);
    if (!ok) {
        closePortalSession(bus, sessionHandle);
        return 1;
    }
    const QList<PipeWireStream> streams = parseStreams(startResults.value(QStringLiteral("streams")));
    if (streams.isEmpty()) {
        err << "portal Start did not return any PipeWire streams\n";
        closePortalSession(bus, sessionHandle);
        return 1;
    }

    const PipeWireStream stream = streams.first();
    if (stream.nodeId == 0) {
        err << "portal returned an invalid PipeWire node id\n";
        closePortalSession(bus, sessionHandle);
        return 1;
    }

    QDBusReply<QDBusUnixFileDescriptor> fdReply = portal.call(
        QStringLiteral("OpenPipeWireRemote"),
        QVariant::fromValue(QDBusObjectPath(sessionHandle)),
        QVariantMap());
    if (!fdReply.isValid()) {
        err << "OpenPipeWireRemote failed: " << fdReply.error().name() << ": " << fdReply.error().message() << "\n";
        closePortalSession(bus, sessionHandle);
        return 1;
    }

    const QDBusUnixFileDescriptor remoteFd = fdReply.value();
    if (!remoteFd.isValid()) {
        err << "OpenPipeWireRemote returned an invalid file descriptor\n";
        closePortalSession(bus, sessionHandle);
        return 1;
    }

    const int status = writePipeWireFrame(remoteFd.fileDescriptor(), stream.nodeId, outputPath, timeoutMs, err);

    closePortalSession(bus, sessionHandle);

    if (status != 0) {
        return status;
    }

    QJsonObject info;
    info.insert(QStringLiteral("node_id"), static_cast<int>(stream.nodeId));
    info.insert(QStringLiteral("session_handle"), sessionHandle);
    const QString newRestoreToken = variantString(startResults, QStringLiteral("restore_token"));
    if (!newRestoreToken.isEmpty()) {
        info.insert(QStringLiteral("restore_token"), newRestoreToken);
    }
    QJsonArray streamArray;
    for (const PipeWireStream &pipeWireStream : streams) {
        QJsonObject streamObject;
        streamObject.insert(QStringLiteral("node_id"), static_cast<int>(pipeWireStream.nodeId));
        QJsonObject properties;
        for (auto it = pipeWireStream.properties.constBegin(); it != pipeWireStream.properties.constEnd(); ++it) {
            properties.insert(it.key(), QJsonValue::fromVariant(it.value()));
        }
        streamObject.insert(QStringLiteral("properties"), properties);
        streamArray.append(streamObject);
    }
    info.insert(QStringLiteral("streams"), streamArray);
    out << QJsonDocument(info).toJson(QJsonDocument::Compact) << "\n";

    return 0;
}

int main(int argc, char *argv[])
{
    QCoreApplication app(argc, argv);
    gst_init(&argc, &argv);
    qDBusRegisterMetaType<PipeWireStream>();
    qDBusRegisterMetaType<QList<PipeWireStream>>();

    QTextStream out(stdout);
    QTextStream err(stderr);
    const QStringList args = QCoreApplication::arguments();

    QString outputPath;
    QString restoreToken;
    bool persist = false;
    int timeoutMs = 60000;

    for (int i = 1; i < args.size(); ++i) {
        const QString arg = args.at(i);
        if (arg == QLatin1String("--window") && i + 1 < args.size()) {
            outputPath = args.at(++i);
        } else if (arg == QLatin1String("--restore-token") && i + 1 < args.size()) {
            restoreToken = args.at(++i);
        } else if (arg == QLatin1String("--persist")) {
            persist = true;
        } else if (arg == QLatin1String("--timeout") && i + 1 < args.size()) {
            bool ok = false;
            const double seconds = args.at(++i).toDouble(&ok);
            if (!ok || seconds <= 0) {
                err << "invalid --timeout value\n";
                return 64;
            }
            timeoutMs = qRound(seconds * 1000.0);
        } else {
            err << "unknown or incomplete argument: " << arg << "\n";
            return 64;
        }
    }

    if (outputPath.isEmpty()) {
        err << "usage: " << args.value(0)
            << " --window OUTPUT.png [--restore-token TOKEN] [--persist] [--timeout SECONDS]\n";
        return 64;
    }

    return captureWindow(app, outputPath, restoreToken, persist, timeoutMs, out, err);
}

#include "portal_capture_frame.moc"