Do not dup dmabuf file descriptors in gpu code

4 files changed, 62 insertions, 42 deletions

diff --git a/capture.c b/capture.c
index 72a641d..f5283db 100644
--- a/capture.c
+++ b/capture.c
@@ -17,6 +17,7 @@
 
 #include "capture.h"
 
+#include <assert.h>
 #include <errno.h>
 #include <stddef.h>
 #include <stdio.h>
@@ -140,8 +141,16 @@ const struct GpuFrame* CaptureContextGetFrame(
     capture_context->gpu_frame = NULL;
   }
 
+  struct GpuFramePlane planes[] = {
+      {.dmabuf_fd = -1},
+      {.dmabuf_fd = -1},
+      {.dmabuf_fd = -1},
+      {.dmabuf_fd = -1},
+  };
+  static_assert(LENGTH(planes) == LENGTH(drm_mode_fb_cmd2.handles),
+                "Suspicious drm_mode_fb_cmd2 structure");
+
   size_t nplanes = 0;
-  struct GpuFramePlane planes[LENGTH(drm_mode_fb_cmd2.handles)];
   for (; nplanes < LENGTH(planes); nplanes++) {
     if (!drm_mode_fb_cmd2.handles[nplanes]) break;
     int status = drmPrimeHandleToFD(capture_context->drm_fd,
@@ -159,10 +168,16 @@ const struct GpuFrame* CaptureContextGetFrame(
   capture_context->gpu_frame = GpuContextCreateFrame(
       capture_context->gpu_context, drm_mode_fb_cmd2.width,
       drm_mode_fb_cmd2.height, drm_mode_fb_cmd2.pixel_format, nplanes, planes);
+  if (!capture_context->gpu_frame) {
+    LOG("Failed to create gpu frame");
+    goto release_planes;
+  }
+  return capture_context->gpu_frame;
 
 release_planes:
-  for (; nplanes; nplanes--) close(planes[nplanes - 1].dmabuf_fd);
-  return capture_context->gpu_frame;
+  CloseUniqueFds((int[]){planes[0].dmabuf_fd, planes[1].dmabuf_fd,
+                         planes[2].dmabuf_fd, planes[3].dmabuf_fd});
+  return NULL;
 }
 
 void CaptureContextDestroy(struct CaptureContext* capture_context) {
diff --git a/encode.c b/encode.c
index 5dc9cdd..cbf9294 100644
--- a/encode.c
+++ b/encode.c
@@ -17,6 +17,7 @@
 
 #include "encode.h"
 
+#include <assert.h>
 #include <drm_fourcc.h>
 #include <errno.h>
 #include <libavcodec/avcodec.h>
@@ -154,10 +155,17 @@ rollback_encode_context:
 
 static struct GpuFrame* PrimeToGpuFrame(
     struct GpuContext* gpu_context, const VADRMPRIMESurfaceDescriptor* prime) {
-  struct GpuFramePlane planes[4];
+  struct GpuFramePlane planes[] = {
+      {.dmabuf_fd = -1},
+      {.dmabuf_fd = -1},
+      {.dmabuf_fd = -1},
+      {.dmabuf_fd = -1},
+  };
+  static_assert(LENGTH(planes) == LENGTH(prime->layers[0].object_index),
+                "Suspicious VADRMPRIMESurfaceDescriptor structure");
+
   for (size_t i = 0; i < prime->layers[0].num_planes; i++) {
     uint32_t object_index = prime->layers[0].object_index[i];
-    if (prime->objects[object_index].fd == -1) break;
     planes[i] = (struct GpuFramePlane){
         .dmabuf_fd = prime->objects[object_index].fd,
         .pitch = prime->layers[0].pitch[i],
@@ -165,11 +173,20 @@ static struct GpuFrame* PrimeToGpuFrame(
         .modifier = prime->objects[object_index].drm_format_modifier,
     };
   }
+
   struct GpuFrame* gpu_frame =
       GpuContextCreateFrame(gpu_context, prime->width, prime->height,
                             prime->fourcc, prime->layers[0].num_planes, planes);
-  for (size_t i = prime->num_objects; i; i--) close(prime->objects[i - 1].fd);
+  if (!gpu_frame) {
+    LOG("Failed to create gpu frame");
+    goto release_planes;
+  }
   return gpu_frame;
+
+release_planes:
+  CloseUniqueFds((int[]){planes[0].dmabuf_fd, planes[1].dmabuf_fd,
+                         planes[2].dmabuf_fd, planes[3].dmabuf_fd});
+  return NULL;
 }
 
 const struct GpuFrame* EncodeContextGetFrame(
diff --git a/gpu.c b/gpu.c
index 9f09323..a00eb53 100644
--- a/gpu.c
+++ b/gpu.c
@@ -651,44 +651,25 @@ struct GpuFrame* GpuContextCreateFrame(struct GpuContext* gpu_context,
       .images = {EGL_NO_IMAGE, EGL_NO_IMAGE},
   };
 
-  for (size_t i = 0; i < nplanes; i++) {
-    gpu_frame_impl->dmabuf_fds[i] = dup(planes[i].dmabuf_fd);
-    if (gpu_frame_impl->dmabuf_fds[i] == -1) {
-      LOG("Failed to dup dmabuf fd (%s)", strerror(errno));
-      goto rollback_dmabuf_fds;
-    }
-  }
-
-  struct GpuFramePlane dummy_planes[4];
-  for (size_t i = 0; i < nplanes; i++) {
-    dummy_planes[i] = (struct GpuFramePlane){
-        .dmabuf_fd = gpu_frame_impl->dmabuf_fds[i],
-        .offset = planes[i].offset,
-        .pitch = planes[i].pitch,
-        .modifier = planes[i].modifier,
-    };
-  }
-
   if (fourcc == DRM_FORMAT_NV12) {
-    gpu_frame_impl->images[0] = CreateEglImage(
-        gpu_context, width, height, DRM_FORMAT_R8, 1, &dummy_planes[0]);
+    gpu_frame_impl->images[0] = CreateEglImage(gpu_context, width, height,
+                                               DRM_FORMAT_R8, 1, &planes[0]);
     if (gpu_frame_impl->images[0] == EGL_NO_IMAGE) {
       LOG("Failed to create luma plane image");
-      goto rollback_dmabuf_fds;
+      goto rollback_gpu_frame;
     }
-    gpu_frame_impl->images[1] =
-        CreateEglImage(gpu_context, width / 2, height / 2, DRM_FORMAT_GR88, 1,
-                       &dummy_planes[1]);
+    gpu_frame_impl->images[1] = CreateEglImage(
+        gpu_context, width / 2, height / 2, DRM_FORMAT_GR88, 1, &planes[1]);
     if (gpu_frame_impl->images[1] == EGL_NO_IMAGE) {
       LOG("Failed to create chroma plane image");
       goto rollback_images;
     }
   } else {
-    gpu_frame_impl->images[0] = CreateEglImage(gpu_context, width, height,
-                                               fourcc, nplanes, dummy_planes);
+    gpu_frame_impl->images[0] =
+        CreateEglImage(gpu_context, width, height, fourcc, nplanes, planes);
     if (gpu_frame_impl->images[0] == EGL_NO_IMAGE) {
       LOG("Failed to create multiplanar image");
-      goto rollback_dmabuf_fds;
+      goto rollback_gpu_frame;
     }
   }
 
@@ -701,6 +682,9 @@ struct GpuFrame* GpuContextCreateFrame(struct GpuContext* gpu_context,
       goto rollback_textures;
     }
   }
+
+  for (size_t i = 0; i < nplanes; i++)
+    gpu_frame_impl->dmabuf_fds[i] = planes[i].dmabuf_fd;
   return (struct GpuFrame*)gpu_frame_impl;
 
 rollback_textures:
@@ -713,11 +697,7 @@ rollback_images:
     if (gpu_frame_impl->images[i - 1] != EGL_NO_IMAGE)
       eglDestroyImage(gpu_context->device, gpu_frame_impl->images[i - 1]);
   }
-rollback_dmabuf_fds:
-  for (size_t i = LENGTH(gpu_frame_impl->dmabuf_fds); i; i--) {
-    if (gpu_frame_impl->dmabuf_fds[i - 1] != -1)
-      close(gpu_frame_impl->dmabuf_fds[i - 1]);
-  }
+rollback_gpu_frame:
   free(gpu_frame_impl);
   return NULL;
 }
@@ -790,10 +770,7 @@ void GpuContextDestroyFrame(struct GpuContext* gpu_context,
     if (gpu_frame_impl->images[i - 1] != EGL_NO_IMAGE)
       eglDestroyImage(gpu_context->device, gpu_frame_impl->images[i - 1]);
   }
-  for (size_t i = LENGTH(gpu_frame_impl->dmabuf_fds); i; i--) {
-    if (gpu_frame_impl->dmabuf_fds[i - 1] != -1)
-      close(gpu_frame_impl->dmabuf_fds[i - 1]);
-  }
+  CloseUniqueFds(gpu_frame_impl->dmabuf_fds);
   free(gpu_frame_impl);
 }
 
@@ -812,3 +789,12 @@ void GpuContextDestroy(struct GpuContext* gpu_context) {
 #endif  // USE_EGL_MESA_PLATFORM_SURFACELESS
   free(gpu_context);
 }
+
+void CloseUniqueFds(int fds[4]) {
+  // TODO(mburakov): Meh, but still better than looping...
+  if (fds[3] != -1 && fds[3] != fds[2] && fds[3] != fds[1] && fds[3] != fds[0])
+    close(fds[3]);
+  if (fds[2] != -1 && fds[2] != fds[1] && fds[2] != fds[0]) close(fds[2]);
+  if (fds[1] != -1 && fds[1] != fds[0]) close(fds[1]);
+  if (fds[0] != -1) close(fds[0]);
+}
diff --git a/gpu.h b/gpu.h
index f8fa886..6d07366 100644
--- a/gpu.h
+++ b/gpu.h
@@ -49,4 +49,6 @@ void GpuContextDestroyFrame(struct GpuContext* gpu_context,
                             struct GpuFrame* gpu_frame);
 void GpuContextDestroy(struct GpuContext* gpu_context);
 
+void CloseUniqueFds(int fds[4]);
+
 #endif  // STREAMER_GPU_H_