[RFC PATCH 6/6] livepatch: Support replace_set in shadow variable API

[Yafang Shao <laoar.shao@gmail.com>]

To support more complex livepatching scenarios where multiple
replacement sets might coexist, extend the klp_shadow API to
include a 'replace_set' identifier.

To maintain compatibility with the existing 64-bit storage in
'struct klp_shadow', the internal @id is now treated as a composite
value. The 64-bit identifier is constructed by packing two 32-bit
values:

  MSB (63-32)          LSB (31-0)
  +--------------------+--------------------+
  |    replace_set     |    original @id    |
  +--------------------+--------------------+

Signed-off-by: Yafang Shao <laoar.shao@gmail.com>
---
 include/linux/livepatch.h | 12 ++++---
 kernel/livepatch/shadow.c | 70 ++++++++++++++++++++++++---------------
 kernel/livepatch/state.c  |  3 +-
 3 files changed, 52 insertions(+), 33 deletions(-)

diff --git a/include/linux/livepatch.h b/include/linux/livepatch.h
index 221f176f1f51..2dd9fca8c01c 100644
--- a/include/linux/livepatch.h
+++ b/include/linux/livepatch.h
@@ -195,15 +195,17 @@ static inline bool klp_have_reliable_stack(void)
 	       IS_ENABLED(CONFIG_HAVE_RELIABLE_STACKTRACE);
 }
 
