summaryrefslogtreecommitdiff
path: root/src/omen.c
diff options
context:
space:
mode:
Diffstat (limited to 'src/omen.c')
-rw-r--r--src/omen.c553
1 files changed, 553 insertions, 0 deletions
diff --git a/src/omen.c b/src/omen.c
new file mode 100644
index 0000000..222e702
--- /dev/null
+++ b/src/omen.c
@@ -0,0 +1,553 @@
+#pragma once
+
+#include <arpa/inet.h>
+#include <errno.h>
+#include <netinet/in.h>
+#include <pthread.h>
+#include <signal.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <strings.h>
+#include <sys/socket.h>
+#include <unistd.h>
+//
+#include "common.c"
+#include "ratelimit.c"
+
+typedef uint8_t byte;
+typedef void (*tome_ptr)(
+ char *ip_addr, size_t addr_size, byte *in, size_t in_len, byte **out, size_t *out_len);
+
+#define MAX_BUFFER 1 << 18
+#define MAX_BODY 1 << 22
+#define MAX_CONNECTIONS 2048
+
+// vars
+
+int worker_count = 64;
+
+pthread_mutex_t queue_mutex = PTHREAD_MUTEX_INITIALIZER;
+pthread_cond_t queue_cond = PTHREAD_COND_INITIALIZER;
+
+static int ring_fds[ MAX_CONNECTIONS ];
+static char ring_ips[ MAX_CONNECTIONS ][ INET6_ADDRSTRLEN ];
+static size_t ring_head = 0;
+static size_t ring_tail = 0;
+size_t queue_size = 0;
+
+tome_ptr global_handler = NULL;
+
+// utils
+
+static ssize_t write_all(int fd, const void *buf, size_t len) {
+ size_t off = 0;
+ const char *p = buf;
+ while (off < len) {
+ ssize_t r = send(fd, p + off, len - off, MSG_NOSIGNAL);
+ if (r < 0) {
+ if (errno == EINTR) continue;
+ if (errno == EAGAIN || errno == EWOULDBLOCK) continue;
+ return -1;
+ }
+ off += (size_t) r;
+ }
+ return (ssize_t) off;
+}
+
+static int read_until_headers(int fd, char **out_buf, size_t *out_len) {
+ size_t total_tries = 0;
+
+ size_t cap = 4096;
+ size_t len = 0;
+ char *buf = malloc(cap);
+ if (!buf) return -1;
+ while (1) {
+ if (total_tries++ > 128) break;
+
+ if (cap - len < 1024) {
+ size_t ncap = cap * 2;
+ char *n = realloc(buf, ncap);
+ if (!n) {
+ free(buf);
+ return -1;
+ }
+ buf = n;
+ cap = ncap;
+ }
+ ssize_t r = recv(fd, buf + len, cap - len, 0);
+ if (r < 0) {
+ if (errno == EINTR) continue;
+ free(buf);
+ return -1;
+ } else if (r == 0) {
+ // TODO: why?
+ // len = len;
+ break;
+ }
+ len += (size_t) r;
+ if (len >= 4 && memmem(buf, len, "\r\n\r\n", 4) != NULL) {
+ if (len >= 19 && memmem(buf, len, "Transfer-Encoding: chunked", 19) != NULL) {
+ const char *err = "HTTP/1.1 501 Not Implemented\r\n"
+ "Content-Length: 31\r\n"
+ "Connection: close\r\n\r\n"
+ "chunked encoding not supported";
+ write_all(fd, err, strlen(err));
+ free(buf);
+ return -1;
+ }
+ break;
+ }
+ if (len > MAX_BUFFER) break;
+ }
+ *out_buf = buf;
+ *out_len = len;
+ return 0;
+}
+
+static int read_exact(int fd, void *buf, size_t count) {
+ size_t off = 0;
+ char *p = buf;
+ while (off < count) {
+ ssize_t r = recv(fd, p + off, count - off, 0);
+ if (r < 0) {
+ if (errno == EINTR) continue;
+ return -1;
+ } else if (r == 0) {
+ return -1;
+ }
+ off += (size_t) r;
+ }
+ return 0;
+}
+
+static ssize_t parse_content_length(const char *hdr, size_t hdr_len) {
+ const char *p = hdr;
+ const char *end = hdr + hdr_len;
+ while (p < end) {
+ const char *line_end = memchr(p, '\n', (size_t) (end - p));
+ if (!line_end) line_end = end;
+ size_t line_len = (size_t) (line_end - p);
+ if (line_len >= 15) {
+ if (strncasecmp(p, "Content-Length:", 15) == 0) {
+ const char *num = p + 15;
+ while (num < line_end && (*num == ' ' || *num == '\t')) num++;
+ char tmp[ 32 ] = { 0 };
+ size_t to_copy = line_len - (num - p);
+ if (to_copy >= sizeof(tmp)) to_copy = sizeof(tmp) - 1;
+ memcpy(tmp, num, to_copy);
+ for (int i = (int) strlen(tmp) - 1; i >= 0
+ && (tmp[ i ] == '\r' || tmp[ i ] == '\n'
+ || tmp[ i ] == ' ' || tmp[ i ] == '\t');
+ --i)
+ tmp[ i ] = 0;
+ long v = strtol(tmp, NULL, 10);
+ if (v < 0) return -1;
+ return (ssize_t) v;
+ }
+ }
+ p = line_end + 1;
+ }
+ return -1;
+}
+
+static int send_http_response(int fd, const byte *payload, size_t payload_len) {
+ const char *template_hdr
+ = "HTTP/1.1 200 OK\r\nContent-Length: %zu\r\nConnection: close\r\n\r\n";
+ char hdr[ 128 ];
+ int n = snprintf(hdr, sizeof(hdr), template_hdr, payload_len);
+ if (n < 0 || (size_t) n >= sizeof(hdr)) return -1;
+ if (write_all(fd, hdr, (size_t) n) < 0) return -1;
+ if (payload_len > 0 && payload) {
+ if (write_all(fd, payload, payload_len) < 0) return -1;
+ }
+ return 0;
+}
+
+// server
+
+void *omen_bound(void *arg) {
+ (void) arg;
+ while (1) {
+ int client_fd = -1;
+ char client_ip[ INET6_ADDRSTRLEN ] = { 0 };
+
+ pthread_mutex_lock(&queue_mutex);
+ while (queue_size == 0) { pthread_cond_wait(&queue_cond, &queue_mutex); }
+
+ client_fd = ring_fds[ ring_head ];
+ strncpy(client_ip, ring_ips[ ring_head ], INET6_ADDRSTRLEN - 1);
+ client_ip[ INET6_ADDRSTRLEN - 1 ] = '\0';
+ ring_head = (ring_head + 1) % MAX_CONNECTIONS;
+ queue_size--;
+ pthread_mutex_unlock(&queue_mutex);
+
+ struct timeval timeout = { 5, 0 };
+ setsockopt(client_fd, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));
+
+ char *hdr_buf = NULL;
+ size_t hdr_len = 0;
+ if (read_until_headers(client_fd, &hdr_buf, &hdr_len) < 0) {
+ close(client_fd);
+ free(hdr_buf);
+ continue;
+ }
+
+ char *sep = memmem(hdr_buf, hdr_len, "\r\n\r\n", 4);
+ size_t body_len = 0;
+ byte *body_ptr = NULL;
+
+ if (sep) {
+ size_t headers_bytes = (size_t) (sep - hdr_buf) + 4;
+ ssize_t cl = parse_content_length(hdr_buf, headers_bytes);
+ if (cl > 0) {
+ size_t already = hdr_len - headers_bytes;
+ if ((size_t) cl > MAX_BODY) cl = MAX_BODY;
+ body_ptr = malloc((size_t) cl);
+ if (!body_ptr) {
+ free(hdr_buf);
+ close(client_fd);
+ continue;
+ }
+ if (already > 0) {
+ size_t copy = already;
+ if (copy > (size_t) cl) copy = (size_t) cl;
+ memcpy(body_ptr, sep + 4, copy);
+ body_len = copy;
+ }
+ if ((size_t) body_len < (size_t) cl) {
+ if (read_exact(client_fd, body_ptr + body_len, (size_t) cl - body_len) < 0) {
+ free(body_ptr);
+ free(hdr_buf);
+ close(client_fd);
+ continue;
+ }
+ body_len = (size_t) cl;
+ }
+ } else {
+ size_t already = hdr_len - (size_t) (sep - hdr_buf) - 4;
+ if (already > 0) {
+ if (already > MAX_BODY) already = MAX_BODY;
+ body_ptr = malloc(already);
+ if (body_ptr) {
+ memcpy(body_ptr, sep + 4, already);
+ body_len = already;
+ }
+ } else {
+ body_ptr = NULL;
+ body_len = 0;
+ }
+ }
+ } else {
+ size_t use = hdr_len;
+ if (use > MAX_BODY) use = MAX_BODY;
+ body_ptr = malloc(use);
+ if (body_ptr) {
+ memcpy(body_ptr, hdr_buf, use);
+ body_len = use;
+ }
+ }
+
+ byte *out = NULL;
+ size_t out_len = 0;
+
+ if (global_handler)
+ global_handler(client_ip, sizeof(client_ip), body_ptr, body_len, &out, &out_len);
+ if (out_len > MAX_BODY) out_len = MAX_BODY;
+
+ send_http_response(client_fd, out, out_len);
+
+ free(out);
+ free(body_ptr);
+ free(hdr_buf);
+ close(client_fd);
+ }
+ return NULL;
+}
+
+int omen(int port, tome_ptr handler) {
+ global_handler = handler;
+ signal(SIGPIPE, SIG_IGN);
+
+ int server_fd = socket(AF_INET6, SOCK_STREAM, 0);
+ if (server_fd < 0) {
+ perror("failed to create IPv6 socket");
+ return -1;
+ }
+
+ int opt = 0;
+ if (setsockopt(server_fd, IPPROTO_IPV6, IPV6_V6ONLY, &opt, sizeof(opt)) < 0) {
+ perror("failed to disable IPV6_V6ONLY");
+ close(server_fd);
+ return -1;
+ }
+
+ opt = 1;
+ if (setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
+ perror("failed to set SO_REUSEADDR");
+ close(server_fd);
+ return -1;
+ }
+
+ struct sockaddr_in6 addr = { 0 };
+ addr.sin6_family = AF_INET6;
+ addr.sin6_addr = in6addr_any;
+ addr.sin6_port = htons(port);
+
+ if (bind(server_fd, (struct sockaddr *) &addr, sizeof(addr)) < 0) {
+ perror("bind failed");
+ close(server_fd);
+ return -1;
+ }
+
+ if (listen(server_fd, MAX_CONNECTIONS) < 0) {
+ perror("listen failed");
+ close(server_fd);
+ return -1;
+ }
+
+ pthread_t *threads = malloc((size_t) worker_count * sizeof(pthread_t));
+ if (!threads) {
+ close(server_fd);
+ return -1;
+ }
+
+ for (int i = 0; i < worker_count; i++) {
+ if (pthread_create(&threads[ i ], NULL, omen_bound, NULL) != 0) {
+ perror("pthread_create failed");
+ free(threads);
+ close(server_fd);
+ return -1;
+ }
+ pthread_detach(threads[ i ]);
+ }
+
+ while (1) {
+ struct sockaddr_storage client_addr;
+ socklen_t len = sizeof(client_addr);
+ int client_fd = accept(server_fd, (struct sockaddr *) &client_addr, &len);
+ if (client_fd < 0) {
+ if (errno != EINTR) { perror("accept failed"); }
+ continue;
+ }
+
+ char client_ip[ INET6_ADDRSTRLEN ] = { 0 };
+ if (client_addr.ss_family == AF_INET) {
+ struct sockaddr_in *addr4 = (struct sockaddr_in *) &client_addr;
+ inet_ntop(AF_INET, &addr4->sin_addr, client_ip, sizeof(client_ip));
+ } else if (client_addr.ss_family == AF_INET6) {
+ struct sockaddr_in6 *addr6 = (struct sockaddr_in6 *) &client_addr;
+ if (IN6_IS_ADDR_V4MAPPED(&addr6->sin6_addr)) {
+ struct in_addr v4;
+ memcpy(&v4.s_addr, &addr6->sin6_addr.s6_addr[ 12 ], sizeof(v4.s_addr));
+ inet_ntop(AF_INET, &v4, client_ip, sizeof(client_ip));
+ } else {
+ inet_ntop(AF_INET6, &addr6->sin6_addr, client_ip, sizeof(client_ip));
+ }
+ } else {
+ strncpy(client_ip, "0.0.0.0", sizeof(client_ip) - 1);
+ client_ip[ sizeof(client_ip) - 1 ] = '\0';
+ }
+
+ if (!ratelimit_check(client_ip)) { usleep(250 * 1000); }
+
+ pthread_mutex_lock(&queue_mutex);
+ if (queue_size >= (size_t) MAX_CONNECTIONS) {
+ char err_resp[] = "HTTP/1.1 418 I'm a teapot\r\n"
+ "Content-Length: 24\r\n"
+ "Connection: close\r\n\r\n"
+ "Max connections reached";
+ write_all(client_fd, err_resp, strlen(err_resp));
+ close(client_fd);
+ pthread_mutex_unlock(&queue_mutex);
+ continue;
+ }
+
+ ring_fds[ ring_tail ] = client_fd;
+ strncpy(ring_ips[ ring_tail ], client_ip, INET6_ADDRSTRLEN - 1);
+ ring_ips[ ring_tail ][ INET6_ADDRSTRLEN - 1 ] = '\0';
+ ring_tail = (ring_tail + 1) % MAX_CONNECTIONS;
+ queue_size++;
+ pthread_cond_signal(&queue_cond);
+ pthread_mutex_unlock(&queue_mutex);
+ }
+
+ free(threads);
+ close(server_fd);
+ return 0;
+}
+
+// client
+
+int study(
+ const char *ip,
+ int port,
+ const byte *data,
+ size_t data_len,
+ byte **response,
+ size_t *resp_len) {
+ bool is_ipv6 = strchr(ip, ':') != NULL;
+ int domain = is_ipv6 ? AF_INET6 : AF_INET;
+ // printf("attempting %s connection to %s:%d\n", is_ipv6 ? "ipv6" : "ipv4", ip, port);
+
+ int sock = socket(domain, SOCK_STREAM, 0);
+ if (sock < 0) {
+ printf("socket creation failed: %s\n", strerror(errno));
+ return -1;
+ }
+
+ struct sockaddr_storage addr = { 0 };
+ addr.ss_family = domain;
+ void *addr_ptr = NULL;
+ size_t addr_len = 0;
+ if (domain == AF_INET) {
+ struct sockaddr_in *addr4 = (struct sockaddr_in *) &addr;
+ addr_ptr = &addr4->sin_addr;
+ addr_len = sizeof(struct sockaddr_in);
+ addr4->sin_port = htons(port);
+ } else {
+ struct sockaddr_in6 *addr6 = (struct sockaddr_in6 *) &addr;
+ addr_ptr = &addr6->sin6_addr;
+ addr_len = sizeof(struct sockaddr_in6);
+ addr6->sin6_port = htons(port);
+ }
+
+ if (inet_pton(domain, ip, addr_ptr) != 1) {
+ printf("invalid %s address: %s\n", is_ipv6 ? "ipv6" : "ipv4", ip);
+ close(sock);
+ return -1;
+ }
+
+ struct timeval timeout = { 5, 0 };
+ if (setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout)) < 0
+ || setsockopt(sock, SOL_SOCKET, SO_SNDTIMEO, &timeout, sizeof(timeout)) < 0) {
+ printf("setsockopt failed: %s\n", strerror(errno));
+ close(sock);
+ return -1;
+ }
+
+ // printf("connecting to %s:%d\n", ip, port);
+ if (connect(sock, (struct sockaddr *) &addr, addr_len) < 0) {
+ printf("connect failed: %s\n", strerror(errno));
+ close(sock);
+ return -1;
+ }
+
+ if (data_len > MAX_BODY) data_len = MAX_BODY;
+ char req_hdr[ 256 ];
+ int hdr_n;
+ if (is_ipv6) {
+ hdr_n = snprintf(
+ req_hdr,
+ sizeof(req_hdr),
+ "POST /data HTTP/1.1\r\nHost: [%s]\r\nContent-Length: %zu\r\nConnection: close\r\n\r\n",
+ ip,
+ data_len);
+ } else {
+ hdr_n = snprintf(
+ req_hdr,
+ sizeof(req_hdr),
+ "POST /data HTTP/1.1\r\nHost: %s\r\nContent-Length: %zu\r\nConnection: close\r\n\r\n",
+ ip,
+ data_len);
+ }
+ if (hdr_n < 0 || (size_t) hdr_n >= sizeof(req_hdr)) {
+ printf("failed to format HTTP header\n");
+ close(sock);
+ return -1;
+ }
+
+ // printf("sending request: %.*s\n", hdr_n, req_hdr);
+ if (write_all(sock, req_hdr, (size_t) hdr_n) < 0
+ || (data_len > 0 && write_all(sock, data, data_len) < 0)) {
+ printf("write failed: %s\n", strerror(errno));
+ close(sock);
+ return -1;
+ }
+
+ char *hdr_buf = NULL;
+ size_t hdr_len = 0;
+ if (read_until_headers(sock, &hdr_buf, &hdr_len) < 0) {
+ printf("failed to read headers: %s\n", strerror(errno));
+ close(sock);
+ free(hdr_buf);
+ return -1;
+ }
+
+ if (memmem(hdr_buf, hdr_len, "HTTP/1.1 418 I'm a teapot", 34) != NULL) {
+ // printf("server responded with 418 I'm a teapot\n");
+ close(sock);
+ free(hdr_buf);
+ return -2;
+ }
+
+ char *sep = memmem(hdr_buf, hdr_len, "\r\n\r\n", 4);
+ size_t body_len = 0;
+ byte *body = NULL;
+ if (sep) {
+ size_t headers_bytes = (size_t) (sep - hdr_buf) + 4;
+ ssize_t cl = parse_content_length(hdr_buf, headers_bytes);
+ size_t already = hdr_len - headers_bytes;
+ if (cl > 0) {
+ if ((size_t) cl > MAX_BODY) cl = MAX_BODY;
+ body = malloc((size_t) cl);
+ if (!body) {
+ printf("failed to allocate body buffer\n");
+ free(hdr_buf);
+ close(sock);
+ return -1;
+ }
+ if (already > 0) {
+ size_t copy = already;
+ if (copy > (size_t) cl) copy = (size_t) cl;
+ memcpy(body, sep + 4, copy);
+ body_len = copy;
+ }
+ if ((size_t) body_len < (size_t) cl) {
+ if (read_exact(sock, body + body_len, (size_t) cl - body_len) < 0) {
+ printf("failed to read body: %s\n", strerror(errno));
+ free(body);
+ free(hdr_buf);
+ close(sock);
+ return -1;
+ }
+ body_len = (size_t) cl;
+ }
+ } else {
+ if (already > 0) {
+ if (already > MAX_BODY) already = MAX_BODY;
+ body = malloc(already);
+ if (body) {
+ memcpy(body, sep + 4, already);
+ body_len = already;
+ }
+ } else {
+ body = NULL;
+ body_len = 0;
+ }
+ }
+ } else {
+ size_t use = hdr_len;
+ if (use > MAX_BODY) use = MAX_BODY;
+ body = malloc(use);
+ if (body) {
+ memcpy(body, hdr_buf, use);
+ body_len = use;
+ }
+ }
+
+ free(hdr_buf);
+ close(sock);
+
+ if (body_len > 0 && body) {
+ *response = body;
+ *resp_len = body_len;
+ } else {
+ *response = NULL;
+ *resp_len = 0;
+ free(body);
+ }
+ // printf("response received, length: %zu\n", body_len);
+ return 0;
+}
generated by cgit v1.2.3 (git 2.39.1) at 2026-10-27 10:10:50 +0000