#define _POSIX_C_SOURCE 200809L
#include <sys/socket.h>
#include <sys/wait.h>
#include <netinet/in.h>
#include <netdb.h>
#include <signal.h>
#include <unistd.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>

#define MAX_HEADERS 16384

static int
write_all(int fd, const char *buf, size_t len)
{
    while (len > 0) {
        ssize_t n = write(fd, buf, len);
        if (n < 0 && errno == EINTR) continue;
        if (n <= 0) return -1;
        buf += n;
        len -= (size_t)n;
    }
    return 0;
}

static void
send_error(int fd, int status)
{
    const char *reason = status == 400 ? "Bad Request" :
                         status == 405 ? "Method Not Allowed" : "Bad Gateway";
    char response[256];
    int n = snprintf(response, sizeof(response),
        "HTTP/1.0 %d %s\r\nContent-Length: 0\r\nConnection: close\r\n%s\r\n",
        status, reason, status == 405 ? "Allow: GET, HEAD\r\n" : "");
    if (n > 0 && (size_t)n < sizeof(response))
        (void)write_all(fd, response, (size_t)n);
}

static int
valid_token(const char *text)
{
    const char *allowed = "!#$%&'*+-.^_`|~0123456789"
                          "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz";
    return *text != '\0' && strspn(text, allowed) == strlen(text);
}

/* Rebuild bounded headers, replacing connection handling and rejecting bodies. */
static int
prepare_request(char *request, char *end, char *forward, size_t *size)
{
    char *line = strstr(request, "\r\n");
    char *method = request, *target, *version, *p;
    int host = 0, n;
    if (line == NULL) return 400;
    *line = '\0';
    line += 2;
    target = strchr(method, ' ');
    if (target == NULL) return 400;
    *target++ = '\0';
    version = strchr(target, ' ');
    if (version == NULL) return 400;
    *version++ = '\0';
    if (!valid_token(method) || *target != '/' ||
        (strcmp(version, "HTTP/1.0") != 0 && strcmp(version, "HTTP/1.1") != 0))
        return 400;
    for (p = target; *p; ++p)
        if ((unsigned char)*p <= 32 || (unsigned char)*p == 127 || *p == '#')
            return 400;
    if (strcmp(method, "GET") != 0 && strcmp(method, "HEAD") != 0)
        return 405;
    n = snprintf(forward, MAX_HEADERS + 64, "%s %s HTTP/1.0\r\n", method, target);
    if (n < 0 || n >= MAX_HEADERS + 64) return 400;
    *size = (size_t)n;
    while (line < end) {
        char *next = strstr(line, "\r\n");
        char *colon, *value, *last;
        char saved;
        size_t len;
        int drop = 0;
        if (next == NULL || next > end) return 400;
        len = (size_t)(next - line);
        *next = '\0';
        colon = strchr(line, ':');
        if (colon == NULL) return 400;
        *colon = '\0';
        if (!valid_token(line)) return 400;
        value = colon + 1;
        for (p = value; *p; ++p)
            if (((unsigned char)*p < 32 && *p != '\t') || (unsigned char)*p == 127)
                return 400;
        while (*value == ' ' || *value == '\t') ++value;
        last = next;
        while (last > value && (last[-1] == ' ' || last[-1] == '\t')) --last;
        saved = *last;
        *last = '\0';
        if (strcasecmp(line, "Host") == 0) {
            if (!*value || ++host > 1) return 400;
        } else if (strcasecmp(line, "Content-Length") == 0) {
            if (!*value || strspn(value, "0") != strlen(value)) return 400;
            drop = 1;
        } else if (strcasecmp(line, "Transfer-Encoding") == 0 ||
                   strcasecmp(line, "Expect") == 0 || strcasecmp(line, "Upgrade") == 0) {
            return 400;
        } else if (strcasecmp(line, "Connection") == 0) {
            /* Other Connection tokens name headers we do not support forwarding. */
            if (strcasecmp(value, "close") != 0 && strcasecmp(value, "keep-alive") != 0)
                return 400;
            drop = 1;
        } else if (strcasecmp(line, "Proxy-Connection") == 0 ||
                   strcasecmp(line, "Keep-Alive") == 0 ||
                   strcasecmp(line, "TE") == 0 || strcasecmp(line, "Trailer") == 0) {
            drop = 1;
        }
        if (!drop) {
            *colon = ':';
            *last = saved;
            if (*size + len + 2 > MAX_HEADERS + 40) return 400;
            memcpy(forward + *size, line, len);
            *size += len;
            memcpy(forward + *size, "\r\n", 2);
            *size += 2;
        }
        line = next + 2;
    }
    if (strcmp(version, "HTTP/1.1") == 0 && !host) return 400;
    memcpy(forward + *size, "Connection: close\r\n\r\n", 21);
    *size += 21;
    return 0;
}

