// SPDX-License-Identifier: GPL-2.0

#include <test_progs.h>

#include "connect4_dropper.skel.h"

#include "cgroup_helpers.h"
#include "network_helpers.h"

static int run_test(int cgroup_fd, int server_fd, bool classid)
{
 struct connect4_dropper *skel;
 int fd, err = 0, port;

 skel = connect4_dropper__open_and_load();
 if (!ASSERT_OK_PTR(skel, "skel_open"))
  return -1;

 port = get_socket_local_port(server_fd);
 if (!ASSERT_GE(port, 0, "get_socket_local_port"))
  return -1;

 skel->bss->port = ntohs(port);

 skel->links.connect_v4_dropper =
  bpf_program__attach_cgroup(skel->progs.connect_v4_dropper,
        cgroup_fd);
 if (!ASSERT_OK_PTR(skel->links.connect_v4_dropper, "prog_attach")) {
  err = -1;
  goto out;
 }

 if (classid && !ASSERT_OK(join_classid(), "join_classid")) {
  err = -1;
  goto out;
 }

 errno = 0;
 fd = connect_to_fd_opts(server_fd, NULL);
 if (fd >= 0) {
  log_err("Unexpected success to connect to server");
  err = -1;
  close(fd);
 } else if (errno != EPERM) {
  log_err("Unexpected errno from connect to server");
  err = -1;
 }
out:
 connect4_dropper__destroy(skel);
 return err;
}

void test_cgroup_v1v2(void)
{
 struct network_helper_opts opts = {};
 int server_fd, client_fd, cgroup_fd;

 /* Step 1: Check base connectivity works without any BPF. */
 server_fd = start_server(AF_INET, SOCK_STREAM, NULL, 0, 0);
 if (!ASSERT_GE(server_fd, 0, "server_fd"))
  return;
 client_fd = connect_to_fd_opts(server_fd, &opts);
 if (!ASSERT_GE(client_fd, 0, "client_fd")) {
  close(server_fd);
  return;
 }
 close(client_fd);
 close(server_fd);

 /* Step 2: Check BPF policy prog attached to cgroups drops connectivity. */
 cgroup_fd = test__join_cgroup("/connect_dropper");
 if (!ASSERT_GE(cgroup_fd, 0, "cgroup_fd"))
  return;
 server_fd = start_server(AF_INET, SOCK_STREAM, NULL, 0, 0);
 if (!ASSERT_GE(server_fd, 0, "server_fd")) {
  close(cgroup_fd);
  return;
 }
 ASSERT_OK(run_test(cgroup_fd, server_fd, false), "cgroup-v2-only");
 setup_classid_environment();
 set_classid();
 ASSERT_OK(run_test(cgroup_fd, server_fd, true), "cgroup-v1v2");
 cleanup_classid_environment();
 close(server_fd);
 close(cgroup_fd);
}