HAMi-core 深度解析 - dlsym Hook 机制

// libvgpu.c 中的核心逻辑
static void find_real_dlsym() {
    // 通过 dlvsym 获取真实的 dlsym，绕过自身的覆盖
    real_dlsym = dlvsym(RTLD_NEXT, "dlsym", "GLIBC_2.2.5");
}

// 版本号数组，按优先级排列
static const char *glibc_versions[] = {
    "GLIBC_2.2.5",   // x86_64 标准版本
    "GLIBC_2.17",    // arm64 (aarch64) 版本
    "GLIBC_2.0",     // 古老的 32 位系统
    NULL
};

static void find_real_dlsym() {
    for (int i = 0; glibc_versions[i] != NULL; i++) {
        real_dlsym = dlvsym(RTLD_NEXT, "dlsym", glibc_versions[i]);
        if (real_dlsym != NULL) break;
    }
}

#define DLMAP_SIZE 128

static struct {
    void* caller_addr[DLMAP_SIZE];
    int index;
} dlmap;

static int check_dlmap(void* addr) {
    // 检查当前调用地址是否已在 dlmap 中（说明是递归）
    for (int i = 0; i < DLMAP_SIZE; i++) {
        if (dlmap.caller_addr[i] == addr)
            return 1;  // 发现递归
    }
    // 记录当前调用地址
    dlmap.caller_addr[dlmap.index % DLMAP_SIZE] = addr;
    dlmap.index++;
    return 0;  // 非递归
}

// 宏定义
#define DLSYM_HOOK_FUNC(name) \
    if (strcmp(symbol, #name) == 0) { \
        return (void*)hook_##name; \
    }

// 在 __dlsym_hook_section 中展开
void* __dlsym_hook_section(const char* symbol) {
    DLSYM_HOOK_FUNC(cuMemAlloc_v2)       // cuMemAlloc_v2 → hook_cuMemAlloc_v2
    DLSYM_HOOK_FUNC(cuMemFree_v2)        // cuMemFree_v2 → hook_cuMemFree_v2
    DLSYM_HOOK_FUNC(cuMemGetInfo_v2)     // cuMemGetInfo_v2 → hook_cuMemGetInfo_v2
    DLSYM_HOOK_FUNC(cuLaunchKernel)      // cuLaunchKernel → hook_cuLaunchKernel
    DLSYM_HOOK_FUNC(cuDeviceGetCount)    // cuDeviceGetCount → hook_cuDeviceGetCount
    DLSYM_HOOK_FUNC(cuDeviceTotalMem_v2) // cuDeviceTotalMem_v2 → hook_cuDeviceTotalMem_v2
    // ... 150+ 个 Hook 函数

    // 未匹配的符号，回退到真实 dlsym
    return real_dlsym(cuda_handle, symbol);
}

void* __dlsym_hook_section_nvml(const char* symbol) {
    DLSYM_HOOK_FUNC(nvmlDeviceGetMemoryInfo)
    DLSYM_HOOK_FUNC(nvmlDeviceGetMemoryInfo_v2)
    DLSYM_HOOK_FUNC(nvmlDeviceGetCount_v2)
    DLSYM_HOOK_FUNC(nvmlDeviceGetHandleByIndex_v2)
    DLSYM_HOOK_FUNC(nvmlDeviceGetProcessUtilization)
    // ... 200+ 个 NVML Hook

    return real_dlsym(nvml_handle, symbol);
}

CUresult hook_cuGetProcAddress(const char *symbol,
                                void **pfn,
                                int cudaVersion,
                                cuuint64_t flags) {
    // 首先调用真实的 cuGetProcAddress 获取原始函数指针
    CUresult result = real_cuGetProcAddress(symbol, pfn, cudaVersion, flags);
    if (result != CUDA_SUCCESS) return result;

    // 检查是否有对应的 Hook
    void *hook = __dlsym_hook_section(symbol);
    if (hook != NULL) {
        *pfn = hook;  // 替换为 Hook 函数指针
    }

    return CUDA_SUCCESS;
}

static pthread_mutex_t dlsym_lock = PTHREAD_MUTEX_INITIALIZER;

void* dlsym(void* handle, const char* symbol) {
    pthread_mutex_lock(&dlsym_lock);

    // ... 路由逻辑 ...

    pthread_mutex_unlock(&dlsym_lock);
    return result;
}

static pthread_once_t dlsym_init_once = PTHREAD_ONCE_INIT;
static pthread_once_t pre_init_once = PTHREAD_ONCE_INIT;
static pthread_once_t post_init_once = PTHREAD_ONCE_INIT;

void* dlsym(void* handle, const char* symbol) {
    pthread_once(&dlsym_init_once, find_real_dlsym);
    // ...
}

CUresult hook_cuInit(unsigned int flags) {
    pthread_once(&pre_init_once, preInit);
    CUresult result = real_cuInit(flags);
    pthread_once(&post_init_once, postInit);
    return result;
}