Merge pull request #21 from dhaern/fix/native-kernel-hardening

fix(native-kernels): harden kernel loading in the bridge and ZLUDA patch
This commit is contained in:
Woutermans authored and GitHub committed 2026-10-06 19:20:31 +02:00
commit 17012ea588
2 files changed
+49 -39

No files matched your search

@@ -1544,7 +1544,7 @@ index 0fbc7f7..ea47914 100644
}); });
&*entry.insert(func) &*entry.insert(func)
} }
@@ -50,6 +61,78 @@ impl ModuleMutable { @@ -50,6 +61,86 @@ impl ModuleMutable {
} }
} }
@@ -1558,6 +1558,14 @@ index 0fbc7f7..ea47914 100644
+ let dir = std::env::var_os("D4R_ZLUDA_NATIVE_DIR")?; + let dir = std::env::var_os("D4R_ZLUDA_NATIVE_DIR")?;
+ let path = std::path::Path::new(&dir).join(format!("{}.hsaco", name.to_str().ok()?)); + let path = std::path::Path::new(&dir).join(format!("{}.hsaco", name.to_str().ok()?));
+ let image = fs::read(&path).ok()?; + let image = fs::read(&path).ok()?;
+ // HIP can crash on a short or truncated code object: an ELF's section headers (e_shoff, e_shentsize, e_shnum) must fit
+ let le = |at: usize, size: usize| image[at..at + size].iter().rev().fold(0u64, |v, &b| (v << 8) | u64::from(b));
+ if image.len() < 64
+ || (image.starts_with(b"\x7fELF") && le(40, 8).saturating_add(le(58, 2) * le(60, 2)) > image.len() as u64)
+ {
+ eprintln!("[d4r-native] {} is truncated ({} bytes)", path.display(), image.len());
+ return None;
+ }
+ // HIP may reference the image after loading; the override lives for the whole process. + // HIP may reference the image after loading; the override lives for the whole process.
+ let image: &'static [u8] = Box::leak(image.into_boxed_slice()); + let image: &'static [u8] = Box::leak(image.into_boxed_slice());
+ let mut native_module = unsafe { mem::zeroed() }; + let mut native_module = unsafe { mem::zeroed() };
@@ -1623,7 +1631,7 @@ index 0fbc7f7..ea47914 100644
impl ZludaObject for Module { impl ZludaObject for Module {
const COOKIE: usize = 0xe9138bd040487d4a; const COOKIE: usize = 0xe9138bd040487d4a;
@@ -263,6 +346,8 @@ fn get_best_ptx_and_compile( @@ -263,6 +354,8 @@ fn get_best_ptx_and_compile(
implicit_max_block: ptx::implicit_max_block(), implicit_max_block: ptx::implicit_max_block(),
sched_strategy: ptx::sched_strategy(), sched_strategy: ptx::sched_strategy(),
wmma: llvm_zluda::compile::wmma_rewrite_requested(), wmma: llvm_zluda::compile::wmma_rewrite_requested(),
@@ -1632,7 +1640,7 @@ index 0fbc7f7..ea47914 100644
}; };
let mut cache_with_key = match (text, global_state.cache_path.as_ref()) { let mut cache_with_key = match (text, global_state.cache_path.as_ref()) {
(Some(text), Some(p)) => (|| { (Some(text), Some(p)) => (|| {
@@ -302,6 +387,8 @@ struct ExtraCacheAttributes { @@ -302,6 +395,8 @@ struct ExtraCacheAttributes {
implicit_max_block: Option<u32>, implicit_max_block: Option<u32>,
sched_strategy: Option<&'static str>, sched_strategy: Option<&'static str>,
wmma: bool, wmma: bool,
@@ -1641,7 +1649,7 @@ index 0fbc7f7..ea47914 100644
} }
fn get_hip_properties<'a>() -> Result<hipDeviceProp_tR0600, CUerror> { fn get_hip_properties<'a>() -> Result<hipDeviceProp_tR0600, CUerror> {
@@ -323,9 +410,27 @@ fn zluda_build_version() -> &'static str { @@ -323,9 +418,27 @@ fn zluda_build_version() -> &'static str {
static VERSION: std::sync::OnceLock<String> = std::sync::OnceLock::new(); static VERSION: std::sync::OnceLock<String> = std::sync::OnceLock::new();
VERSION.get_or_init(|| { VERSION.get_or_init(|| {
let sha = env!("VERGEN_GIT_SHA"); let sha = env!("VERGEN_GIT_SHA");
+37 -35
View File
@@ -475,7 +475,7 @@ static void prepare_native_kernels(const char* cache_home)
return; return;
} }
snprintf(path, sizeof(path), "%s/d4r-kernels.txt", source); snprintf(path, sizeof(path), "%s/d4r-kernels.txt", source);
manifest = fopen(path, "r"); manifest = cache_home != NULL ? fopen(path, "r") : NULL; /* a manifest set is served from the cache */
if (manifest == NULL) if (manifest == NULL)
{ {
unsetenv("D4R_ZLUDA_NATIVE_DIR"); unsetenv("D4R_ZLUDA_NATIVE_DIR");
@@ -511,7 +511,7 @@ static void prepare_native_kernels(const char* cache_home)
/* <cache>/d4r-native/<pid>; directories of processes that no longer exist are removed */ /* <cache>/d4r-native/<pid>; directories of processes that no longer exist are removed */
char base[1024]; char base[1024];
snprintf(base, sizeof(base), "%s/d4r-native", cache_home != NULL ? cache_home : "/tmp"); snprintf(base, sizeof(base), "%s/d4r-native", cache_home);
make_directories(base); make_directories(base);
DIR* directory = opendir(base); DIR* directory = opendir(base);
if (directory != NULL) if (directory != NULL)
@@ -531,15 +531,18 @@ static void prepare_native_kernels(const char* cache_home)
} }
snprintf(native_served, sizeof(native_served), "%s/%ld", base, (long)getpid()); snprintf(native_served, sizeof(native_served), "%s/%ld", base, (long)getpid());
remove_directory(native_served); remove_directory(native_served);
if (mkdir(native_served, 0755) != 0) char* resolved = realpath(source, NULL);
if (resolved == NULL || strlen(resolved) >= sizeof(native_source) || mkdir(native_served, 0755) != 0)
{ {
set_load_error("cannot create %s; native kernels disabled", native_served); free(resolved);
set_load_error("cannot create %s from %s; native kernels disabled", native_served, source);
native_served[0] = '\0'; native_served[0] = '\0';
native_kernel_count = 0; native_kernel_count = 0;
unsetenv("D4R_ZLUDA_NATIVE_DIR"); unsetenv("D4R_ZLUDA_NATIVE_DIR");
return; return;
} }
snprintf(native_source, sizeof(native_source), "%s", source); snprintf(native_source, sizeof(native_source), "%s", resolved);
free(resolved);
setenv("D4R_ZLUDA_NATIVE_DIR", native_served, 1); setenv("D4R_ZLUDA_NATIVE_DIR", native_served, 1);
atexit(native_cleanup); atexit(native_cleanup);
tracef("native kernels: %zu in %s, each served from %s once DLSS's PTX for it matches", tracef("native kernels: %zu in %s, each served from %s once DLSS's PTX for it matches",
@@ -696,7 +699,7 @@ static void verify_native_kernels(const void* image)
memcpy(&flags, bytes + offset + 40, sizeof(flags)); memcpy(&flags, bytes + offset + 40, sizeof(flags));
if (entry_header >= 64) if (entry_header >= 64)
memcpy(&decompressed, bytes + offset + 56, sizeof(decompressed)); memcpy(&decompressed, bytes + offset + 56, sizeof(decompressed));
if (entry_header < 16 || entry_size > end - offset - entry_header) if (entry_header < 16 || entry_header > end - offset || entry_size > end - offset - entry_header)
break; break;
const unsigned char* payload = bytes + offset + entry_header; const unsigned char* payload = bytes + offset + entry_header;
if (kind == 1) /* PTX */ if (kind == 1) /* PTX */
@@ -727,8 +730,8 @@ static struct
{ {
CUfunction function; CUfunction function;
int native; int native;
} output_kernels[OUTPUT_KERNEL_CAPACITY]; } output_kernels[OUTPUT_KERNEL_CAPACITY]; /* the latest lookups, newest last */
static int output_kernel_count; static pthread_mutex_t output_kernel_lock = PTHREAD_MUTEX_INITIALIZER;
static int last_output_native = -1; static int last_output_native = -1;
static void note_function_lookup(CUfunction function, const char* name) static void note_function_lookup(CUfunction function, const char* name)
@@ -745,25 +748,24 @@ static void note_function_lookup(CUfunction function, const char* name)
snprintf(path, sizeof(path), "%s/%s.hsaco", directory, name); snprintf(path, sizeof(path), "%s/%s.hsaco", directory, name);
native = access(path, R_OK) == 0; native = access(path, R_OK) == 0;
} }
pthread_mutex_lock(&native_lock); pthread_mutex_lock(&output_kernel_lock);
if (output_kernel_count < OUTPUT_KERNEL_CAPACITY) memmove(output_kernels, output_kernels + 1, sizeof(output_kernels) - sizeof(output_kernels[0]));
{ output_kernels[OUTPUT_KERNEL_CAPACITY - 1].function = function;
output_kernels[output_kernel_count].function = function; output_kernels[OUTPUT_KERNEL_CAPACITY - 1].native = native;
output_kernels[output_kernel_count].native = native; pthread_mutex_unlock(&output_kernel_lock);
__atomic_store_n(&output_kernel_count, output_kernel_count + 1, __ATOMIC_RELEASE);
}
pthread_mutex_unlock(&native_lock);
} }
static void note_launch(CUfunction function) static void note_launch(CUfunction function)
{ {
const int count = __atomic_load_n(&output_kernel_count, __ATOMIC_ACQUIRE); pthread_mutex_lock(&output_kernel_lock);
for (int i = 0; i < count; ++i) /* newest first, so a recycled handle resolves with its latest lookup */
for (int i = OUTPUT_KERNEL_CAPACITY - 1; i >= 0 && output_kernels[i].function != NULL; --i)
if (output_kernels[i].function == function) if (output_kernels[i].function == function)
{ {
__atomic_store_n(&last_output_native, output_kernels[i].native, __ATOMIC_RELAXED); __atomic_store_n(&last_output_native, output_kernels[i].native, __ATOMIC_RELAXED);
return; break;
} }
pthread_mutex_unlock(&output_kernel_lock);
} }
/* 1 when the last output kernel launched was d4r's native one, 0 when it was ZLUDA's compile of /* 1 when the last output kernel launched was d4r's native one, 0 when it was ZLUDA's compile of
@@ -805,35 +807,33 @@ static void load_zluda(void)
} }
load_error[0] = '\0'; load_error[0] = '\0';
const char* cache = getenv("D4R_ZLUDA_CACHE_HOME"); const char* cache = getenv("D4R_ZLUDA_CACHE_HOME");
char cache_home[1024] = {0}; const char* xdg = getenv("XDG_CACHE_HOME");
const char* home = getenv("HOME");
char cache_home[1024] = {0}, root[1024] = {0};
if (cache != NULL && cache[0] != '\0') if (cache != NULL && cache[0] != '\0')
{ {
expand_home(cache, cache_home, sizeof(cache_home)); expand_home(cache, cache_home, sizeof(cache_home));
make_directories(cache_home); make_directories(cache_home);
} }
/* the cache root; no /tmp fallback, since other users can plant links there */
if (cache_home[0] != '\0')
snprintf(root, sizeof(root), "%s", cache_home);
else if (xdg != NULL && xdg[0] != '\0')
snprintf(root, sizeof(root), "%s", xdg);
else if (home != NULL)
snprintf(root, sizeof(root), "%s/.cache", home);
if (accuracy) if (accuracy)
{ {
/* FAST_MATH is not fingerprinted by older ZLUDA runtimes. Use a separate cache even /* FAST_MATH is not fingerprinted by older ZLUDA runtimes. Use a separate cache even
when a previous run explicitly enabled that experimental compiler switch. */ when a previous run explicitly enabled that experimental compiler switch. */
char base[1024]; if (snprintf(cache_home, sizeof(cache_home), "%s/d4r-accuracy", root[0] != '\0' ? root : "/tmp/.cache") >= (int)sizeof(cache_home))
const char* xdg = getenv("XDG_CACHE_HOME");
if (cache_home[0] != '\0')
snprintf(base, sizeof(base), "%s", cache_home);
else if (xdg != NULL && xdg[0] != '\0')
snprintf(base, sizeof(base), "%s", xdg);
else
{
const char* home = getenv("HOME");
snprintf(base, sizeof(base), "%s/.cache", home != NULL ? home : "/tmp");
}
if (snprintf(cache_home, sizeof(cache_home), "%s/d4r-accuracy", base) >= (int)sizeof(cache_home))
{ {
set_load_error("accuracy cache path is too long"); set_load_error("accuracy cache path is too long");
return; return;
} }
make_directories(cache_home); make_directories(cache_home);
} }
prepare_native_kernels(cache_home[0] != '\0' ? cache_home : NULL); prepare_native_kernels(root[0] == '\0' ? NULL : accuracy ? cache_home : root);
if (cache_home[0] != '\0') if (cache_home[0] != '\0')
{ {
/* ZLUDA picks its cache directory once, in cuInit */ /* ZLUDA picks its cache directory once, in cuInit */
@@ -2482,6 +2482,8 @@ CUresult WINAPI cuDestroyExternalMemory(CUexternalMemory memory)
return result; return result;
} }
CUresult WINAPI cuSurfObjectDestroy(CUsurfObject object);
CUresult WINAPI cuSurfObjectCreate(CUsurfObject* object, const void* descriptor) CUresult WINAPI cuSurfObjectCreate(CUsurfObject* object, const void* descriptor)
{ {
CUSURFOBJECTCREATE_FN function = (CUSURFOBJECTCREATE_FN)find_zluda_symbol("cuSurfObjectCreate"); CUSURFOBJECTCREATE_FN function = (CUSURFOBJECTCREATE_FN)find_zluda_symbol("cuSurfObjectCreate");
@@ -2495,8 +2497,8 @@ CUresult WINAPI cuSurfObjectCreate(CUsurfObject* object, const void* descriptor)
const CUdeviceptr pointer = redirect_pointer; const CUdeviceptr pointer = redirect_pointer;
const uint32_t pitch = redirect_pitch; const uint32_t pitch = redirect_pitch;
pthread_mutex_unlock(&instrumentation_lock); pthread_mutex_unlock(&instrumentation_lock);
if (redirected) if (redirected && (result = write_redirect_tail(*object, pointer, pitch, 0)) != CUDA_SUCCESS)
write_redirect_tail(*object, pointer, pitch, 0); cuSurfObjectDestroy(*object);
} }
TRACE_CALL(result, "cuSurfObjectCreate descriptor=%p result=%d object=0x%llx", descriptor, result, TRACE_CALL(result, "cuSurfObjectCreate descriptor=%p result=%d object=0x%llx", descriptor, result,
(unsigned long long)(object != NULL ? *object : 0)); (unsigned long long)(object != NULL ? *object : 0));