Mergepatch bugfix

mergepatch.c

// compile with `gcc mergepatch.c -Wall -Werror -fPIC -ldl -shared -o mergepatch.so`
// usage: `LD_PRELOAD=./mergepatch.so nvidia-smi`

#define _GNU_SOURCE // for RTLD_NEXT

#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <unistd.h>
#include <dlfcn.h>
#include <link.h>
#include <sys/mman.h>

#define ARRAY_SIZE(a) (sizeof(a) / sizeof(*a))

typedef struct driver_patch_loc_st {
	const char* version;
	size_t offset;
} driver_patch_loc_t;

static void*(*dlsymReal)(void*, const char*);
static int(*nvmlInitWithFlagsReal)(unsigned int);
static int(*nvmlSystemGetNVMLVersionReal)(char*, unsigned int);
static struct link_map* nvmlLinkMap;

// -- ADD TO THIS ON DRIVER UPDATE --
static driver_patch_loc_t driver_patch_locs[] = {
	{ "11.510.85.01", 0xCCC84 },
	{ "11.510.85.02", 0xCCC84 },
	{ "12.525.60.11", 0xBD684 },
	{ "12.525.60.12", 0xBD684 },
	{ "12.525.60.13", 0xBD684 }
};
// -- ADD TO THIS ON DRIVER UPDATE --

int setProtect(void* ptr, bool writable) {
	return mprotect((void*)((uintptr_t)ptr & ~(getpagesize() - 1)), getpagesize(), (writable ? PROT_WRITE : PROT_EXEC) | PROT_READ);
}

int nvmlInitWithFlagsHook(unsigned int flags) {
	unsigned int ret = nvmlInitWithFlagsReal(flags);

	char buf[16];
	nvmlSystemGetNVMLVersionReal(buf, sizeof(buf));

	for(size_t i = 0; i < ARRAY_SIZE(driver_patch_locs); i++) {
		printf("%s %s\n", buf, driver_patch_locs[i].version);
		if(!strcmp(buf, driver_patch_locs[i].version)) {
			uint8_t* byte = (uint8_t*)(nvmlLinkMap->l_addr + driver_patch_locs[i].offset);
			setProtect((void*)byte, true);
			*byte = 0xEB; // JL -> JMP
			setProtect((void*)byte, false);
			return ret;
		}
	}

	fprintf(stderr, "failed to patch nvml: version %s\n", buf);
	abort();

	return 999; // NVML_ERROR_UNKNOWN
}

void* dlsym(void* handle, const char* name) {
	if(!dlsymReal) {
		// https://stackoverflow.com/questions/15599026/how-can-i-intercept-dlsym-calls-using-ld-preload
		dlsymReal = dlvsym(RTLD_NEXT, "dlsym", "GLIBC_2.2.5");
		dlsymReal = dlsymReal(RTLD_NEXT, "dlsym");
		if(!dlsymReal) {
			fprintf(stderr, "dlsym: %s\n", dlerror());
			abort();
		}
	}

	if(!strcmp(name, "nvmlInitWithFlags")) {
		dlinfo(handle, RTLD_DI_LINKMAP, &nvmlLinkMap);
		nvmlInitWithFlagsReal = dlsymReal(handle, "nvmlInitWithFlags");
		nvmlSystemGetNVMLVersionReal = dlsymReal(handle, "nvmlSystemGetNVMLVersion");
		return nvmlInitWithFlagsHook;
	}
	
	return dlsymReal(handle, name);
}