Introduce a per-socket bit SK_BPF_CB_TIMESTAMPING_V2 for BPF Timestamping
V2, that allows a socket opt into recording a "start time" on every
skb it produces or receives, which should be reflected in the rest of
patches.

Add a static key to avoid normal traffic suffered from the performance
affect. It works exactly like netstamp_needed_key.

Use bpf_setsockopt(SK_BPF_CB_FLAGS) to toggle the feature that is the
same mechanism used by SK_BPF_CB_TX_TIMESTAMPING which is actually
version 1.

The whole feature is only reachable from BPF, not from user-space
setsockopt().

Signed-off-by: Jason Xing <kerneljasonxing@gmail.com>
---
 include/linux/skbuff.h         |  4 +++
 include/uapi/linux/bpf.h       |  5 ++--
 net/core/dev.c                 | 53 ++++++++++++++++++++++++++++++++++
 net/core/filter.c              |  7 +++++
 net/core/sock.c                |  5 ++++
 tools/include/uapi/linux/bpf.h |  5 ++--
 6 files changed, 75 insertions(+), 4 deletions(-)

diff --git a/include/linux/skbuff.h b/include/linux/skbuff.h
index 671c13494566..55ae1653a1cf 100644
--- a/include/linux/skbuff.h
+++ b/include/linux/skbuff.h
@@ -4514,6 +4514,10 @@ static inline void skb_set_delivery_type_by_clockid(struct sk_buff *skb,
 
 DECLARE_STATIC_KEY_FALSE(netstamp_needed_key);
 
+DECLARE_STATIC_KEY_FALSE(bpfts_v2_needed_key);
+void bpfts_v2_enable(void);
+void bpfts_v2_disable(void);
+
 /* It is used in the ingress path to clear the delivery_time.
  * If needed, set the skb->tstamp to the (rcv) timestamp.
  */
diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
index 732b35cc08d1..40feea175a10 100644
--- a/include/uapi/linux/bpf.h
+++ b/include/uapi/linux/bpf.h
@@ -7143,8 +7143,9 @@ enum {
 
 enum {
 	SK_BPF_CB_TX_TIMESTAMPING	= 1<<0,
-	SK_BPF_CB_MASK			= (SK_BPF_CB_TX_TIMESTAMPING - 1) |
-					   SK_BPF_CB_TX_TIMESTAMPING
+	SK_BPF_CB_TIMESTAMPING_V2		= 1<<1,
+	SK_BPF_CB_MASK			= (SK_BPF_CB_TIMESTAMPING_V2 - 1) |
+					   SK_BPF_CB_TIMESTAMPING_V2
 };
 
 /* List of known BPF sock_ops operators.
diff --git a/net/core/dev.c b/net/core/dev.c
index 38336858c168..4479fc87b789 100644
--- a/net/core/dev.c
+++ b/net/core/dev.c
@@ -2440,6 +2440,59 @@ void net_disable_timestamp(void)
 }
 EXPORT_SYMBOL(net_disable_timestamp);
 
+DEFINE_STATIC_KEY_FALSE(bpfts_v2_needed_key);
+EXPORT_SYMBOL(bpfts_v2_needed_key);
+#ifdef CONFIG_JUMP_LABEL
+static atomic_t bpfts_v2_needed_deferred;
+static atomic_t bpfts_v2_wanted;
+static void bpfts_v2_clear(struct work_struct *work)
+{
+	int deferred = atomic_xchg(&bpfts_v2_needed_deferred, 0);
+	int wanted;
+
+	wanted = atomic_add_return(deferred, &bpfts_v2_wanted);
+	if (wanted > 0)
+		static_branch_enable(&bpfts_v2_needed_key);
+	else
+		static_branch_disable(&bpfts_v2_needed_key);
+}
+static DECLARE_WORK(bpfts_v2_work, bpfts_v2_clear);
+#endif
+
+void bpfts_v2_enable(void)
+{
+#ifdef CONFIG_JUMP_LABEL
+	int wanted = atomic_read(&bpfts_v2_wanted);
+
+	while (wanted > 0) {
+		if (atomic_try_cmpxchg(&bpfts_v2_wanted, &wanted, wanted + 1))
+			return;
+	}
+	atomic_inc(&bpfts_v2_needed_deferred);
+	schedule_work(&bpfts_v2_work);
+#else
+	static_branch_inc(&bpfts_v2_needed_key);
+#endif
+}
+EXPORT_SYMBOL(bpfts_v2_enable);
+
+void bpfts_v2_disable(void)
+{
+#ifdef CONFIG_JUMP_LABEL
+	int wanted = atomic_read(&bpfts_v2_wanted);
+
+	while (wanted > 1) {
+		if (atomic_try_cmpxchg(&bpfts_v2_wanted, &wanted, wanted - 1))
+			return;
+	}
+	atomic_dec(&bpfts_v2_needed_deferred);
+	schedule_work(&bpfts_v2_work);
+#else
+	static_branch_dec(&bpfts_v2_needed_key);
+#endif
+}
+EXPORT_SYMBOL(bpfts_v2_disable);
+
 static inline void net_timestamp_set(struct sk_buff *skb)
 {
 	skb->tstamp = 0;
diff --git a/net/core/filter.c b/net/core/filter.c
index 61940e753552..3486c5c2d170 100644
--- a/net/core/filter.c
+++ b/net/core/filter.c
@@ -5468,6 +5468,13 @@ static int sk_bpf_set_get_cb_flags(struct sock *sk, char *optval, bool getopt)
 	if (sk_bpf_cb_flags & ~SK_BPF_CB_MASK)
 		return -EINVAL;
 
+	if ((sk_bpf_cb_flags ^ sk->sk_bpf_cb_flags) & SK_BPF_CB_TIMESTAMPING_V2) {
+		if (sk_bpf_cb_flags & SK_BPF_CB_TIMESTAMPING_V2)
+			bpfts_v2_enable();
+		else
+			bpfts_v2_disable();
+	}
+
 	sk->sk_bpf_cb_flags = sk_bpf_cb_flags;
 
 	return 0;
diff --git a/net/core/sock.c b/net/core/sock.c
index 1ad41904db25..251c88e2ee45 100644
--- a/net/core/sock.c
+++ b/net/core/sock.c
@@ -2364,6 +2364,9 @@ static void __sk_destruct(struct rcu_head *head)
 
 	sock_disable_timestamp(sk, SK_FLAGS_TIMESTAMP);
 
+	if (sk->sk_bpf_cb_flags & SK_BPF_CB_TIMESTAMPING_V2)
+		bpfts_v2_disable();
+
 #ifdef CONFIG_BPF_SYSCALL
 	bpf_sk_storage_free(sk);
 #endif
@@ -2550,6 +2553,8 @@ struct sock *sk_clone(const struct sock *sk, const gfp_t priority,
 
 	if (sock_needs_netstamp(sk) && newsk->sk_flags & SK_FLAGS_TIMESTAMP)
 		net_enable_timestamp();
+	if (newsk->sk_bpf_cb_flags & SK_BPF_CB_TIMESTAMPING_V2)
+		bpfts_v2_enable();
 
 	rcu_read_lock();
 	filter = rcu_dereference(sk->sk_filter);
diff --git a/tools/include/uapi/linux/bpf.h b/tools/include/uapi/linux/bpf.h
index 732b35cc08d1..40feea175a10 100644
--- a/tools/include/uapi/linux/bpf.h
+++ b/tools/include/uapi/linux/bpf.h
@@ -7143,8 +7143,9 @@ enum {
 
 enum {
 	SK_BPF_CB_TX_TIMESTAMPING	= 1<<0,
-	SK_BPF_CB_MASK			= (SK_BPF_CB_TX_TIMESTAMPING - 1) |
-					   SK_BPF_CB_TX_TIMESTAMPING
+	SK_BPF_CB_TIMESTAMPING_V2		= 1<<1,
+	SK_BPF_CB_MASK			= (SK_BPF_CB_TIMESTAMPING_V2 - 1) |
+					   SK_BPF_CB_TIMESTAMPING_V2
 };
 
 /* List of known BPF sock_ops operators.
-- 
2.43.7