-void *klp_shadow_get(void *obj, unsigned long id);
-void *klp_shadow_alloc(void *obj, unsigned long id,
+void *klp_shadow_get(void *obj, unsigned int replace_set, unsigned int id);
+void *klp_shadow_alloc(void *obj, unsigned int replace_set, unsigned int id,
 		       size_t size, gfp_t gfp_flags,
 		       klp_shadow_ctor_t ctor, void *ctor_data);
-void *klp_shadow_get_or_alloc(void *obj, unsigned long id,
+void *klp_shadow_get_or_alloc(void *obj, unsigned int replace_set, unsigned int id,
 			      size_t size, gfp_t gfp_flags,
 			      klp_shadow_ctor_t ctor, void *ctor_data);
-void klp_shadow_free(void *obj, unsigned long id, klp_shadow_dtor_t dtor);
-void klp_shadow_free_all(unsigned long id, klp_shadow_dtor_t dtor);
+void klp_shadow_free(void *obj, unsigned int replace_set, unsigned int id,
+		     klp_shadow_dtor_t dtor);
+void klp_shadow_free_all(unsigned int replace_set, unsigned int id,
+			 klp_shadow_dtor_t dtor);
 
 struct klp_state *klp_get_state(struct klp_patch *patch, unsigned long id);
 struct klp_state *klp_get_prev_state(unsigned long id);
diff --git a/kernel/livepatch/shadow.c b/kernel/livepatch/shadow.c
index c2e724d97ddf..35e507fae445 100644
--- a/kernel/livepatch/shadow.c
+++ b/kernel/livepatch/shadow.c
@@ -48,7 +48,8 @@ static DEFINE_SPINLOCK(klp_shadow_lock);
  * @node:	klp_shadow_hash hash table node
  * @rcu_head:	RCU is used to safely free this structure
  * @obj:	pointer to parent object
- * @id:		data identifier
+ * @id:		combined data identifier
+ *		higher 32 bits: replace_set, lower 32 bits: resource ID
  * @data:	data area
  */
 struct klp_shadow {
@@ -59,6 +60,11 @@ struct klp_shadow {
 	char data[];
 };
 
+static unsigned long klp_shadow_combined_id(unsigned int set, unsigned int id)
+{
+	return ((unsigned long)set << 32) | id;
+}
+
 /**
  * klp_shadow_match() - verify a shadow variable matches given <obj, id>
  * @shadow:	shadow variable to match
@@ -76,11 +82,12 @@ static inline bool klp_shadow_match(struct klp_shadow *shadow, void *obj,
 /**
  * klp_shadow_get() - retrieve a shadow variable data pointer
  * @obj:	pointer to parent object
+ * @replace_set:identifier for the livepatch replacement set
  * @id:		data identifier
  *
  * Return: the shadow variable data element, NULL on failure.
  */
-void *klp_shadow_get(void *obj, unsigned long id)
+void *klp_shadow_get(void *obj, unsigned int replace_set, unsigned int id)
 {
 	struct klp_shadow *shadow;
 
@@ -89,7 +96,8 @@ void *klp_shadow_get(void *obj, unsigned long id)
 	hash_for_each_possible_rcu(klp_shadow_hash, shadow, node,
 				   (unsigned long)obj) {
 
-		if (klp_shadow_match(shadow, obj, id)) {
+		if (klp_shadow_match(shadow, obj,
+				     klp_shadow_combined_id(replace_set, id))) {
 			rcu_read_unlock();
 			return shadow->data;
 		}
@@ -101,7 +109,7 @@ void *klp_shadow_get(void *obj, unsigned long id)
 }
 EXPORT_SYMBOL_GPL(klp_shadow_get);
 
-static void *__klp_shadow_get_or_alloc(void *obj, unsigned long id,
+static void *__klp_shadow_get_or_alloc(void *obj, unsigned int set, unsigned int id,
 				       size_t size, gfp_t gfp_flags,
 				       klp_shadow_ctor_t ctor, void *ctor_data,
 				       bool warn_on_exist)
@@ -111,7 +119,7 @@ static void *__klp_shadow_get_or_alloc(void *obj, unsigned long id,
 	unsigned long flags;
 
 	/* Check if the shadow variable already exists */
-	shadow_data = klp_shadow_get(obj, id);
+	shadow_data = klp_shadow_get(obj, set, id);
 	if (shadow_data)
 		goto exists;
 
@@ -126,7 +134,7 @@ static void *__klp_shadow_get_or_alloc(void *obj, unsigned long id,
 
 	/* Look for <obj, id> again under the lock */
 	spin_lock_irqsave(&klp_shadow_lock, flags);
-	shadow_data = klp_shadow_get(obj, id);
+	shadow_data = klp_shadow_get(obj, set, id);
 	if (unlikely(shadow_data)) {
 		/*
 		 * Shadow variable was found, throw away speculative
@@ -147,8 +155,8 @@ static void *__klp_shadow_get_or_alloc(void *obj, unsigned long id,
 		if (err) {
 			spin_unlock_irqrestore(&klp_shadow_lock, flags);
 			kfree(new_shadow);
-			pr_err("Failed to construct shadow variable <%p, %lx> (%d)\n",
-			       obj, id, err);
+			pr_err("Failed to construct shadow variable <%p, %x, %x> (%d)\n",
+			       obj, set, id, err);
 			return NULL;
 		}
 	}
@@ -162,7 +170,7 @@ static void *__klp_shadow_get_or_alloc(void *obj, unsigned long id,
 
 exists:
 	if (warn_on_exist) {
-		WARN(1, "Duplicate shadow variable <%p, %lx>\n", obj, id);
+		WARN(1, "Duplicate shadow variable <%p, %x, %x>\n", obj, set, id);
 		return NULL;
 	}
 
@@ -172,6 +180,7 @@ static void *__klp_shadow_get_or_alloc(void *obj, unsigned long id,
 /**
  * klp_shadow_alloc() - allocate and add a new shadow variable
  * @obj:	pointer to parent object
+ * @replace_set:identifier for the livepatch replacement set
  * @id:		data identifier
  * @size:	size of attached data
  * @gfp_flags:	GFP mask for allocation
@@ -183,8 +192,8 @@ static void *__klp_shadow_get_or_alloc(void *obj, unsigned long id,
  * function if it is not NULL.  The new shadow variable is then added
  * to the global hashtable.
  *
- * If an existing <obj, id> shadow variable can be found, this routine will
- * issue a WARN, exit early and return NULL.
+ * If an existing <obj, replace_set, id> shadow variable can be found, this
+ * routine will issue a WARN, exit early and return NULL.
  *
  * This function guarantees that the constructor function is called only when
  * the variable did not exist before.  The cost is that @ctor is called
@@ -193,11 +202,11 @@ static void *__klp_shadow_get_or_alloc(void *obj, unsigned long id,
  * Return: the shadow variable data element, NULL on duplicate or
  * failure.
  */
-void *klp_shadow_alloc(void *obj, unsigned long id,
+void *klp_shadow_alloc(void *obj, unsigned int replace_set, unsigned int id,
 		       size_t size, gfp_t gfp_flags,
 		       klp_shadow_ctor_t ctor, void *ctor_data)
 {
-	return __klp_shadow_get_or_alloc(obj, id, size, gfp_flags,
+	return __klp_shadow_get_or_alloc(obj, replace_set, id, size, gfp_flags,
 					 ctor, ctor_data, true);
 }
 EXPORT_SYMBOL_GPL(klp_shadow_alloc);
@@ -205,28 +214,29 @@ EXPORT_SYMBOL_GPL(klp_shadow_alloc);
 /**
  * klp_shadow_get_or_alloc() - get existing or allocate a new shadow variable
  * @obj:	pointer to parent object
+ * @replace_set:identifier for the livepatch replacement set
  * @id:		data identifier
  * @size:	size of attached data
  * @gfp_flags:	GFP mask for allocation
  * @ctor:	custom constructor to initialize the shadow data (optional)
  * @ctor_data:	pointer to any data needed by @ctor (optional)
  *
- * Returns a pointer to existing shadow data if an <obj, id> shadow
+ * Returns a pointer to existing shadow data if an <obj, replace_set, id> shadow
  * variable is already present.  Otherwise, it creates a new shadow
  * variable like klp_shadow_alloc().
  *
  * This function guarantees that only one shadow variable exists with the given
- * @id for the given @obj.  It also guarantees that the constructor function
- * will be called only when the variable did not exist before.  The cost is
- * that @ctor is called in atomic context under a spin lock.
+ * @id for the given @obj within the same replace_set.  It also guarantees that
+ * the constructor function will be called only when the variable did not exist
+ * before.  The cost is that @ctor is called in atomic context under a spin lock.
  *
  * Return: the shadow variable data element, NULL on failure.
  */
-void *klp_shadow_get_or_alloc(void *obj, unsigned long id,
-			      size_t size, gfp_t gfp_flags,
+void *klp_shadow_get_or_alloc(void *obj, unsigned int replace_set,
+			      unsigned int id, size_t size, gfp_t gfp_flags,
 			      klp_shadow_ctor_t ctor, void *ctor_data)
 {
-	return __klp_shadow_get_or_alloc(obj, id, size, gfp_flags,
+	return __klp_shadow_get_or_alloc(obj, replace_set, id, size, gfp_flags,
 					 ctor, ctor_data, false);
 }
 EXPORT_SYMBOL_GPL(klp_shadow_get_or_alloc);
@@ -243,14 +253,16 @@ static void klp_shadow_free_struct(struct klp_shadow *shadow,
 /**
  * klp_shadow_free() - detach and free a <obj, id> shadow variable
  * @obj:	pointer to parent object
+ * @replace_set:identifier for the livepatch replacement set
  * @id:		data identifier
  * @dtor:	custom callback that can be used to unregister the variable
  *		and/or free data that the shadow variable points to (optional)
  *
- * This function releases the memory for this <obj, id> shadow variable
+ * This function releases the memory for this <obj, replace_set, id> shadow variable
  * instance, callers should stop referencing it accordingly.
  */
-void klp_shadow_free(void *obj, unsigned long id, klp_shadow_dtor_t dtor)
+void klp_shadow_free(void *obj, unsigned int replace_set, unsigned int id,
+		     klp_shadow_dtor_t dtor)
 {
 	struct klp_shadow *shadow;
 	unsigned long flags;
@@ -261,7 +273,8 @@ void klp_shadow_free(void *obj, unsigned long id, klp_shadow_dtor_t dtor)
 	hash_for_each_possible(klp_shadow_hash, shadow, node,
 			       (unsigned long)obj) {
 
-		if (klp_shadow_match(shadow, obj, id)) {
+		if (klp_shadow_match(shadow, obj,
+				     klp_shadow_combined_id(replace_set, id))) {
 			klp_shadow_free_struct(shadow, dtor);
 			break;
 		}
@@ -272,15 +285,17 @@ void klp_shadow_free(void *obj, unsigned long id, klp_shadow_dtor_t dtor)
 EXPORT_SYMBOL_GPL(klp_shadow_free);
 
 /**
- * klp_shadow_free_all() - detach and free all <_, id> shadow variables
+ * klp_shadow_free_all() - detach and free all <_, replace_set, id> shadow variables
+ * @replace_set:identifier for the livepatch replacement set
  * @id:		data identifier
  * @dtor:	custom callback that can be used to unregister the variable
  *		and/or free data that the shadow variable points to (optional)
  *
- * This function releases the memory for all <_, id> shadow variable
+ * This function releases the memory for all <_, replace_set, id> shadow variable
  * instances, callers should stop referencing them accordingly.
  */
-void klp_shadow_free_all(unsigned long id, klp_shadow_dtor_t dtor)
+void klp_shadow_free_all(unsigned int replace_set, unsigned int id,
+			 klp_shadow_dtor_t dtor)
 {
 	struct klp_shadow *shadow;
 	unsigned long flags;
@@ -290,7 +305,8 @@ void klp_shadow_free_all(unsigned long id, klp_shadow_dtor_t dtor)
 
 	/* Delete all <_, id> from hash */
 	hash_for_each(klp_shadow_hash, i, shadow, node) {
-		if (klp_shadow_match(shadow, shadow->obj, id))
+		if (klp_shadow_match(shadow, shadow->obj,
+				     klp_shadow_combined_id(replace_set, id)))
 			klp_shadow_free_struct(shadow, dtor);
 	}
 
diff --git a/kernel/livepatch/state.c b/kernel/livepatch/state.c
index 43115e8e8453..6e3d6fb92e64 100644
--- a/kernel/livepatch/state.c
+++ b/kernel/livepatch/state.c
@@ -203,7 +203,8 @@ void klp_states_post_unpatch(struct klp_patch *patch)
 			state->callbacks.post_unpatch(patch, state);
 
 		if (state->is_shadow)
-			klp_shadow_free_all(state->id, state->callbacks.shadow_dtor);
+			klp_shadow_free_all(patch->replace_set, state->id,
+					    state->callbacks.shadow_dtor);
 
 		state->callbacks.pre_patch_succeeded = 0;
 	}
-- 
2.47.3