diff --git a/linuxkm/linuxkm_wc_port.h b/linuxkm/linuxkm_wc_port.h
index 892f323a8a..83e35dc692 100644
--- a/linuxkm/linuxkm_wc_port.h
+++ b/linuxkm/linuxkm_wc_port.h
@@ -167,13 +167,14 @@
     #ifndef WC_LINUXKM_INTR_SIGNALS
         #define WC_LINUXKM_INTR_SIGNALS { SIGKILL, SIGABRT, SIGHUP, SIGINT }
     #endif
-    extern int wc_linuxkm_sig_ignore_begin(void);
-    extern int wc_linuxkm_sig_ignore_end(void);
-    extern int wc_linuxkm_check_for_intr_signals(void);
+    WOLFSSL_API int wc_linuxkm_can_block(void);
+    WOLFSSL_API int wc_linuxkm_sig_ignore_begin(void);
+    WOLFSSL_API int wc_linuxkm_sig_ignore_end(void);
+    WOLFSSL_API int wc_linuxkm_check_for_intr_signals(void);
     #ifndef WC_LINUXKM_MAX_NS_WITHOUT_YIELD
         #define WC_LINUXKM_MAX_NS_WITHOUT_YIELD 1000000000
     #endif
-    extern void wc_linuxkm_relax_long_loop(void);
+    WOLFSSL_API void wc_linuxkm_relax_long_loop(void);
 
     #ifndef WC_SIG_IGNORE_BEGIN
         #define WC_SIG_IGNORE_BEGIN() wc_linuxkm_sig_ignore_begin()
@@ -1221,6 +1222,7 @@
         typeof(wc_lkm_LockMutex) *wc_lkm_LockMutex;
         #endif
 
+        typeof(wc_linuxkm_can_block) *wc_linuxkm_can_block;
         typeof(wc_linuxkm_sig_ignore_begin) *wc_linuxkm_sig_ignore_begin;
         typeof(wc_linuxkm_sig_ignore_end) *wc_linuxkm_sig_ignore_end;
         typeof(wc_linuxkm_check_for_intr_signals) *wc_linuxkm_check_for_intr_signals;
@@ -1521,8 +1523,9 @@
         wolfssl_spin_unlock_irqrestore_rt((lock), (flags))
     #endif
 
-    #define wc_linuxkm_sig_ignore_begin WC_PIE_INDIRECT_SYM(wc_linuxkm_sig_ignore_begin);
-    #define wc_linuxkm_sig_ignore_end WC_PIE_INDIRECT_SYM(wc_linuxkm_sig_ignore_end);
+    #define wc_linuxkm_can_block WC_PIE_INDIRECT_SYM(wc_linuxkm_can_block)
+    #define wc_linuxkm_sig_ignore_begin WC_PIE_INDIRECT_SYM(wc_linuxkm_sig_ignore_begin)
+    #define wc_linuxkm_sig_ignore_end WC_PIE_INDIRECT_SYM(wc_linuxkm_sig_ignore_end)
     #define wc_linuxkm_check_for_intr_signals WC_PIE_INDIRECT_SYM(wc_linuxkm_check_for_intr_signals)
     #define wc_linuxkm_relax_long_loop WC_PIE_INDIRECT_SYM(wc_linuxkm_relax_long_loop)
 
@@ -1905,21 +1908,21 @@
         #endif
     #else /* !WC_LINUXKM_USE_HEAP_WRAPPERS */
     #ifdef USE_KVMALLOC
-        #define malloc(size) kvmalloc_node(WC_LINUXKM_ROUND_UP_P_OF_2(size), (preempt_count() == 0 ? GFP_KERNEL : GFP_ATOMIC), NUMA_NO_NODE)
+        #define malloc(size) kvmalloc_node(WC_LINUXKM_ROUND_UP_P_OF_2(size), (wc_linuxkm_can_block() ? GFP_KERNEL : GFP_ATOMIC), NUMA_NO_NODE)
         #if LINUX_VERSION_CODE >= KERNEL_VERSION(7, 2, 0)
-            #define free(ptr) (preempt_count() == 0 ? kvfree(ptr) : kvfree_atomic(ptr))
+            #define free(ptr) (wc_linuxkm_can_block() ? kvfree(ptr) : kvfree_atomic(ptr))
         #else
             #define free(ptr) kvfree(ptr)
         #endif
         #ifdef USE_KVREALLOC
