VWeb/Source/Router.cpp

#include "StringUtils.h"

#include <VWeb.h>
#include <type_traits>
#include <utility>

namespace VWeb {

template <typename E> constexpr auto to_underlying(E e) noexcept {
  return static_cast<std::underlying_type_t<E>>(e);
}

class ErrorRoute : public Route {
public:
  bool Execute(Request &request, Response &response) override {
    response.Reset();
    response << "Unhandled Error: Status "
             << std::to_string(to_underlying(response.Status));
    response.SetType("text/plain");
    return true;
  }
};

class InstanceHandleRoute : public Route {
public:
  bool Get(Request &request, Response &response) override {
    return GetFunc && GetFunc(request, response);
  }
  bool Post(Request &request, Response &response) override {
    return PostFunc && PostFunc(request, response);
  }
  bool Put(Request &request, Response &response) override {
    return PutFunc && PutFunc(request, response);
  }
  bool Patch(Request &request, Response &response) override {
    return PatchFunc && PatchFunc(request, response);
  }
  bool Delete(Request &request, Response &response) override {
    return DeleteFunc && DeleteFunc(request, response);
  }

protected:
  RouteFunction GetFunc{nullptr};
  RouteFunction PostFunc{nullptr};
  RouteFunction PutFunc{nullptr};
  RouteFunction PatchFunc{nullptr};
  RouteFunction DeleteFunc{nullptr};
  friend Router;
};

Router::Router() { m_Routes["@"] = CreateRef<ErrorRoute>(); }

void Router::AddRoute(const std::string &name, const Ref<Route> &route) {
  m_Routes[name] = route;
}

Ref<Route> &Router::GetRoute(const std::string &name) { return m_Routes[name]; }

void Router::DeleteRoute(const std::string &name) {
  if (m_Routes.contains(name)) {
    m_Routes.erase(name);
  }
}

Ref<Response> Router::HandleRoute(Ref<Request> &request) {
  auto response = CreateRef<Response>();
  auto route = FindRoute(request);
  response->CookieData = request->CookieData;
  response->SessionData = request->SessionData;
  response->Method = request->Method;
  if (!route) {
    response->SetStatus(HttpStatusCode::NotFound);
    m_Routes["@"]->Execute(*request, *response);
    return response;
  }

  if (!route->IsAllowed(*request)) {
    response->SetStatus(HttpStatusCode::Forbidden);
    m_Routes["@"]->Execute(*request, *response);
    return response;
  }

  if (!route->Execute(*request, *response)) {
    std::string rKey = "@" + std::to_string(to_underlying(response->Status));
    m_Routes.contains(rKey) ? m_Routes[rKey]->Execute(*request, *response)
                            : m_Routes["@"]->Execute(*request, *response);
  }
  return response;
}

Ref<Route> Router::FindRoute(Ref<Request> &request) {
  auto &url = request->URI;
  if (m_Routes.contains(url.data())) {
    auto &route = m_Routes.at(url.data());
    if (route->SupportsMethod(*request))
      return route;
  }
  if (url.starts_with("@"))
    return nullptr;

  auto split = String::Split(url.data(), "/");
  if (split.size() > 1) {
    AddToArgs(request, split);
    while (split.size() > 1) {
      std::string nUrl = String::Join(split, "/");
      if (m_Routes.contains(nUrl)) {
        auto &route = m_Routes[nUrl];
        if (route->SupportsMethod(*request))
          return route;
      }
      AddToArgs(request, split);
    }
  }
  if (m_Routes.contains("/")) {
    auto &route = m_Routes["/"];
    if (route->SupportsMethod(*request))
      return route;
  }
  return nullptr;
}

void Router::AddToArgs(Ref<Request> &request, Vector<std::string> &items) {
  request->URLParameters.push_back(items[items.size() - 1]);
  items.pop_back();
}

InstanceHandleRoute* GetOrNull(const std::string& path, Router* router) {
  InstanceHandleRoute* route;
  if (!router->m_Routes.contains(path)) {
    router->AddRoute(path, CreateRef<InstanceHandleRoute>());
  }
  route = dynamic_cast<InstanceHandleRoute*>(router->m_Routes[path].get());
  return route;
}
void Router::Get(const std::string &path, RouteFunction func) {
  auto* route = GetOrNull(path, this);
  if (route)
    route->GetFunc = std::move(func);
}
void Router::Post(const std::string &path, RouteFunction func) {
  auto* route = GetOrNull(path, this);
  if (route)
    route->PostFunc = std::move(func);
}
void Router::Put(const std::string &path, RouteFunction func) {
  auto* route = GetOrNull(path, this);
  if (route)
    route->PutFunc = std::move(func);
}
void Router::Patch(const std::string &path, RouteFunction func) {
  auto* route = GetOrNull(path, this);
  if (route)
    route->PatchFunc = std::move(func);
}
void Router::Delete(const std::string &path, RouteFunction func) {
  auto* route = GetOrNull(path, this);
  if (route)
    route->DeleteFunc = std::move(func);
}
} // namespace VWeb