Initial commit

udp-proxy.c

@@ -0,0 +1,166 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <arpa/inet.h>
+#include <netinet/ip.h>
+#include <sys/epoll.h>
+#include <sys/socket.h>
+#include <time.h>
+#include <unistd.h>
+#include <errno.h>
+
+#define SESSION_TIMEOUT 120
+#define ENTRIES_SIZE 16
+
+struct relay_entry {
+    int relay_fd;
+    struct sockaddr_in client_addr;
+    time_t last_used;
+};
+
+int server_fd;
+
+int epoll_fd;
+
+int key;
+
+struct sockaddr_in dest_addr;
+
+struct relay_entry relay_entries[ENTRIES_SIZE];
+
+void recv_server_msg() {
+    unsigned char buf[65536];
+    struct sockaddr_in client_addr;
+    socklen_t client_addr_len = sizeof(client_addr);
+    ssize_t msg_len = recvfrom(server_fd, buf, sizeof(buf), 0, (struct sockaddr*) &client_addr, &client_addr_len);
+
+    struct relay_entry *entry = NULL;
+    for (int i = 0; i < ENTRIES_SIZE; i ++) {
+        if (relay_entries[i].relay_fd != -1
+                && client_addr.sin_addr.s_addr == relay_entries[i].client_addr.sin_addr.s_addr
+                && client_addr.sin_port == relay_entries[i].client_addr.sin_port) {
+            entry = &relay_entries[i];
+            break;
+        }
+    }
+    if (!entry) {
+        int relay_fd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
+        if (relay_fd == -1) {
+            perror("socket");
+            exit(1);
+        }
+
+        int i;
+        for (i = 0; i < ENTRIES_SIZE; i ++) {
+            if (relay_entries[i].relay_fd == -1) {
+                break;
+            }
+        }
+        if (i == ENTRIES_SIZE) {
+            printf("%s:%d no more space.\n", inet_ntoa(client_addr.sin_addr), ntohs(client_addr.sin_port));
+            return;
+        }
+
+        entry = &relay_entries[i];
+        memcpy(&entry->client_addr, &client_addr, sizeof(client_addr));
+        entry->relay_fd = relay_fd;
+
+        struct epoll_event event = {
+            .events = EPOLLIN,
+            .data = { .ptr = entry }
+        };
+        epoll_ctl(epoll_fd, EPOLL_CTL_ADD, relay_fd, &event);
+
+        printf("%s:%d connected.\n", inet_ntoa(client_addr.sin_addr), ntohs(client_addr.sin_port));
+    }
+    for (int i = 0; i < msg_len; i ++) {
+        buf[i] ^= key;
+    }
+    sendto(entry->relay_fd, buf, msg_len, 0, (struct sockaddr *) &dest_addr, sizeof(dest_addr));
+    entry->last_used = time(NULL);
+}
+
+void recv_relay_msg(struct relay_entry *entry) {
+    unsigned char buf[65536];
+    ssize_t msg_len = recv(entry->relay_fd, buf, sizeof(buf), 0);
+    for (int i = 0; i < msg_len; i ++) {
+        buf[i] ^= key;
+    }
+    sendto(server_fd, buf, msg_len, 0, (struct sockaddr *) &entry->client_addr, sizeof(entry->client_addr));
+    entry->last_used = time(NULL);
+}
+
+void cleanup() {
+    static int last_clean = 0;
+    int cur_time = time(NULL);
+    if (cur_time - last_clean >= SESSION_TIMEOUT) {
+        for (int i = 0; i < ENTRIES_SIZE; i ++) {
+            struct relay_entry *entry = &relay_entries[i];
+            if (entry->relay_fd != -1 && cur_time - entry->last_used >= SESSION_TIMEOUT) {
+                printf("%s:%d destroyed.\n", inet_ntoa(entry->client_addr.sin_addr), ntohs(entry->client_addr.sin_port));
+                close(entry->relay_fd);
+                entry->relay_fd = -1;
+            }
+        }
+        last_clean = cur_time;
+    }
+}
+
+int main(int argc, char *argv[]) {
+	if (argc < 6) {
+		puts("usage: udp-proxy <listen-address> <listen-port> <remote-address> <remote-port> <key>");
+		exit(1);
+	}
+
+	key = atoi(argv[5]);
+
+    for (int i = 0; i < ENTRIES_SIZE; i ++) {
+        relay_entries[i].relay_fd = -1;
+    }
+
+    server_fd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
+    if (server_fd == -1) {
+        perror("socket");
+        exit(1);
+    }
+
+    struct sockaddr_in server_addr = { .sin_family = AF_INET };
+    inet_aton(argv[1], &server_addr.sin_addr);
+    server_addr.sin_port = htons(atoi(argv[2]));
+
+    dest_addr.sin_family = AF_INET;
+    inet_aton(argv[3], &dest_addr.sin_addr);
+    dest_addr.sin_port = htons(atoi(argv[4]));
+
+    int res = bind(server_fd, (struct sockaddr*) &server_addr, sizeof(server_addr));
+    if (res == -1) {
+        perror("bind");
+        exit(1);
+    }
+
+    epoll_fd = epoll_create1(0);
+    if (epoll_fd == -1) {
+        perror("epoll_create1");
+        exit(1);
+    }
+
+    struct epoll_event event_filter = { .events = EPOLLIN };
+    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, server_fd, &event_filter);
+
+    struct epoll_event event;
+    while (1) {
+        res = epoll_wait(epoll_fd, &event, 1, SESSION_TIMEOUT * 1000);
+        if (res == -1 && errno != EINTR) {
+            perror("epoll_wait");
+            exit(1);
+        }
+        if (res > 0) {
+            if (event.data.ptr == NULL) {
+                recv_server_msg();
+            } else {
+                recv_relay_msg(event.data.ptr);
+            }
+        }
+        cleanup();
+    }
+}