[PATCH net-next v2 09/12] net-timestamp: add tx OPT_ID_TCP support for bpf case

[Jason Xing <kerneljasonxing@gmail.com>]

From: Jason Xing <kernelxing@tencent.com>

We can set OPT_ID|OPT_ID_TCP before we initialize the last skb
from each sendmsg. We only set the socket once like how we use
setsockopt() with OPT_ID|OPT_ID_TCP flags.

Note: we will check if non-bpf _and_ bpf sk_tsflags have OPT_ID
flag. If either of them has been set before, we will not initialize
the key any more, or else it will affect the existing printing
from applications or BPF program behaviour.

Signed-off-by: Jason Xing <kernelxing@tencent.com>
---
 include/net/sock.h |  1 +
 net/core/filter.c  |  5 +++++
 net/core/skbuff.c  | 14 ++++++++++----
 net/core/sock.c    | 29 +++++++++++++++++++++--------
 4 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/include/net/sock.h b/include/net/sock.h
index b7c51b95c92d..2b4ac289c8fa 100644
--- a/include/net/sock.h
+++ b/include/net/sock.h
@@ -2893,6 +2893,7 @@ 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);
+int sock_set_tskey(struct sock *sk, int val, int type);
 int sock_set_timestamping(struct sock *sk, int optname,
 			  struct so_timestamping timestamping);
 
diff --git a/net/core/filter.c b/net/core/filter.c
index 08135f538c99..3b4afaa273d9 100644
--- a/net/core/filter.c
+++ b/net/core/filter.c
@@ -5210,6 +5210,7 @@ static int bpf_sock_set_timestamping(struct sock *sk,
 				     struct so_timestamping *timestamping)
 {
 	u32 flags = timestamping->flags;
+	int ret;
 
 	if (flags & ~SOF_TIMESTAMPING_MASK)
 		return -EINVAL;
@@ -5218,6 +5219,10 @@ static int bpf_sock_set_timestamping(struct sock *sk,
 	      SOF_TIMESTAMPING_TX_ACK)))
 		return -EINVAL;
 
+	ret = sock_set_tskey(sk, flags, BPFPROG_TS_REQUESTOR);
+	if (ret)
+		return ret;
+
 	WRITE_ONCE(sk->sk_tsflags[BPFPROG_TS_REQUESTOR], flags);
 	static_branch_enable(&bpf_tstamp_control);
 
diff --git a/net/core/skbuff.c b/net/core/skbuff.c
index e18305b03a01..1ef379a87f88 100644
--- a/net/core/skbuff.c
+++ b/net/core/skbuff.c
@@ -5619,7 +5619,7 @@ 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,
+static void bpf_skb_tstamp_tx_output(struct sock *sk, struct sk_buff *skb, int tstype,
 				     struct skb_shared_hwtstamps *hwtstamps)
 {
 	struct tcp_sock *tp;
@@ -5635,7 +5635,7 @@ static void bpf_skb_tstamp_tx_output(struct sock *sk, int tstype,
 	tp = tcp_sk(sk);
 	if (BPF_SOCK_OPS_TEST_FLAG(tp, BPF_SOCK_OPS_TX_TIMESTAMPING_OPT_CB_FLAG)) {
 		struct timespec64 tstamp;
-		u32 cb_flag;
+		u32 cb_flag, key = 0;
 
 		switch (tstype) {
 		case SCM_TSTAMP_SCHED:
@@ -5651,11 +5651,17 @@ static void bpf_skb_tstamp_tx_output(struct sock *sk, int tstype,
 			return;
 		}
 
+		if (sk_is_tcp(sk)) {
+			key = skb_shinfo(skb)->tskey;
+			key -= atomic_read(&sk->sk_tskey);
+		}
+
 		if (hwtstamps)
 			tstamp = ktime_to_timespec64(hwtstamps->hwtstamp);
 		else
 			tstamp = ktime_to_timespec64(ktime_get_real());
-		tcp_call_bpf_2arg(sk, cb_flag, tstamp.tv_sec, tstamp.tv_nsec);
+
+		tcp_call_bpf_3arg(sk, cb_flag, key, tstamp.tv_sec, tstamp.tv_nsec);
 	}
 }
 
@@ -5668,7 +5674,7 @@ void __skb_tstamp_tx(struct sk_buff *orig_skb,
 		return;
 
 	if (static_branch_unlikely(&bpf_tstamp_control))
-		bpf_skb_tstamp_tx_output(sk, tstype, hwtstamps);
+		bpf_skb_tstamp_tx_output(sk, orig_skb, tstype, hwtstamps);
 
 	skb_tstamp_tx_output(orig_skb, ack_skb, hwtstamps, sk, tstype);
 }
diff --git a/net/core/sock.c b/net/core/sock.c
index a6e0d51a5f72..c15edbd382d5 100644
--- a/net/core/sock.c
+++ b/net/core/sock.c
@@ -915,21 +915,18 @@ int sock_get_timestamping(struct so_timestamping *timestamping,
 	return 0;
 }
 
-int sock_set_timestamping(struct sock *sk, int optname,
-			  struct so_timestamping timestamping)
+int sock_set_tskey(struct sock *sk, int val, int type)
 {
-	int val = timestamping.flags;
-	int ret;
-
-	if (val & ~SOF_TIMESTAMPING_MASK)
-		return -EINVAL;
+	u32 tsflags;
 
 	if (val & SOF_TIMESTAMPING_OPT_ID_TCP &&
 	    !(val & SOF_TIMESTAMPING_OPT_ID))
 		return -EINVAL;
 
+	tsflags |= (sk->sk_tsflags[SOCKETOPT_TS_REQUESTOR] |
+		    sk->sk_tsflags[BPFPROG_TS_REQUESTOR]);
 	if (val & SOF_TIMESTAMPING_OPT_ID &&
-	    !(sk->sk_tsflags[SOCKETOPT_TS_REQUESTOR] & SOF_TIMESTAMPING_OPT_ID)) {
+	    !(tsflags & SOF_TIMESTAMPING_OPT_ID)) {
 		if (sk_is_tcp(sk)) {
 			if ((1 << sk->sk_state) &
 			    (TCPF_CLOSE | TCPF_LISTEN))
@@ -943,6 +940,22 @@ int sock_set_timestamping(struct sock *sk, int optname,
 		}
 	}
 
+	return 0;
+}
+
+int sock_set_timestamping(struct sock *sk, int optname,
+			  struct so_timestamping timestamping)
+{
+	int val = timestamping.flags;
+	int ret;
+
+	if (val & ~SOF_TIMESTAMPING_MASK)
+		return -EINVAL;
+
+	ret = sock_set_tskey(sk, val, SOCKETOPT_TS_REQUESTOR);
+	if (ret)
+		return ret;
+
 	if (val & SOF_TIMESTAMPING_OPT_STATS &&
 	    !(val & SOF_TIMESTAMPING_OPT_TSONLY))
 		return -EINVAL;
-- 
2.37.3