Merge branch 'socket-lookup-cg_sock'

Andrey Ignatov says:

====================
This patch set makes bpf_sk_lookup_tcp, bpf_sk_lookup_udp and
bpf_sk_release helpers available in programs of type
BPF_PROG_TYPE_CGROUP_SOCK_ADDR.

Patch 1 is a fix for bpf_sk_lookup_udp that was already merged to bpf
(stable) tree. Here it's prerequisite for patch 3.

Patch 2 is the main patch in the set, it makes the helpers available for
BPF_PROG_TYPE_CGROUP_SOCK_ADDR and provides more details about use-case.

Patch 3 adds selftest for new functionality.

v1->v2:
- remove "Split bpf_sk_lookup" patch since it was already split by:
  commit c8123ead13a5 ("bpf: Extend the sk_lookup() helper to XDP
  hookpoint.");
- avoid unnecessary bpf_sock_addr_sk_lookup function.
====================

Signed-off-by: Alexei Starovoitov <ast@kernel.org>

3 files changed, 125 insertions, 24 deletions

diff --git a/net/core/filter.c b/net/core/filter.c
index 53d50fb75ea1..f6ca38a7d433 100644
--- a/net/core/filter.c
+++ b/net/core/filter.c
@@ -4867,17 +4867,16 @@ static struct sock *sk_lookup(struct net *net, struct bpf_sock_tuple *tuple,
 	} else {
 		struct in6_addr *src6 = (struct in6_addr *)&tuple->ipv6.saddr;
 		struct in6_addr *dst6 = (struct in6_addr *)&tuple->ipv6.daddr;
-		u16 hnum = ntohs(tuple->ipv6.dport);
 
 		if (proto == IPPROTO_TCP)
 			sk = __inet6_lookup(net, &tcp_hashinfo, NULL, 0,
 					    src6, tuple->ipv6.sport,
-					    dst6, hnum,
+					    dst6, ntohs(tuple->ipv6.dport),
 					    dif, sdif, &refcounted);
 		else if (likely(ipv6_bpf_stub))
 			sk = ipv6_bpf_stub->udp6_lib_lookup(net,
 							    src6, tuple->ipv6.sport,
-							    dst6, hnum,
+							    dst6, tuple->ipv6.dport,
 							    dif, sdif,
 							    &udp_table, NULL);
 #endif
@@ -5043,6 +5042,43 @@ static const struct bpf_func_proto bpf_xdp_sk_lookup_tcp_proto = {
 	.arg4_type      = ARG_ANYTHING,
 	.arg5_type      = ARG_ANYTHING,
 };
+
+BPF_CALL_5(bpf_sock_addr_sk_lookup_tcp, struct bpf_sock_addr_kern *, ctx,
+	   struct bpf_sock_tuple *, tuple, u32, len, u64, netns_id, u64, flags)
+{
+	return __bpf_sk_lookup(NULL, tuple, len, sock_net(ctx->sk), 0,
+			       IPPROTO_TCP, netns_id, flags);
+}
+
+static const struct bpf_func_proto bpf_sock_addr_sk_lookup_tcp_proto = {
+	.func		= bpf_sock_addr_sk_lookup_tcp,
+	.gpl_only	= false,
+	.ret_type	= RET_PTR_TO_SOCKET_OR_NULL,
+	.arg1_type	= ARG_PTR_TO_CTX,
+	.arg2_type	= ARG_PTR_TO_MEM,
+	.arg3_type	= ARG_CONST_SIZE,
+	.arg4_type	= ARG_ANYTHING,
+	.arg5_type	= ARG_ANYTHING,
+};
+
+BPF_CALL_5(bpf_sock_addr_sk_lookup_udp, struct bpf_sock_addr_kern *, ctx,
+	   struct bpf_sock_tuple *, tuple, u32, len, u64, netns_id, u64, flags)
+{
+	return __bpf_sk_lookup(NULL, tuple, len, sock_net(ctx->sk), 0,
+			       IPPROTO_UDP, netns_id, flags);
+}
+
+static const struct bpf_func_proto bpf_sock_addr_sk_lookup_udp_proto = {
+	.func		= bpf_sock_addr_sk_lookup_udp,
+	.gpl_only	= false,
+	.ret_type	= RET_PTR_TO_SOCKET_OR_NULL,
+	.arg1_type	= ARG_PTR_TO_CTX,
+	.arg2_type	= ARG_PTR_TO_MEM,
+	.arg3_type	= ARG_CONST_SIZE,
+	.arg4_type	= ARG_ANYTHING,
+	.arg5_type	= ARG_ANYTHING,
+};
+
 #endif /* CONFIG_INET */
 
 bool bpf_helper_changes_pkt_data(void *func)
@@ -5149,6 +5185,14 @@ sock_addr_func_proto(enum bpf_func_id func_id, const struct bpf_prog *prog)
 		return &bpf_get_socket_cookie_sock_addr_proto;
 	case BPF_FUNC_get_local_storage:
 		return &bpf_get_local_storage_proto;
+#ifdef CONFIG_INET
+	case BPF_FUNC_sk_lookup_tcp:
+		return &bpf_sock_addr_sk_lookup_tcp_proto;
+	case BPF_FUNC_sk_lookup_udp:
+		return &bpf_sock_addr_sk_lookup_udp_proto;
+	case BPF_FUNC_sk_release:
+		return &bpf_sk_release_proto;
+#endif /* CONFIG_INET */
 	default:
 		return bpf_base_func_proto(func_id);
 	}
diff --git a/tools/testing/selftests/bpf/connect4_prog.c b/tools/testing/selftests/bpf/connect4_prog.c
index 5a88a681d2ab..b8395f3c43e9 100644
--- a/tools/testing/selftests/bpf/connect4_prog.c
+++ b/tools/testing/selftests/bpf/connect4_prog.c
@@ -21,23 +21,48 @@ int _version SEC("version") = 1;
 SEC("cgroup/connect4")
 int connect_v4_prog(struct bpf_sock_addr *ctx)
 {
+	struct bpf_sock_tuple tuple = {};
 	struct sockaddr_in sa;
+	struct bpf_sock *sk;
+
+	/* Verify that new destination is available. */
+	memset(&tuple.ipv4.saddr, 0, sizeof(tuple.ipv4.saddr));
+	memset(&tuple.ipv4.sport, 0, sizeof(tuple.ipv4.sport));
+
+	tuple.ipv4.daddr = bpf_htonl(DST_REWRITE_IP4);
+	tuple.ipv4.dport = bpf_htons(DST_REWRITE_PORT4);
+
+	if (ctx->type != SOCK_STREAM && ctx->type != SOCK_DGRAM)
+		return 0;
+	else if (ctx->type == SOCK_STREAM)
+		sk = bpf_sk_lookup_tcp(ctx, &tuple, sizeof(tuple.ipv4), 0, 0);
+	else
+		sk = bpf_sk_lookup_udp(ctx, &tuple, sizeof(tuple.ipv4), 0, 0);
+
+	if (!sk)
+		return 0;
+
+	if (sk->src_ip4 != tuple.ipv4.daddr ||
+	    sk->src_port != DST_REWRITE_PORT4) {
+		bpf_sk_release(sk);
+		return 0;
+	}
+
+	bpf_sk_release(sk);
 
 	/* Rewrite destination. */
 	ctx->user_ip4 = bpf_htonl(DST_REWRITE_IP4);
 	ctx->user_port = bpf_htons(DST_REWRITE_PORT4);
 
-	if (ctx->type == SOCK_DGRAM || ctx->type == SOCK_STREAM) {
-		///* Rewrite source. */
-		memset(&sa, 0, sizeof(sa));
+	/* Rewrite source. */
+	memset(&sa, 0, sizeof(sa));
 
-		sa.sin_family = AF_INET;
-		sa.sin_port = bpf_htons(0);
-		sa.sin_addr.s_addr = bpf_htonl(SRC_REWRITE_IP4);
+	sa.sin_family = AF_INET;
+	sa.sin_port = bpf_htons(0);
+	sa.sin_addr.s_addr = bpf_htonl(SRC_REWRITE_IP4);
 
-		if (bpf_bind(ctx, (struct sockaddr *)&sa, sizeof(sa)) != 0)
-			return 0;
-	}
+	if (bpf_bind(ctx, (struct sockaddr *)&sa, sizeof(sa)) != 0)
+		return 0;
 
 	return 1;
 }
diff --git a/tools/testing/selftests/bpf/connect6_prog.c b/tools/testing/selftests/bpf/connect6_prog.c
index 8ea3f7d12dee..25f5dc7b7aa0 100644
--- a/tools/testing/selftests/bpf/connect6_prog.c
+++ b/tools/testing/selftests/bpf/connect6_prog.c
@@ -29,7 +29,41 @@ int _version SEC("version") = 1;
 SEC("cgroup/connect6")
 int connect_v6_prog(struct bpf_sock_addr *ctx)
 {
+	struct bpf_sock_tuple tuple = {};
 	struct sockaddr_in6 sa;
+	struct bpf_sock *sk;
+
+	/* Verify that new destination is available. */
+	memset(&tuple.ipv6.saddr, 0, sizeof(tuple.ipv6.saddr));
+	memset(&tuple.ipv6.sport, 0, sizeof(tuple.ipv6.sport));
+
+	tuple.ipv6.daddr[0] = bpf_htonl(DST_REWRITE_IP6_0);
+	tuple.ipv6.daddr[1] = bpf_htonl(DST_REWRITE_IP6_1);
+	tuple.ipv6.daddr[2] = bpf_htonl(DST_REWRITE_IP6_2);
+	tuple.ipv6.daddr[3] = bpf_htonl(DST_REWRITE_IP6_3);
+
+	tuple.ipv6.dport = bpf_htons(DST_REWRITE_PORT6);
+
+	if (ctx->type != SOCK_STREAM && ctx->type != SOCK_DGRAM)
+		return 0;
+	else if (ctx->type == SOCK_STREAM)
+		sk = bpf_sk_lookup_tcp(ctx, &tuple, sizeof(tuple.ipv6), 0, 0);
+	else
+		sk = bpf_sk_lookup_udp(ctx, &tuple, sizeof(tuple.ipv6), 0, 0);
+
+	if (!sk)
+		return 0;
+
+	if (sk->src_ip6[0] != tuple.ipv6.daddr[0] ||
+	    sk->src_ip6[1] != tuple.ipv6.daddr[1] ||
+	    sk->src_ip6[2] != tuple.ipv6.daddr[2] ||
+	    sk->src_ip6[3] != tuple.ipv6.daddr[3] ||
+	    sk->src_port != DST_REWRITE_PORT6) {
+		bpf_sk_release(sk);
+		return 0;
+	}
+
+	bpf_sk_release(sk);
 
 	/* Rewrite destination. */
 	ctx->user_ip6[0] = bpf_htonl(DST_REWRITE_IP6_0);
@@ -39,21 +73,19 @@ int connect_v6_prog(struct bpf_sock_addr *ctx)
 
 	ctx->user_port = bpf_htons(DST_REWRITE_PORT6);
 
-	if (ctx->type == SOCK_DGRAM || ctx->type == SOCK_STREAM) {
-		/* Rewrite source. */
-		memset(&sa, 0, sizeof(sa));
+	/* Rewrite source. */
+	memset(&sa, 0, sizeof(sa));
 
-		sa.sin6_family = AF_INET6;
-		sa.sin6_port = bpf_htons(0);
+	sa.sin6_family = AF_INET6;
+	sa.sin6_port = bpf_htons(0);
 
-		sa.sin6_addr.s6_addr32[0] = bpf_htonl(SRC_REWRITE_IP6_0);
-		sa.sin6_addr.s6_addr32[1] = bpf_htonl(SRC_REWRITE_IP6_1);
-		sa.sin6_addr.s6_addr32[2] = bpf_htonl(SRC_REWRITE_IP6_2);
-		sa.sin6_addr.s6_addr32[3] = bpf_htonl(SRC_REWRITE_IP6_3);
+	sa.sin6_addr.s6_addr32[0] = bpf_htonl(SRC_REWRITE_IP6_0);
+	sa.sin6_addr.s6_addr32[1] = bpf_htonl(SRC_REWRITE_IP6_1);
+	sa.sin6_addr.s6_addr32[2] = bpf_htonl(SRC_REWRITE_IP6_2);
+	sa.sin6_addr.s6_addr32[3] = bpf_htonl(SRC_REWRITE_IP6_3);
 
-		if (bpf_bind(ctx, (struct sockaddr *)&sa, sizeof(sa)) != 0)
-			return 0;
-	}
+	if (bpf_bind(ctx, (struct sockaddr *)&sa, sizeof(sa)) != 0)
+		return 0;
 
 	return 1;
 }