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;
