Add the BPF_SOCK_OPS_UDP_CONNECTED_CB callback as a sockops hook where
connected UDP sockets can be inserted into a socket map. This is
invoked on calls to connect() for UDP sockets right after the socket is
hashed. Together with the next patch, this provides the missing piece
allowing us to fully manage the contents of a socket hash in an
environment where we want to keep track of all UDP and TCP sockets
connected to some backend.

is_locked_tcp_sock was recently introduced in [1] to prevent access to
TCP-specific socket fields in contexts where the socket lock isn't held.
This patch extends the use of this field to prevent access to these
fields in UDP socket contexts.

Note: Technically, there should be nothing preventing the use of
      bpf_sock_ops_setsockopt() and bpf_sock_ops_getsockopt() in this
      context, but I've avoided removing the is_locked_tcp_sock_ops()
      guard from these helpers for now to keep the changes in this patch
      series more focused.

[1]: https://lore.kernel.org/all/20250220072940.99994-4-kerneljasonxing@gmail.com/

Signed-off-by: Jordan Rife <jordan@jrife.io>
---
 include/net/udp.h              | 43 ++++++++++++++++++++++++++++++++++
 include/uapi/linux/bpf.h       |  3 +++
 net/ipv4/udp.c                 |  1 +
 net/ipv6/udp.c                 |  1 +
 tools/include/uapi/linux/bpf.h |  3 +++
 5 files changed, 51 insertions(+)

diff --git a/include/net/udp.h b/include/net/udp.h
index e2af3bda90c9..0f55c489e90f 100644
--- a/include/net/udp.h
+++ b/include/net/udp.h
@@ -18,6 +18,7 @@
 #ifndef _UDP_H
 #define _UDP_H
 
+#include <linux/filter.h>
 #include <linux/list.h>
 #include <linux/bug.h>
 #include <net/inet_sock.h>
@@ -25,6 +26,7 @@
 #include <net/sock.h>
 #include <net/snmp.h>
 #include <net/ip.h>
+#include <linux/bpf-cgroup.h>
 #include <linux/ipv6.h>
 #include <linux/seq_file.h>
 #include <linux/poll.h>
@@ -661,4 +663,45 @@ struct sk_psock;
 int udp_bpf_update_proto(struct sock *sk, struct sk_psock *psock, bool restore);
 #endif
 
+#ifdef CONFIG_BPF
+
+/* Call BPF_SOCK_OPS program that returns an int. If the return value
+ * is < 0, then the BPF op failed (for example if the loaded BPF
+ * program does not support the chosen operation or there is no BPF
+ * program loaded).
+ */
+static inline int udp_call_bpf(struct sock *sk, int op)
+{
+	struct bpf_sock_ops_kern sock_ops;
+	int ret;
+
+	memset(&sock_ops, 0, offsetof(struct bpf_sock_ops_kern, temp));
+	if (sk_fullsock(sk)) {
+		sock_ops.is_fullsock = 1;
+		/* sock_ops.is_locked_tcp_sock not set. This prevents
+		 * access to TCP-specific fields.
+		 */
+		sock_owned_by_me(sk);
+	}
+
+	sock_ops.sk = sk;
+	sock_ops.op = op;
+
+	ret = BPF_CGROUP_RUN_PROG_SOCK_OPS(&sock_ops);
+	if (ret == 0)
+		ret = sock_ops.reply;
+	else
+		ret = -1;
+	return ret;
+}
+
+#else
+
+static inline int udp_call_bpf(struct sock *sk, int op, u32 nargs, u32 *args)
+{
+	return -EPERM;
+}
+
+#endif /* CONFIG_BPF */
+
 #endif	/* _UDP_H */
diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
index 22761dea4635..e30515af1f27 100644
--- a/include/uapi/linux/bpf.h
+++ b/include/uapi/linux/bpf.h
@@ -7122,6 +7122,9 @@ enum {
 					 * sendmsg timestamp with corresponding
 					 * tskey.
 					 */
+	BPF_SOCK_OPS_UDP_CONNECTED_CB,	/* Called on connect() for UDP sockets
+					 * right after the socket is hashed.
+					 */
 };
 
 /* List of TCP states. There is a build check in net/ipv4/tcp.c to detect
diff --git a/net/ipv4/udp.c b/net/ipv4/udp.c
index cc3ce0f762ec..2d51d0ead70d 100644
--- a/net/ipv4/udp.c
+++ b/net/ipv4/udp.c
@@ -2153,6 +2153,7 @@ static int udp_connect(struct sock *sk, struct sockaddr *uaddr, int addr_len)
 	res = __ip4_datagram_connect(sk, uaddr, addr_len);
 	if (!res)
 		udp4_hash4(sk);
+	udp_call_bpf(sk, BPF_SOCK_OPS_UDP_CONNECTED_CB);
 	release_sock(sk);
 	return res;
 }
diff --git a/net/ipv6/udp.c b/net/ipv6/udp.c
index 6a68f77da44b..304b43851e16 100644
--- a/net/ipv6/udp.c
+++ b/net/ipv6/udp.c
@@ -1310,6 +1310,7 @@ static int udpv6_connect(struct sock *sk, struct sockaddr *uaddr, int addr_len)
 	res = __ip6_datagram_connect(sk, uaddr, addr_len);
 	if (!res)
 		udp6_hash4(sk);
+	udp_call_bpf(sk, BPF_SOCK_OPS_UDP_CONNECTED_CB);
 	release_sock(sk);
 	return res;
 }
diff --git a/tools/include/uapi/linux/bpf.h b/tools/include/uapi/linux/bpf.h
index 22761dea4635..e30515af1f27 100644
--- a/tools/include/uapi/linux/bpf.h
+++ b/tools/include/uapi/linux/bpf.h
@@ -7122,6 +7122,9 @@ enum {
 					 * sendmsg timestamp with corresponding
 					 * tskey.
 					 */
+	BPF_SOCK_OPS_UDP_CONNECTED_CB,	/* Called on connect() for UDP sockets
+					 * right after the socket is hashed.
+					 */
 };
 
 /* List of TCP states. There is a build check in net/ipv4/tcp.c to detect
-- 
2.43.0