// SPDX-License-Identifier: GPL-2.0
 // Copyright (C) 2020 ARM Limited
 
+#define _GNU_SOURCE
+
 #include <sys/auxv.h>
 #include <sys/types.h>
 #include <sys/wait.h>
 #include <signal.h>
 #include <setjmp.h>
+#include <sched.h>
 
 #include "../../kselftest_harness.h"
 #include "helper.h"
  * The VA space size is 48 bits. Bigger is opt-in.
  */
 #define PAC_MASK (~0xff80ffffffffffff)
+#define ARBITRARY_VALUE (0x1234)
 #define ASSERT_PAUTH_ENABLED() \
 do { \
        unsigned long hwcaps = getauxval(AT_HWCAP); \
        return res;
 }
 
+int n_same_single_set(struct signatures *sign, int nkeys)
+{
+       size_t vals[nkeys];
+       int same = 0;
+
+       vals[0] = sign->keyia & PAC_MASK;
+       vals[1] = sign->keyib & PAC_MASK;
+       vals[2] = sign->keyda & PAC_MASK;
+       vals[3] = sign->keydb & PAC_MASK;
+
+       if (nkeys >= 4)
+               vals[4] = sign->keyg & PAC_MASK;
+
+       for (int i = 0; i < nkeys - 1; i++) {
+               for (int j = i + 1; j < nkeys; j++) {
+                       if (vals[i] == vals[j])
+                               same += 1;
+               }
+       }
+       return same;
+}
+
 int exec_sign_all(struct signatures *signed_vals, size_t val)
 {
        int new_stdin[2];
        int new_stdout[2];
        int status;
+       int i;
        ssize_t ret;
        pid_t pid;
+       cpu_set_t mask;
 
        ret = pipe(new_stdin);
        if (ret == -1) {
                return -1;
        }
 
+       /*
+        * pin this process and all its children to a single CPU, so it can also
+        * guarantee a context switch with its child
+        */
+       sched_getaffinity(0, sizeof(mask), &mask);
+
+       for (i = 0; i < sizeof(cpu_set_t); i++)
+               if (CPU_ISSET(i, &mask))
+                       break;
+
+       CPU_ZERO(&mask);
+       CPU_SET(i, &mask);
+       sched_setaffinity(0, sizeof(mask), &mask);
+
        pid = fork();
        // child
        if (pid == 0) {
        ASSERT_NE(0, keyg)  TH_LOG("keyg instructions did nothing");
 }
 
+TEST(single_thread_different_keys)
+{
+       int same = 10;
+       int nkeys = NKEYS;
+       int tmp;
+       struct signatures signed_vals;
+       unsigned long hwcaps = getauxval(AT_HWCAP);
+
+       /* generic and data key instructions are not in NOP space. This prevents a SIGILL */
+       ASSERT_NE(0, hwcaps & HWCAP_PACA) TH_LOG("PAUTH not enabled");
+       if (!(hwcaps & HWCAP_PACG)) {
+               TH_LOG("WARNING: Generic PAUTH not enabled. Skipping generic key checks");
+               nkeys = NKEYS - 1;
+       }
+
+       /*
+        * In Linux the PAC field can be up to 7 bits wide. Even if keys are
+        * different, there is about 5% chance for PACs to collide with
+        * different addresses. This chance rapidly increases with fewer bits
+        * allocated for the PAC (e.g. wider address). A comparison of the keys
+        * directly will be more reliable.
+        * All signed values need to be different at least once out of n
+        * attempts to be certain that the keys are different
+        */
+       for (int i = 0; i < PAC_COLLISION_ATTEMPTS; i++) {
+               if (nkeys == NKEYS)
+                       sign_all(&signed_vals, i);
+               else
+                       sign_specific(&signed_vals, i);
+
+               tmp = n_same_single_set(&signed_vals, nkeys);
+               if (tmp < same)
+                       same = tmp;
+       }
+
+       ASSERT_EQ(0, same) TH_LOG("%d keys clashed every time", same);
+}
+
 /*
  * fork() does not change keys. Only exec() does so call a worker program.
  * Its only job is to sign a value and report back the resutls
        ASSERT_EQ(0, same) TH_LOG("exec() did not change %d keys", same);
 }
 
+TEST(context_switch_keep_keys)
+{
+       int ret;
+       struct signatures trash;
+       struct signatures before;
+       struct signatures after;
+
+       ASSERT_PAUTH_ENABLED();
+
+       sign_specific(&before, ARBITRARY_VALUE);
+
+       /* will context switch with a process with different keys at least once */
+       ret = exec_sign_all(&trash, ARBITRARY_VALUE);
+       ASSERT_EQ(0, ret) TH_LOG("failed to run worker");
+
+       sign_specific(&after, ARBITRARY_VALUE);
+
+       ASSERT_EQ(before.keyia, after.keyia) TH_LOG("keyia changed after context switching");
+       ASSERT_EQ(before.keyib, after.keyib) TH_LOG("keyib changed after context switching");
+       ASSERT_EQ(before.keyda, after.keyda) TH_LOG("keyda changed after context switching");
+       ASSERT_EQ(before.keydb, after.keydb) TH_LOG("keydb changed after context switching");
+}
+
+TEST(context_switch_keep_keys_generic)
+{
+       int ret;
+       struct signatures trash;
+       size_t before;
+       size_t after;
+
+       ASSERT_GENERIC_PAUTH_ENABLED();
+
+       before = keyg_sign(ARBITRARY_VALUE);
+
+       /* will context switch with a process with different keys at least once */
+       ret = exec_sign_all(&trash, ARBITRARY_VALUE);
+       ASSERT_EQ(0, ret) TH_LOG("failed to run worker");
+
+       after = keyg_sign(ARBITRARY_VALUE);
+
+       ASSERT_EQ(before, after) TH_LOG("keyg changed after context switching");
+}
+
 TEST_HARNESS_MAIN