test

#include "WebSocketProxy.h"

#include <algorithm>

#include <atomic>

#include <cctype>

#include <cstring>

#include <iostream>

#include <map>

#include <set>

#include <sstream>

#include <thread>

#include <vector>

#include <netplus/connection.h>

#include <netplus/exception.h>

namespace proxyplus::httpproxy {

namespace {

std::string toLower(const std::string &s) {

    std::string out = s;

    std::transform(out.begin(), out.end(), out.begin(),

                    [](unsigned char c) { return static_cast<char>(std::tolower(c)); });

    return out;

}

// "Connection: keep-alive, Upgrade" -- a comma-separated token list (RFC

// 7230 §3.2), any one of which may be "upgrade" case-insensitively.

bool connectionHeaderRequestsUpgrade(const std::string &value) {

    std::stringstream ss(value);

    std::string token;

    while (std::getline(ss, token, ',')) {

        size_t b = token.find_first_not_of(" \t");

        size_t e = token.find_last_not_of(" \t");

        if (b == std::string::npos) continue;

        if (toLower(token.substr(b, e - b + 1)) == "upgrade") return true;

    }

    return false;

}

std::string headerValue(libhttppp::HttpRequest &req, const std::string &key) {

    if (auto *hd = req.getHeaderData(key)) {

        if (auto *v = hd->getfirstValue()) return v->getvalue();

    }

    return {};

}

// Bounds how many concurrent WebSocket tunnels this process will open --

// each owns a dedicated relay thread for its whole (potentially very long)

// lifetime, unlike the bounded, short-task-oriented dispatch pools used

// elsewhere (see quicDispatchPool()/_h2DispatchPool/_h3StreamPool). Home

// Assistant-style deployments run a handful of concurrent sessions, not

// thousands, so this is generous headroom, not a real limit in practice --

// its job is purely to guarantee an *unbounded* number of upgrade requests

// can never turn into unbounded thread creation (see

// project_proxyplus_h3_thread_starvation in memory for exactly what that

// class of bug looks like in production).

constexpr int kMaxConcurrentWsTunnels = 256;

std::atomic<int> g_activeWsTunnels{0};

struct ScopedTunnelSlot {

    bool acquired = false;

    ScopedTunnelSlot() {

        int cur = g_activeWsTunnels.fetch_add(1, std::memory_order_relaxed);

        if (cur < kMaxConcurrentWsTunnels) {

            acquired = true;

        } else {

            g_activeWsTunnels.fetch_sub(1, std::memory_order_relaxed);

        }

    }

    ~ScopedTunnelSlot() {

        if (acquired) g_activeWsTunnels.fetch_sub(1, std::memory_order_relaxed);

    }

