[PATCH bpf-next 16/25] ftrace: supporting replace direct ftrace_ops

[Menglong Dong <menglong8.dong@gmail.com>]

Introduce the function replace_ftrace_direct(). This is used to replace
the direct ftrace_ops for a function, and will be used in the next patch.

Let's call the origin ftrace_ops A, and the new ftrace_ops B. First, we
register B directly, and the callback of the functions in A and B will
fallback to the ftrace_ops_list case.

Then, we modify the address of the entry in the direct_functions to
B->direct_call, and remove it from A. This will update the dyn_rec and
make the functions call b->direct_call directly. If no function in
A->filter_hash, just unregister it.

So a record can have more than one direct ftrace_ops, and we need check
if there is any direct ops for the record before remove the
FTRACE_OPS_FL_DIRECT in __ftrace_hash_rec_update().

Signed-off-by: Menglong Dong <dongml2@chinatelecom.cn>
---
 include/linux/ftrace.h |  8 ++++
 kernel/trace/ftrace.c  | 87 +++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 93 insertions(+), 2 deletions(-)

diff --git a/include/linux/ftrace.h b/include/linux/ftrace.h
index 40727d3f125d..1d162e331e99 100644
--- a/include/linux/ftrace.h
+++ b/include/linux/ftrace.h
@@ -528,6 +528,9 @@ void ftrace_stub_direct_tramp(void);
 
 int reset_ftrace_direct_ips(struct ftrace_ops *ops, unsigned long *ips,
 			    unsigned int cnt);
+int replace_ftrace_direct(struct ftrace_ops *ops, struct ftrace_ops *src_ops,
+			  unsigned long addr);
+
 #else
 struct ftrace_ops;
 static inline unsigned long ftrace_find_rec_direct(unsigned long ip)
@@ -556,6 +559,11 @@ static inline int reset_ftrace_direct_ips(struct ftrace_ops *ops, unsigned long
 {
 	return -ENODEV;
 }
+static inline int replace_ftrace_direct(struct ftrace_ops *ops, struct ftrace_ops *src_ops,
+					unsigned long addr)
+{
+	return -ENODEV;
+}
 
 /*
  * This must be implemented by the architecture.
diff --git a/kernel/trace/ftrace.c b/kernel/trace/ftrace.c
index 5b6b74ea4c20..7f2313e4c3d9 100644
--- a/kernel/trace/ftrace.c
+++ b/kernel/trace/ftrace.c
@@ -1727,6 +1727,24 @@ static bool skip_record(struct dyn_ftrace *rec)
 		!(rec->flags & FTRACE_FL_ENABLED);
 }
 
+static struct ftrace_ops *
+ftrace_find_direct_ops_any_other(struct dyn_ftrace *rec, struct ftrace_ops *op_exclude)
+{
+	struct ftrace_ops *op;
+	unsigned long ip = rec->ip;
+
+	do_for_each_ftrace_op(op, ftrace_ops_list) {
+
+		if (op == op_exclude || !(op->flags & FTRACE_OPS_FL_DIRECT))
+			continue;
+
+		if (hash_contains_ip(ip, op->func_hash))
+			return op;
+	} while_for_each_ftrace_op(op);
+
+	return NULL;
+}
+
 /*
  * This is the main engine to the ftrace updates to the dyn_ftrace records.
  *
@@ -1831,8 +1849,10 @@ static bool __ftrace_hash_rec_update(struct ftrace_ops *ops,
 			 * function, then that function should no longer
 			 * be direct.
 			 */
-			if (ops->flags & FTRACE_OPS_FL_DIRECT)
-				rec->flags &= ~FTRACE_FL_DIRECT;
+			if (ops->flags & FTRACE_OPS_FL_DIRECT) {
+				if (!ftrace_find_direct_ops_any_other(rec, ops))
+					rec->flags &= ~FTRACE_FL_DIRECT;
+			}
 
 			/*
 			 * If the rec had REGS enabled and the ops that is
@@ -6033,6 +6053,69 @@ int register_ftrace_direct(struct ftrace_ops *ops, unsigned long addr)
 }
 EXPORT_SYMBOL_GPL(register_ftrace_direct);
 
+int replace_ftrace_direct(struct ftrace_ops *ops, struct ftrace_ops *src_ops,
+			  unsigned long addr)
+{
+	struct ftrace_hash *hash;
+	struct ftrace_func_entry *entry, *iter;
+	int err = -EBUSY, size, count;
+
+	if (ops->func || ops->trampoline)
+		return -EINVAL;
+	if (!(ops->flags & FTRACE_OPS_FL_INITIALIZED))
+		return -EINVAL;
+	if (ops->flags & FTRACE_OPS_FL_ENABLED)
+		return -EINVAL;
+
+	hash = ops->func_hash->filter_hash;
+	if (ftrace_hash_empty(hash))
+		return -EINVAL;
+
+	mutex_lock(&direct_mutex);
+
+	ops->func = call_direct_funcs;
+	ops->flags = MULTI_FLAGS;
+	ops->trampoline = FTRACE_REGS_ADDR;
+	ops->direct_call = addr;
+
+	err = register_ftrace_function_nolock(ops);
+	if (err)
+		goto out_unlock;
+
+	hash = ops->func_hash->filter_hash;
+	size = 1 << hash->size_bits;
+	for (int i = 0; i < size; i++) {
+		hlist_for_each_entry(iter, &hash->buckets[i], hlist) {
+			entry = __ftrace_lookup_ip(direct_functions, iter->ip);
+			if (!entry) {
+				err = -ENOENT;
+				goto out_unlock;
+			}
+			WRITE_ONCE(entry->direct, addr);
+			/* remove the ip from the hash, and this will make the trampoline
+			 * be called directly.
+			 */
+			count = src_ops->func_hash->filter_hash->count;
+			if (count <= 1) {
+				if (WARN_ON_ONCE(!count))
+					continue;
+				err = __unregister_ftrace_direct(src_ops, src_ops->direct_call,
+								 true);
+			} else {
+				err = ftrace_set_filter_ip(src_ops, iter->ip, 1, 0);
+			}
+			if (err)
+				goto out_unlock;
+		}
+	}
+
+out_unlock:
+	mutex_unlock(&direct_mutex);
+
+	return err;
+}
+EXPORT_SYMBOL_GPL(replace_ftrace_direct);
+
 /**
  * unregister_ftrace_direct - Remove calls to custom trampoline
  * previously registered by register_ftrace_direct for @ops object.
-- 
2.39.5