-            #define realloc(ptr, newsize) kvrealloc(ptr, WC_LINUXKM_ROUND_UP_P_OF_2(newsize), (preempt_count() == 0 ? GFP_KERNEL : GFP_ATOMIC))
+            #define realloc(ptr, newsize) kvrealloc(ptr, WC_LINUXKM_ROUND_UP_P_OF_2(newsize), (wc_linuxkm_can_block() ? GFP_KERNEL : GFP_ATOMIC))
         #else
             #define realloc(ptr, newsize) ((void)(ptr), (void)(newsize), NULL)
         #endif
     #else
-        #define malloc(size) kmalloc(WC_LINUXKM_ROUND_UP_P_OF_2(size), (preempt_count() == 0 ? GFP_KERNEL : GFP_ATOMIC))
+        #define malloc(size) kmalloc(WC_LINUXKM_ROUND_UP_P_OF_2(size), (wc_linuxkm_can_block() ? GFP_KERNEL : GFP_ATOMIC))
         #define free(ptr) kfree(ptr)
-        #define realloc(ptr, newsize) krealloc(ptr, WC_LINUXKM_ROUND_UP_P_OF_2(newsize), (preempt_count() == 0 ? GFP_KERNEL : GFP_ATOMIC))
+        #define realloc(ptr, newsize) krealloc(ptr, WC_LINUXKM_ROUND_UP_P_OF_2(newsize), (wc_linuxkm_can_block() ? GFP_KERNEL : GFP_ATOMIC))
     #endif
     #endif /* !WC_LINUXKM_USE_HEAP_WRAPPERS */
 
