Created
December 11, 2022 15:43
-
-
Save 0xNULLderef/e4a757567734345c3be23683e73db7c5 to your computer and use it in GitHub Desktop.
Patch for nvidia-smi to display process information with a merged driver
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#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.02", 0xCCC84 }, | |
{ "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++) { | |
if(!strcmp(buf, driver_patch_locs[i].version)) { | |
uint8_t* byte = (uint8_t*)(nvmlLinkMap->l_addr + driver_patch_locs->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); | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Updated version with bugfix:
https://gist.github.com/KonradCzerw/80d3b907946b5c940f50da7795364bc1