1 files changed, 11 insertions, 3 deletions

diff --git a/mm/kasan/kasan.c b/mm/kasan/kasan.c
index 78fee632a7ee..936d81661c47 100644
--- a/mm/kasan/kasan.c
+++ b/mm/kasan/kasan.c
@@ -29,6 +29,7 @@
 #include <linux/stacktrace.h>
 #include <linux/string.h>
 #include <linux/types.h>
+#include <linux/vmalloc.h>
 #include <linux/kasan.h>
 
 #include "kasan.h"
@@ -414,12 +415,19 @@ int kasan_module_alloc(void *addr, size_t size)
                         GFP_KERNEL | __GFP_HIGHMEM | __GFP_ZERO,
                         PAGE_KERNEL, VM_NO_GUARD, NUMA_NO_NODE,
                         __builtin_return_address(0));
-        return ret ? 0 : -ENOMEM;
+
+        if (ret) {
+                find_vm_area(addr)->flags |= VM_KASAN;
+                return 0;
+        }
+
+        return -ENOMEM;
 }
 
-void kasan_module_free(void *addr)
+void kasan_free_shadow(const struct vm_struct *vm)
 {
-        vfree(kasan_mem_to_shadow(addr));
+        if (vm->flags & VM_KASAN)
+                vfree(kasan_mem_to_shadow(vm->addr));
 }
 
 static void register_global(struct kasan_global *global)