Set uid/gid inside the GPU jail based on ownership of wayland socket

Fixes running as root connecting to a users wayland socket when using nvidia
drivers.

src/crosvm/sys/linux/gpu.rs

@@ -8,8 +8,11 @@ static_assertions::assert_cfg!(feature = "gpu");
 
 use std::collections::HashMap;
 use std::env;
+use std::os::unix::fs::MetadataExt;
+use std::path::Path;
 use std::path::PathBuf;
 
+use base::geteuid;
 use base::linux::move_proc_to_cgroup;
 use jail::*;
 use serde::Deserialize;
@@ -19,6 +22,38 @@ use serde_keyvalue::FromKeyValues;
 use super::*;
 use crate::crosvm::config::Config;
 
+/// Get the uid/gid of a wayland socket's owner.
+/// Falls back to the socket's parent directory owner if the socket doesn't exist yet.
+fn get_wayland_socket_owner(socket_path: &Path) -> Option<(u32, u32)> {
+    // Try the socket itself first, then its parent directory
+    let path_to_check = if socket_path.exists() {
+        socket_path.to_path_buf()
+    } else {
+        socket_path.parent()?.to_path_buf()
+    };
+
+    let metadata = std::fs::metadata(&path_to_check).ok()?;
+    Some((metadata.uid(), metadata.gid()))
+}
+
+/// Get the group owner of the first available DRM render node or card.
+fn get_drm_device_group() -> Option<u32> {
+    // Try render nodes first, then card devices
+    for i in 128..192 {
+        let path = PathBuf::from(format!("/dev/dri/renderD{}", i));
+        if let Ok(metadata) = std::fs::metadata(&path) {
+            return Some(metadata.gid());
+        }
+    }
+    for i in 0..64 {
+        let path = PathBuf::from(format!("/dev/dri/card{}", i));
+        if let Ok(metadata) = std::fs::metadata(&path) {
+            return Some(metadata.gid());
+        }
+    }
+    None
+}
+
 pub struct GpuCacheInfo<'a> {
     directory: Option<&'a str>,
     environment: Vec<(&'a str, &'a str)>,
@@ -165,6 +200,37 @@ pub fn create_gpu_device(
         // Allow changes made externally take effect immediately to allow shaders to be dynamically
         // added by external processes.
         config.remount_mode = Some(libc::MS_SLAVE);
+
+        // When running as root, run the GPU jail as the wayland socket owner. This is needed
+        // for the nVidia driver when connecting to a user socket.
+        if geteuid() == 0 {
+            // Try to get uid/gid from wayland socket owner
+            if let Some((uid, _wayland_gid)) = cfg
+                .wayland_socket_paths
+                .values()
+                .next()
+                .and_then(|p| get_wayland_socket_owner(p))
+            {
+                if uid != 0 {
+                    // Get the group owner of the DRM device (typically "video" group)
+                    // so the jail can access GPU devices
+                    let drm_gid = get_drm_device_group().unwrap_or(0);
+
+                    config.run_as = RunAsUser::Specified(uid, drm_gid);
+                    // Map uid/gid inside the jail to the same ids outside.
+                    // The uid must match the wayland socket owner for nvidia driver compatibility
+                    // and wayland socket access.
+                    // The gid must match the DRM device group for GPU access.
+                    // Format: "inside_id outside_id count"
+                    let uid_map = format!("{uid} {uid} 1");
+                    let gid_map = format!("{drm_gid} {drm_gid} 1");
+                    // Leak the strings since ugid_map expects &'a str with lifetime of config
+                    config.ugid_map =
+                        Some((Box::leak(uid_map.into_boxed_str()), Box::leak(gid_map.into_boxed_str())));
+                }
+            }
+        }
+
         let mut jail = create_gpu_minijail(
             &jail_config.pivot_root,
             &config,