From 4d06da0f3abd8100fa72a072aa1af2a8d69b2e34 Mon Sep 17 00:00:00 2001
From: Jai Luthra <me@jailuthra.in>
Date: Thu, 23 Nov 2023 23:36:36 +0530
Subject: [PATCH] linux: autopatch: Find bytecode and update patch.sh

Have a two-script workflow where one focuses on finding the bytecode in
the driver files, and the other focuses on updating patch.sh using the
newly found bytecode or a previously existing one.

Signed-off-by: Jai Luthra <me@jailuthra.in>
---
 .gitignore                       |  3 ++
 tools/autopatch/autopatch.sh     | 30 -----------
 tools/autopatch/find_bytecode.sh | 45 +++++++++++++++++
 tools/autopatch/update_patch.sh  | 87 ++++++++++++++++++++++++++++++++
 4 files changed, 135 insertions(+), 30 deletions(-)
 delete mode 100755 tools/autopatch/autopatch.sh
 create mode 100755 tools/autopatch/find_bytecode.sh
 create mode 100755 tools/autopatch/update_patch.sh

diff --git a/.gitignore b/.gitignore
index 9dc51cb..620c105 100644
--- a/.gitignore
+++ b/.gitignore
@@ -160,3 +160,6 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 .idea/
 
+# Temp folder
+tools/autopatch/temp
+win/tools/autopatch/temp
diff --git a/tools/autopatch/autopatch.sh b/tools/autopatch/autopatch.sh
deleted file mode 100755
index 111093b..0000000
--- a/tools/autopatch/autopatch.sh
+++ /dev/null
@@ -1,30 +0,0 @@
-#!/bin/bash
-
-set -euo pipefail
-
-(( $# == 2 )) || {
-    >&2 echo "Usage: $0 <patch-file.sh> <new-version>"
-    exit 2
-}
-
-patch_file=$1
-new_version=$2
-
-if [[ ! -e $patch_file ]]; then
-    >&2 echo "Patch file $patch_file not found"
-    exit 2
-fi
-
-# Find the latest patch line
-latest=$(grep -n "\s.*\[\".*\"\]='.*/g\?'" $patch_file | tail -1)
-
-# Find the line number to insert at
-line=$(cut -d : -f 1 <<<"$latest")
-line=$((line + 1))
-
-# Use the same bytecode, and escape it
-bytecode=$(cut -d = -f 2 <<<"$latest")
-bytecode=$(printf '%s\n' "$bytecode" | sed -e 's/[]\/$*.^[]/\\&/g');
-
-# Insert bytecode
-sed -i "${line} i \ \ \ \ [\"${new_version}\"]=${bytecode}" $patch_file
diff --git a/tools/autopatch/find_bytecode.sh b/tools/autopatch/find_bytecode.sh
new file mode 100755
index 0000000..66b0c9d
--- /dev/null
+++ b/tools/autopatch/find_bytecode.sh
@@ -0,0 +1,45 @@
+#!/bin/bash
+
+set -euo pipefail
+
+(( $# == 2 )) || {
+    >&2 echo "Usage: $0 <version> <URL>"
+    exit 2
+}
+
+MATCH_STR=feff85c04189c4
+driver_version=$1
+driver_url=$2
+driver_file=NVIDIA-Linux-x86_64-$driver_version.run
+
+download_driver() {
+    wget -c $driver_url -O $driver_file 1>&2
+    chmod +x $driver_file
+}
+
+extract_driver() {
+    if [[ ! -e ${driver_file%".run"} ]]; then
+        ./$driver_file -x
+    fi
+}
+
+search_bytecode() {
+    nvenc_file=${driver_file%".run"}/libnvidia-encode.so.$driver_version
+    bytecode=$(xxd -c0 -ps $nvenc_file | grep -oP ".{0,6}$MATCH_STR")
+    echo $bytecode
+}
+
+get_patch_str() {
+    bytecode=$1
+    fixed=${bytecode:0:10}29${bytecode:(-8):8}
+    bytecode=$(echo "$bytecode" | sed 's/../\\x&/g')
+    fixed=$(echo "$fixed" | sed 's/../\\x&/g')
+    echo "[\"$driver_version\"]='s/$bytecode/$fixed/g'"
+}
+
+mkdir -p temp
+cd temp
+download_driver
+extract_driver
+get_patch_str $(search_bytecode)
+cd ..
diff --git a/tools/autopatch/update_patch.sh b/tools/autopatch/update_patch.sh
new file mode 100755
index 0000000..7e0d7b3
--- /dev/null
+++ b/tools/autopatch/update_patch.sh
@@ -0,0 +1,87 @@
+#!/bin/bash
+
+set -euo pipefail
+
+print_usage() { printf '
+SYNOPSIS
+   update_patch.sh -f PATCH_FILE [-v VERSION [-o OLD_VERSION] | -b PATCHSTR | -h]
+
+DESCRIPTION
+    Update the patch for Nvidia NVENC or NVFBC drivers for a new version
+
+    -f PATCH_FILE               The file (patch.sh/patch-fbc.sh) that should be updated
+    -v VERSION                  Driver version (by default copies latest existing patch)
+    -o OLD_VERSION              Copy patch string from this older driver version
+    -b PATCHSTR                 Append PATCHSTR to the patch_list
+    -h                          Print help
+'
+}
+
+opmode="copy"
+
+while getopts 'hf:v:o:b:' flag; do
+    case "${flag}" in
+        f) patch_file="$OPTARG" ;;
+        v) new_version="$OPTARG" ;;
+        o) old_version="$OPTARG" ;;
+        b) opmode="new" ; patch_string="$OPTARG" ;;
+        h) opmode="help" ;;
+        *) echo "Incorrect option specified in command line" ; exit 2 ;;
+    esac
+done
+
+
+get_last_line() {
+    if [[ -v old_version ]]; then
+        # Find old patch line
+        last=$(grep -n "\s.*\[\"$old_version\"\]='.*/g\?'" $patch_file | tail -1)
+    else
+        # Find the latest patch line
+        last=$(grep -n "\s.*\[\".*\"\]='.*/g\?'" $patch_file | tail -1)
+    fi
+    echo $last
+}
+
+get_last_line_number() {
+    last=$1
+
+    # Find the line number to insert at
+    line=$(cut -d : -f 1 <<<"$last")
+    line=$((line + 1))
+
+    echo $line
+}
+
+copy_patch() {
+    last=$(get_last_line)
+    line=$(get_last_line_number $last)
+
+    # Use the same bytecode, and escape it
+    bytecode=$(cut -d = -f 2 <<<"$last")
+    bytecode=$(printf '%s\n' "$bytecode" | sed -e 's/[]\/$*.^[]/\\&/g');
+
+    # Insert bytecode
+    sed -i "${line} i \ \ \ \ [\"${new_version}\"]=${bytecode}" $patch_file
+
+    echo "Successfully inserted bytecode for $new_version"
+}
+
+apply_new_patch() {
+    line=$(get_last_line_number $(get_last_line))
+
+    # Escape the patch string
+    bytecode=$(printf '%s\n' "$patch_string" | sed -e 's/[]\/$*.^[]/\\&/g');
+
+    # Insert it at the end
+    sed -i "${line} i \ \ \ \ ${bytecode}" $patch_file
+
+    echo "Successfully inserted $patch_string"
+}
+
+case "${opmode}" in
+    help) print_usage ; exit 2 ;;
+    copy) copy_patch ;;
+    new) apply_new_patch ;;
+    *) echo "Incorrect combination of flags. Use option -h to get help."
+       exit 2 ;;
+esac