static int
connect_upstream(const char *host, const char *port)
{
    struct addrinfo hints = {0}, *addresses, *addr;
    int fd = -1, result, saved = ECONNREFUSED;
    hints.ai_family = AF_INET;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_flags = AI_NUMERICSERV;
    result = getaddrinfo(host, port, &hints, &addresses);
    if (result != 0) {
        fprintf(stderr, "getaddrinfo: %s\n", gai_strerror(result));
        return -1;
    }
    for (addr = addresses; addr != NULL; addr = addr->ai_next) {
        do {
            fd = socket(addr->ai_family, addr->ai_socktype, addr->ai_protocol);
            if (fd < 0) { saved = errno; break; }
            if (connect(fd, addr->ai_addr, addr->ai_addrlen) == 0) break;
            saved = errno;
            close(fd);
            fd = -1;
        } while (saved == EINTR);
        if (fd >= 0) break;
    }
    freeaddrinfo(addresses);
    if (fd < 0) { errno = saved; perror("connect"); }
    return fd;
}

static void
handle_client(int client, const char *host, const char *port)
{
    char request[MAX_HEADERS + 1], forward[MAX_HEADERS + 64], buf[16384];
    char *end = NULL;
    size_t used = 0, size;
    int upstream, status, replied = 0;
    while (used < MAX_HEADERS) {
        ssize_t n = recv(client, request + used, MAX_HEADERS - used, 0);
        if (n < 0 && errno == EINTR) continue;
        if (n <= 0) {
            if (n < 0) perror("recv client");
            else if (used) send_error(client, 400);
            return;
        }
        if (memchr(request + used, '\0', (size_t)n) != NULL) {
            send_error(client, 400);
            return;
        }
        used += (size_t)n;
        request[used] = '\0';
        end = strstr(request, "\r\n\r\n");
        if (end != NULL) break;
    }
    if (end == NULL) { send_error(client, 400); return; }
    status = prepare_request(request, end, forward, &size);
    if (status) { send_error(client, status); return; }
    upstream = connect_upstream(host, port);
    if (upstream < 0) { send_error(client, 502); return; }
    if (write_all(upstream, forward, size) < 0) {
        perror("write upstream");
        send_error(client, 502);
        close(upstream);
        return;
    }
    for (;;) {
        ssize_t n = recv(upstream, buf, sizeof(buf), 0);
        if (n < 0 && errno == EINTR) continue;
        if (n <= 0) {
            if (n < 0) perror("recv upstream");
            if (!replied) send_error(client, 502);
            break;
        }
        if (write_all(client, buf, (size_t)n) < 0) break;
        replied = 1;
    }
    close(upstream);
}

static void
reap_children(int sig)
{
    int saved = errno;
    pid_t pid;
    (void)sig;
    do {
        pid = waitpid(-1, NULL, WNOHANG);
    } while (pid > 0 || (pid < 0 && errno == EINTR));
    errno = saved;
}

int
main(int argc, char **argv)
{
    long ports[2];
    int listener, one = 1, i;
    struct sockaddr_in address = {0};
    struct sigaction action = {0};
    if (argc != 4) {
        fprintf(stderr, "usage: %s listen-port upstream-host upstream-port\n", argv[0]);
        return 1;
    }
    for (i = 0; i < 2; ++i) {
        char *end;
        const char *text = argv[i == 0 ? 1 : 3];
        errno = 0;
        ports[i] = strtol(text, &end, 10);
        if (errno || !*text || *end || ports[i] < 1 || ports[i] > 65535) {
            fprintf(stderr, "invalid port: %s\n", text);
            return 1;
        }
    }
    sigemptyset(&action.sa_mask);
    action.sa_handler = SIG_IGN;
    if (sigaction(SIGPIPE, &action, NULL) < 0) { perror("sigaction"); return 1; }
    action.sa_handler = reap_children;
    action.sa_flags = SA_RESTART | SA_NOCLDSTOP;
    if (sigaction(SIGCHLD, &action, NULL) < 0) { perror("sigaction"); return 1; }
    listener = socket(AF_INET, SOCK_STREAM, 0);
    if (listener < 0) { perror("socket"); return 1; }
    address.sin_family = AF_INET;
    address.sin_addr.s_addr = htonl(INADDR_ANY);
    address.sin_port = htons((unsigned short)ports[0]);
    if (setsockopt(listener, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one)) < 0 ||
        bind(listener, (struct sockaddr *)&address, sizeof(address)) < 0 ||
        listen(listener, 32) < 0) {
        perror("proxy");
        close(listener);
        return 1;
    }
    printf("listening on port %ld\n", ports[0]);
    fflush(stdout);
    for (;;) {
        int client = accept(listener, NULL, NULL);
        pid_t pid;
        if (client < 0) {
            if (errno == EINTR || errno == ECONNABORTED) continue;
            perror("accept");
            break;
        }
        pid = fork();
        if (pid == 0) {
            close(listener);
            handle_client(client, argv[2], argv[3]);
            close(client);
            _exit(0);
        }
        if (pid < 0) perror("fork");
        close(client);
    }
    close(listener);
    return 1;
}