portproxy.cpp

#include <iostream>
#include <string>
#include <memory>
#include <thread>
#include <filesystem>
#include <stdexcept>
#include <boost/asio.hpp>
#include <gsl/narrow>

namespace ip = boost::asio::ip;
namespace fs = std::filesystem;

constexpr auto BUFFER_SIZE = 4096;

class Proxy {
public:
	Proxy(boost::asio::io_context& io_context, ip::port_type listen_port, const std::string& remote_host, ip::port_type remote_port)
		: acceptor_(io_context, ip::tcp::endpoint(ip::tcp::v4(), listen_port))
	{
		// Resolve the remote host.
		ip::tcp::resolver resolver(io_context);
		auto endpoints = resolver.resolve(remote_host, std::to_string(remote_port));
		if (endpoints.empty())
			throw std::runtime_error("Host " + remote_host + " not found.");
		remote_endpoint_ = *endpoints.begin(); // Get the first resolved endpoint
		std::cout << "Resolved " << remote_host << ":" << remote_port << " to " << remote_endpoint_ << "\n";

		// Start listening.
		startAccept();
	}

private:
	void startAccept() {
		std::cout << "Listening for clients...\n";
		auto client_socket = std::make_shared<ip::tcp::socket>(acceptor_.get_executor());
		acceptor_.async_accept(*client_socket, [this, client_socket](const boost::system::error_code& error) {
			if (!error) {
				startProxy(client_socket);
			} else {
				std::cerr << "Failed to accept connection from " << client_socket << ". err=" << error.what() << "\n";
			}
			startAccept(); //< Keep listening for new clients.
		});
	}

	void startProxy(std::shared_ptr<ip::tcp::socket> client_socket) {
		std::cout << "Client connected from " << client_socket->remote_endpoint() << "\n";
		client_socket->set_option(ip::tcp::no_delay(true)); //< Disable nagle's algorithm for lower latency.
		auto remote_socket = std::make_shared<ip::tcp::socket>(client_socket->get_executor());

		remote_socket->async_connect(remote_endpoint_,
			[this, client_socket, remote_socket](const boost::system::error_code& error) {
				if (error) {
					std::cerr << "Could not connect to remote socket " << remote_socket << ". err=" << error.what() << "\n";
					return;
				}

				remote_socket->set_option(ip::tcp::no_delay(true)); //< Disable nagle's algorithm for lower latency.

				// Start proxying traffic in both directions.
				std::cout << "Proxying client from " << client_socket->remote_endpoint() << " to " << remote_socket->remote_endpoint() << "\n";
				startReadWriteProxy(client_socket, remote_socket);
				startReadWriteProxy(remote_socket, client_socket);
		});
	}

	void startReadWriteProxy(std::shared_ptr<ip::tcp::socket> read_socket, std::shared_ptr<ip::tcp::socket> write_socket) {
		auto buffer = std::make_shared<std::array<char, BUFFER_SIZE>>();
		read_socket->async_read_some(boost::asio::buffer(*buffer),
			[this, read_socket, write_socket, buffer](const boost::system::error_code& error, std::size_t bytes_transferred) {
				if (error) {
					std::cerr << "Failed to read from socket " << read_socket << ". err=" << error.what() << "\n";
					return;
				}

				boost::asio::async_write(*write_socket, boost::asio::buffer(*buffer, bytes_transferred),
					[this, read_socket, write_socket, buffer](const boost::system::error_code& error, std::size_t) {
						if (error) {
							std::cerr << "Failed to write to socket " << write_socket << ". err=" << error.what() << "\n";
							return;
						}

						startReadWriteProxy(read_socket, write_socket); //< Keep proxying...
					});
			});
	}

	ip::tcp::acceptor acceptor_;
	ip::tcp::endpoint remote_endpoint_;
};

static void printUsage(const std::string& arg0) {
	std::cout << "Usage: " << fs::path(arg0).stem() << " <listen_port> <remote_host> <remote_port>\n";
}

int main(int argc, char* argv[]) {
	try {
		if (argc != 4) {
			printUsage(argc > 0 ? argv[0] : "PortProxy");
			return 1;
		}

		// Decode command line parameters.
		auto listen_port = gsl::narrow<ip::port_type>(std::stoi(argv[1]));
		std::string remote_host = argv[2];
		auto remote_port = gsl::narrow<ip::port_type>(std::stoi(argv[3]));

		// Run the proxy.
		boost::asio::io_context io_context;
		Proxy proxy(io_context, listen_port, remote_host, remote_port);
		io_context.run();
	} catch (std::exception& e) {
		std::cerr << "Exception: " << e.what() << std::endl;
		printUsage(argv[0]);
	}
}