moloapp/example/fw/Server.cpp

#include "Server.hpp"

#include <arpa/inet.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <unistd.h>

#include <cstring>
#include <sstream>

#include "Caps.hpp"

namespace molo::example::fw {

Server::Server(std::shared_ptr<Application> app) : app_(std::move(app)) {
  stop();
}

Server::~Server() { stop(); }

bool Server::start(const std::string& host, int port) {
  if (running_) return false;
  listen_fd_ = ::socket(AF_INET, SOCK_STREAM, 0);
  if (listen_fd_ < 0) return false;
  int one = 1;
  ::setsockopt(listen_fd_, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
  sockaddr_in addr{};
  addr.sin_family = AF_INET;
  addr.sin_port = htons((uint16_t)port);
  if (::inet_pton(AF_INET, host.c_str(), &addr.sin_addr) != 1) {
    ::close(listen_fd_);
    listen_fd_ = -1;
    return false;
  }
  if (::bind(listen_fd_, (sockaddr*)&addr, sizeof(addr)) != 0) {
    ::close(listen_fd_);
    listen_fd_ = -1;
    return false;
  }
  if (port == 0) {
    sockaddr_in bound{};
    socklen_t len = sizeof(bound);
    if (::getsockname(listen_fd_, (sockaddr*)&bound, &len) == 0)
      port_ = ntohs(bound.sin_port);
  } else {
    port_ = port;
  }
  if (::listen(listen_fd_, caps::kServerBacklog) != 0) {
    ::close(listen_fd_);
    listen_fd_ = -1;
    return false;
  }
  running_ = true;
  worker_ = std::thread([this]() { loop(); });
  return true;
}

void Server::stop() {
  bool was = running_.exchange(false);
  if (listen_fd_ >= 0) {
    ::shutdown(listen_fd_, SHUT_RDWR);
    ::close(listen_fd_);
    listen_fd_ = -1;
  }
  if (was && worker_.joinable() &&
      worker_.get_id() != std::this_thread::get_id())
    worker_.join();
  else if (worker_.joinable() &&
           worker_.get_id() == std::this_thread::get_id())
    worker_.detach();
}

int Server::port() const { return port_; }

bool Server::running() const { return running_; }

void Server::loop() {
  int served = 0;
  while (running_ && served < caps::kMaxConnections) {
    sockaddr_in peer{};
    socklen_t len = sizeof(peer);
    int fd = ::accept(listen_fd_, (sockaddr*)&peer, &len);
    if (fd < 0) {
      if (!running_) break;
      continue;
    }
    struct timeval tv{};
    tv.tv_sec = caps::kServerTimeoutSec;
    ::setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
    handle_one(fd);
    ::close(fd);
    ++served;
  }
}

void Server::handle_one(int fd) {
  std::string raw;
  char buf[4096];
  while (raw.size() < caps::kMaxHeaderBytes) {
    ssize_t n = ::recv(fd, buf, sizeof(buf), 0);
    if (n <= 0) break;
    raw.append(buf, (std::size_t)n);
    if (raw.find("\r\n\r\n") != std::string::npos) break;
  }
  std::size_t hdr_end = raw.find("\r\n\r\n");
  if (hdr_end == std::string::npos) return;
  std::string head = raw.substr(0, hdr_end);
  std::string rest = raw.substr(hdr_end + 4);
  std::istringstream lines(head);
  std::string line;
  if (!std::getline(lines, line)) return;
  if (!line.empty() && line.back() == '\r') line.pop_back();
  std::istringstream rl(line);
  std::string method_s, target, version;
  rl >> method_s >> target >> version;
  Request req;
  req.method = parse_method(method_s);
  std::size_t q = target.find('?');
  if (q == std::string::npos) {
    req.path = target.empty() ? "/" : target;
  } else {
    req.path = target.substr(0, q);
    req.query = parse_query(target.substr(q + 1));
  }
  std::size_t header_count = 0;
  std::size_t content_length = 0;
  while (std::getline(lines, line)) {
    if (!line.empty() && line.back() == '\r') line.pop_back();
    if (line.empty()) continue;
    if (++header_count > caps::kMaxHeaderCount) break;
    std::size_t colon = line.find(':');
    if (colon == std::string::npos) continue;
    std::string name = line.substr(0, colon);
    std::string value = line.substr(colon + 1);
    while (!value.empty() && value.front() == ' ') value.erase(value.begin());
    req.headers.set(name, value);
    if (name == "Content-Length" || name == "content-length") {
      try {
        content_length = (std::size_t)std::stoul(value);
      } catch (...) {
        content_length = 0;
      }
    }
  }
  std::string cookie_h = req.headers.get("Cookie");
  if (!cookie_h.empty()) req.cookies = parse_cookies(cookie_h);
  req.body = rest;
  while (req.body.size() < content_length &&
         req.body.size() < caps::kMaxRequestBody) {
    ssize_t n = ::recv(fd, buf, sizeof(buf), 0);
    if (n <= 0) break;
    req.body.append(buf, (std::size_t)n);
  }
  Response res = app_->dispatch(std::move(req));
  if (res.body.size() > caps::kMaxTestBody) {
    res.status = 500;
    res.body = "response over cap";
  }
  std::ostringstream out;
  out << "HTTP/1.1 " << res.status << " " << status_text(res.status) << "\r\n";
  out << "Content-Length: " << res.body.size() << "\r\n";
  out << "Connection: close\r\n";
  for (const auto& [k, v] : res.headers.entries()) {
    if (k == "Content-Length" || k == "content-length") continue;
    if (k == "Transfer-Encoding" || k == "transfer-encoding") continue;
    out << k << ": " << v << "\r\n";
  }
  out << "\r\n" << res.body;
  std::string wire = out.str();
  std::size_t sent = 0;
  while (sent < wire.size()) {
    ssize_t n = ::send(fd, wire.data() + sent, wire.size() - sent, 0);
    if (n <= 0) break;
    sent += (std::size_t)n;
  }
}

std::string http_roundtrip(const std::string& host, int port,
                           const std::string& raw_request) {
  int fd = ::socket(AF_INET, SOCK_STREAM, 0);
  if (fd < 0) return "";
  sockaddr_in addr{};
  addr.sin_family = AF_INET;
  addr.sin_port = htons((uint16_t)port);
  if (::inet_pton(AF_INET, host.c_str(), &addr.sin_addr) != 1) {
    ::close(fd);
    return "";
  }
  if (::connect(fd, (sockaddr*)&addr, sizeof(addr)) != 0) {
    ::close(fd);
    return "";
  }
  ::send(fd, raw_request.data(), raw_request.size(), 0);
  std::string out;
  char buf[4096];
  while (out.size() < caps::kMaxTestBody) {
    ssize_t n = ::recv(fd, buf, sizeof(buf), 0);
    if (n <= 0) break;
    out.append(buf, (std::size_t)n);
  }
  ::close(fd);
  return out;
}

}  // namespace molo::example::fw