diff --git a/cet/Makefile b/cet/Makefile
index 3a7c22b..5d891ea 100644
--- a/cet/Makefile
+++ b/cet/Makefile
@@ -31,7 +31,7 @@ endif
 all: $(BIN)
 
 shstk_alloc: shstk_alloc.c
-	gcc $(CETFLAGS) $^ -o $@
+	gcc $(NOCETFLAGS) $^ -o $@
 
 test_shadow_stack: test_shadow_stack.c
 	gcc -pthread $(NOCETFLAGS) $^ -o $@
diff --git a/cet/shstk_alloc.c b/cet/shstk_alloc.c
index 52ee623..74ed097 100644
--- a/cet/shstk_alloc.c
+++ b/cet/shstk_alloc.c
@@ -12,6 +12,7 @@
 
 #define _GNU_SOURCE
 
+#include <sys/syscall.h>
 #include <sys/mman.h>
 #include <sys/stat.h>
 #include <stdio.h>
@@ -23,6 +24,16 @@
 #include <stdbool.h>
 #include <x86intrin.h>
 
+/* It's from arch/x86/include/uapi/asm/prctl.h file. */
+#define ARCH_SHSTK_ENABLE	0x5001
+#define ARCH_SHSTK_DISABLE	0x5002
+#define ARCH_SHSTK_LOCK		0x5003
+#define ARCH_SHSTK_UNLOCK	0x5004
+#define ARCH_SHSTK_STATUS	0x5005
+/* ARCH_SHSTK_ features bits */
+#define ARCH_SHSTK_SHSTK		(1ULL <<  0)
+#define ARCH_SHSTK_WRSS			(1ULL <<  1)
+
 #define SHADOW_STACK_SET_TOKEN	0x1
 #ifndef __NR_map_shadow_stack
 #define __NR_map_shadow_stack 451
@@ -30,6 +41,33 @@
 
 size_t shstk_size = 0x200000;
 
+/*
+ * For use in inline enablement of shadow stack.
+ *
+ * The program cannot return from the point where the shadow stack is enabled
+ * because there will be no address on the shadow stack. So it can't use
+ * syscall() for enablement, since it is a function.
+ *
+ * Based on code from nolibc.h. Keep a copy here because this can't pull in all
+ * of nolibc.h.
+ */
+#define ARCH_PRCTL(arg1, arg2)					\
+({								\
+	long _ret;						\
+	register long _num  asm("eax") = __NR_arch_prctl;	\
+	register long _arg1 asm("rdi") = (long)(arg1);		\
+	register long _arg2 asm("rsi") = (long)(arg2);		\
+								\
+	asm volatile (						\
+		"syscall\n"					\
+		: "=a"(_ret)					\
+		: "r"(_arg1), "r"(_arg2),			\
+		  "0"(_num)					\
+		: "rcx", "r11", "memory", "cc"			\
+	);							\
+	_ret;							\
+})
+
 #define get_ssp()						\
 ({								\
 	unsigned long _ret;					\
@@ -87,25 +125,12 @@ int main(int argc, char *argv[])
 	void *shstk;
 	unsigned long ssp, *asm_ssp, *bp;
 
-	if (!get_ssp()) {
-		printf("[SKIP]\tshadow stack was disabled.\n");
-		printf("[INFO]\tMaybe glibc doesn't support SHSTK!\n");
-		return 2;
+	if (ARCH_PRCTL(ARCH_SHSTK_ENABLE, ARCH_SHSTK_SHSTK)) {
+		printf("[FAIL]\tParent process could not enable SHSTK!\n");
+		return 1;
 	}
+	printf("[PASS]\tParent process enable SHSTK.\n");
 
-	ssp = get_ssp();
-	#ifdef __x86_64__
-		asm volatile ("rdsspq %rbx");
-		asm("movq %%rbx,%0" : "=r"(asm_ssp));
-		asm("movq %%rbp,%0" : "=r"(bp));
-	#else
-		asm volatile ("rdsspd %ebx");
-		asm("mov %%ebx,%0" : "=r"(asm_ssp));
-		asm("mov %%ebp,%0" : "=r"(bp));
-	#endif
-	printf("Libc get_ssp -> ssp:%lx, content:%lx\n", ssp, *(long *)ssp);
-	printf("Show: bp+1:%p, *(bp+1):0x%lx, asm_ssp:%p, *asm_ssp:0x%lx\n",
-	       bp + 1, *(bp + 1), asm_ssp, *asm_ssp);
 	shstk = create_shstk();
 	if (shstk == MAP_FAILED) {
 		printf("[FAIL]\tError creating shadow stack: %d\n", errno);
@@ -116,6 +141,14 @@ int main(int argc, char *argv[])
 	try_shstk((unsigned long)shstk + shstk_size - 8);
 
 	printf("[PASS]\tCreated and allocated new shstk addr test.\n");
+
+	/* Disable SHSTK in parent process to avoid segfault issue. */
+	if (ARCH_PRCTL(ARCH_SHSTK_DISABLE, ARCH_SHSTK_SHSTK)) {
+		printf("[FAIL]\tParent process disable shadow stack failed.\n");
+		return 1;
+	} else {
+		printf("[PASS]\tParent process disable shadow stack successfully.\n");
+	}
 	return 0;
 }
 #endif
diff --git a/cet/shstk_unlock_test.c b/cet/shstk_unlock_test.c
index 34d404d..2d4c643 100644
--- a/cet/shstk_unlock_test.c
+++ b/cet/shstk_unlock_test.c
@@ -129,11 +129,19 @@ long unlock_shstk(pid_t pid)
 	}
 
 	if (ptrace(PTRACE_ARCH_PRCTL, pid, ARCH_SHSTK_SHSTK, ARCH_SHSTK_UNLOCK)) {
-		printf("[FAIL]\tCan't unlock CET for %d task", pid);
+		printf("[FAIL]\tCan't unlock user SHSTK for child pid:%d", pid);
 		err_num++;
 		goto detach;
 	} else {
-		printf("[PASS]\tUnlock CET successfully for pid:%d\n", pid);
+		printf("[PASS]\tUnlock SHSTK for child pid:%d\n", pid);
+	}
+
+	if (ptrace(PTRACE_ARCH_PRCTL, pid, ARCH_SHSTK_WRSS, ARCH_SHSTK_UNLOCK)) {
+		printf("[FAIL]\tCan't unlock WRSS for child pid:%d", pid);
+		err_num++;
+		goto detach;
+	} else {
+		printf("[PASS]\tUnlock WRSS for child pid:%d\n", pid);
 	}
 
 	ret = ptrace(PTRACE_GETREGSET, pid, NT_X86_SHSTK, &iov_cet);
@@ -279,7 +287,7 @@ int main(void)
 			printf("[FAIL]\tCould not re-enable Shadow stack.\n");
 			err_num++;
 		} else {
-			printf("[PASS]\tChild process re-enable ssp\n");
+			printf("[PASS]\tChild process re-enable SHSTK\n");
 		}
 
 		ret = ARCH_PRCTL(ARCH_SHSTK_STATUS, &feature);