/** @file networkproxy.c  Network proxy (SOCKS v5).

@authors Copyright (c) 2026 Jaakko Keränen <jaakko.keranen@iki.fi>

@par License

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice, this
   list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright notice,
   this list of conditions and the following disclaimer in the documentation
   and/or other materials provided with the distribution.

<small>THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.</small>
*/

#include "the_Foundation/networkproxy.h"
#include "the_Foundation/mutex.h"
#include "the_Foundation/socket.h"

#include "networkproxy.h" /* internal params */

#if defined(iPlatformWindows)
#define WIN32_LEAN_AND_MEAN
#   include <WinSock2.h>
#   include <Windows.h>
#   include <ws2tcpip.h>
typedef int iSockResult_;
#else
#   include <sys/socket.h>
#   include <netinet/in.h>
#   include <unistd.h>
typedef ssize_t iSockResult_;
#endif

struct Impl_NetworkProxy {
    iNetworkProxyParams params;
};

static iNetworkProxy *networkProxy_ = NULL;
static iMutex         networkProxyMutex_;

/* Called from init_Foundation(). */
void init_NetworkProxy_(void) {
    init_Mutex(&networkProxyMutex_);
}

/* Called from deinit_Foundation(). */
void deinit_NetworkProxy_(void) {
    lock_Mutex(&networkProxyMutex_);
    iNetworkProxy *proxy = networkProxy_;
    networkProxy_        = NULL;
    unlock_Mutex(&networkProxyMutex_);
    delete_NetworkProxy(proxy);
    deinit_Mutex(&networkProxyMutex_);
}

iDefineTypeConstructionArgs(NetworkProxy, (const iString *host, uint16_t port), host, port)

void init_NetworkProxy(iNetworkProxy *d, const iString *host, uint16_t port) {
    d->params.port = port;
    if (host) {
        initCopy_String(&d->params.host, host);
    }
    else {
        init_String(&d->params.host);
    }
    init_String(&d->params.username);
    init_String(&d->params.password);
}

void deinit_NetworkProxy(iNetworkProxy *d) {
    deinit_NetworkProxyParams(&d->params);
}

iBool isUsed_NetworkProxy_(void) {
    iBool result;
    iGuardMutex(&networkProxyMutex_, result = (networkProxy_ != NULL));
    return result;
}

void init_NetworkProxyParams(iNetworkProxyParams *d) {
    /* The parameters are taken from the active global NetworkProxy instance,
       if there is one. */
    lock_Mutex(&networkProxyMutex_);
    if (!networkProxy_) {
        unlock_Mutex(&networkProxyMutex_);
        d->port = 0;
        init_String(&d->host);
        init_String(&d->username);
        init_String(&d->password);
        return;
    }
    const iNetworkProxyParams *params = &networkProxy_->params;
    d->port = params->port;
    initCopy_String(&d->host,     &params->host);
    initCopy_String(&d->username, &params->username);
    initCopy_String(&d->password, &params->password);
    unlock_Mutex(&networkProxyMutex_);
}

void deinit_NetworkProxyParams(iNetworkProxyParams *d) {
    deinit_String(&d->password);
    deinit_String(&d->username);
    deinit_String(&d->host);
}

iDefineTypeConstruction(NetworkProxyParams)

void setCredentials_NetworkProxy(iNetworkProxy *d, const iString *username,
                                 const iString *password) {
    /* In case the proxy is already in use, lock the mutex while applying changes. */
    iGuardMutex(&networkProxyMutex_, {
        iString empty;
        init_String(&empty);
        set_String(&d->params.username, username ? username : &empty);
        set_String(&d->params.password, password ? password : &empty);
        deinit_String(&empty);
    });
}

void use_NetworkProxy(iNetworkProxy *d) {
    iNetworkProxy *old;
    iGuardMutex(&networkProxyMutex_, {
        old           = networkProxy_;
        networkProxy_ = d;
    });
    delete_NetworkProxy(old);
}

static iBool sendAll_(iNativeSocket fd, const uint8_t *data, size_t size) {
    while (size > 0) {
        iSockResult_ sent = send((iNativeSocket) fd, (const char *) data, (int) size, 0);
        if (sent <= 0) return iFalse;
        data += sent;
        size -= (size_t) sent;
    }
    return iTrue;
}

static iBool recvAll_(iNativeSocket fd, uint8_t *buf, size_t size) {
    while (size > 0) {
        iSockResult_ n = recv((iNativeSocket) fd, (char *) buf, (int) size, 0);
        if (n <= 0) return iFalse;
        buf += n;
        size -= (size_t) n;
    }
    return iTrue;
}

