#endif /* CONFIG_KASAN_SW_TAGS */
 
 #ifdef CONFIG_KASAN_VMALLOC
-int kasan_populate_vmalloc(unsigned long requested_size,
-                          struct vm_struct *area);
-void kasan_poison_vmalloc(void *start, unsigned long size);
+int kasan_populate_vmalloc(unsigned long addr, unsigned long size);
+void kasan_poison_vmalloc(const void *start, unsigned long size);
+void kasan_unpoison_vmalloc(const void *start, unsigned long size);
 void kasan_release_vmalloc(unsigned long start, unsigned long end,
                           unsigned long free_region_start,
                           unsigned long free_region_end);
 #else
-static inline int kasan_populate_vmalloc(unsigned long requested_size,
-                                        struct vm_struct *area)
+static inline int kasan_populate_vmalloc(unsigned long start,
+                                       unsigned long size)
 {
        return 0;
 }
 
-static inline void kasan_poison_vmalloc(void *start, unsigned long size) {}
+static inline void kasan_poison_vmalloc(const void *start, unsigned long size)
+{ }
+static inline void kasan_unpoison_vmalloc(const void *start, unsigned long size)
+{ }
 static inline void kasan_release_vmalloc(unsigned long start,
                                         unsigned long end,
                                         unsigned long free_region_start,
 
        return 0;
 }
 
