1 files changed, 50 insertions, 30 deletions

diff --git a/samples/bpf/test_cgrp2_attach.c b/samples/bpf/test_cgrp2_attach.c
index 63ef2083f766..a19484c45b79 100644
--- a/samples/bpf/test_cgrp2_attach.c
+++ b/samples/bpf/test_cgrp2_attach.c
@@ -10,8 +10,6 @@
  *   incremented on each iteration by the number of bytes stored in
  *   the skb.
  *
- * - Detaches any eBPF program previously attached to the cgroup
- *
  * - Attaches the new program to a cgroup using BPF_PROG_ATTACH
  *
  * - Every second, reads map[0] and map[1] to see how many bytes and
@@ -75,35 +73,16 @@ static int prog_load(int map_fd, int verdict)
 
 static int usage(const char *argv0)
 {
-	printf("Usage: %s <cg-path> <egress|ingress> [drop]\n", argv0);
+	printf("Usage: %s [-d] [-D] <cg-path> <egress|ingress>\n", argv0);
+	printf("	-d	Drop Traffic\n");
+	printf("	-D	Detach filter, and exit\n");
 	return EXIT_FAILURE;
 }
 
-int main(int argc, char **argv)
+static int attach_filter(int cg_fd, int type, int verdict)
 {
-	int cg_fd, map_fd, prog_fd, key, ret;
+	int prog_fd, map_fd, ret, key;
 	long long pkt_cnt, byte_cnt;
-	enum bpf_attach_type type;
-	int verdict = 1;
-
-	if (argc < 3)
-		return usage(argv[0]);
-
-	if (strcmp(argv[2], "ingress") == 0)
-		type = BPF_CGROUP_INET_INGRESS;
-	else if (strcmp(argv[2], "egress") == 0)
-		type = BPF_CGROUP_INET_EGRESS;
-	else
-		return usage(argv[0]);
-
-	if (argc > 3 && strcmp(argv[3], "drop") == 0)
-		verdict = 0;
-
-	cg_fd = open(argv[1], O_DIRECTORY | O_RDONLY);
-	if (cg_fd < 0) {
-		printf("Failed to open cgroup path: '%s'\n", strerror(errno));
-		return EXIT_FAILURE;
-	}
 
 	map_fd = bpf_create_map(BPF_MAP_TYPE_ARRAY,
 				sizeof(key), sizeof(byte_cnt),
@@ -121,16 +100,12 @@ int main(int argc, char **argv)
 		return EXIT_FAILURE;
 	}
 
-	ret = bpf_prog_detach(cg_fd, type);
-	printf("bpf_prog_detach() returned '%s' (%d)\n", strerror(errno), errno);
-
 	ret = bpf_prog_attach(prog_fd, cg_fd, type);
 	if (ret < 0) {
 		printf("Failed to attach prog to cgroup: '%s'\n",
 		       strerror(errno));
 		return EXIT_FAILURE;
 	}
-
 	while (1) {
 		key = MAP_KEY_PACKETS;
 		assert(bpf_lookup_elem(map_fd, &key, &pkt_cnt) == 0);
@@ -145,3 +120,48 @@ int main(int argc, char **argv)
 
 	return EXIT_SUCCESS;
 }
+
+int main(int argc, char **argv)
+{
+	int detach_only = 0, verdict = 1;
+	enum bpf_attach_type type;
+	int opt, cg_fd, ret;
+
+	while ((opt = getopt(argc, argv, "Dd")) != -1) {
+		switch (opt) {
+		case 'd':
+			verdict = 0;
+			break;
+		case 'D':
+			detach_only = 1;
+			break;
+		default:
+			return usage(argv[0]);
+		}
+	}
+
+	if (argc - optind < 2)
+		return usage(argv[0]);
+
+	if (strcmp(argv[optind + 1], "ingress") == 0)
+		type = BPF_CGROUP_INET_INGRESS;
+	else if (strcmp(argv[optind + 1], "egress") == 0)
+		type = BPF_CGROUP_INET_EGRESS;
+	else
+		return usage(argv[0]);
+
+	cg_fd = open(argv[optind], O_DIRECTORY | O_RDONLY);
+	if (cg_fd < 0) {
+		printf("Failed to open cgroup path: '%s'\n", strerror(errno));
+		return EXIT_FAILURE;
+	}
+
+	if (detach_only) {
+		ret = bpf_prog_detach(cg_fd, type);
+		printf("bpf_prog_detach() returned '%s' (%d)\n",
+		       strerror(errno), errno);
+	} else
+		ret = attach_filter(cg_fd, type, verdict);
+
+	return ret;
+}