[PATCH 2/2] liveupdate: Reference count incoming FLB data

[David Matlack <dmatlack@google.com>]

Increment the incoming FLB refcount in liveupdate_flb_get_incoming() so
that the FLB structure cannot be freed while the caller is actively using
it. Add an additional liveupdate_flb_put_incoming() function so the
caller can explicitly indicate when it is done using the FLB data.

During a Live Update, a subsystem might need to hold onto the incoming
File-Lifecycle-Bound (FLB) data for an extended period, such as during
device enumeration. Incrementing the reference count guarantees that the
data remains valid and accessible until the subsystem releases it,
preventing future use-after-free bugs.

Fixes: cab056f2aae7 ("liveupdate: luo_flb: introduce File-Lifecycle-Bound global state")
Signed-off-by: David Matlack <dmatlack@google.com>
---
 include/linux/liveupdate.h  |  6 ++++++
 kernel/liveupdate/luo_flb.c | 32 +++++++++++++++++---------------
 lib/tests/liveupdate.c      |  3 +++
 3 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/include/linux/liveupdate.h b/include/linux/liveupdate.h
index 8d3bbc35c828..88722e5caf02 100644
--- a/include/linux/liveupdate.h
+++ b/include/linux/liveupdate.h
@@ -240,6 +240,8 @@ void liveupdate_unregister_flb(struct liveupdate_file_handler *fh,
 			       struct liveupdate_flb *flb);
 
 int liveupdate_flb_get_incoming(struct liveupdate_flb *flb, void **objp);
+void liveupdate_flb_put_incoming(struct liveupdate_flb *flb);
+
 int liveupdate_flb_get_outgoing(struct liveupdate_flb *flb, void **objp);
 
 #else /* CONFIG_LIVEUPDATE */
@@ -280,6 +282,10 @@ static inline int liveupdate_flb_get_incoming(struct liveupdate_flb *flb,
 	return -EOPNOTSUPP;
 }
 
+static inline void liveupdate_flb_put_incoming(struct liveupdate_flb *flb)
+{
+}
+
 static inline int liveupdate_flb_get_outgoing(struct liveupdate_flb *flb,
 					      void **objp)
 {
diff --git a/kernel/liveupdate/luo_flb.c b/kernel/liveupdate/luo_flb.c
index 59c5f31ab767..8f5c5dd01cd0 100644
--- a/kernel/liveupdate/luo_flb.c
+++ b/kernel/liveupdate/luo_flb.c
@@ -165,7 +165,7 @@ static int luo_flb_retrieve_one(struct liveupdate_flb *flb)
 	bool found = false;
 	int err;
 
-	guard(mutex)(&private->incoming.lock);
+	lockdep_assert_held(&private->incoming.lock);
 
 	if (private->incoming.finished)
 		return -ENODATA;
@@ -206,12 +206,14 @@ static int luo_flb_retrieve_one(struct liveupdate_flb *flb)
 	return 0;
 }
 
-static void luo_flb_file_finish_one(struct liveupdate_flb *flb)
+void liveupdate_flb_put_incoming(struct liveupdate_flb *flb)
 {
 	struct luo_flb_private *private = luo_flb_get_private(flb);
+	struct liveupdate_flb_op_args args = {0};
 
-	if (refcount_dec_and_test(&private->incoming.count)) {
-		struct liveupdate_flb_op_args args = {0};
+	scoped_guard(mutex, &private->incoming.lock) {
+		if (!refcount_dec_and_test(&private->incoming.count))
+			return;
 
 		if (!private->incoming.retrieved) {
 			int err = luo_flb_retrieve_one(flb);
@@ -220,16 +222,14 @@ static void luo_flb_file_finish_one(struct liveupdate_flb *flb)
 				return;
 		}
 
-		scoped_guard(mutex, &private->incoming.lock) {
-			args.flb = flb;
-			args.obj = private->incoming.obj;
-			flb->ops->finish(&args);
+		args.flb = flb;
+		args.obj = private->incoming.obj;
+		flb->ops->finish(&args);
 
-			private->incoming.data = 0;
-			private->incoming.obj = NULL;
-			private->incoming.finished = true;
-			module_put(flb->ops->owner);
-		}
+		private->incoming.data = 0;
+		private->incoming.obj = NULL;
+		private->incoming.finished = true;
+		module_put(flb->ops->owner);
 	}
 }
 
@@ -312,7 +312,7 @@ void luo_flb_file_finish(struct liveupdate_file_handler *fh)
 
 	guard(rwsem_read)(&luo_register_rwlock);
 	list_for_each_entry_reverse(iter, flb_list, list)
-		luo_flb_file_finish_one(iter->flb);
+		liveupdate_flb_put_incoming(iter->flb);
 }
 
 static void luo_flb_unregister_one(struct liveupdate_file_handler *fh,
@@ -509,6 +509,8 @@ int liveupdate_flb_get_incoming(struct liveupdate_flb *flb, void **objp)
 	if (!liveupdate_enabled())
 		return -EOPNOTSUPP;
 
+	guard(mutex)(&private->incoming.lock);
+
 	if (!private->incoming.obj) {
 		int err = luo_flb_retrieve_one(flb);
 
@@ -516,7 +518,7 @@ int liveupdate_flb_get_incoming(struct liveupdate_flb *flb, void **objp)
 			return err;
 	}
 
-	guard(mutex)(&private->incoming.lock);
+	refcount_inc(&private->incoming.count);
 	*objp = private->incoming.obj;
 
 	return 0;
diff --git a/lib/tests/liveupdate.c b/lib/tests/liveupdate.c
index e4b0ecbee32f..4c08a7c6fb78 100644
--- a/lib/tests/liveupdate.c
+++ b/lib/tests/liveupdate.c
@@ -105,6 +105,9 @@ static void liveupdate_test_init(void)
 			pr_err("liveupdate_flb_get_incoming for %s failed: %pe\n",
 			       flb->compatible, ERR_PTR(err));
 		}
+
+		if (!err)
+			liveupdate_flb_put_incoming(flb);
 	}
 	initialized = true;
 }
-- 
2.54.0.rc2.544.gc7ae2d5bb8-goog