[PATCH net-next v3 12/14] net-timestamp: add OPT_ID for UDP proto

[Jason Xing <kerneljasonxing@xxxxxxxxx>]

From: Jason Xing <kernelxing@xxxxxxxxxxx>

Let it work for UDP proto.

Signed-off-by: Jason Xing <kernelxing@xxxxxxxxxxx>
---
 net/ipv4/ip_output.c  | 16 +++++++++++-----
 net/ipv6/ip6_output.c | 16 +++++++++++-----
 2 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/net/ipv4/ip_output.c b/net/ipv4/ip_output.c
index 9d94a209057b..45033105b34c 100644
--- a/net/ipv4/ip_output.c
+++ b/net/ipv4/ip_output.c
@@ -1049,11 +1049,17 @@ static int __ip_append_data(struct sock *sk,
 
 	cork->length += length;
 
-	if (cork->tx_flags & SKBTX_ANY_TSTAMP &&
-	    READ_ONCE(sk->sk_tsflags) & SOF_TIMESTAMPING_OPT_ID) {
-		if (cork->flags & IPCORK_TS_OPT_ID) {
-			tskey = cork->ts_opt_id;
-		} else {
+	if (cork->tx_flags & SKBTX_ANY_TSTAMP) {
+		if (READ_ONCE(sk->sk_tsflags) & SOF_TIMESTAMPING_OPT_ID) {
+			if (cork->flags & IPCORK_TS_OPT_ID) {
+				tskey = cork->ts_opt_id;
+			} else {
+				tskey = atomic_inc_return(&sk->sk_tskey) - 1;
+				hold_tskey = true;
+			}
+		}
+		if (!hold_tskey &&
+		    READ_ONCE(sk->sk_tsflags_bpf) & SOF_TIMESTAMPING_OPT_ID) {
 			tskey = atomic_inc_return(&sk->sk_tskey) - 1;
 			hold_tskey = true;
 		}
diff --git a/net/ipv6/ip6_output.c b/net/ipv6/ip6_output.c
index 230e8d5a792c..ec956ada7179 100644
--- a/net/ipv6/ip6_output.c
+++ b/net/ipv6/ip6_output.c
@@ -1547,11 +1547,17 @@ static int __ip6_append_data(struct sock *sk,
 			flags &= ~MSG_SPLICE_PAGES;
 	}
 
-	if (cork->tx_flags & SKBTX_ANY_TSTAMP &&
-	    READ_ONCE(sk->sk_tsflags) & SOF_TIMESTAMPING_OPT_ID) {
-		if (cork->flags & IPCORK_TS_OPT_ID) {
-			tskey = cork->ts_opt_id;
-		} else {
+	if (cork->tx_flags & SKBTX_ANY_TSTAMP) {
+		if (READ_ONCE(sk->sk_tsflags) & SOF_TIMESTAMPING_OPT_ID) {
+			if (cork->flags & IPCORK_TS_OPT_ID) {
+				tskey = cork->ts_opt_id;
+			} else {
+				tskey = atomic_inc_return(&sk->sk_tskey) - 1;
+				hold_tskey = true;
+			}
+		}
+		if (!hold_tskey &&
+		    READ_ONCE(sk->sk_tsflags_bpf) & SOF_TIMESTAMPING_OPT_ID) {
 			tskey = atomic_inc_return(&sk->sk_tskey) - 1;
 			hold_tskey = true;
 		}
-- 
2.37.3