[PATCH] bpf: Add bpf_skb_get_sock_comm() helper

[PATCH] bpf: Add bpf_skb_get_sock_comm() helper

[Jiang Yu]

skb distinguished by uid can only recorded to user who consume them.
in many case, skb should been recorded more specific to process who
consume them. E.g, the unexpected large data traffic of illegal process
in metered network.

this helper is used in tracing task comm of the sock to which a skb
belongs.

Signed-off-by: Jiang Yu <jyu.jiang@vivo.com>
---
 include/net/sock.h             |  1 +
 include/uapi/linux/bpf.h       |  1 +
 net/core/filter.c              | 32 ++++++++++++++++++++++++++++++++
 net/core/sock.c                | 20 ++++++++++++++++++++
 tools/include/uapi/linux/bpf.h |  1 +
 5 files changed, 55 insertions(+)

diff --git a/include/net/sock.h b/include/net/sock.h
index 064637d1ddf6..9c6e8e61940f 100644
--- a/include/net/sock.h
+++ b/include/net/sock.h
@@ -519,6 +519,7 @@ struct sock {
 #ifdef CONFIG_BPF_SYSCALL
 	struct bpf_sk_storage __rcu	*sk_bpf_storage;
 #endif
+	char sk_task_com[TASK_COMM_LEN];
 	struct rcu_head		sk_rcu;
 };
 
