Re: [PATCH net-next 1/2] Add a helper function to get socket cookie in eBPF

From: Chenbo Feng <redacted>

Retrieve the socket cookie generated by sock_gen_cookie() from a sk_buff
with a known socket. Generates a new cookie if one was not yet set.If
the socket pointer inside sk_buff is NULL, 0 is returned. The helper
function coud be useful in monitoring per socket networking traffic
statistics and provide a unique socket identifier per namespace.

Signed-off-by: Chenbo Feng <redacted>
---
  include/linux/bpf.h       |  1 +
  include/linux/sock_diag.h |  1 +
  include/uapi/linux/bpf.h  |  9 ++++++++-
  net/core/filter.c         | 15 +++++++++++++++
  net/core/sock_diag.c      |  3 ++-
  5 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/include/linux/bpf.h b/include/linux/bpf.h
index 3ed1f3b1d594..3f2e0af28c6e 100644
--- a/include/linux/bpf.h
+++ b/include/linux/bpf.h

@@ -349,6 +349,7 @@ extern const struct bpf_func_proto bpf_get_current_comm_proto;
  extern const struct bpf_func_proto bpf_skb_vlan_push_proto;
  extern const struct bpf_func_proto bpf_skb_vlan_pop_proto;
  extern const struct bpf_func_proto bpf_get_stackid_proto;
+extern const struct bpf_func_proto bpf_get_socket_cookie_proto;

  /* Shared helpers among cBPF and eBPF. */
  void bpf_user_rnd_init_once(void);

diff --git a/include/linux/sock_diag.h b/include/linux/sock_diag.h
index a0596ca0e80a..a2f8109bb215 100644
--- a/include/linux/sock_diag.h
+++ b/include/linux/sock_diag.h

@@ -24,6 +24,7 @@ void sock_diag_unregister(const struct sock_diag_handler *h);
  void sock_diag_register_inet_compat(int (*fn)(struct sk_buff *skb, struct nlmsghdr *nlh));
  void sock_diag_unregister_inet_compat(int (*fn)(struct sk_buff *skb, struct nlmsghdr *nlh));

+u64 sock_gen_cookie(struct sock *sk);
  int sock_diag_check_cookie(struct sock *sk, const __u32 *cookie);
  void sock_diag_save_cookie(struct sock *sk, __u32 *cookie);

diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
index 0eb0e87dbe9f..62ee5fab08e5 100644
--- a/include/uapi/linux/bpf.h
+++ b/include/uapi/linux/bpf.h

@@ -430,6 +430,12 @@ union bpf_attr {
   *     @xdp_md: pointer to xdp_md
   *     @delta: An positive/negative integer to be added to xdp_md.data
   *     Return: 0 on success or negative on error
+ *
+ * u64 bpf_bpf_get_socket_cookie(skb)
+ *     Get the cookie for the socket stored inside sk_buff.
+ *     @skb: pointer to skb
+ *     Return: 8 Bytes non-decreasing number on success or 0 if the socket
+ *     field is missing inside sk_buff
   */
  #define __BPF_FUNC_MAPPER(FN)		\
  	FN(unspec),			\

@@ -476,7 +482,8 @@ union bpf_attr {
  	FN(set_hash_invalid),		\
  	FN(get_numa_node_id),		\
  	FN(skb_change_head),		\
-	FN(xdp_adjust_head),
+	FN(xdp_adjust_head),		\
+	FN(get_socket_cookie),

  /* integer value in 'imm' field of BPF_CALL instruction selects which helper
   * function eBPF program intends to call

diff --git a/net/core/filter.c b/net/core/filter.c
index 1969b3f118c1..913b14d3b484 100644
--- a/net/core/filter.c
+++ b/net/core/filter.c

@@ -26,6 +26,7 @@
  #include <linux/mm.h>
  #include <linux/fcntl.h>
  #include <linux/socket.h>
+#include <linux/sock_diag.h>
  #include <linux/in.h>
  #include <linux/inet.h>
  #include <linux/netdevice.h>

@@ -2597,6 +2598,18 @@ static const struct bpf_func_proto bpf_xdp_event_output_proto = {
  	.arg5_type	= ARG_CONST_STACK_SIZE,
  };

+BPF_CALL_1(bpf_get_socket_cookie, struct sk_buff *, skb)
+{
+	return skb->sk ? sock_gen_cookie(skb->sk) : 0;
+}
+
+const struct bpf_func_proto bpf_get_socket_cookie_proto = {

+	.func           = bpf_get_socket_cookie,
+	.gpl_only       = false,
+	.ret_type       = RET_INTEGER,
+	.arg1_type      = ARG_PTR_TO_CTX,
+};
+
  static const struct bpf_func_proto *
  sk_filter_func_proto(enum bpf_func_id func_id)
  {

@@ -2620,6 +2633,8 @@ sk_filter_func_proto(enum bpf_func_id func_id)
  	case BPF_FUNC_trace_printk:
  		if (capable(CAP_SYS_ADMIN))
  			return bpf_get_trace_printk_proto();
+	case BPF_FUNC_get_socket_cookie:
+		return &bpf_get_socket_cookie_proto;
  	default:
  		return NULL;
  	}

diff --git a/net/core/sock_diag.c b/net/core/sock_diag.c
index 6b10573cc9fa..b07b0a784f33 100644
--- a/net/core/sock_diag.c
+++ b/net/core/sock_diag.c

@@ -19,7 +19,7 @@ static int (*inet_rcv_compat)(struct sk_buff *skb, struct nlmsghdr *nlh);
  static DEFINE_MUTEX(sock_diag_table_mutex);
  static struct workqueue_struct *broadcast_wq;

-static u64 sock_gen_cookie(struct sock *sk)
+u64 sock_gen_cookie(struct sock *sk)
  {
  	while (1) {
  		u64 res = atomic64_read(&sk->sk_cookie);

@@ -30,6 +30,7 @@ static u64 sock_gen_cookie(struct sock *sk)
  		atomic64_cmpxchg(&sk->sk_cookie, 0, res);
  	}
  }
+EXPORT_SYMBOL_GPL(sock_gen_cookie);