// A proxy, in this process, that does exactly what a test tells it to.

//

// The proxy tests need the three things a real proxy does between accepting a

// connection and relaying bytes: read what the client asked for, answer it, and

// then splice the client to the target. A handler scripts the first two and

// `Peer::relay_to` does the third, so one test can make the proxy refuse, stall

// or hang up, and another can let a request through to the TLS server from

// tls_test_server.hpp and check what arrived.

//

// Plain TCP on 127.0.0.1, port chosen by the OS. It uses mbedtls's net layer

// for the same reason tls_test_server.hpp does: it is already a dependency and

// it is the same on every platform the library builds for.

#pragma once

#include <mbedtls/net_sockets.h>

#include <atomic>

#include <chrono>

#include <functional>

#include <mutex>

#include <string>

#include <string_view>

#include <thread>

#include <utility>

#include <vector>

// Which socket interface this is built against; see socket.cppm.

#if defined(_WIN32) && !defined(TINYHTTPS_POSIX_SOCKETS)

#define TINYHTTPS_WINSOCK 1

#endif

#ifdef TINYHTTPS_WINSOCK

#include <winsock2.h>

#include <ws2tcpip.h>

#else

#include <sys/socket.h>

#include <netinet/in.h>

#endif

namespace proxy_test {

class Server;

// The accepted connection, as the handler sees it.

class Peer {

public:

Peer(Server& server, mbedtls_net_context net) : server_(server), net_(net) {}

~Peer() { mbedtls_net_free(&net_); }

Peer(const Peer&) = delete;

Peer& operator=(const Peer&) = delete;

// Everything up to and including the blank line that ends an HTTP head,

// read a byte at a time so nothing past it is taken. Empty on a hang-up.

std::string read_head();

// Exactly `n` bytes, or empty when the client hangs up first.

std::string read_exact(std::size_t n);

// Whatever arrives first, up to 4 KiB. Empty on a hang-up.

std::string read_some();

bool write(std::string_view data) {

std::size_t sent = 0;

while (sent < data.size()) {

int ret = mbedtls_net_send(&net_,

reinterpret_cast<const unsigned char*>(data.data() + sent),

data.size() - sent);

if (ret <= 0) return false;

sent += static_cast<std::size_t>(ret);

}

return true;

}

void note(std::string text);

// Splices this connection to 127.0.0.1:`port` until either side ends.

void relay_to(int port);

private:

bool wait_readable(int ms) { return mbedtls_net_poll(&net_, MBEDTLS_NET_POLL_READ, ms) > 0; }

Server& server_;

mbedtls_net_context net_;

};

class Server {

public:

using Handler = std::function<void(Peer&)>;

explicit Server(Handler handler) : handler_(std::move(handler)) {

mbedtls_net_init(&listener_);

if (mbedtls_net_bind(&listener_, "127.0.0.1", "0", MBEDTLS_NET_PROTO_TCP) != 0) {

failed_ = true;

return;

}

sockaddr_in addr {};

socklen_t len = sizeof addr;

if (::getsockname(listener_.fd, reinterpret_cast<sockaddr*>(&addr), &len) != 0) {

failed_ = true;

return;

}

port_ = ntohs(addr.sin_port);

acceptor_ = std::thread([this] { accept_loop(); });

}

~Server() { stop(); }

Server(const Server&) = delete;

Server& operator=(const Server&) = delete;

void stop() {

if (stopping_.exchange(true)) return;

if (acceptor_.joinable()) acceptor_.join();

for (auto& t : connections_) {

if (t.joinable()) t.join();

}

connections_.clear();

mbedtls_net_free(&listener_);

}

[[nodiscard]] bool failed() const { return failed_; }

[[nodiscard]] int port() const { return port_; }

[[nodiscard]] bool stopping() const { return stopping_.load(); }

[[nodiscard]] int accepts() const { return accepts_.load(); }

// The text a handler stored with `note`, for the test to assert on once the

// request has completed.

void note(std::string text) {

std::lock_guard<std::mutex> lock(mutex_);

notes_.push_back(std::move(text));

}

[[nodiscard]] std::vector<std::string> notes() {

std::lock_guard<std::mutex> lock(mutex_);

return notes_;

}

private:

void accept_loop() {

while (!stopping_.load()) {

int ready = mbedtls_net_poll(&listener_, MBEDTLS_NET_POLL_READ, 50);

if (ready < 0) break;

if (ready == 0) continue;

mbedtls_net_context client {};

mbedtls_net_init(&client);

if (mbedtls_net_accept(&listener_, &client, nullptr, 0, nullptr) != 0) {

mbedtls_net_free(&client);

continue;

}

accepts_.fetch_add(1);

connections_.emplace_back([this, client] {

Peer peer(*this, client);

handler_(peer);

});

}

}

Handler handler_;

mbedtls_net_context listener_ {};

int port_ { 0 };

bool failed_ { false };

std::atomic<bool> stopping_ { false };

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

std::thread acceptor_;

std::vector<std::thread> connections_;

std::mutex mutex_;

std::vector<std::string> notes_;

};

inline void Peer::note(std::string text) { server_.note(std::move(text)); }

inline std::string Peer::read_head() {

std::string head;

unsigned char c {};

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

if (head.size() > 65536) return {};

// 20 s is far beyond any test; the loop exists so `stop` is noticed.

bool ready = false;

for (int waited = 0; waited < 20000 && !ready; waited += 50) {

if (server_.stopping()) return {};

ready = wait_readable(50);

}

if (!ready) return {};

if (mbedtls_net_recv(&net_, &c, 1) <= 0) return {};

head.push_back(static_cast<char>(c));

}

return head;

}

inline std::string Peer::read_exact(std::size_t n) {

std::string out;

unsigned char c {};

while (out.size() < n) {

bool ready = false;

for (int waited = 0; waited < 20000 && !ready; waited += 50) {

if (server_.stopping()) return {};

ready = wait_readable(50);

}

if (!ready || mbedtls_net_recv(&net_, &c, 1) <= 0) return {};

out.push_back(static_cast<char>(c));

}

return out;

}

inline std::string Peer::read_some() {

unsigned char buf[4096];

for (int waited = 0; waited < 20000; waited += 50) {

if (server_.stopping()) return {};

if (!wait_readable(50)) continue;

int n = mbedtls_net_recv(&net_, buf, sizeof buf);

return n > 0 ? std::string(reinterpret_cast<char*>(buf), static_cast<std::size_t>(n))

: std::string();

}

return {};

}

inline void Peer::relay_to(int port) {

mbedtls_net_context target {};

mbedtls_net_init(&target);

if (mbedtls_net_connect(&target, "127.0.0.1", std::to_string(port).c_str(),

MBEDTLS_NET_PROTO_TCP) != 0) {

return;

}

// Round-robin over two sockets with a short poll each: slower than a

// select over both, and portable without a second code path.

unsigned char buf[4096];

bool open = true;

while (open && !server_.stopping()) {

for (auto [from, to] : {std::pair{&net_, &target}, std::pair{&target, &net_}}) {

if (mbedtls_net_poll(from, MBEDTLS_NET_POLL_READ, 10) <= 0) continue;

int n = mbedtls_net_recv(from, buf, sizeof buf);

if (n <= 0) { open = false; break; }

for (int sent = 0; sent < n;) {

int w = mbedtls_net_send(to, buf + sent, static_cast<std::size_t>(n - sent));

if (w <= 0) { open = false; break; }

sent += w;

}

}

}

mbedtls_net_free(&target);

}

} // namespace proxy_test