diff --git a/include/uapi/linux/bpf.h b/include/uapi/linux/bpf.h
index b134e679e9db..c7f62215a483 100644
--- a/include/uapi/linux/bpf.h
+++ b/include/uapi/linux/bpf.h
@@ -3538,6 +3538,7 @@ union bpf_attr {
 	FN(skc_to_tcp_request_sock),	\
 	FN(skc_to_udp6_sock),		\
 	FN(get_task_stack),		\
+	FN(skb_get_sock_comm),		\
 	/* */
 
 /* integer value in 'imm' field of BPF_CALL instruction selects which helper
diff --git a/net/core/filter.c b/net/core/filter.c
index 7124f0fe6974..972c0bf8e7ca 100644
--- a/net/core/filter.c
+++ b/net/core/filter.c
@@ -4313,6 +4313,36 @@ static const struct bpf_func_proto bpf_get_socket_uid_proto = {
 	.arg1_type      = ARG_PTR_TO_CTX,
 };
 
+BPF_CALL_3(bpf_skb_get_sock_comm,     struct sk_buff *, skb, char *, buf, u32, size)
+{
+	struct sock *sk;
+
+	if (!buf || 0 == size)
+		return -EINVAL;
+
+	sk = sk_to_full_sk(skb->sk);
+	if (!sk || !sk_fullsock(sk))
+		goto err_clear;
+
+	memcpy(buf, sk->sk_task_com, size);
+	buf[size - 1] = 0;
+	return 0;
+
+err_clear:
+	memset(buf, 0, size);
+	buf[size - 1] = 0;
+	return -ENOENT;
+}
+
+const struct bpf_func_proto bpf_skb_get_sock_comm_proto = {
+	.func           = bpf_skb_get_sock_comm,
+	.gpl_only       = false,
+	.ret_type       = RET_INTEGER,
+	.arg1_type      = ARG_PTR_TO_CTX,
+	.arg2_type      = ARG_PTR_TO_MEM,
+	.arg3_type      = ARG_CONST_SIZE,
+};
+
 #define SOCKOPT_CC_REINIT (1 << 0)
 
 static int _bpf_setsockopt(struct sock *sk, int level, int optname,
@@ -6313,6 +6343,8 @@ sk_filter_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
 		return &bpf_get_socket_cookie_proto;
 	case BPF_FUNC_get_socket_uid:
 		return &bpf_get_socket_uid_proto;
+	case BPF_FUNC_skb_get_sock_comm:
+		return &bpf_skb_get_sock_comm_proto;
 	case BPF_FUNC_perf_event_output:
 		return &bpf_skb_event_output_proto;
 	default:
diff --git a/net/core/sock.c b/net/core/sock.c
index d29709e0790d..79d81afa048f 100644
--- a/net/core/sock.c
+++ b/net/core/sock.c
@@ -2961,6 +2961,24 @@ void sk_stop_timer(struct sock *sk, struct timer_list* timer)
 }
 EXPORT_SYMBOL(sk_stop_timer);
 
+void sock_init_task_comm(struct sock *sk)
+{
+	struct pid *pid = NULL;
+	struct task_struct *tgid_task = NULL;
+
+	pid = find_get_pid(current->tgid);
+	if (pid) {
+		tgid_task = get_pid_task(pid, PIDTYPE_PID);
+
+		if (tgid_task) {
+			snprintf(sk->sk_task_com, TASK_COMM_LEN, tgid_task->comm);
+			put_task_struct(tgid_task);
+		}
+
+		put_pid(pid);
+	}
+}
+
 void sock_init_data(struct socket *sock, struct sock *sk)
 {
 	sk_init_common(sk);
@@ -3031,6 +3049,8 @@ void sock_init_data(struct socket *sock, struct sock *sk)
 	WRITE_ONCE(sk->sk_pacing_shift, 10);
 	sk->sk_incoming_cpu = -1;
 
+	sock_init_task_comm(sk);
+
 	sk_rx_queue_clear(sk);
 	/*
 	 * Before updating sk_refcnt, we must commit prior changes to memory
diff --git a/tools/include/uapi/linux/bpf.h b/tools/include/uapi/linux/bpf.h
index b134e679e9db..c7f62215a483 100644
--- a/tools/include/uapi/linux/bpf.h
+++ b/tools/include/uapi/linux/bpf.h
@@ -3538,6 +3538,7 @@ union bpf_attr {
 	FN(skc_to_tcp_request_sock),	\
 	FN(skc_to_udp6_sock),		\
 	FN(get_task_stack),		\
+	FN(skb_get_sock_comm),		\
 	/* */
 
 /* integer value in 'imm' field of BPF_CALL instruction selects which helper
-- 
2.25.1

Re: [PATCH] bpf: Add bpf_skb_get_sock_comm() helper

[Eric Dumazet]

On 8/10/20 6:09 AM, Jiang Yu wrote:
> skb distinguished by uid can only recorded to user who consume them.
> in many case, skb should been recorded more specific to process who
> consume them. E.g, the unexpected large data traffic of illegal process
> in metered network.
> 
> this helper is used in tracing task comm of the sock to which a skb
> belongs.
> 
> Signed-off-by: Jiang Yu <jyu.jiang@vivo.com>

fd can be passed from one process to another.

Re: [PATCH] bpf: Add bpf_skb_get_sock_comm() helper

[Martin KaFai Lau]

On Mon, Aug 10, 2020 at 06:09:48AM -0700, Jiang Yu wrote:
> skb distinguished by uid can only recorded to user who consume them.
> in many case, skb should been recorded more specific to process who
> consume them. E.g, the unexpected large data traffic of illegal process
> in metered network.
> 
> this helper is used in tracing task comm of the sock to which a skb
> belongs.
> 
> Signed-off-by: Jiang Yu <jyu.jiang@vivo.com>
> ---
>  include/net/sock.h             |  1 +
>  include/uapi/linux/bpf.h       |  1 +
>  net/core/filter.c              | 32 ++++++++++++++++++++++++++++++++
>  net/core/sock.c                | 20 ++++++++++++++++++++
>  tools/include/uapi/linux/bpf.h |  1 +
>  5 files changed, 55 insertions(+)
> 
> diff --git a/include/net/sock.h b/include/net/sock.h
> index 064637d1ddf6..9c6e8e61940f 100644
> --- a/include/net/sock.h
> +++ b/include/net/sock.h
> @@ -519,6 +519,7 @@ struct sock {
>  #ifdef CONFIG_BPF_SYSCALL
>  	struct bpf_sk_storage __rcu	*sk_bpf_storage;
>  #endif
> +	char sk_task_com[TASK_COMM_LEN];
One possibility is to use the "sk_bpf_storage" member immediately above
instead of adding "sk_task_com[]".

It is an extensible sk storage for bpf.  There are examples in selftests,
e.g tools/testing/selftests/bpf/progs/udp_limits.c which creates sk storage
at socket creation time.  Another hook point option could be "connect()"
for tcp, i.e. "cgroup/connect[46]".

Search "BPF_MAP_TYPE_SK_STORAGE" under the selftests/bpf for other examples.

It seems there is already a "bpf_get_current_comm()" helper which
could be used to initialize the task comm string in the bpf sk storage.

btw, bpf-next is still closed.

Re: [PATCH] bpf: Add bpf_skb_get_sock_comm() helper

[Martin KaFai Lau]

On Wed, Aug 19, 2020 at 02:10:56PM +0800, 江禹 wrote:
> Dear Martin,
> 
> 
> > One possibility is to use the "sk_bpf_storage" member immediately above
> > instead of adding "sk_task_com[]".
> > 
> > It is an extensible sk storage for bpf.  There are examples in selftests,
> > e.g tools/testing/selftests/bpf/progs/udp_limits.c which creates sk storage
> > at socket creation time.  Another hook point option could be "connect()"
> > for tcp, i.e. "cgroup/connect[46]".
> > 
> > Search "BPF_MAP_TYPE_SK_STORAGE" under the selftests/bpf for other examples.
> > 
> > It seems there is already a "bpf_get_current_comm()" helper which
> > could be used to initialize the task comm string in the bpf sk storage.
> > 
> > btw, bpf-next is still closed.
> 
> 
> I have rewrite my code according to your suggestion.  In general,it works as designed.
> 
> 
> But the task comm string got from "bpf_get_current_comm()" helper is belong to specific thread.
> It is not a obvious label for skb tracing. More reasonable tracing key is the task comm of process
> which this skb belongs.
>  
> It seems a new bpf helper is still needed.
May be.  It is not clear to me whether it is better to account this skb
to its process or to its thread.  I think this depends on the
use-case/assumption.

If bpf_get_current_comm() does not get what you need,
then another helper may be added to get the process comm.
This new helper may be useful for other use cases also.

btw, Have your thought about tracing at the sendmsg time?
e.g. tcp_sendmsg()?