    ScopedTunnelSlot(const ScopedTunnelSlot &) = delete;

};

// Opens a raw (not HttpClient) connection to the upstream, blocking until

// connected (and, for TLS, until the handshake completes) -- mirrors

// libhttppp::HttpClient::resetConnection()'s HTTPS/plain-TCP branches

// exactly, minus the H2/H3 ALPN negotiation: a WebSocket upgrade is always

// spoken as HTTP/1.1, so there's nothing to negotiate.

std::unique_ptr<netplus::socket> connectUpstream(const libhttppp::HttpUrl &upstream,

                                                   const netplus::TlsTrustPolicy &trustPolicy,

                                                   int timeoutSec) {

    if (upstream.getOriginalProtocol() == libhttppp::HttpUrl::HTTPS) {

        std::map<std::string, netplus::ssl::CertificateBundle> certs;

        auto sslsock = std::make_unique<netplus::ssl>(certs, -1);

        sslsock->getTls().trust_policy = trustPolicy;

        sslsock->getTls().client_alpn_protocols = std::string("\x08http/1.1", 9);

        sslsock->connect(upstream.getHost(), upstream.getPort(), false);

        return sslsock;

    }

    auto tcpsock = std::make_unique<netplus::tcp>(-1);

    tcpsock->connectTimeout(upstream.getHost(), upstream.getPort(), timeoutSec * 1000);

    return tcpsock;

}

void sendAll(netplus::socket &sock, const std::string &data) {

    size_t sent = 0;

    while (sent < data.size()) {

        netplus::buffer buf(data.data() + sent, data.size() - sent);

        size_t n = sock.sendData(buf, 0);

        if (n == 0) {

            netplus::NetException e;

            e[netplus::NetException::Error] << "WebSocket handshake: upstream write failed";

            throw e;

        }

        sent += n;

    }

}

// Reads from sock until "\r\n\r\n" is seen (i.e. the full status line +

// headers), capped so a misbehaving upstream can't grow this unbounded.

// Returns the raw bytes read, which may include some response body /

// leftover bytes past the header terminator -- callers that care split on

// the terminator themselves.

constexpr size_t kMaxHandshakeResponseBytes = 64 * 1024;

std::string readUntilHeadersEnd(netplus::socket &sock) {

    std::string acc;

    while (acc.find("\r\n\r\n") == std::string::npos) {

        if (acc.size() > kMaxHandshakeResponseBytes) {

            netplus::NetException e;

            e[netplus::NetException::Error] << "WebSocket handshake: upstream response headers too large";

            throw e;

        }

        netplus::buffer buf(4096);

        size_t n = sock.recvData(buf, 0);

        if (n == 0) {

            netplus::NetException e;

            e[netplus::NetException::Error] << "WebSocket handshake: upstream closed before responding";

            throw e;

        }

        acc.append(buf.data.buf, n);

    }

    return acc;

}

int parseStatusCode(const std::string &statusLine) {

    // "HTTP/1.1 101 Switching Protocols" -- code is the token after the

    // first space.

    size_t sp1 = statusLine.find(' ');

    if (sp1 == std::string::npos) return 0;

    size_t sp2 = statusLine.find(' ', sp1 + 1);

    std::string code = statusLine.substr(sp1 + 1, sp2 == std::string::npos

                                                        ? std::string::npos

                                                        : sp2 - sp1 - 1);

    try {

        return std::stoi(code);

    } catch (...) {

        return 0;

    }

}

// Relays raw bytes both ways between the client and the upstream until

// either side closes or errors, then closes both. Runs on its own

// dedicated thread for the tunnel's entire lifetime -- see

// kMaxConcurrentWsTunnels for why that's an acceptable tradeoff here

// (unlike the per-request dispatch pools, which must never block on

// arbitrarily long-lived work).

void relayLoop(std::unique_ptr<netplus::socket> client,

               std::unique_ptr<netplus::socket> upstream,

               std::string clientLeftover,

               std::string upstreamLeftover,

               ScopedTunnelSlot slot) {

    try {

        // This thread owns both sockets exclusively for its whole lifetime

        // (nothing else will ever touch client's fd again -- see

        // detachConnection()'s contract), so there's no reason to deal with

        // EAGAIN/Note here: plain blocking reads/writes are simplest and

        // correct. client arrives non-blocking (the event loop always runs

        // non-blocking sockets); upstream is fresh from connectUpstream() and

        // likely already blocking, but set explicitly rather than assume.

        client->setBlock();

        upstream->setBlock();

        if (!upstreamLeftover.empty()) sendAll(*client, upstreamLeftover);

        if (!clientLeftover.empty()) sendAll(*upstream, clientLeftover);

        netplus::socketwait sw;

        std::vector<netplus::socket *> socks{client.get(), upstream.get()};

        std::vector<bool> ready;

        for (;;) {

            int n = sw.waitReadMulti(socks, ready, /*timeout_ms=*/60000);

            if (n <= 0) continue; // timeout -- keep waiting, WebSocket pings keep this alive

            if (ready[0]) {

                netplus::buffer buf(16384);

                size_t r = client->recvData(buf, 0);

                if (r == 0) break; // client closed

                sendAll(*upstream, std::string(buf.data.buf, r));

            }

            if (ready[1]) {

                netplus::buffer buf(16384);

                size_t r = upstream->recvData(buf, 0);

                if (r == 0) break; // upstream closed

                sendAll(*client, std::string(buf.data.buf, r));

            }

        }

    } catch (const netplus::NetException &) {

        // Either side closed/reset mid-relay -- normal end of a WebSocket

        // session, nothing to report.

    } catch (const std::exception &e) {

        std::cerr << "[ws] relay error: " << e.what() << std::endl;

    }

    client->close();

    upstream->close();

}

} // namespace

bool isWebSocketUpgrade(libhttppp::HttpRequest &curreq) {

    std::string connection = headerValue(curreq, "connection");

    std::string upgrade = headerValue(curreq, "upgrade");

    return connectionHeaderRequestsUpgrade(connection) && toLower(upgrade) == "websocket";

}

void handleWebSocketUpgrade(libhttppp::HttpRequest &curreq,

                             const std::string &sourceIp,

                             const std::string &upstreamBase,

                             const std::string &forwardPath,

                             const netplus::TlsTrustPolicy &trustPolicy) {

    try {

        libhttppp::HttpUrl upstream(upstreamBase + forwardPath);

        auto upstreamSock = connectUpstream(upstream, trustPolicy, /*timeoutSec=*/30);

        // Build the raw HTTP/1.1 request text. Unlike RequestEvent's normal

        // proxying path, Connection/Upgrade (and the Sec-WebSocket-* set)

        // must survive to the upstream -- only genuinely hop-by-hop /

        // recomputed-here headers are dropped.

        static const std::set<std::string> kSkip = {

            "host", "content-length", "transfer-encoding", "keep-alive",

            "proxy-authenticate", "proxy-authorization", "te", "trailer",

        };

        std::ostringstream req;

        req << "GET " << upstream.getPath() << " HTTP/1.1\r\n";

        req << "Host: " << upstream.getHost();

        if (upstream.getPort() != (upstream.getOriginalProtocol() == libhttppp::HttpUrl::HTTPS ? 443 : 80)) {

            req << ":" << upstream.getPort();

        }

        req << "\r\n";

        for (auto *hd = curreq.getfirstHeaderData(); hd; hd = hd->nextHeaderData()) {

            const std::string &key = hd->getkey();

            if (key.empty() || key[0] == ':' || kSkip.count(key)) continue;

            for (auto *v = hd->getfirstValue(); v; v = v->nextvalue()) {

                req << key << ": " << v->getvalue() << "\r\n";

            }

        }

        std::string xff = sourceIp;

        if (auto *hd = curreq.getHeaderData("x-forwarded-for")) {

            if (auto *v = hd->getfirstValue()) xff = v->getvalue() + ", " + sourceIp;

        }

        req << "X-Forwarded-For: " << xff << "\r\n";

        if (auto *hostHdr = curreq.getHeaderData("host")) {

            if (auto *v = hostHdr->getfirstValue()) req << "X-Forwarded-Host: " << v->getvalue() << "\r\n";

        }

        req << "\r\n";

        sendAll(*upstreamSock, req.str());

        std::string resp = readUntilHeadersEnd(*upstreamSock);

        size_t termPos = resp.find("\r\n\r\n");

        std::string headerPart = resp.substr(0, termPos + 4);

        std::string upstreamLeftover = resp.substr(termPos + 4);

        size_t lineEnd = headerPart.find("\r\n");

        std::string statusLine = headerPart.substr(0, lineEnd);

        int status = parseStatusCode(statusLine);

        if (status != 101) {

            // Not a WebSocket response (auth failure, wrong path, upstream

            // just doesn't support WS here, ...) -- forward it back to the

            // client as a normal completed response instead of tunneling.

            std::string reason;

            size_t sp1 = statusLine.find(' ');

            size_t sp2 = (sp1 == std::string::npos) ? std::string::npos

                                                      : statusLine.find(' ', sp1 + 1);

            if (sp2 != std::string::npos) reason = statusLine.substr(sp2 + 1);

            libhttppp::HttpResponse curres;

            curres.setState(std::to_string(status) + (reason.empty() ? "" : " " + reason));

            std::vector<char> body(upstreamLeftover.begin(), upstreamLeftover.end());

            curres.send(curreq, body, static_cast<int>(body.size()));

            upstreamSock->close();

            return;

        }

        auto clientSock = netplus::detachConnection(curreq);

        if (!clientSock) {

            // Backend doesn't support detaching (see connection.h), or the

            // connection is already gone -- nothing safe left to do.

            upstreamSock->close();

            return;

        }

        std::vector<char> leftoverVec(curreq.RecvData.begin(), curreq.RecvData.end());

        std::string clientLeftover(leftoverVec.begin(), leftoverVec.end());

        ScopedTunnelSlot slot;

        if (!slot.acquired) {

            std::cerr << "[ws] too many concurrent WebSocket tunnels (" << kMaxConcurrentWsTunnels

                       << "), rejecting new upgrade" << std::endl;

            clientSock->close();

            upstreamSock->close();

            return;

        }

        sendAll(*clientSock, headerPart);

        std::thread(relayLoop, std::move(clientSock), std::move(upstreamSock),

                    std::move(clientLeftover), std::move(upstreamLeftover), std::move(slot))

            .detach();

    } catch (const std::exception &e) {

        std::cerr << "[ws] upgrade failed: " << e.what() << std::endl;

        libhttppp::HttpResponse curres;

        curres.setState(HTTP502);

        std::vector<char> body;

        curres.send(curreq, body, 0);

    }

}

} // namespace proxyplus::httpproxy

