4 files changed, 84 insertions, 12 deletions

diff --git a/sys/amd64/amd64/pmap.c b/sys/amd64/amd64/pmap.c
index 9fd6c4091c95..bae3735eeafe 100644
--- a/sys/amd64/amd64/pmap.c
+++ b/sys/amd64/amd64/pmap.c
@@ -11598,7 +11598,7 @@ pmap_pkru_update_range(pmap_t pmap, vm_offset_t sva, vm_offset_t eva,
     u_int keyidx)
 {
 	pml4_entry_t *pml4e;
-	pdp_entry_t *pdpe;
+	pdp_entry_t newpdpe, *pdpe;
 	pd_entry_t newpde, ptpaddr, *pde;
 	pt_entry_t newpte, *ptep, pte;
 	vm_offset_t va, va_next;
@@ -11624,6 +11624,22 @@ pmap_pkru_update_range(pmap_t pmap, vm_offset_t sva, vm_offset_t eva,
 				va_next = eva;
 			continue;
 		}
+		if ((*pdpe & PG_PS) != 0) {
+			va_next = (va + NBPDP) & ~PDPMASK;
+			if (va_next < va)
+				va_next = eva;
+			KASSERT(va_next <= eva,
+			    ("partial update of non-transparent 1G mapping "
+			    "pdpe %#lx va %#lx eva %#lx va_next %#lx",
+			    *pdpe, va, eva, va_next));
+			newpdpe = (*pdpe & ~X86_PG_PKU_MASK) |
+			    X86_PG_PKU(keyidx);
+			if (newpdpe != *pdpe) {
+				*pdpe = newpdpe;
+				changed = true;
+			}
+			continue;
+		}
 
 		va_next = (va + NBPDR) & ~PDRMASK;
 		if (va_next < va)
@@ -11676,8 +11692,6 @@ pmap_pkru_check_uargs(pmap_t pmap, vm_offset_t sva, vm_offset_t eva,
 	if (pmap->pm_type != PT_X86 || keyidx > PMAP_MAX_PKRU_IDX ||
 	    (flags & ~(AMD64_PKRU_PERSIST | AMD64_PKRU_EXCL)) != 0)
 		return (EINVAL);
-	if (eva <= sva || eva > VM_MAXUSER_ADDRESS)
-		return (EFAULT);
 	if ((cpu_stdext_feature2 & CPUID_STDEXT2_PKU) == 0)
 		return (ENOTSUP);
 	return (0);
diff --git a/sys/amd64/amd64/sys_machdep.c b/sys/amd64/amd64/sys_machdep.c
index a7ee1f49887c..74b182e73dc2 100644
--- a/sys/amd64/amd64/sys_machdep.c
+++ b/sys/amd64/amd64/sys_machdep.c
@@ -32,7 +32,6 @@
  *	from: @(#)sys_machdep.c	5.5 (Berkeley) 1/19/91
  */
 
-#include <sys/cdefs.h>
 #include "opt_capsicum.h"
 #include "opt_ktrace.h"
 
@@ -371,32 +370,58 @@ sysarch(struct thread *td, struct sysarch_args *uap)
 		break;
 
 	case I386_SET_PKRU:
-	case AMD64_SET_PKRU:
+	case AMD64_SET_PKRU: {
+		vm_offset_t addr, start, end;
+		vm_size_t len;
+
+		addr = (uintptr_t)a64pkru.addr;
+		len = a64pkru.len;
+
 		/*
 		 * Read-lock the map to synchronize with parallel
 		 * pmap_vmspace_copy() on fork.
 		 */
 		map = &td->td_proc->p_vmspace->vm_map;
 		vm_map_lock_read(map);
-		error = pmap_pkru_set(PCPU_GET(curpmap),
-		    (vm_offset_t)a64pkru.addr, (vm_offset_t)a64pkru.addr +
-		    a64pkru.len, a64pkru.keyidx, a64pkru.flags);
+		if (len == 0 || !vm_map_check_boundary(map, addr, addr + len)) {
+			vm_map_unlock_read(map);
+			error = EINVAL;
+			break;
+		}
+		start = trunc_page(addr);
+		end = round_page(addr + len);
+		error = pmap_pkru_set(PCPU_GET(curpmap), start, end,
+		    a64pkru.keyidx, a64pkru.flags);
 		vm_map_unlock_read(map);
 		break;
+	}
 
 	case I386_CLEAR_PKRU:
-	case AMD64_CLEAR_PKRU:
+	case AMD64_CLEAR_PKRU: {
+		vm_offset_t addr, start, end;
+		vm_size_t len;
+
 		if (a64pkru.flags != 0 || a64pkru.keyidx != 0) {
 			error = EINVAL;
 			break;
 		}
+
+		addr = (uintptr_t)a64pkru.addr;
+		len = a64pkru.len;
+
 		map = &td->td_proc->p_vmspace->vm_map;
 		vm_map_lock_read(map);
-		error = pmap_pkru_clear(PCPU_GET(curpmap),
-		    (vm_offset_t)a64pkru.addr,
-		    (vm_offset_t)a64pkru.addr + a64pkru.len);
+		if (len == 0 || !vm_map_check_boundary(map, addr, addr + len)) {
+			vm_map_unlock_read(map);
+			error = EINVAL;
+			break;
+		}
+		start = trunc_page(addr);
+		end = round_page(addr + len);
+		error = pmap_pkru_clear(PCPU_GET(curpmap), start, end);
 		vm_map_unlock_read(map);
 		break;
+	}
 
 	case AMD64_DISABLE_TLSBASE:
 		clear_pcb_flags(pcb, PCB_TLSBASE);
diff --git a/sys/vm/vm_map.c b/sys/vm/vm_map.c
index cf028c0c47ef..631a5e302879 100644
--- a/sys/vm/vm_map.c
+++ b/sys/vm/vm_map.c
@@ -4141,6 +4141,38 @@ vm_map_check_protection(vm_map_t map, vm_offset_t start, vm_offset_t end,
 }
 
 /*
+ * Check whether the specified range partially overlaps a map entry with
+ * fixed boundaries, and return false if so.
+ *
+ * The map must be locked.
+ */
+bool
+vm_map_check_boundary(vm_map_t map, vm_offset_t start, vm_offset_t end)
+{
+	vm_map_entry_t entry;
+	int bdry_idx;
+
+	if (!vm_map_range_valid(map, start, end))
+		return (false);
+	if (start == end)
+		return (true);
+
+	if (vm_map_lookup_entry(map, start, &entry)) {
+		bdry_idx = MAP_ENTRY_SPLIT_BOUNDARY_INDEX(entry);
+		if (bdry_idx != 0 &&
+		    (start & (pagesizes[bdry_idx] - 1)) != 0)
+			return (false);
+	}
+	if (vm_map_lookup_entry(map, end - 1, &entry)) {
+		bdry_idx = MAP_ENTRY_SPLIT_BOUNDARY_INDEX(entry);
+		if (bdry_idx != 0 &&
+		    (end & (pagesizes[bdry_idx] - 1)) != 0)
+			return (false);
+	}
+	return (true);
+}
+
+/*
  *
  *	vm_map_copy_swap_object:
  *
diff --git a/sys/vm/vm_map.h b/sys/vm/vm_map.h
index 5d41f892e83d..66a787b0d2d8 100644
--- a/sys/vm/vm_map.h
+++ b/sys/vm/vm_map.h
@@ -471,6 +471,7 @@ vm_map_entry_read_succ(void *token, struct vm_map_entry *const clone,
 #endif				/* ! _KERNEL */
 
 #ifdef _KERNEL
+bool vm_map_check_boundary(vm_map_t, vm_offset_t, vm_offset_t);
 boolean_t vm_map_check_protection (vm_map_t, vm_offset_t, vm_offset_t, vm_prot_t);
 int vm_map_delete(vm_map_t, vm_offset_t, vm_offset_t);
 int vm_map_find(vm_map_t, vm_object_t, vm_ooffset_t, vm_offset_t *, vm_size_t,