int set_memory_ro(unsigned long addr, int numpages);
 int set_memory_rw(unsigned long addr, int numpages);
 int set_memory_nx(unsigned long addr, int numpages);
+int set_memory_x(unsigned long addr, int numpages);
 
 #endif /* _S390_CACHEFLUSH_H */
 
        }
 }
 
-/* Setting memory back to RW+NX before releasing it */
+/* Setting memory back to W+X before releasing it */
 void unset_section_ro_nx(struct module *mod, void *module_region)
 {
-       unsigned long total_pages;
-
        if (mod->module_core == module_region) {
-               /* Set core as NX+RW */
-               total_pages = MOD_NUMBER_OF_PAGES(mod->module_core, mod->core_size);
-               set_memory_nx((unsigned long)mod->module_core, total_pages);
-               set_memory_rw((unsigned long)mod->module_core, total_pages);
-
+               set_page_attributes(mod->module_core + mod->core_text_size,
+                       mod->module_core + mod->core_size,
+                       set_memory_x);
+               set_page_attributes(mod->module_core,
+                       mod->module_core + mod->core_ro_size,
+                       set_memory_rw);
        } else if (mod->module_init == module_region) {
-               /* Set init as NX+RW */
-               total_pages = MOD_NUMBER_OF_PAGES(mod->module_init, mod->init_size);
-               set_memory_nx((unsigned long)mod->module_init, total_pages);
-               set_memory_rw((unsigned long)mod->module_init, total_pages);
+               set_page_attributes(mod->module_init + mod->init_text_size,
+                       mod->module_init + mod->init_size,
+                       set_memory_x);
+               set_page_attributes(mod->module_init,
+                       mod->module_init + mod->init_ro_size,
+                       set_memory_rw);
        }
 }