2 files changed, 51 insertions, 12 deletions

diff --git a/sys/amd64/amd64/pmap.c b/sys/amd64/amd64/pmap.c
index 956e9c5e78d2..497c85c3f0c2 100644
--- a/sys/amd64/amd64/pmap.c
+++ b/sys/amd64/amd64/pmap.c
@@ -11551,7 +11551,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;
@@ -11577,6 +11577,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)
@@ -11629,8 +11645,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 51f55687bbcf..1df73a25c05e 100644
--- a/sys/amd64/amd64/sys_machdep.c
+++ b/sys/amd64/amd64/sys_machdep.c
@@ -30,7 +30,6 @@
  * SUCH DAMAGE.
  */
 
-#include <sys/cdefs.h>
 #include "opt_capsicum.h"
 #include "opt_ktrace.h"
 
@@ -369,32 +368,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);