// Copyright (C) 2026 Douglas Quigg (dstroy0) <dquigg123@gmail.com>
// SPDX-License-Identifier: AGPL-3.0-or-later

/**
 * @file proxy_protocol.cpp
 * @brief HAProxy PROXY protocol v1 / v2 parser + builder (pure, host-tested).
 */

#include "services/proxy_protocol/proxy_protocol.h"

#if DETWS_ENABLE_PROXY_PROTOCOL

#include <stdio.h> // snprintf for the v1 dotted-quad text
#include <string.h>

static const uint8_t kV2Sig[PROXY_V2_SIG_LEN] = {0x0D, 0x0A, 0x0D, 0x0A, 0x00, 0x0D,
                                                 0x0A, 0x51, 0x55, 0x49, 0x54, 0x0A};

static uint16_t rd16(const uint8_t *p)
{
    return (uint16_t)(((uint16_t)p[0] << 8) | p[1]);
}

static uint32_t rd32(const uint8_t *p)
{
    return ((uint32_t)p[0] << 24) | ((uint32_t)p[1] << 16) | ((uint32_t)p[2] << 8) | p[3];
}

// Parse a dotted-quad IPv4 in [s, s+n) into a host-order uint32; false on malformed.
static bool parse_ipv4(const char *s, size_t n, uint32_t *out)
{
    uint32_t v = 0;
    int octets = 0;
    size_t i = 0;
    while (octets < 4)
    {
        if (i >= n || s[i] < '0' || s[i] > '9')
            return false;
        uint32_t o = 0;
        size_t digits = 0;
        while (i < n && s[i] >= '0' && s[i] <= '9')
        {
            o = o * 10 + (uint32_t)(s[i] - '0');
            i++;
            if (++digits > 3 || o > 255)
                return false;
        }
        v = (v << 8) | o;
        octets++;
        if (octets < 4)
        {
            if (i >= n || s[i] != '.')
                return false;
            i++;
        }
    }
    if (i != n) // trailing junk
        return false;
    *out = v;
    return true;
}

static bool parse_u16(const char *s, size_t n, uint16_t *out)
{
    if (n == 0 || n > 5)
        return false;
    uint32_t v = 0;
    for (size_t i = 0; i < n; i++)
    {
        if (s[i] < '0' || s[i] > '9')
            return false;
        v = v * 10 + (uint32_t)(s[i] - '0');
    }
    if (v > 0xFFFF)
        return false;
    *out = (uint16_t)v;
    return true;
}

// Parse the v1 text header (already known to start with "PROXY ").
static bool parse_v1(const uint8_t *buf, size_t len, ProxyInfo *out, size_t *consumed)
{
    // Find the terminating CRLF (the line is bounded at 107 octets).
    size_t crlf = len;
    size_t scan = len < 108 ? len : 108;
    for (size_t i = 0; i + 1 < scan; i++)
        if (buf[i] == '\r' && buf[i + 1] == '\n')
        {
            crlf = i;
            break;
        }
    if (crlf == len)
        return false; // line not complete

    const char *s = (const char *)buf;
    // Tokenize the line by single spaces.
    const char *tok[6];
    size_t tlen[6], ntok = 0, i = 0;
    while (i < crlf && ntok < 6)
    {
        while (i < crlf && s[i] == ' ')
            i++;
        if (i >= crlf)
            break;
        size_t start = i;
        while (i < crlf && s[i] != ' ')
            i++;
        tok[ntok] = s + start;
        tlen[ntok] = i - start;
        ntok++;
    }
    out->version = 1;
    out->has_addr = false;
    out->src_addr = out->dst_addr = 0;
    out->src_port = out->dst_port = 0;
    *consumed = crlf + 2;
    // "PROXY TCP4 <src> <dst> <sport> <dport>"; anything else (UNKNOWN/TCP6) yields no addr.
    if (ntok == 6 && tlen[1] == 4 && memcmp(tok[1], "TCP4", 4) == 0)
    {
        if (parse_ipv4(tok[2], tlen[2], &out->src_addr) && parse_ipv4(tok[3], tlen[3], &out->dst_addr) &&
            parse_u16(tok[4], tlen[4], &out->src_port) && parse_u16(tok[5], tlen[5], &out->dst_port))
            out->has_addr = true;
    }
    return true;
}

