fix(server): prevent a segfault on nvidia

Caused by use-after free on the rate control layer array.
colinmarc · May 7, 2024 · 8b331b5 · 8b331b5
1 parent b3f7353
commit 8b331b5
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 14 deletions.
diff --git a/mm-server/src/compositor/video/vulkan_encode/h265.rs b/mm-server/src/compositor/video/vulkan_encode/h265.rs
@@ -405,24 +405,23 @@ impl H265Encoder {
                     }),
             );
 
-            rc_layers = h265_rc_layers
-                .iter_mut()
-                .map(|h265_layer| {
+            for h265_layer in h265_rc_layers.iter_mut() {
+                rc_layers.push(
                     vk::VideoEncodeRateControlLayerInfoKHR::default()
                         .max_bitrate(settings.peak_bitrate)
                         .average_bitrate(settings.average_bitrate)
                         .frame_rate_numerator(self.inner.framerate)
                         .frame_rate_denominator(1)
-                        .push_next(h265_layer)
-                })
-                .collect::<Vec<_>>();
+                        .push_next(h265_layer),
+                );
+            }
         }
 
         let mut h265_rc_info = vk::VideoEncodeH265RateControlInfoEXT::default()
             .gop_frame_count(self.structure.gop_size)
             .idr_period(self.structure.gop_size)
             .consecutive_b_frame_count(0)
-            .sub_layer_count(self.structure.layers)
+            .sub_layer_count(1)
             .flags(vk::VideoEncodeH265RateControlFlagsEXT::REGULAR_GOP | pattern);
 
         let vbv_size = match self.rc_mode {
@@ -432,8 +431,11 @@ impl H265Encoder {
 
         let mut rc_info = vk::VideoEncodeRateControlInfoKHR::default()
             .rate_control_mode(self.rc_mode.as_vk_flags())
-            .virtual_buffer_size_in_ms(vbv_size)
-            .layers(&rc_layers);
+            .virtual_buffer_size_in_ms(vbv_size);
+
+        if !rc_layers.is_empty() {
+            rc_info = rc_info.layers(&rc_layers);
+        }
 
         // Doesn't have a push_next method, because we're supposed to call it on
         // the parent struct.

diff --git a/mm-server/src/compositor/video/vulkan_encode/rate_control.rs b/mm-server/src/compositor/video/vulkan_encode/rate_control.rs
@@ -77,11 +77,8 @@ pub fn select_rc_mode(
             min_qp: 17,
             max_qp: target_qp,
         })
-    } else if caps
-        .rate_control_modes
-        .contains(vk::VideoEncodeRateControlModeFlagsKHR::DISABLED)
-    {
-        // Fall back to CRF with a high bitrate.
+    } else if supports_crf    {
+        // Fall back to CRF with a high QP.
         RateControlMode::ConstantQp(target_qp)
     } else {
         warn!("no rate control modes available, using driver defaults!");