From 9ed2141fc852a2e143bd671605771de6ebe89b42 Mon Sep 17 00:00:00 2001
From: loki <loki@fakeemail.com>
Date: Sat, 7 Aug 2021 14:39:18 +0200
Subject: [PATCH] Fix X11 screengrabbing with vaapi

---
 sunshine/platform/linux/graphics.cpp |  6 ------
 sunshine/platform/linux/kmsgrab.cpp  | 17 +++++++++++++----
 sunshine/platform/linux/vaapi.cpp    |  6 ++++++
 3 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/sunshine/platform/linux/graphics.cpp b/sunshine/platform/linux/graphics.cpp
index 5093255c..bfc10b97 100644
--- a/sunshine/platform/linux/graphics.cpp
+++ b/sunshine/platform/linux/graphics.cpp
@@ -635,12 +635,6 @@ int egl_t::_set_frame(AVFrame *frame) {
   this->hwframe.reset(frame);
   this->frame = frame;
 
-  if(av_hwframe_get_buffer(frame->hw_frames_ctx, frame, 0)) {
-    BOOST_LOG(error) << "Couldn't get hwframe for VAAPI"sv;
-
-    return -1;
-  }
-
   // Ensure aspect ratio is maintained
   auto scalar       = std::fminf(frame->width / (float)in_width, frame->height / (float)in_height);
   auto out_width_f  = in_width * scalar;
diff --git a/sunshine/platform/linux/kmsgrab.cpp b/sunshine/platform/linux/kmsgrab.cpp
index b04f83f4..e4cc7127 100644
--- a/sunshine/platform/linux/kmsgrab.cpp
+++ b/sunshine/platform/linux/kmsgrab.cpp
@@ -11,6 +11,7 @@
 #include "sunshine/utility.h"
 
 #include "graphics.h"
+#include "vaapi.h"
 
 using namespace std::literals;
 
@@ -154,7 +155,7 @@ void print(plane_t::pointer plane, fb_t::pointer fb) {
 
 class display_t : public platf::display_t {
 public:
-  display_t() : platf::display_t() {}
+  display_t(mem_type_e mem_type) : platf::display_t(), mem_type { mem_type } {}
 
   int init(const std::string &display_name, int framerate) {
     if(!gbm::create_device) {
@@ -287,10 +288,16 @@ public:
     return capture_e::ok;
   }
 
+  std::shared_ptr<hwdevice_t> make_hwdevice(pix_fmt_e pix_fmt) override {
+    if(mem_type == mem_type_e::vaapi) {
+      return va::make_hwdevice(width, height);
+    }
+
+    return std::make_shared<hwdevice_t>();
+  }
+
   capture_e snapshot(img_t *img_out_base, std::chrono::milliseconds timeout, bool cursor) {
-
     gl::ctx.BindTexture(GL_TEXTURE_2D, rgb->tex[0]);
-
     gl::ctx.GetTexImage(GL_TEXTURE_2D, 0, GL_BGRA, GL_UNSIGNED_BYTE, img_out_base->data);
 
     return capture_e::ok;
@@ -311,6 +318,8 @@ public:
     return 0;
   }
 
+  mem_type_e mem_type;
+
   std::chrono::nanoseconds delay;
 
   card_t card;
@@ -325,7 +334,7 @@ public:
 } // namespace kms
 
 std::shared_ptr<display_t> kms_display(mem_type_e hwdevice_type, const std::string &display_name, int framerate) {
-  auto disp = std::make_shared<kms::display_t>();
+  auto disp = std::make_shared<kms::display_t>(hwdevice_type);
 
   if(disp->init(display_name, framerate)) {
     return nullptr;
diff --git a/sunshine/platform/linux/vaapi.cpp b/sunshine/platform/linux/vaapi.cpp
index 6b1a123a..22031e44 100644
--- a/sunshine/platform/linux/vaapi.cpp
+++ b/sunshine/platform/linux/vaapi.cpp
@@ -184,6 +184,12 @@ public:
   int set_frame(AVFrame *frame) override {
     // No deallocation necessary
 
+    if(av_hwframe_get_buffer(frame->hw_frames_ctx, frame, 0)) {
+      BOOST_LOG(error) << "Couldn't get hwframe for VAAPI"sv;
+
+      return -1;
+    }
+
     va::DRMPRIMESurfaceDescriptor prime;
     va::VASurfaceID surface = (std::uintptr_t)frame->data[3];