summaryrefslogtreecommitdiff
path: root/omen.c
diff options
context:
space:
mode:
Diffstat (limited to 'omen.c')
-rw-r--r--omen.c467
1 files changed, 467 insertions, 0 deletions
diff --git a/omen.c b/omen.c
new file mode 100644
index 0000000..2f60a8c
--- /dev/null
+++ b/omen.c
@@ -0,0 +1,467 @@
+#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"
+
+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);
+struct Task {
+ int fd;
+ struct Task *next;
+};
+
+#define MAX_BUFFER 1 << 18
+#define MAX_BODY 1 << 17
+#define MAX_CONNECTIONS 2048
+
+// vars
+
+int worker_count = 16;
+
+pthread_mutex_t queue_mutex = PTHREAD_MUTEX_INITIALIZER;
+pthread_cond_t queue_cond = PTHREAD_COND_INITIALIZER;
+struct Task *head = NULL;
+struct Task *tail = NULL;
+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 cap = 4096;
+ size_t len = 0;
+ char *buf = malloc(cap);
+ if (!buf) return -1;
+ while (1) {
+ 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) {
+ len = len;
+ break;
+ }
+ len += (size_t) r;
+ if (len >= 4 && memmem(buf, len, "\r\n\r\n", 4) != NULL) 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;
+ struct sockaddr_in client_addr;
+ socklen_t client_addr_len = sizeof(client_addr);
+
+ pthread_mutex_lock(&queue_mutex);
+ while (head == NULL) { pthread_cond_wait(&queue_cond, &queue_mutex); }
+ struct Task *task = head;
+ head = head->next;
+ if (head == NULL) tail = NULL;
+ queue_size--;
+ client_fd = task->fd;
+
+ if (getpeername(client_fd, (struct sockaddr *) &client_addr, &client_addr_len) < 0) {
+ memset(&client_addr, 0, sizeof(client_addr));
+ }
+
+ free(task);
+ pthread_mutex_unlock(&queue_mutex);
+
+ char client_ip[ INET_ADDRSTRLEN ] = { 0 };
+ if (client_addr.sin_family == AF_INET) {
+ inet_ntop(AF_INET, &client_addr.sin_addr, client_ip, INET_ADDRSTRLEN);
+ }
+
+ 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_INET, SOCK_STREAM, 0);
+ if (server_fd < 0) return -1;
+
+ int opt = 1;
+ setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
+
+ struct sockaddr_in addr = { 0 };
+ addr.sin_family = AF_INET;
+ addr.sin_addr.s_addr = INADDR_ANY;
+ addr.sin_port = htons(port);
+
+ if (bind(server_fd, (struct sockaddr *) &addr, sizeof(addr)) < 0
+ || listen(server_fd, MAX_CONNECTIONS) < 0) {
+ 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");
+ } else {
+ pthread_detach(threads[ i ]);
+ }
+ }
+
+ while (1) {
+ struct sockaddr_in client_addr;
+ socklen_t len = sizeof(client_addr);
+ int client_fd = accept(server_fd, (struct sockaddr *) &client_addr, &len);
+ if (client_fd < 0) continue;
+
+ pthread_mutex_lock(&queue_mutex);
+ if (queue_size >= (size_t) MAX_CONNECTIONS) {
+ // note: The server converts into a teapot if the number of concurrent connections
+ // reaches the max amount. this can be mitigated by disconnecting some clients,
+ // which will turn the teapot back to an actual server.
+ 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, 92);
+ close(client_fd);
+ pthread_mutex_unlock(&queue_mutex);
+ continue;
+ }
+
+ struct Task *new_task = malloc(sizeof(struct Task));
+ if (!new_task) {
+ close(client_fd);
+ pthread_mutex_unlock(&queue_mutex);
+ continue;
+ }
+ new_task->fd = client_fd;
+ new_task->next = NULL;
+ if (tail) tail->next = new_task;
+ else
+ head = new_task;
+ tail = new_task;
+ queue_size++;
+ pthread_cond_signal(&queue_cond);
+ pthread_mutex_unlock(&queue_mutex);
+ }
+
+ 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) {
+ int sock = socket(AF_INET, SOCK_STREAM, 0);
+ if (sock < 0) return -1;
+
+ struct sockaddr_in addr = { 0 };
+ addr.sin_family = AF_INET;
+ if (inet_pton(AF_INET, ip, &addr.sin_addr) != 1) {
+ close(sock);
+ return -1;
+ }
+ addr.sin_port = htons(port);
+
+ struct timeval timeout = { 5, 0 };
+ setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));
+ setsockopt(sock, SOL_SOCKET, SO_SNDTIMEO, &timeout, sizeof(timeout));
+
+ if (connect(sock, (struct sockaddr *) &addr, sizeof(addr)) < 0) {
+ close(sock);
+ return -1;
+ }
+
+ if (data_len > MAX_BODY) data_len = MAX_BODY;
+ char req_hdr[ 256 ];
+ int 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)) {
+ close(sock);
+ return -1;
+ }
+
+ if (write_all(sock, req_hdr, (size_t) hdr_n) < 0
+ || (data_len > 0 && write_all(sock, data, data_len) < 0)) {
+ close(sock);
+ return -1;
+ }
+
+ char *hdr_buf = NULL;
+ size_t hdr_len = 0;
+ if (read_until_headers(sock, &hdr_buf, &hdr_len) < 0) {
+ close(sock);
+ free(hdr_buf);
+ return -1;
+ }
+
+ if (memmem(hdr_buf, hdr_len, "HTTP/1.1 418 I'm a teapot", 34) != NULL) {
+ 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) {
+ 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) {
+ 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);
+ }
+ return 0;
+} \ No newline at end of file
generated by cgit v1.2.3 (git 2.39.1) at 2026-10-27 10:36:59 +0000