// Copyright (c) 2026 Bitflash developers
// Distributed under the MIT/X11 software license.
//
// Transparent anonymous tunnel. See btftunnel.h.

#include "btftunnel.h"
#include "btfchan.h"

#ifdef _WIN32
#include <winsock2.h>
#include <ws2tcpip.h>
#else
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
typedef int SOCKET;
#ifndef INVALID_SOCKET
#define INVALID_SOCKET (-1)
#endif
#define closesocket(s) ::close(s)
#ifndef SD_BOTH
#define SD_BOTH SHUT_RDWR
#endif
#endif
#include "sockcount.h"
#include <cstring>
#include <cstdio>
#include <memory>
#include <new>
#include <vector>
#include <thread>
#include <system_error>

namespace btf
{

bool TunnelInit()
{
    return RvInit() && ChanInit();
}

// Create a connected loopback socket pair (no socketpair() on Windows).
// Returns app end and pump end; both are ordinary full-duplex TCP sockets.
static bool MakeLocalPair(btf_socket_t& appEnd, btf_socket_t& pumpEnd)
{
    SOCKET listener = BtfSocketTag(socket(AF_INET, SOCK_STREAM, IPPROTO_TCP), SOCK_PAIR_LISTEN);
    if (listener == INVALID_SOCKET) return false;

    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
    addr.sin_port = 0; // ephemeral
    if (bind(listener, (struct sockaddr*)&addr, sizeof(addr)) != 0) { BtfCloseSocket(listener); return false; }
    if (listen(listener, 1) != 0) { BtfCloseSocket(listener); return false; }

#ifdef _WIN32
    int len = sizeof(addr);
#else
    socklen_t len = sizeof(addr);
#endif
    if (getsockname(listener, (struct sockaddr*)&addr, &len) != 0) { BtfCloseSocket(listener); return false; }

    SOCKET a = BtfSocketTag(socket(AF_INET, SOCK_STREAM, IPPROTO_TCP), SOCK_PAIR_APP);
    if (a == INVALID_SOCKET) { BtfCloseSocket(listener); return false; }
    if (connect(a, (struct sockaddr*)&addr, sizeof(addr)) != 0) { BtfCloseSocket(a); BtfCloseSocket(listener); return false; }

    SOCKET b = BtfSocketTag(accept(listener, NULL, NULL), SOCK_PAIR_PUMP);
    BtfCloseSocket(listener);
    if (b == INVALID_SOCKET) { BtfCloseSocket(a); return false; }

    appEnd = a;
    pumpEnd = b;
    return true;
}

static bool SockReadN(SOCKET s, void* buf, int n)
{
    char* p = (char*)buf; int off = 0;
    while (off < n) { int r = recv(s, p + off, n - off, 0); if (r <= 0) return false; off += r; }
    return true;
}
static bool SockWriteN(SOCKET s, const void* buf, int n)
{
    const char* p = (const char*)buf; int off = 0;
    while (off < n) { int r = send(s, p + off, n - off, 0); if (r <= 0) return false; off += r; }
    return true;
}

// Outbound pump: plaintext from the app -> echan frame -> rendezvous socket.
static void PumpOut(SOCKET pumpEnd, RvSocket rv, std::vector<unsigned char> key)
{
    char buf[16 * 1024];
    for (;;)
    {
        int r = recv(pumpEnd, buf, sizeof(buf), 0);
        if (r <= 0) break;
        std::vector<unsigned char> frame = ChanEncryptFrame(key.data(), (const unsigned char*)buf, r);
        if (!RvWriteN(rv, frame.data(), (int)frame.size())) break;
    }
    // Tear down: unblock the inbound pump and the app.
    shutdown((SOCKET)rv, SD_BOTH);
    shutdown(pumpEnd, SD_BOTH);
}

// Inbound pump: echan frames from the rendezvous socket -> plaintext to the app.
static void PumpIn(SOCKET pumpEnd, RvSocket rv, std::vector<unsigned char> key)
{
    for (;;)
    {
        unsigned char len[4];
        if (!RvReadN(rv, len, 4)) break;
        size_t n = ((size_t)len[0]<<24)|((size_t)len[1]<<16)|((size_t)len[2]<<8)|len[3];
        if (n < (size_t)(CHAN_NONCEBYTES + CHAN_MACBYTES) || n > (1u<<20)) break;
        std::vector<unsigned char> body(n);
        if (!RvReadN(rv, body.data(), (int)n)) break;
        std::vector<unsigned char> pt;
        if (!ChanDecryptBody(key.data(), body.data(), (int)n, pt)) break; // auth fail -> drop
        if (!pt.empty() && !SockWriteN(pumpEnd, pt.data(), (int)pt.size())) break;
    }
    shutdown(pumpEnd, SD_BOTH);
    shutdown((SOCKET)rv, SD_BOTH);
}

// The pump end and the rendezvous socket, owned jointly by whoever is still
// using them. Closed once, by the last reference to go.
//
// Joint ownership rather than "the inbound pump owns the cleanup", because
// either pump can reach its exit first and BOTH touch BOTH descriptors on the
// way out -- PumpOut ends in shutdown(rv), shutdown(pumpEnd) and PumpIn ends in
// the mirror image. A single owner therefore has no exit at which the other is
// guaranteed to be finished. Holding a reference for the whole body of each
// pump gives it one: the close cannot happen until the last pump has stopped
// touching the descriptors, so no shutdown() can land on a number that has
// already been closed and handed back out by the kernel to an unrelated
// connection.
struct TunnelEnds
{
    SOCKET   pumpEnd;
    RvSocket rv;
    TunnelEnds(SOCKET p, RvSocket r) : pumpEnd(p), rv(r) {}
    ~TunnelEnds() { BtfCloseSocket(pumpEnd); RvClose(rv); }
private:
    TunnelEnds(const TunnelEnds&);
    TunnelEnds& operator=(const TunnelEnds&);
};

// Start the two pumps and hand back the app-facing socket. When both pumps
// finish, close the pump end and the rendezvous socket.
//
// That sentence used to be a comment describing something the code did not do.
// The inbound pump closed both descriptors as soon as it exited, without
// waiting for the outbound pump, which then called shutdown() on both at its
// own exit. Whether that shutdown hit a live descriptor or a closed one came
// down to which pump got there first, and the rendezvous socket dying before
// the app socket -- a relay restart, the far side hanging up -- is precisely
// the order that puts the inbound pump ahead. Measured over 25 such teardowns
// against the real code: 29 shutdowns landed on an already-closed descriptor.
//
// appEnd is not part of this. It is handed to the caller on success and closed
// here on every failure path; no pump ever touches it.
//
// Every tunnel is two threads, and a tunnel is one peer connection (a .btf
// dial or accept, a pool miner, a participant's pool session) -- so the
// thread count scales with peer count, which is attacker-influenced. When
// pthread_create is refused std::thread throws std::system_error; uncaught it
// reached std::terminate and took the whole process down over one connection
// it could not serve. Both spawns are guarded, and the failure paths return
// INVALID_SOCKET, which every caller already handles by dropping just that
// connection.
static btf_socket_t StartTunnel(RvSocket rv, const unsigned char shared[32])
{
    btf_socket_t appEnd, pumpEnd;
    if (!MakeLocalPair(appEnd, pumpEnd)) { RvClose(rv); return INVALID_SOCKET; }

    std::vector<unsigned char> key(shared, shared + 32);

    // make_shared, not shared_ptr(new ...): it allocates before it constructs,
    // so it either comes back owning both descriptors or throws having taken
    // ownership of neither. The two-step form can run the deleter and then
    // throw, which would close both here and again below.
    std::shared_ptr<TunnelEnds> ends;
    try
    {
        ends = std::make_shared<TunnelEnds>(pumpEnd, rv);
    }
    catch (const std::bad_alloc&)
    {
        fprintf(stderr, "btf tunnel: out of memory starting tunnel, tunnel dropped\n");
        BtfCloseSocket(pumpEnd);
        BtfCloseSocket(appEnd);
        RvClose(rv);
        return INVALID_SOCKET;
    }

    try
    {
        std::thread([ends, key]() { PumpIn(ends->pumpEnd, ends->rv, key); }).detach();
    }
    catch (const std::system_error& e)
    {
        // Nothing is running yet: `ends` is the only reference, and returning
        // drops it, which is what closes pumpEnd and rv.
        fprintf(stderr, "btf tunnel: cannot spawn inbound pump (%s), tunnel dropped\n",
                e.what());
        BtfCloseSocket(appEnd);
        return INVALID_SOCKET;
    }

    try
    {
        std::thread([ends, key]() { PumpOut(ends->pumpEnd, ends->rv, key); }).detach();
    }
    catch (const std::system_error& e)
    {
        // The inbound pump may be parked in RvReadN(rv) or may already have
        // exited -- when the spawn is refused because the process is loaded,
        // the rendezvous socket is often dead for the same reason, and then
        // PumpIn returns on its first length read. Either way this thread
        // still holds a reference, so both descriptors are still open: the
        // shutdown wakes the pump if it is parked and is harmless if it is
        // not. The close belongs to whichever reference goes last.
        fprintf(stderr, "btf tunnel: cannot spawn outbound pump (%s), tunnel dropped\n",
                e.what());
        shutdown((SOCKET)ends->rv, SD_BOTH);
        shutdown(ends->pumpEnd, SD_BOTH);
        BtfCloseSocket(appEnd);
        return INVALID_SOCKET;
    }
    return appEnd;
}

const char* TunnelFailText(TunnelFail f)
{
    switch (f)
    {
    case TUNNEL_FAIL_NONE:      return "no failure";
    case TUNNEL_FAIL_INIT:      return "local socket/crypto init failed";
    case TUNNEL_FAIL_DIAL:      return "meeting node unreachable (tcp dial failed)";
    case TUNNEL_FAIL_REFUSED:   return "relay reachable, but no service registered there for this address";
    case TUNNEL_FAIL_PAIR:      return "relay connection died before pairing completed";
    case TUNNEL_FAIL_HANDSHAKE: return "end-to-end handshake failed";
    case TUNNEL_FAIL_LOCAL:     return "local tunnel setup failed";
    }
    return "unknown failure";
}

btf_socket_t BtfClientTunnel(const char* meeting_host, unsigned short port,
                             const unsigned char target_pubkey[32],
                             const unsigned char service_enc_pub[32],
                             TunnelFail* pFail)
{
    if (pFail) *pFail = TUNNEL_FAIL_INIT;
    if (!TunnelInit()) return INVALID_SOCKET;
    RvClientFail rvFail = RVC_FAIL_NONE;
    RvSocket rv = RvClientConnect(meeting_host, port, target_pubkey, &rvFail);
    if (rv == RV_INVALID)
    {
        if (pFail)
            *pFail = rvFail == RVC_FAIL_DIAL    ? TUNNEL_FAIL_DIAL
                   : rvFail == RVC_FAIL_REFUSED ? TUNNEL_FAIL_REFUSED
                                                : TUNNEL_FAIL_PAIR;
        return INVALID_SOCKET;
    }

    unsigned char eph_pub[32], shared[32];
    if (pFail) *pFail = TUNNEL_FAIL_HANDSHAKE;
    if (!ChanClientHandshake(service_enc_pub, eph_pub, shared)) { RvClose(rv); return INVALID_SOCKET; }
    if (!RvWriteN(rv, eph_pub, 32)) { RvClose(rv); return INVALID_SOCKET; }
    if (pFail) *pFail = TUNNEL_FAIL_LOCAL;
    btf_socket_t app = StartTunnel(rv, shared);
    if (app != INVALID_SOCKET && pFail) *pFail = TUNNEL_FAIL_NONE;
    return app;
}

btf_socket_t BtfServiceWrap(RvSocket rv, const unsigned char my_enc_sk[32])
{
    if (!TunnelInit()) { RvClose(rv); return INVALID_SOCKET; }
    unsigned char eph_pub[32], shared[32];
    if (!RvReadN(rv, eph_pub, 32)) { RvClose(rv); return INVALID_SOCKET; }
    if (!ChanEndpointHandshake(eph_pub, my_enc_sk, shared)) { RvClose(rv); return INVALID_SOCKET; }
    return StartTunnel(rv, shared);
}

} // namespace btf