[PATCH net-next v2 04/12] net-timestamp: add static key to control the whole bpf extension

[Jason Xing <kerneljasonxing@gmail.com>]

From: Jason Xing <kernelxing@tencent.com>

Willem suggested that we use a static key to control. The advantage
is that we will not affect the existing applications at all if we
don't load BPF program.

In this patch, except the static key, I also add one logic that is
used to test if the socket has enabled its tsflags in order to
support bpf logic to allow both cases to happen at the same time.
Or else, the skb carring related timestamp flag doesn't know which
way of printing is desirable.

One thing important is this patch allows print from both applications
and bpf program at the same time. Now we have three kinds of print:
1) only BPF program prints
2) only application program prints
3) both can print without side effect

Signed-off-by: Jason Xing <kernelxing@tencent.com>
---
 include/net/sock.h |  1 +
 net/core/filter.c  |  3 +++
 net/core/skbuff.c  | 38 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 42 insertions(+)

diff --git a/include/net/sock.h b/include/net/sock.h
index 66ecd78f1dfe..b7c51b95c92d 100644
--- a/include/net/sock.h
+++ b/include/net/sock.h
@@ -2889,6 +2889,7 @@ static inline bool sk_dev_equal_l3scope(struct sock *sk, int dif)
 void sock_def_readable(struct sock *sk);
 
 int sock_bindtoindex(struct sock *sk, int ifindex, bool lock_sk);
+DECLARE_STATIC_KEY_FALSE(bpf_tstamp_control);
 void sock_set_timestamp(struct sock *sk, int optname, bool valbool);
 int sock_get_timestamping(struct so_timestamping *timestamping,
 			  sockptr_t optval, unsigned int optlen);
diff --git a/net/core/filter.c b/net/core/filter.c
index 996426095bd9..08135f538c99 100644
--- a/net/core/filter.c
+++ b/net/core/filter.c
@@ -5204,6 +5204,8 @@ static const struct bpf_func_proto bpf_get_socket_uid_proto = {
 	.arg1_type      = ARG_PTR_TO_CTX,
 };
 
+DEFINE_STATIC_KEY_FALSE(bpf_tstamp_control);
+
 static int bpf_sock_set_timestamping(struct sock *sk,
 				     struct so_timestamping *timestamping)
 {
@@ -5217,6 +5219,7 @@ static int bpf_sock_set_timestamping(struct sock *sk,
 		return -EINVAL;
 
 	WRITE_ONCE(sk->sk_tsflags[BPFPROG_TS_REQUESTOR], flags);
+	static_branch_enable(&bpf_tstamp_control);
 
 	return 0;
 }
diff --git a/net/core/skbuff.c b/net/core/skbuff.c
index f36eb9daa31a..d0f912f1ff7b 100644
--- a/net/core/skbuff.c
+++ b/net/core/skbuff.c
@@ -5540,6 +5540,29 @@ void skb_complete_tx_timestamp(struct sk_buff *skb,
 }
 EXPORT_SYMBOL_GPL(skb_complete_tx_timestamp);
 
+static bool sk_tstamp_tx_flags(struct sock *sk, u32 tsflags, int tstype)
+{
+	u32 testflag;
+
+	switch (tstype) {
+	case SCM_TSTAMP_SCHED:
+		testflag = SOF_TIMESTAMPING_TX_SCHED;
+		break;
+	case SCM_TSTAMP_SND:
+		testflag = SOF_TIMESTAMPING_TX_SOFTWARE;
+		break;
+	case SCM_TSTAMP_ACK:
+		testflag = SOF_TIMESTAMPING_TX_ACK;
+		break;
+	default:
+		return false;
+	}
+	if (tsflags & testflag)
+		return true;
+
+	return false;
+}
+
 static void skb_tstamp_tx_output(struct sk_buff *orig_skb,
 				 const struct sk_buff *ack_skb,
 				 struct skb_shared_hwtstamps *hwtstamps,
@@ -5558,6 +5581,9 @@ static void skb_tstamp_tx_output(struct sk_buff *orig_skb,
 	if (!skb_may_tx_timestamp(sk, tsonly))
 		return;
 
+	if (!sk_tstamp_tx_flags(sk, tsflags, tstype))
+		return;
+
 	if (tsonly) {
 #ifdef CONFIG_INET
 		if ((tsflags & SOF_TIMESTAMPING_OPT_STATS) &&
@@ -5593,6 +5619,15 @@ static void skb_tstamp_tx_output(struct sk_buff *orig_skb,
 	__skb_complete_tx_timestamp(skb, sk, tstype, opt_stats);
 }
 
+static void bpf_skb_tstamp_tx_output(struct sock *sk, int tstype)
+{
+	u32 tsflags;
+
+	tsflags = READ_ONCE(sk->sk_tsflags[BPFPROG_TS_REQUESTOR]);
+	if (!sk_tstamp_tx_flags(sk, tsflags, tstype))
+		return;
+}
+
 void __skb_tstamp_tx(struct sk_buff *orig_skb,
 		     const struct sk_buff *ack_skb,
 		     struct skb_shared_hwtstamps *hwtstamps,
@@ -5601,6 +5636,9 @@ void __skb_tstamp_tx(struct sk_buff *orig_skb,
 	if (!sk)
 		return;
 
+	if (static_branch_unlikely(&bpf_tstamp_control))
+		bpf_skb_tstamp_tx_output(sk, tstype);
+
 	skb_tstamp_tx_output(orig_skb, ack_skb, hwtstamps, sk, tstype);
 }
 EXPORT_SYMBOL_GPL(__skb_tstamp_tx);
-- 
2.37.3