From e262303c2a9f537714071551b3c90e49b5dd0292 Mon Sep 17 00:00:00 2001
From: Arjun <agadiyar@nvidia.com>
Date: Tue, 16 Jul 2024 16:47:14 +0000
Subject: [PATCH] new branches with changes

---
 ubuntu22.04/empty         |  0
 ubuntu22.04/nvidia-driver | 41 ++++++++++++++++++++-------------------
 2 files changed, 21 insertions(+), 20 deletions(-)
 delete mode 100644 ubuntu22.04/empty

diff --git a/ubuntu22.04/empty b/ubuntu22.04/empty
deleted file mode 100644
index e69de29b..00000000
diff --git a/ubuntu22.04/nvidia-driver b/ubuntu22.04/nvidia-driver
index aedeeea2..9ec95a03 100755
--- a/ubuntu22.04/nvidia-driver
+++ b/ubuntu22.04/nvidia-driver
@@ -63,39 +63,40 @@ _resolve_kernel_version() {
 }
 
 # Install the kernel modules header/builtin/order files and generate the kernel version string.
-_install_prerequisites() (
+_install_prerequisites() { 
     local tmp_dir=$(mktemp -d)
 
     trap "rm -rf ${tmp_dir}" EXIT
     cd ${tmp_dir}
+    
+    if [ "${PERSIST_DRIVER}" = false ]; then
+        rm -rf /lib/modules/${KERNEL_VERSION}
+    fi
+    
+    if [ ! -d "/usr/src/linux-headers-$(uname -r)/" ]; then
+        echo "Installing Linux kernel headers..."
+        apt-get -qq install --no-install-recommends linux-headers-${KERNEL_VERSION} > /dev/null
+    fi
+    
+    echo ${KERNEL_VERSION} > version
 
-    rm -rf /lib/modules/${KERNEL_VERSION}
+    echo "Generating Linux kernel version string..."
+    
     mkdir -p /lib/modules/${KERNEL_VERSION}/proc
 
-    echo "Installing Linux kernel headers..."
-    apt-get -qq install --no-install-recommends linux-headers-${KERNEL_VERSION} > /dev/null
-
-    echo "Installing Linux kernel module files..."
-    apt-get -qq download linux-image-${KERNEL_VERSION} && dpkg -x linux-image*.deb .
-    { apt-get -qq download linux-modules-${KERNEL_VERSION} && dpkg -x linux-modules*.deb . || true; } 2> /dev/null
-    mv lib/modules/${KERNEL_VERSION}/modules.* /lib/modules/${KERNEL_VERSION}
-    mv lib/modules/${KERNEL_VERSION}/kernel /lib/modules/${KERNEL_VERSION}
-    depmod ${KERNEL_VERSION}
-
-    echo "Generating Linux kernel version string..."
+    mv version /lib/modules/${KERNEL_VERSION}/proc
 
-    ls -1 boot/vmlinuz-* | sed  's/\/boot\/vmlinuz-//g' - > version
-    if [ -z "$(<version)" ]; then
-        echo "Could not locate Linux kernel version string" >&2
-        return 1
+    if [ "${PERSIST_DRIVER}" = false ]; then
+        mv lib/modules/${KERNEL_VERSION}/modules.* /lib/modules/${KERNEL_VERSION}
+	mv lib/modules/${KERNEL_VERSION}/kernel /lib/modules/${KERNEL_VERSION}
+	depmod ${KERNEL_VERSION}
     fi
-    mv version /lib/modules/${KERNEL_VERSION}/proc
-)
+}
 
 # Cleanup the prerequisites installed above.
 _remove_prerequisites() {
     if [ "${PACKAGE_TAG:-}" != "builtin" ]; then
-        apt-get -qq purge linux-headers-${KERNEL_VERSION} > /dev/null
+        apt-get -qq purge linux-headers-${KERNEL_VERSION} > /dev/null || true
         # TODO remove module files not matching an existing driver package.
     fi
 }