From e796f36996c2a05e6a916b03910ef8fb9e272eb8 Mon Sep 17 00:00:00 2001
From: Lucas Alvares Gomes <lucasagomes@gmail.com>
Date: Thu, 18 Jun 2026 20:54:50 +0100
Subject: [PATCH] Add GPU device mounts for NVIDIA and DRI

Detect /dev/dri and /dev/nvidia* on the host and bind them into the
bubblewrap sandbox so GPU workloads work inside bwai.

Signed-off-by: Lucas Alvares Gomes <lucasagomes@gmail.com>
---
 README.md               |  3 ++-
 cmd/bwai/config.go      |  1 +
 cmd/bwai/main.go        |  1 +
 cmd/bwai/mounts.go      | 18 ++++++++++++++++++
 cmd/bwai/mounts_test.go | 24 ++++++++++++++++++++++++
 5 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index e04b031..11c7769 100644
--- a/README.md
+++ b/README.md
@@ -112,7 +112,8 @@ Example `~/.bwai.json`:
     ".bashrc.d",
     ".password-store",
     ".bash_history*",
-    ".config/Bitwarden"
+    ".config/Bitwarden",
+    ".cache/nvidia"
   ],
   "env_allow": [
     "TERM",
diff --git a/cmd/bwai/config.go b/cmd/bwai/config.go
index 2b4c53e..38ca9f6 100644
--- a/cmd/bwai/config.go
+++ b/cmd/bwai/config.go
@@ -92,6 +92,7 @@ func defaultConfig() Config {
 			".password-store",
 			".bash_history*",
 			".config/Bitwarden",
+			".cache/nvidia",
 		},
 	}
 }
diff --git a/cmd/bwai/main.go b/cmd/bwai/main.go
index 48fa370..9e5e3cd 100644
--- a/cmd/bwai/main.go
+++ b/cmd/bwai/main.go
@@ -85,6 +85,7 @@ func main() {
 		// Device nodes
 		"--dev", "/dev",
 	)
+	args = append(args, gpuMounts()...)
 	args = append(args, shmMount()...)
 	args = append(args,
 		// Virtual filesystems
diff --git a/cmd/bwai/mounts.go b/cmd/bwai/mounts.go
index 6d8e186..e959a7a 100644
--- a/cmd/bwai/mounts.go
+++ b/cmd/bwai/mounts.go
@@ -3,6 +3,7 @@ package main
 import (
 	"os"
 	"path/filepath"
+	"sort"
 	"strings"
 )
 
@@ -82,3 +83,20 @@ func dnsMounts() []string {
 	}
 	return nil
 }
+
+// Bind GPU device nodes into the sandbox if they exist on this host.
+// This includes the DRI subsystem and any NVIDIA character devices.
+func gpuMounts() []string {
+	var args []string
+	if info, err := os.Stat("/dev/dri"); err == nil && info.IsDir() {
+		args = append(args, devBind("/dev/dri")...)
+	}
+	matches, err := filepath.Glob("/dev/nvidia*")
+	if err == nil {
+		sort.Strings(matches)
+		for _, p := range matches {
+			args = append(args, devBind(p)...)
+		}
+	}
+	return args
+}
diff --git a/cmd/bwai/mounts_test.go b/cmd/bwai/mounts_test.go
index e56e03b..8c387f5 100644
--- a/cmd/bwai/mounts_test.go
+++ b/cmd/bwai/mounts_test.go
@@ -3,6 +3,7 @@ package main
 import (
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 )
 
@@ -63,6 +64,29 @@ func TestSubPathMounts(t *testing.T) {
 	}
 }
 
+func TestGPUMounts(t *testing.T) {
+	// gpuMounts() reads from the real /dev, so we can only assert stable
+	// structural properties: arguments come in --dev-bind triplets and each
+	// source path belongs to the expected /dev/dri or /dev/nvidia* families.
+	args := gpuMounts()
+
+	for i := 0; i < len(args); i += 3 {
+		if args[i] != "--dev-bind" {
+			t.Errorf("expected --dev-bind at index %d, got %q", i, args[i])
+		}
+		if i+2 >= len(args) {
+			t.Fatalf("incomplete --dev-bind triplet at index %d: %v", i, args[i:])
+		}
+		src := args[i+1]
+		if _, err := os.Stat(src); err != nil {
+			t.Errorf("gpuMounts returned non-existent path %q: %v", src, err)
+		}
+		if !strings.HasPrefix(src, "/dev/dri") && !strings.HasPrefix(src, "/dev/nvidia") {
+			t.Errorf("gpuMounts returned unexpected path %q", src)
+		}
+	}
+}
+
 func TestHomeMounts(t *testing.T) {
 	home := t.TempDir()