diff --git a/fastfetch_server.c b/fastfetch_server.c index 6d0c8a0..8c2f639 100644 --- a/fastfetch_server.c +++ b/fastfetch_server.c @@ -24,6 +24,7 @@ static struct { typedef struct pthread_input_t { int client_fd; float window; + int rate_limit; const char* command; } pthread_input_t; @@ -33,7 +34,7 @@ static double now(void) { return ts.tv_sec + ts.tv_nsec * 1e-9; } -static int rate_limited(const uint8_t addr[16], float window) { +static int rate_limited(const uint8_t addr[16], float window, int rate_limit) { double n = now(); pthread_mutex_lock(&ratelimit_table.lock); @@ -58,7 +59,7 @@ static int rate_limited(const uint8_t addr[16], float window) { e->count = 0; } - const int allowed = e->count < RATE_LIMIT; + const int allowed = e->count < rate_limit; if (allowed) e->count++; pthread_mutex_unlock(&ratelimit_table.lock); @@ -69,13 +70,14 @@ static void* handle_client(void* _arg) { pthread_input_t* arg = _arg; int client_fd = arg->client_fd; float window = arg->window; + int rate_limit = arg->rate_limit; const char* command = arg->command; struct sockaddr_in6 peer; socklen_t peerlen = sizeof(peer); getpeername(client_fd, (struct sockaddr*)&peer, &peerlen); - if (rate_limited(peer.sin6_addr.s6_addr, window)) { + if (rate_limited(peer.sin6_addr.s6_addr, window, rate_limit)) { dprintf(client_fd, "HTTP/1.1 429 Too Many Requests\r\n" "Content-Type: text/plain\r\n" @@ -83,7 +85,6 @@ static void* handle_client(void* _arg) { "Connection: close\r\n" "\r\n" "Too Many Requests\r\n" - "HTTP/1.1 200 OK\r\n" ); goto cleanup; } @@ -125,6 +126,7 @@ cleanup: int main(int argc, char *argv[]) { int port = 80; float window = 60.0f; + int rate_limit = 5; const char* command = "fastfetch --pipe false 2>/dev/null"; struct in6_addr bind_addr = in6addr_any; @@ -133,10 +135,12 @@ int main(int argc, char *argv[]) { puts("usage:\n" "\tfastfetch_server --addr 127.0.01 (optional)\n" "\t --window 60 (optional)\n" + "\t --ratelimit 5 (optional)" "\t --hyfetch (optional)\n" "\t port (optional)\n" "\t --addr expects an address in an IPv4 or IPv6 format\n" - "\t --window sets a ratelimit delay in seconds\n" + "\t --window sets a rate limit delay in seconds\n" + "\t --ratelimit sets the rate limit amount" "\t --hyfetch uses hyfetch with fastfetch backend instead of fastfetch\n" "\t port specifies the port, 80 by default, might need superuser for\n" "\t ports below 1024" @@ -158,7 +162,10 @@ int main(int argc, char *argv[]) { } } else if (streq(argv[i], "--window") && i + 1 < argc) { - window = atof(argv[i++]); + window = atof(argv[++i]); + } + else if (streq(argv[i], "--ratelimit") && i + 1 < argc) { + rate_limit = atoi(argv[++i]); } else if (streq(argv[i], "--hyfetch")) { command = "hyfetch --backend=fastfetch --args=\"--pipe false\" 2>/dev/null"; @@ -188,7 +195,7 @@ int main(int argc, char *argv[]) { char addr_str[INET6_ADDRSTRLEN]; inet_ntop(AF_INET6, &bind_addr, addr_str, sizeof(addr_str)); printf("listening on [%s]:%d (limit: %d req / %fs per IP)\n", - addr_str, port, RATE_LIMIT, window); + addr_str, port, rate_limit, window); for (;;) { struct sockaddr_in6 client_addr; @@ -196,7 +203,7 @@ int main(int argc, char *argv[]) { int client_fd = accept(server_fd, (struct sockaddr*)&client_addr, &addrlen); if (client_fd < 0) continue; - pthread_input_t pi = {.client_fd = client_fd, .window = window, .command = command}; + pthread_input_t pi = {.client_fd = client_fd, .window = window, .rate_limit = rate_limit, .command = command}; pthread_t tid; if (pthread_create(&tid, NULL, handle_client, (void*)(intptr_t)&pi) == 0)