diff --git a/ebpf/aw-bpfctl.c b/ebpf/aw-bpfctl.c index 022ef4d6..913372e9 100644 --- a/ebpf/aw-bpfctl.c +++ b/ebpf/aw-bpfctl.c @@ -69,6 +69,8 @@ print_stats_ipv4(__be32 ip, struct traffic_stats *stats) stats->incoming.total_bytes, stats->incoming.total_packets, calc_rate_estimator(stats, true)); printf(" Outgoing: total_bytes=%llu, total_packets=%llu, rate=%u\n", stats->outgoing.total_bytes, stats->outgoing.total_packets, calc_rate_estimator(stats, false)); + printf(" Rate Limits: incoming=%u bps, outgoing=%u bps\n", + stats->incoming_rate_limit, stats->outgoing_rate_limit); } static void @@ -82,6 +84,8 @@ print_stats_ipv6(struct in6_addr ip, struct traffic_stats *stats) stats->incoming.total_bytes, stats->incoming.total_packets, calc_rate_estimator(stats, true)); printf(" Outgoing: total_bytes=%llu, total_packets=%llu, rate=%u\n", stats->outgoing.total_bytes, stats->outgoing.total_packets, calc_rate_estimator(stats, false)); + printf(" Rate Limits: incoming=%u bps, outgoing=%u bps\n", + stats->incoming_rate_limit, stats->outgoing_rate_limit); } static void @@ -94,6 +98,8 @@ print_stats_mac(struct mac_addr mac, struct traffic_stats *stats) stats->incoming.total_bytes, stats->incoming.total_packets, calc_rate_estimator(stats, true)); printf(" Outgoing: total_bytes=%llu, total_packets=%llu, rate=%u\n", stats->outgoing.total_bytes, stats->outgoing.total_packets, calc_rate_estimator(stats, false)); + printf(" Rate Limits: incoming=%u bps, outgoing=%u bps\n", + stats->incoming_rate_limit, stats->outgoing_rate_limit); } static struct json_object* @@ -239,6 +245,7 @@ main(int argc, char **argv) { fprintf(stderr, " %s del \n", argv[0]); fprintf(stderr, " %s flush\n", argv[0]); fprintf(stderr, " %s json\n", argv[0]); + fprintf(stderr, " %s update downrate uprate \n", argv[0]); return EXIT_FAILURE; } @@ -487,6 +494,73 @@ main(int argc, char **argv) { return EXIT_FAILURE; } printf("Flushed all entries in the map.\n"); + } else if (strcmp(cmd, "update") == 0) { + if (argc != 8) { + fprintf(stderr, "Usage: %s %s update downrate uprate \n", argv[0], map_type); + return EXIT_FAILURE; + } + + const char *ip_str = argv[3]; + const char *downrate_str = argv[5]; + const char *uprate_str = argv[7]; + uint32_t downrate = atoi(downrate_str); + uint32_t uprate = atoi(uprate_str); + + if (downrate == 0 || uprate == 0) { + fprintf(stderr, "Invalid rate value\n"); + return EXIT_FAILURE; + } + + void *key = NULL; + int key_size = 0; + bool success = false; + + if (strcmp(map_type, "ipv4") == 0) { + __be32 ipv4_key = 0; + if (inet_pton(AF_INET, ip_str, &ipv4_key) != 1) { + perror("inet_pton (IPv4)"); + return EXIT_FAILURE; + } + key = &ipv4_key; + key_size = sizeof(ipv4_key); + success = true; + } else if (strcmp(map_type, "mac") == 0) { + if (!is_valid_mac_addr(ip_str)) { + fprintf(stderr, "Invalid MAC address format\n"); + return EXIT_FAILURE; + } + struct mac_addr mac_key = {0}; + parse_mac_address(&mac_key, ip_str); + key = &mac_key; + key_size = sizeof(mac_key); + success = true; + } else if (strcmp(map_type, "ipv6") == 0) { + struct in6_addr ipv6_key = {0}; + if (inet_pton(AF_INET6, ip_str, &ipv6_key) != 1) { + perror("inet_pton (IPv6)"); + return EXIT_FAILURE; + } + key = &ipv6_key; + key_size = sizeof(ipv6_key); + success = true; + } + + if (!success || !key) { + fprintf(stderr, "Invalid address type\n"); + return EXIT_FAILURE; + } + + struct traffic_stats stats = {0}; + if (bpf_map_lookup_elem(map_fd, key, &stats) < 0) { + // add new entry if not found + if (bpf_map_update_elem(map_fd, key, &stats, BPF_NOEXIST) < 0) { + perror("bpf_map_update_elem"); + return EXIT_FAILURE; + } + } + + stats.incoming_rate_limit = downrate; + stats.outgoing_rate_limit = uprate; } else { fprintf(stderr, "Invalid command. Use 'add ' or 'list'.\n"); return EXIT_FAILURE; diff --git a/ebpf/aw-loader.c b/ebpf/aw-loader.c new file mode 100644 index 00000000..397868ec --- /dev/null +++ b/ebpf/aw-loader.c @@ -0,0 +1,59 @@ +#include +#include +#include +#include +#include +#include + +static volatile bool running = true; + +static void sig_handler(int sig) +{ + running = false; +} + +int main(int argc, char **argv) +{ + struct bpf_object *obj; + struct bpf_program *prog; + int err; + + // Set up signal handler + signal(SIGINT, sig_handler); + signal(SIGTERM, sig_handler); + + // Open and load the BPF object file + obj = bpf_object__open_file("aw-bpf.o", NULL); + if (libbpf_get_error(obj)) { + fprintf(stderr, "Error opening BPF object file\n"); + return 1; + } + + prog = bpf_object__find_program_by_name(obj, "tc_ingress"); + if (!prog) { + fprintf(stderr, "Error finding BPF program\n"); + bpf_object__close(obj); + return 1; + } + + bpf_program__set_type(prog, BPF_PROG_TYPE_SCHED_CLS); + + // Load the BPF program into the kernel + err = bpf_object__load(obj); + if (err) { + fprintf(stderr, "Error loading BPF object: %d\n", err); + bpf_object__close(obj); + return 1; + } + + printf("BPF program loaded successfully\n"); + + // Keep the program running until signal is received + while (running) { + sleep(1); + } + + // Cleanup + bpf_object__close(obj); + return 0; +} \ No newline at end of file