From 0d0f0c5020bc425c9a51c8d17b16ca831c2598fb Mon Sep 17 00:00:00 2001
From: Namkyu Kim <namkyu78.kim@samsung.com>
Date: Thu, 4 Mar 2021 09:31:56 +0900
Subject: [PATCH] ANDROID: Add a vendor hook that allow a module to modify the
 wake flag

 android_vh_do_wake_up_sync:
  To modify the mode value of __wake_up_sync_key

 android_vh_set_wake_flags:
  To modify the wake flag from a module

Bug: 181743516
Signed-off-by: Namkyu Kim <namkyu78.kim@samsung.com>
Change-Id: I972e2469c3f139373d21f1e8c85974763388a693
(cherry picked from commit 97368fc2dcc29777e8d3d637d0afdef90e611763)
---
 drivers/android/vendor_hooks.c |  2 ++
 include/trace/hooks/sched.h    | 10 ++++++++++
 kernel/sched/wait.c            |  6 +++++-
 net/core/sock.c                | 13 ++++++++++++-
 4 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/drivers/android/vendor_hooks.c b/drivers/android/vendor_hooks.c
index 90c1ebd65429..b41558070b4c 100644
--- a/drivers/android/vendor_hooks.c
+++ b/drivers/android/vendor_hooks.c
@@ -197,3 +197,5 @@ EXPORT_TRACEPOINT_SYMBOL_GPL(android_vh_check_file_open);
 EXPORT_TRACEPOINT_SYMBOL_GPL(android_vh_check_bpf_syscall);
 EXPORT_TRACEPOINT_SYMBOL_GPL(android_rvh_post_init_entity_util_avg);
 EXPORT_TRACEPOINT_SYMBOL_GPL(android_rvh_find_new_ilb);
+EXPORT_TRACEPOINT_SYMBOL_GPL(android_vh_do_wake_up_sync);
+EXPORT_TRACEPOINT_SYMBOL_GPL(android_vh_set_wake_flags);
diff --git a/include/trace/hooks/sched.h b/include/trace/hooks/sched.h
index f845da877389..1a2640552695 100644
--- a/include/trace/hooks/sched.h
+++ b/include/trace/hooks/sched.h
@@ -314,6 +314,16 @@ DECLARE_RESTRICTED_HOOK(android_rvh_find_new_ilb,
 	TP_PROTO(struct cpumask *nohz_idle_cpus_mask, int *ilb),
 	TP_ARGS(nohz_idle_cpus_mask, ilb), 1);
 
+DECLARE_HOOK(android_vh_do_wake_up_sync,
+	TP_PROTO(struct wait_queue_head *wq_head, int *done),
+	TP_ARGS(wq_head, done));
+
+DECLARE_HOOK(android_vh_set_wake_flags,
+	TP_PROTO(int *wake_flags, unsigned int *mode),
+	TP_ARGS(wake_flags, mode));
+
+/* macro versions of hooks are no longer required */
+
 #endif /* _TRACE_HOOK_SCHED_H */
 /* This part must be outside protection */
 #include <trace/define_trace.h>
diff --git a/kernel/sched/wait.c b/kernel/sched/wait.c
index eca38107b32f..636f8e4e86fa 100644
--- a/kernel/sched/wait.c
+++ b/kernel/sched/wait.c
@@ -5,6 +5,7 @@
  * (C) 2004 Nadia Yvette Chambers, Oracle
  */
 #include "sched.h"
+#include <trace/hooks/sched.h>
 
 void __init_waitqueue_head(struct wait_queue_head *wq_head, const char *name, struct lock_class_key *key)
 {
@@ -199,10 +200,13 @@ EXPORT_SYMBOL_GPL(__wake_up_locked_key_bookmark);
 void __wake_up_sync_key(struct wait_queue_head *wq_head, unsigned int mode,
 			void *key)
 {
+	int wake_flags = WF_SYNC;
+
 	if (unlikely(!wq_head))
 		return;
 
-	__wake_up_common_lock(wq_head, mode, 1, WF_SYNC, key);
+	trace_android_vh_set_wake_flags(&wake_flags, &mode);
+	__wake_up_common_lock(wq_head, mode, 1, wake_flags, key);
 }
 EXPORT_SYMBOL_GPL(__wake_up_sync_key);
 
diff --git a/net/core/sock.c b/net/core/sock.c
index 4ff806d71921..a665e0d32cab 100644
--- a/net/core/sock.c
+++ b/net/core/sock.c
@@ -135,6 +135,7 @@
 #include <net/bpf_sk_storage.h>
 
 #include <trace/events/sock.h>
+#include <trace/hooks/sched.h>
 
 #include <net/tcp.h>
 #include <net/busy_poll.h>
@@ -3143,9 +3144,19 @@ void sock_def_readable(struct sock *sk)
 
 	rcu_read_lock();
 	wq = rcu_dereference(sk->sk_wq);
-	if (skwq_has_sleeper(wq))
+
+	if (skwq_has_sleeper(wq)) {
+		int done = 0;
+
+		trace_android_vh_do_wake_up_sync(&wq->wait, &done);
+		if (done)
+			goto out;
+
 		wake_up_interruptible_sync_poll(&wq->wait, EPOLLIN | EPOLLPRI |
 						EPOLLRDNORM | EPOLLRDBAND);
+	}
+
+out:
 	sk_wake_async(sk, SOCK_WAKE_WAITD, POLL_IN);
 	rcu_read_unlock();
 }