proxy_v4_fixed.cpp

// proxy_v4_fixed.cpp
// Note: Fixes the v3 data race using a mutex-protected round-robin routing
// implementation. Usage: ./proxy_v4_fixed <listen_port> <backend_host>
// <backend_port_1> [backend_port_2 ...] [--threads N]

#include <arpa/inet.h>
#include <cstring>
#include <iostream>
#include <mutex>
#include <netdb.h>
#include <netinet/in.h>
#include <string>
#include <sys/socket.h>
#include <unistd.h>
#include <vector>

#include "thread_pool.h"

namespace {

constexpr int kListenBacklog = 128;
constexpr size_t kBufferSize = 4096;
constexpr size_t kNumThreads = 4;

int connect_to_backend(const std::string &host, int port) {
  addrinfo hints{};
  hints.ai_family = AF_INET;
  hints.ai_socktype = SOCK_STREAM;

  addrinfo *result = nullptr;
  std::string port_str = std::to_string(port);
  int rc = getaddrinfo(host.c_str(), port_str.c_str(), &hints, &result);
  if (rc != 0) {
    std::cerr << "getaddrinfo failed for " << host << ":" << port << " - "
              << gai_strerror(rc) << "\n";
    return -1;
  }

  int fd = -1;
  for (addrinfo *p = result; p != nullptr; p = p->ai_next) {
    fd = socket(p->ai_family, p->ai_socktype, p->ai_protocol);
    if (fd < 0)
      continue;
    if (connect(fd, p->ai_addr, p->ai_addrlen) == 0) {
      break;
    }
    close(fd);
    fd = -1;
  }
  freeaddrinfo(result);
  return fd;
}

ssize_t forward_once(int src_fd, int dst_fd) {
  char buffer[kBufferSize];
  ssize_t n_read = recv(src_fd, buffer, sizeof(buffer), 0);
  if (n_read <= 0) {
    return n_read;
  }

  ssize_t total_sent = 0;
  while (total_sent < n_read) {
    ssize_t n_sent = send(dst_fd, buffer + total_sent, n_read - total_sent, 0);
    if (n_sent < 0) {
      perror("send");
      return -1;
    }
    total_sent += n_sent;
  }
  return total_sent;
}

void handle_client(int client_fd, const std::string &backend_host,
                   int backend_port) {
  int backend_fd = connect_to_backend(backend_host, backend_port);
  if (backend_fd < 0) {
    std::cerr << "[proxy] failed to connect to backend " << backend_host << ":"
              << backend_port << "\n";
    close(client_fd);
    return;
  }

  ssize_t forwarded = forward_once(client_fd, backend_fd);
  if (forwarded <= 0) {
    close(client_fd);
    close(backend_fd);
    return;
  }

  forward_once(backend_fd, client_fd);

  close(backend_fd);
  close(client_fd);
}

std::vector<int> g_backend_ports;
int g_next_backend = 0;
std::mutex g_backend_mutex;

int pick_backend_port() {
  std::lock_guard<std::mutex> lock(g_backend_mutex);
  int chosen_index = g_next_backend;
  g_next_backend =
      (g_next_backend + 1) % static_cast<int>(g_backend_ports.size());
  return g_backend_ports[chosen_index];
}

} // namespace

int main(int argc, char *argv[]) {
  if (argc < 4) {
    std::cerr << "Usage: " << argv[0]
              << " <listen_port> <backend_host> <backend_port_1> "
                 "[backend_port_2 ...]\n";
    return 1;
  }

  int listen_port = std::stoi(argv[1]);
  std::string backend_host = argv[2];
  size_t num_threads = kNumThreads;

  for (int i = 3; i < argc; ++i) {
    std::string arg = argv[i];
    if (arg == "--threads" && i + 1 < argc) {
      num_threads = static_cast<size_t>(std::stoi(argv[++i]));
    } else {
      g_backend_ports.push_back(std::stoi(arg));
    }
  }

  int listen_fd = socket(AF_INET, SOCK_STREAM, 0);
  if (listen_fd < 0) {
    perror("socket");
    return 1;
  }

  int opt = 1;
  if (setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
    perror("setsockopt");
  }

  sockaddr_in addr{};
  addr.sin_family = AF_INET;
  addr.sin_addr.s_addr = INADDR_ANY;
  addr.sin_port = htons(listen_port);

  if (bind(listen_fd, reinterpret_cast<sockaddr *>(&addr), sizeof(addr)) < 0) {
    perror("bind");
    close(listen_fd);
    return 1;
  }

  if (listen(listen_fd, kListenBacklog) < 0) {
    perror("listen");
    close(listen_fd);
    return 1;
  }

  std::cout << "[proxy v4 - FIXED round-robin, " << g_backend_ports.size()
            << " backends, " << num_threads << " threads] listening on port "
            << listen_port << std::endl;

  ThreadPool pool(num_threads);

  while (true) {
    sockaddr_in client_addr{};
    socklen_t client_len = sizeof(client_addr);
    int client_fd = accept(
        listen_fd, reinterpret_cast<sockaddr *>(&client_addr), &client_len);
    if (client_fd < 0) {
      perror("accept");
      continue;
    }

    pool.submit([client_fd, &backend_host] {
      int backend_port = pick_backend_port();
      handle_client(client_fd, backend_host, backend_port);
    });
  }

  close(listen_fd);
  return 0;
}