Sometimes a 64-bit task might need to have a shadow stack allocated from
within 32-bit address range.  One example is selftests/x86/sigreturn.

Currently arch_prctl(ARCH_X86_CET_ALLOC_SHSTK) takes a input parameter for
the desired shadow stack size.  Modify it and use bit[0] of the parameter
to indicate the desire to allocate from 32-bit address range.

Signed-off-by: Yu-cheng Yu <redacted>
---
 arch/x86/include/asm/cet.h        |  2 +-
 arch/x86/include/uapi/asm/prctl.h |  2 ++
 arch/x86/kernel/cet.c             | 19 ++++++++++++-------
 arch/x86/kernel/cet_prctl.c       |  6 +++++-
 4 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/arch/x86/include/asm/cet.h b/arch/x86/include/asm/cet.h
index f163c805a559..534b02785a39 100644
--- a/arch/x86/include/asm/cet.h
+++ b/arch/x86/include/asm/cet.h

@@ -22,7 +22,7 @@ struct cet_status {
 int prctl_cet(int option, u64 arg2);
 int cet_setup_shstk(void);
 int cet_setup_thread_shstk(struct task_struct *p);
-int cet_alloc_shstk(unsigned long *arg);
+int cet_alloc_shstk(unsigned long *arg, int map_32bit);
 void cet_disable_free_shstk(struct task_struct *p);
 int cet_verify_rstor_token(bool ia32, unsigned long ssp, unsigned long *new_ssp);
 void cet_restore_signal(struct sc_ext *sc);

diff --git a/arch/x86/include/uapi/asm/prctl.h b/arch/x86/include/uapi/asm/prctl.h
index d962f0ec9ccf..e254c6a21475 100644
--- a/arch/x86/include/uapi/asm/prctl.h
+++ b/arch/x86/include/uapi/asm/prctl.h

@@ -19,4 +19,6 @@
 #define ARCH_X86_CET_LOCK		0x3003
 #define ARCH_X86_CET_ALLOC_SHSTK	0x3004
 
+#define ARCH_X86_CET_ALLOC_SHSTK_32BIT	0x1UL

+
 #endif /* _ASM_X86_PRCTL_H */

diff --git a/arch/x86/kernel/cet.c b/arch/x86/kernel/cet.c
index 92b8730c0b08..d6f93e1864b2 100644
--- a/arch/x86/kernel/cet.c
+++ b/arch/x86/kernel/cet.c

@@ -57,14 +57,19 @@ static unsigned long cet_get_shstk_addr(void)
 	return ssp;
 }
 
-static unsigned long alloc_shstk(unsigned long size)
+static unsigned long alloc_shstk(unsigned long size, int map_32bit)
 {
 	struct mm_struct *mm = current->mm;
 	unsigned long addr, populate;
+	unsigned long map_flags;
+
+	map_flags = MAP_ANONYMOUS | MAP_PRIVATE;
+	if (map_32bit)
+		map_flags |= MAP_32BIT;
 
 	down_write(&mm->mmap_sem);
-	addr = do_mmap(NULL, 0, size, PROT_READ, MAP_ANONYMOUS | MAP_PRIVATE,
-		       VM_SHSTK, 0, &populate, NULL);
+	addr = do_mmap(NULL, 0, size, PROT_READ, map_flags, VM_SHSTK, 0,
+		       &populate, NULL);
 	up_write(&mm->mmap_sem);
 
 	if (populate)

@@ -147,14 +152,14 @@ static int create_rstor_token(bool ia32, unsigned long ssp,
 	return 0;
 }
 
-int cet_alloc_shstk(unsigned long *arg)
+int cet_alloc_shstk(unsigned long *arg, int map_32bit)
 {
 	unsigned long len = *arg;
 	unsigned long addr;
 	unsigned long token;
 	unsigned long ssp;
 
-	addr = alloc_shstk(round_up(len, PAGE_SIZE));
+	addr = alloc_shstk(round_up(len, PAGE_SIZE), map_32bit);
 
 	if (IS_ERR((void *)addr))
 		return PTR_ERR((void *)addr);

@@ -185,7 +190,7 @@ int cet_setup_shstk(void)
 		return -EOPNOTSUPP;
 
 	size = round_up(min(rlimit(RLIMIT_STACK), 1UL << 32), PAGE_SIZE);
-	addr = alloc_shstk(size);
+	addr = alloc_shstk(size, 0);
 
 	if (IS_ERR((void *)addr))
 		return PTR_ERR((void *)addr);

@@ -226,7 +231,7 @@ int cet_setup_thread_shstk(struct task_struct *tsk)
 	if (in_compat_syscall())
 		size /= 4;
 	size = round_up(size, PAGE_SIZE);
-	addr = alloc_shstk(size);
+	addr = alloc_shstk(size, 0);
 
 	if (IS_ERR((void *)addr)) {
 		cet->shstk_base = 0;

diff --git a/arch/x86/kernel/cet_prctl.c b/arch/x86/kernel/cet_prctl.c
index a8e68fefd524..364ed2420202 100644
--- a/arch/x86/kernel/cet_prctl.c
+++ b/arch/x86/kernel/cet_prctl.c

@@ -35,12 +35,16 @@ static int handle_alloc_shstk(u64 arg2)
 	unsigned long arg;
 	unsigned long addr = 0;
 	unsigned long size = 0;
+	int map_32bit;
 
 	if (get_user(arg, (unsigned long __user *)arg2))
 		return -EFAULT;