#include <linux/jump_label.h>
 
 DECLARE_STATIC_KEY_TRUE(pkey_disabled);
-#define ARCH_VM_PKEY_FLAGS 0
+extern int pkeys_total; /* total pkeys as per device tree */
+extern u32 initial_allocation_mask; /* bits set for reserved keys */
+
+/*
+ * Define these here temporarily so we're not dependent on patching linux/mm.h.
+ * Once it's updated we can drop these.
+ */
+#ifndef VM_PKEY_BIT0
+# define VM_PKEY_SHIFT VM_HIGH_ARCH_BIT_0
+# define VM_PKEY_BIT0  VM_HIGH_ARCH_0
+# define VM_PKEY_BIT1  VM_HIGH_ARCH_1
+# define VM_PKEY_BIT2  VM_HIGH_ARCH_2
+# define VM_PKEY_BIT3  VM_HIGH_ARCH_3
+# define VM_PKEY_BIT4  VM_HIGH_ARCH_4
+#endif
+
+#define ARCH_VM_PKEY_FLAGS (VM_PKEY_BIT0 | VM_PKEY_BIT1 | VM_PKEY_BIT2 | \
+                           VM_PKEY_BIT3 | VM_PKEY_BIT4)
+
+#define arch_max_pkey() pkeys_total
+
+#define pkey_alloc_mask(pkey) (0x1 << pkey)
+
+#define mm_pkey_allocation_map(mm) (mm->context.pkey_allocation_map)
+
+#define __mm_pkey_allocated(mm, pkey) {        \
+       mm_pkey_allocation_map(mm) |= pkey_alloc_mask(pkey); \
+}
+
+#define __mm_pkey_free(mm, pkey) {     \
+       mm_pkey_allocation_map(mm) &= ~pkey_alloc_mask(pkey);   \
+}
+
+#define __mm_pkey_is_allocated(mm, pkey)       \
+       (mm_pkey_allocation_map(mm) & pkey_alloc_mask(pkey))
+
+#define __mm_pkey_is_reserved(pkey) (initial_allocation_mask & \
+                                      pkey_alloc_mask(pkey))
 
 static inline bool mm_pkey_is_allocated(struct mm_struct *mm, int pkey)
 {
-       return false;
+       /* A reserved key is never considered as 'explicitly allocated' */
+       return ((pkey < arch_max_pkey()) &&
+               !__mm_pkey_is_reserved(pkey) &&
+               __mm_pkey_is_allocated(mm, pkey));
 }
 
+/*
+ * Returns a positive, 5-bit key on success, or -1 on failure.
+ * Relies on the mmap_sem to protect against concurrency in mm_pkey_alloc() and
+ * mm_pkey_free().
+ */
 static inline int mm_pkey_alloc(struct mm_struct *mm)
 {
-       return -1;
+       /*
+        * Note: this is the one and only place we make sure that the pkey is
+        * valid as far as the hardware is concerned. The rest of the kernel
+        * trusts that only good, valid pkeys come out of here.
+        */
+       u32 all_pkeys_mask = (u32)(~(0x0));
+       int ret;
+
+       if (static_branch_likely(&pkey_disabled))
+               return -1;
+
+       /*
+        * Are we out of pkeys? We must handle this specially because ffz()
+        * behavior is undefined if there are no zeros.
+        */
+       if (mm_pkey_allocation_map(mm) == all_pkeys_mask)
+               return -1;
+
+       ret = ffz((u32)mm_pkey_allocation_map(mm));
+       __mm_pkey_allocated(mm, ret);
+       return ret;
 }
 
 static inline int mm_pkey_free(struct mm_struct *mm, int pkey)
 {
-       return -EINVAL;
+       if (static_branch_likely(&pkey_disabled))
+               return -1;
+
+       if (!mm_pkey_is_allocated(mm, pkey))
+               return -EINVAL;
+
+       __mm_pkey_free(mm, pkey);
+
+       return 0;
 }
 
 /*
 {
        return 0;
 }
+
+extern void pkey_mm_init(struct mm_struct *mm);
 #endif /*_ASM_POWERPC_KEYS_H */
 
 
 DEFINE_STATIC_KEY_TRUE(pkey_disabled);
 bool pkey_execute_disable_supported;
+int  pkeys_total;              /* Total pkeys as per device tree */
+u32  initial_allocation_mask;  /* Bits set for reserved keys */
 
 int pkey_initialize(void)
 {
+       int os_reserved, i;
+
        /*
         * Disable the pkey system till everything is in place. A subsequent
         * patch will enable it.
         */
        static_branch_enable(&pkey_disabled);
 
+       /* Lets assume 32 keys */
+       pkeys_total = 32;
+
+       /*
+        * Adjust the upper limit, based on the number of bits supported by
+        * arch-neutral code.
+        */
+       pkeys_total = min_t(int, pkeys_total,
+                       (ARCH_VM_PKEY_FLAGS >> VM_PKEY_SHIFT));
+
        /*
         * Disable execute_disable support for now. A subsequent patch will
         * enable it.
         */
        pkey_execute_disable_supported = false;
+
+#ifdef CONFIG_PPC_4K_PAGES
+       /*
+        * The OS can manage only 8 pkeys due to its inability to represent them
+        * in the Linux 4K PTE.
+        */
+       os_reserved = pkeys_total - 8;
+#else
+       os_reserved = 0;
+#endif
+       /*
+        * Bits are in LE format. NOTE: 1, 0 are reserved.
+        * key 0 is the default key, which allows read/write/execute.
+        * key 1 is recommended not to be used. PowerISA(3.0) page 1015,
+        * programming note.
+        */
+       initial_allocation_mask = ~0x0;
+       for (i = 2; i < (pkeys_total - os_reserved); i++)
+               initial_allocation_mask &= ~(0x1 << i);
        return 0;
 }
 
 arch_initcall(pkey_initialize);
+
+void pkey_mm_init(struct mm_struct *mm)
+{
+       if (static_branch_likely(&pkey_disabled))
+               return;
+       mm_pkey_allocation_map(mm) = initial_allocation_mask;
+}