-int kasan_populate_vmalloc(unsigned long requested_size, struct vm_struct *area)
+int kasan_populate_vmalloc(unsigned long addr, unsigned long size)
 {
        unsigned long shadow_start, shadow_end;
        int ret;
 
-       shadow_start = (unsigned long)kasan_mem_to_shadow(area->addr);
+       if (!is_vmalloc_or_module_addr((void *)addr))
+               return 0;
+
+       shadow_start = (unsigned long)kasan_mem_to_shadow((void *)addr);
        shadow_start = ALIGN_DOWN(shadow_start, PAGE_SIZE);
-       shadow_end = (unsigned long)kasan_mem_to_shadow(area->addr +
-                                                       area->size);
+       shadow_end = (unsigned long)kasan_mem_to_shadow((void *)addr + size);
        shadow_end = ALIGN(shadow_end, PAGE_SIZE);
 
        ret = apply_to_page_range(&init_mm, shadow_start,
 
        flush_cache_vmap(shadow_start, shadow_end);
 
-       kasan_unpoison_shadow(area->addr, requested_size);
-
-       area->flags |= VM_KASAN;
-
        /*
         * We need to be careful about inter-cpu effects here. Consider:
         *
  * Poison the shadow for a vmalloc region. Called as part of the
  * freeing process at the time the region is freed.
  */
-void kasan_poison_vmalloc(void *start, unsigned long size)
+void kasan_poison_vmalloc(const void *start, unsigned long size)
 {
+       if (!is_vmalloc_or_module_addr(start))
+               return;
+
        size = round_up(size, KASAN_SHADOW_SCALE_SIZE);
        kasan_poison_shadow(start, size, KASAN_VMALLOC_INVALID);
 }
 
+void kasan_unpoison_vmalloc(const void *start, unsigned long size)
+{
+       if (!is_vmalloc_or_module_addr(start))
+               return;
+
+       kasan_unpoison_shadow(start, size);
+}
+
 static int kasan_depopulate_vmalloc_pte(pte_t *ptep, unsigned long addr,
                                        void *unused)
 {
 
        return nva_start_addr;
 }
 
+/*
+ * Free a region of KVA allocated by alloc_vmap_area
+ */
+static void free_vmap_area(struct vmap_area *va)
+{
+       /*
+        * Remove from the busy tree/list.
+        */
+       spin_lock(&vmap_area_lock);
+       unlink_va(va, &vmap_area_root);
+       spin_unlock(&vmap_area_lock);
+
+       /*
+        * Insert/Merge it back to the free tree/list.
+        */
+       spin_lock(&free_vmap_area_lock);
+       merge_or_add_vmap_area(va, &free_vmap_area_root, &free_vmap_area_list);
+       spin_unlock(&free_vmap_area_lock);
+}
+
 /*
  * Allocate a region of KVA of the specified size and alignment, within the
  * vstart and vend.
        struct vmap_area *va, *pva;
        unsigned long addr;
        int purged = 0;
+       int ret;
 
        BUG_ON(!size);
        BUG_ON(offset_in_page(size));
        va->va_end = addr + size;
        va->vm = NULL;
 
+
        spin_lock(&vmap_area_lock);
        insert_vmap_area(va, &vmap_area_root, &vmap_area_list);
        spin_unlock(&vmap_area_lock);
        BUG_ON(va->va_start < vstart);
        BUG_ON(va->va_end > vend);
 
+       ret = kasan_populate_vmalloc(addr, size);
+       if (ret) {
+               free_vmap_area(va);
+               return ERR_PTR(ret);
+       }
+
        return va;
 
 overflow:
 }
 EXPORT_SYMBOL_GPL(unregister_vmap_purge_notifier);
 
-/*
- * Free a region of KVA allocated by alloc_vmap_area
- */
-static void free_vmap_area(struct vmap_area *va)
-{
-       /*
-        * Remove from the busy tree/list.
-        */
-       spin_lock(&vmap_area_lock);
-       unlink_va(va, &vmap_area_root);
-       spin_unlock(&vmap_area_lock);
-
-       /*
-        * Insert/Merge it back to the free tree/list.
-        */
-       spin_lock(&free_vmap_area_lock);
-       merge_or_add_vmap_area(va, &free_vmap_area_root, &free_vmap_area_list);
-       spin_unlock(&free_vmap_area_lock);
-}
-
 /*
  * Clear the pagetable entries of a given vmap_area
  */
        BUG_ON(addr > VMALLOC_END);
        BUG_ON(!PAGE_ALIGNED(addr));
 
+       kasan_poison_vmalloc(mem, size);
+
        if (likely(count <= VMAP_MAX_ALLOC)) {
                debug_check_no_locks_freed(mem, size);
                vb_free(mem, size);
                addr = va->va_start;
                mem = (void *)addr;
        }
+
+       kasan_unpoison_vmalloc(mem, size);
+
        if (vmap_page_range(addr, addr + size, prot, pages) < 0) {
                vm_unmap_ram(mem, count);
                return NULL;
 {
        struct vmap_area *va;
        struct vm_struct *area;
+       unsigned long requested_size = size;
 
        BUG_ON(in_interrupt());
        size = PAGE_ALIGN(size);
                return NULL;
        }
 
-       setup_vmalloc_vm(area, va, flags, caller);
+       kasan_unpoison_vmalloc((void *)va->va_start, requested_size);
 
-       /*
-        * For KASAN, if we are in vmalloc space, we need to cover the shadow
-        * area with real memory. If we come here through VM_ALLOC, this is
-        * done by a higher level function that has access to the true size,
-        * which might not be a full page.
-        *
-        * We assume module space comes via VM_ALLOC path.
-        */
-       if (is_vmalloc_addr(area->addr) && !(area->flags & VM_ALLOC)) {
-               if (kasan_populate_vmalloc(area->size, area)) {
-                       unmap_vmap_area(va);
-                       kfree(area);
-                       return NULL;
-               }
-       }
+       setup_vmalloc_vm(area, va, flags, caller);
 
        return area;
 }
        debug_check_no_locks_freed(area->addr, get_vm_area_size(area));
        debug_check_no_obj_freed(area->addr, get_vm_area_size(area));
 
-       if (area->flags & VM_KASAN)
-               kasan_poison_vmalloc(area->addr, area->size);
+       kasan_poison_vmalloc(area->addr, area->size);
 
        vm_remove_mappings(area, deallocate_pages);
 
        if (!size || (size >> PAGE_SHIFT) > totalram_pages())
                goto fail;
 
-       area = __get_vm_area_node(size, align, VM_ALLOC | VM_UNINITIALIZED |
+       area = __get_vm_area_node(real_size, align, VM_ALLOC | VM_UNINITIALIZED |
                                vm_flags, start, end, node, gfp_mask, caller);
        if (!area)
                goto fail;
        if (!addr)
                return NULL;
 
-       if (is_vmalloc_or_module_addr(area->addr)) {
-               if (kasan_populate_vmalloc(real_size, area))
-                       return NULL;
-       }
-
        /*
         * In this function, newly allocated vm_struct has VM_UNINITIALIZED
         * flag. It means that vm_struct is not fully initialized.
        /* populate the shadow space outside of the lock */
        for (area = 0; area < nr_vms; area++) {
                /* assume success here */
-               kasan_populate_vmalloc(sizes[area], vms[area]);
+               kasan_populate_vmalloc(vas[area]->va_start, sizes[area]);
+               kasan_unpoison_vmalloc((void *)vms[area]->addr, sizes[area]);
        }
 
        kfree(vas);