#pragma once

#include <string>

#include <httppp/http.h>

#include <netplus/socket.h>

namespace proxyplus::httpproxy {

// True if this request is asking to upgrade to a WebSocket connection --

// RFC 6455 §4.1: a "Connection" header whose token list includes "Upgrade"

// (case-insensitively, RFC 7230 §3.2) and an "Upgrade" header naming

// "websocket".

bool isWebSocketUpgrade(libhttppp::HttpRequest &curreq);

// Handles a WebSocket upgrade request end to end. WebSocket is universally

// negotiated over plain HTTP/1.1 even when the surrounding page loaded over

// H2/H3 -- browsers open a fresh H1.1 connection for it whenever the server

// doesn't advertise RFC 8441 extended CONNECT support (which proxyplus

// doesn't), so this always speaks H1.1 to the upstream regardless of what

// protocol curreq itself arrived over.

//

// Connects to the upstream, replays the original request line + headers

// (this time keeping Connection/Upgrade, which the normal proxying path in

// RequestEvent deliberately strips as hop-by-hop -- see its kHopByHop set),

// and:

//   - on a 101 response: hands curreq's socket off to a dedicated relay

//     thread (via netplus::detachConnection) that splices raw bytes both

//     ways for the connection's lifetime. curreq must not be touched again

//     by the caller after this returns -- its socket is gone.

//   - on any other response: forwards the upstream's status back to the

//     client as a normal completed HTTP response (e.g. upstream auth

//     failure, or a path that just isn't a WebSocket endpoint); curreq

//     completes normally, exactly like the non-WebSocket proxying path.

//

// upstreamBase/forwardPath/trustPolicy mirror exactly what RequestEvent's

// normal proxying path already computed for this request -- passed in

// rather than recomputed, so route/pool selection stays single-sourced.

//

// Returns whether the upstream was reachable and answered at all (101 or

// any other status) -- mirrors the normal path's "a real HTTP response, even

// a non-2xx one, counts as up" convention (RequestEvent's `status < 500`),

// for callers that report this to a load-balancer pool. False only for a

// transport-level failure (connect refused, timeout, ...).

bool handleWebSocketUpgrade(libhttppp::HttpRequest &curreq,

                             const std::string &sourceIp,

                             const std::string &upstreamBase,

                             const std::string &forwardPath,

                             const netplus::TlsTrustPolicy &trustPolicy);

} // namespace proxyplus::httpproxy