diff --git a/linuxkm/lkcapi_sha_glue.c b/linuxkm/lkcapi_sha_glue.c
index 39754a6424..9a80247bc3 100644
--- a/linuxkm/lkcapi_sha_glue.c
+++ b/linuxkm/lkcapi_sha_glue.c
@@ -1009,7 +1009,7 @@ static volatile int wc_linuxkm_rng_initing_default_bank_flag = 0;
 
 static int linuxkm_affinity_lock(void *arg) {
     (void)arg;
-    if (preempt_count() != 0)
+    if (! wc_linuxkm_can_block())
         return ALREADY_E;
 #if defined(CONFIG_SMP) && (LINUX_VERSION_CODE >= KERNEL_VERSION(5, 7, 0))
     migrate_disable(); /* this actually makes irq_count() nonzero, so that
@@ -1118,7 +1118,7 @@ static struct wc_rng_bank_inst *linuxkm_get_drbg(struct wc_rng_bank *ctx) {
         WC_RNG_BANK_FLAG_CAN_WAIT |
         WC_RNG_BANK_FLAG_PREFER_AFFINITY_INST;
 
-    if (preempt_count() == 0)
+    if (wc_linuxkm_can_block())
         flags |= WC_RNG_BANK_FLAG_AFFINITY_LOCK;
     else
         flags |= WC_RNG_BANK_FLAG_NO_VECTOR_OPS;
@@ -1487,7 +1487,7 @@ static int wc_mix_pool_bytes(const void *buf, size_t len) {
     struct wc_rng_bank *ctx;
     size_t i;
     int n;
-    int can_sleep = (preempt_count() == 0);
+    int can_sleep = wc_linuxkm_can_block();
 
     if (len == 0)
         return 0;
@@ -1545,7 +1545,7 @@ static int wc_mix_pool_bytes(const void *buf, size_t len) {
 
 static int wc_crng_reseed(void) {
     struct wc_rng_bank *ctx;
-    int can_sleep = (preempt_count() == 0);
+    int can_sleep = wc_linuxkm_can_block();
     int ret = wc_rng_bank_default_checkout(&ctx);
 
     if (ret) {
diff --git a/linuxkm/module_hooks.c b/linuxkm/module_hooks.c
index be1fd01d7b..e5b94ba278 100644
--- a/linuxkm/module_hooks.c
+++ b/linuxkm/module_hooks.c
@@ -316,6 +316,13 @@ MODULE_PARM_DESC(rodata_dump_path,
     #include "linuxkm/lkcapi_glue.c"
 #endif
 
+int wc_linuxkm_can_block(void) {
+    /* We can't use preemptible() for this, because we need an accurate test
+     * even in !CONFIG_PREEMPT_COUNT configs where preemptible() is always 0.
+     */
+    return (preempt_count() == 0) && (! irqs_disabled());
+}
+
 /* for simplicity, we use a global count to suspend signal processing while any
  * thread is running fipsEntry(), wolfCrypt_IntegrityTest_fips(),
  * linuxkm_lkcapi_register(), or linuxkm_lkcapi_unregister().  This only affects
@@ -336,7 +343,7 @@ int wc_linuxkm_sig_ignore_end(void) {
 
 int wc_linuxkm_check_for_intr_signals(void) {
     static const int intr_signals[] = WC_LINUXKM_INTR_SIGNALS;
-    if (preempt_count() != 0)
+    if (! wc_linuxkm_can_block())
         return 0;
     if (signal_pending(current)) {
         int i;
@@ -365,7 +372,7 @@ int wc_linuxkm_check_for_intr_signals(void) {
 
 void wc_linuxkm_relax_long_loop(void) {
     #if WC_LINUXKM_MAX_NS_WITHOUT_YIELD >= 0
-    if (preempt_count() == 0) {
+    if (wc_linuxkm_can_block()) {
         #if (WC_LINUXKM_MAX_NS_WITHOUT_YIELD == 0) || !defined(CONFIG_SCHED_INFO)
         cond_resched();
         #else
@@ -662,6 +669,11 @@ static int wolfssl_init(void)
     }
 #endif
 
+#if defined(WC_LINUXKM_USE_HEAP_WRAPPERS) && defined(CONFIG_HAVE_KPROBES)
+    /* cache the function pointer to find_vm_area(). */
+    (void)wc_linuxkm_malloc_usable_size(NULL);
+#endif
+
 #ifdef WOLFCRYPT_FIPS_CORE_DYNAMIC_HASH_VALUE
 #ifdef CONFIG_MODULE_SIG
     if (THIS_MODULE->sig_ok == false) {
@@ -1283,13 +1295,13 @@ static const struct wc_reloc_table_segments seg_map = {
 
 void *wc_linuxkm_malloc(size_t size)
 {
-    return kvmalloc_node(WC_LINUXKM_ROUND_UP_P_OF_2(size), (preempt_count() == 0 ? GFP_KERNEL : GFP_ATOMIC), NUMA_NO_NODE);
+    return kvmalloc_node(WC_LINUXKM_ROUND_UP_P_OF_2(size), (wc_linuxkm_can_block() ? GFP_KERNEL : GFP_ATOMIC), NUMA_NO_NODE);
 }
 
 void wc_linuxkm_free(void *ptr)
 {
 #if LINUX_VERSION_CODE >= KERNEL_VERSION(7, 2, 0)
-    if (preempt_count() == 0)
+    if (wc_linuxkm_can_block())
         kvfree(ptr);
     else
         kvfree_atomic(ptr);
@@ -1300,12 +1312,41 @@ void wc_linuxkm_free(void *ptr)
 
 void *wc_linuxkm_realloc(void *ptr, size_t newsize)
 {
-    return kvrealloc(ptr, WC_LINUXKM_ROUND_UP_P_OF_2(newsize), (preempt_count() == 0 ? GFP_KERNEL : GFP_ATOMIC));
+    return kvrealloc(ptr, WC_LINUXKM_ROUND_UP_P_OF_2(newsize), (wc_linuxkm_can_block() ? GFP_KERNEL : GFP_ATOMIC));
 }
 
+#ifdef CONFIG_HAVE_KPROBES
+#include <linux/vmalloc.h>
+#endif
+
 size_t wc_linuxkm_malloc_usable_size(void *ptr)
 {
-    return ksize(ptr);
+    if ((ptr == NULL) || is_vmalloc_addr(ptr)) {
+#ifdef CONFIG_HAVE_KPROBES
+        static typeof(find_vm_area) *find_vm_area_ptr = NULL;
+        if (find_vm_area_ptr == NULL) {
+            if (! wc_linuxkm_can_block())
+                return 0;
+            find_vm_area_ptr = my_kallsyms_lookup_name("find_vm_area");
+        }
+        if (find_vm_area_ptr == NULL)
+            return 0;
+        else if (ptr == NULL)
+            return 0;
+        else {
+            struct vm_struct *vm = find_vm_area_ptr(ptr);
+            if (vm)
+                return get_vm_area_size(vm);
+            else
+                return 0;
+        }
+#else
+        return 0;
+#endif
+    }
+    else {
+        return ksize(ptr);
+    }
 }
 
 #endif /* WC_LINUXKM_USE_HEAP_WRAPPERS */
@@ -1678,6 +1719,7 @@ static int set_up_wolfssl_linuxkm_pie_redirect_table(void) {
 #endif
 #endif
 
+    wolfssl_linuxkm_pie_redirect_table.wc_linuxkm_can_block = wc_linuxkm_can_block;
     wolfssl_linuxkm_pie_redirect_table.wc_linuxkm_sig_ignore_begin = wc_linuxkm_sig_ignore_begin;
     wolfssl_linuxkm_pie_redirect_table.wc_linuxkm_sig_ignore_end = wc_linuxkm_sig_ignore_end;
     wolfssl_linuxkm_pie_redirect_table.wc_linuxkm_check_for_intr_signals = wc_linuxkm_check_for_intr_signals;