bool proxy_parse(const uint8_t *buf, size_t len, ProxyInfo *out, size_t *consumed)
{
    if (!buf || !out || !consumed)
        return false;

    // v2: the 12-octet binary signature.
    if (len >= PROXY_V2_SIG_LEN && memcmp(buf, kV2Sig, PROXY_V2_SIG_LEN) == 0)
    {
        if (len < 16) // signature + ver_cmd + fam + 2-octet length
            return false;
        uint8_t ver_cmd = buf[12];
        uint8_t fam = buf[13];
        uint16_t addr_len = rd16(buf + 14);
        size_t total = 16 + (size_t)addr_len;
        if (total > len)
            return false;             // address block not fully buffered
        if ((ver_cmd & 0xF0) != 0x20) // must be version 2
            return false;
        out->version = 2;
        out->has_addr = false;
        out->src_addr = out->dst_addr = 0;
        out->src_port = out->dst_port = 0;
        if (ver_cmd == PROXY_V2_VER_CMD_PROXY && fam == PROXY_V2_FAM_TCP4 && addr_len >= 12)
        {
            out->src_addr = rd32(buf + 16);
            out->dst_addr = rd32(buf + 20);
            out->src_port = rd16(buf + 24);
            out->dst_port = rd16(buf + 26);
            out->has_addr = true;
        }
        *consumed = total;
        return true;
    }

    // v1: the "PROXY " text prefix.
    if (len >= 6 && memcmp(buf, "PROXY ", 6) == 0)
        return parse_v1(buf, len, out, consumed);

    return false; // no PROXY header present
}

size_t proxy_v1_build(char *buf, size_t cap, uint32_t src_addr, uint32_t dst_addr, uint16_t src_port, uint16_t dst_port)
{
    if (!buf)
        return 0;
    int n = snprintf(buf, cap, "PROXY TCP4 %u.%u.%u.%u %u.%u.%u.%u %u %u\r\n", (unsigned)((src_addr >> 24) & 0xFF),
                     (unsigned)((src_addr >> 16) & 0xFF), (unsigned)((src_addr >> 8) & 0xFF),
                     (unsigned)(src_addr & 0xFF), (unsigned)((dst_addr >> 24) & 0xFF),
                     (unsigned)((dst_addr >> 16) & 0xFF), (unsigned)((dst_addr >> 8) & 0xFF),
                     (unsigned)(dst_addr & 0xFF), (unsigned)src_port, (unsigned)dst_port);
    if (n < 0 || (size_t)n >= cap) // snprintf truncated (no room for the content + NUL)
        return 0;
    return (size_t)n;
}

size_t proxy_v2_build(uint8_t *buf, size_t cap, uint32_t src_addr, uint32_t dst_addr, uint16_t src_port,
                      uint16_t dst_port)
{
    const size_t total = 16 + 12; // header + TCP/IPv4 address block
    if (!buf || cap < total)
        return 0;
    memcpy(buf, kV2Sig, PROXY_V2_SIG_LEN);
    buf[12] = PROXY_V2_VER_CMD_PROXY;
    buf[13] = PROXY_V2_FAM_TCP4;
    buf[14] = 0x00; // address-block length (12), big-endian
    buf[15] = 0x0C;
    buf[16] = (uint8_t)(src_addr >> 24);
    buf[17] = (uint8_t)(src_addr >> 16);
    buf[18] = (uint8_t)(src_addr >> 8);
    buf[19] = (uint8_t)(src_addr);
    buf[20] = (uint8_t)(dst_addr >> 24);
    buf[21] = (uint8_t)(dst_addr >> 16);
    buf[22] = (uint8_t)(dst_addr >> 8);
    buf[23] = (uint8_t)(dst_addr);
    buf[24] = (uint8_t)(src_port >> 8);
    buf[25] = (uint8_t)(src_port);
    buf[26] = (uint8_t)(dst_port >> 8);
    buf[27] = (uint8_t)(dst_port);
    return total;
}

#endif // DETWS_ENABLE_PROXY_PROTOCOL
