From d43af98345abd527bd730e814ffff8d44849846e Mon Sep 17 00:00:00 2001
From: Malte Poll <1780588+malt3@users.noreply.github.com>
Date: Mon, 15 Apr 2024 14:49:57 +0200
Subject: [PATCH] image: install nvidia driver

Co-authored-by: derpsteb <ob@edgeless.systems>
---
 WORKSPACE.bazel                   |  4 +-
 bazel/toolchains/linux_kernel.bzl | 61 +++++++++++++++++++++++++++++++
 image/BUILD.bazel                 | 23 ++++++++++++
 image/base/BUILD.bazel            | 10 ++++-
 4 files changed, 96 insertions(+), 2 deletions(-)

diff --git a/WORKSPACE.bazel b/WORKSPACE.bazel
index 21c6e1982..34bf99f97 100644
--- a/WORKSPACE.bazel
+++ b/WORKSPACE.bazel
@@ -331,10 +331,12 @@ load("//bazel/toolchains:k8s.bzl", "k8s_deps")
 k8s_deps()
 
 # kernel rpms
-load("//bazel/toolchains:linux_kernel.bzl", "kernel_rpms")
+load("//bazel/toolchains:linux_kernel.bzl", "kernel_rpms", "nvidia_kos")
 
 kernel_rpms()
 
+nvidia_kos()
+
 # mkosi rpms
 load("//bazel/rpm:package_manifest.bzl", "rpm_repository")
 
diff --git a/bazel/toolchains/linux_kernel.bzl b/bazel/toolchains/linux_kernel.bzl
index 7b7617563..4d65ce3ab 100644
--- a/bazel/toolchains/linux_kernel.bzl
+++ b/bazel/toolchains/linux_kernel.bzl
@@ -2,6 +2,67 @@
 
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_file")
 
+def nvidia_kos():
+    """nvidia ko"""
+
+    # nvidia kernel module
+    http_file(
+        name = "nvidia_ko",
+        urls = [
+            "https://cdn.confidential.cloud/constellation/kernel/nvidia/6.2.0/535.129.03/nvidia.ko",
+        ],
+        sha256 = "ee2555a032cf17f2756312ae2004673acc461c3d6fbc4b3021b2d4735034fb11",
+        downloaded_file_path = "nvidia.ko",
+    )
+    http_file(
+        name = "nvidia_drm_ko",
+        urls = [
+            "https://cdn.confidential.cloud/constellation/kernel/nvidia/6.2.0/535.129.03/nvidia-drm.ko",
+        ],
+        sha256 = "78e08dce97ba7306bbd5658183dbcb1221bdf9cb8fe2fc5528797b0fa0e9e31d",
+        downloaded_file_path = "nvidia-drm.ko",
+    )
+    http_file(
+        name = "nvidia_modeset_ko",
+        urls = [
+            "https://cdn.confidential.cloud/constellation/kernel/nvidia/6.2.0/535.129.03/nvidia-modeset.ko",
+        ],
+        sha256 = "18d669cc4c089f896457560b69cc6eb30344a434de5a35ab5846ac65b88dde5e",
+        downloaded_file_path = "nvidia-modeset.ko",
+    )
+    http_file(
+        name = "nvidia_peermem_ko",
+        urls = [
+            "https://cdn.confidential.cloud/constellation/kernel/nvidia/6.2.0/535.129.03/nvidia-peermem.ko",
+        ],
+        sha256 = "52ce0116713a35a4db6b36a6d029a3d1a4ae1d30c032c8c71281545e878e5923",
+        downloaded_file_path = "nvidia-peermem.ko",
+    )
+    http_file(
+        name = "nvidia_uvm_ko",
+        urls = [
+            "https://cdn.confidential.cloud/constellation/kernel/nvidia/6.2.0/535.129.03/nvidia-uvm.ko",
+        ],
+        sha256 = "c0c4e044f2bbaa939d9e1bfb6e11ce1b441aeaac683defc2aa33bfb7b2b6c217",
+        downloaded_file_path = "nvidia-uvm.ko",
+    )
+    http_file(
+        name = "gsp_ga10x",
+        urls = [
+            "https://cdn.confidential.cloud/constellation/kernel/nvidia/6.2.0/535.129.03/gsp_ga10x.bin",
+        ],
+        sha256 = "1f6d303a192388b3ccd97f468fa4ed64b5921a8b76ae1d4660f2caed5568dc17",
+        downloaded_file_path = "gsp_ga10x.bin",
+    )
+    http_file(
+        name = "gsp_tu10x",
+        urls = [
+            "https://cdn.confidential.cloud/constellation/kernel/nvidia/6.2.0/535.129.03/gsp_tu10x.bin",
+        ],
+        sha256 = "6ada90fdfbfa134ab02c588be09441a9c64670a79d0cbb200106d0f3d3f672fe",
+        downloaded_file_path = "gsp_tu10x.bin",
+    )
+
 def kernel_rpms():
     """kernel rpms"""
 
diff --git a/image/BUILD.bazel b/image/BUILD.bazel
index 681c6f385..f9be715fd 100644
--- a/image/BUILD.bazel
+++ b/image/BUILD.bazel
@@ -27,3 +27,26 @@ copy_file(
     allow_symlink = True,
     visibility = ["//visibility:public"],
 )
+
+pkg_tar(
+    name = "nvidia_kernel_modules",
+    srcs = [
+        "@nvidia_drm_ko//file",
+        "@nvidia_ko//file",
+        "@nvidia_modeset_ko//file",
+        "@nvidia_peermem_ko//file",
+        "@nvidia_uvm_ko//file",
+    ],
+    package_dir = "lib/modules/6.2.0-100.constellation.fc38.x86_64/kernel/nvidia",
+    visibility = ["//visibility:public"],
+)
+
+pkg_tar(
+    name = "nvidia_gsp_firmware",
+    srcs = [
+        "@gsp_ga10x//file",
+        "@gsp_tu10x//file",
+    ],
+    package_dir = "usr/lib/firmware/nvidia/535.129.03",
+    visibility = ["//visibility:public"],
+)
diff --git a/image/base/BUILD.bazel b/image/base/BUILD.bazel
index 9028b8376..69460bc66 100644
--- a/image/base/BUILD.bazel
+++ b/image/base/BUILD.bazel
@@ -26,6 +26,14 @@ copy_to_directory(
     replace_prefixes = {"file": ""},
 )
 
+# TODO(malt3): only install nvidia kernel modules for specific nvidia image
+extra_trees = {
+    "lts": [
+        "//image:nvidia_kernel_modules",
+    ],
+    "mainline": [],
+}
+
 [
     mkosi_image(
         name = "base_" + kernel_variant,
@@ -47,7 +55,7 @@ copy_to_directory(
         extra_trees = [
             "//image:sysroot_tar",
             "//image:cryptsetup_closure",
-        ],
+        ] + extra_trees[kernel_variant],
         local_mirror = ["@mkosi_rpms//:repo"],
         mkosi_conf = "mkosi.conf",
         output = kernel_variant,