iBool doSocks5Handshake_NetworkProxy_(iNativeSocket fd, const char *host, uint16_t port,
                                      const iNetworkProxyParams *params) {
    const iBool  hasAuth = !isEmpty_String(&params->username);
    uint8_t      response[4];
    const size_t hlen = strlen(host);
    if (hlen > 255) {
        iWarning("[NetworkProxy] SOCKS5: target hostname too long (>255)\n");
        return iFalse;
    }
    /* Step 1: Send greeting with supported auth methods. */
    uint8_t greeting[4];
    size_t  greetSize;
    greeting[0] = 0x05; /* SOCKS version */
    if (hasAuth) {
        greeting[1] = 0x02; /* 2 methods */
        greeting[2] = 0x00; /* no authentication */
        greeting[3] = 0x02; /* username/password (RFC 1929) */
        greetSize   = 4;
    }
    else {
        greeting[1] = 0x01; /* 1 method */
        greeting[2] = 0x00; /* no authentication */
        greetSize   = 3;
    }
    if (!sendAll_(fd, greeting, greetSize)) {
        iWarning("[NetworkProxy] SOCKS5: failed to send greeting\n");
        return iFalse;
    }
    /* Step 2: Receive server's chosen method. */
    if (!recvAll_(fd, response, 2) || response[0] != 0x05) {
        iWarning("[NetworkProxy] SOCKS5: invalid greeting response\n");
        return iFalse;
    }
    if (response[1] == 0xFF) {
        iWarning("[NetworkProxy] SOCKS5: no acceptable authentication method\n");
        return iFalse;
    }
    /* Step 3: Username/password authentication if required (RFC 1929). */
    if (response[1] == 0x02) {
        if (!hasAuth) {
            iWarning("[NetworkProxy] SOCKS5: server requires authentication but none provided\n");
            return iFalse;
        }
        const size_t ulen = size_String(&params->username);
        const size_t plen = size_String(&params->password);
        if (ulen > 255 || plen > 255) {
            iWarning("[NetworkProxy] SOCKS5: username or password exceeds 255 characters\n");
            return iFalse;
        }
        uint8_t authBuf[3 + 255 + 255];
        size_t  pos    = 0;
        authBuf[pos++] = 0x01; /* sub-negotiation version */
        authBuf[pos++] = (uint8_t) ulen;
        memcpy(authBuf + pos, cstr_String(&params->username), ulen);
        pos += ulen;
        authBuf[pos++] = (uint8_t) plen;
        if (plen > 0) {
            memcpy(authBuf + pos, cstr_String(&params->password), plen);
            pos += plen;
        }
        if (!sendAll_(fd, authBuf, pos)) {
            iWarning("[NetworkProxy] SOCKS5: failed to send credentials\n");
            return iFalse;
        }
        if (!recvAll_(fd, response, 2) || response[0] != 0x01 || response[1] != 0x00) {
            iWarning("[NetworkProxy] SOCKS5: authentication failed\n");
            return iFalse;
        }
    }
    /* Step 4: Send CONNECT request using DOMAINNAME address type. */
    uint8_t request[4 + 1 + 255 + 2];
    size_t  pos = 0;
    request[pos++]  = 0x05; /* SOCKS version */
    request[pos++]  = 0x01; /* CONNECT command */
    request[pos++]  = 0x00; /* reserved */
    request[pos++]  = 0x03; /* address type: DOMAINNAME */
    request[pos++]  = (uint8_t) hlen;
    memcpy(request + pos, host, hlen);
    pos += hlen;
    request[pos++] = (uint8_t) (port >> 8);
    request[pos++] = (uint8_t) (port & 0xFF);
    if (!sendAll_(fd, request, pos)) {
        iWarning("[NetworkProxy] SOCKS5: failed to send CONNECT request\n");
        return iFalse;
    }
    /* Step 5: Receive server response. */
    if (!recvAll_(fd, response, 4) || response[0] != 5) {
        iWarning("[NetworkProxy] SOCKS5: invalid CONNECT response\n");
        return iFalse;
    }
    if (response[1] != 0) {
        iWarning("[NetworkProxy] SOCKS5: CONNECT failed (server reply: 0x%02x)\n", response[1]);
        return iFalse;
    }
    /* Parse the bound address from the response to complete the handshake. */
    uint8_t addrBytes[16];
    uint8_t portBytes[2];
    switch (response[3]) {
        case 0x01: { /* IPv4: 4 bytes + 2 port */
            if (!recvAll_(fd, addrBytes, 4) || !recvAll_(fd, portBytes, 2)) return iFalse;
            break;
        }
        case 0x04: { /* IPv6: 16 bytes + 2 port */
            if (!recvAll_(fd, addrBytes, 16) || !recvAll_(fd, portBytes, 2)) return iFalse;
            break;
        }
        case 0x03: { /* Domain name: 1-byte length + N bytes + 2 port */
            uint8_t dlen;
            char    domain[256];
            if (!recvAll_(fd, &dlen, 1)) return iFalse;
            if (!recvAll_(fd, (uint8_t *) domain, dlen) || !recvAll_(fd, portBytes, 2)) {
                return iFalse;
            }
            break;
        }
        default:
            iWarning("[NetworkProxy] SOCKS5: unknown address type in response: 0x%02x\n",
                     response[3]);
            return iFalse;
    }
    iDebug("[NetworkProxy] SOCKS5 handshake successful: %s:%u\n", host, port);
    return iTrue;
}
