request_response.h

#ifndef THIRD_PARTY_MILOTIC_INTERNAL_CC_PROXY_REQUEST_RESPONSE_H_
#define THIRD_PARTY_MILOTIC_INTERNAL_CC_PROXY_REQUEST_RESPONSE_H_

#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/container/flat_hash_map.h"
#include "absl/log/log.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "redfish_query_engine/http/client.h"
#include "nlohmann/json_fwd.hpp"
#include "proxy_config.pb.h"

namespace milotic {

using HttpHeaders = absl::flat_hash_map<std::string, std::string>;

constexpr absl::string_view SchemeToString(
    milotic_grpc_proxy::TargetScheme scheme) {
  switch (scheme) {
    case milotic_grpc_proxy::TargetScheme::TARGET_SCHEME_HTTPS:
      return "https://";
    case milotic_grpc_proxy::TargetScheme::TARGET_SCHEME_HTTP:
      return "http://";
    case milotic_grpc_proxy::TargetScheme::TARGET_SCHEME_GRPC:
      return "dns:///";
    default:
      return "";
  }
}

class Host {
 public:
  explicit Host(milotic_grpc_proxy::TargetScheme scheme =
                    milotic_grpc_proxy::TargetScheme::TARGET_SCHEME_UNSPECIFIED,
                std::string hostname = "", uint16_t port = 0)
      : scheme_(scheme), hostname_(std::move(hostname)), port_(port) {
    endpoint_ = absl::StrCat(SchemeToString(scheme_), hostname_);
    if (port_ != 0) {
      absl::StrAppend(&endpoint_, ":", port_);
    }
  }

  explicit Host(const milotic_grpc_proxy::Target& target) {
    scheme_ = target.scheme();
    if (target.has_hostname()) {
      hostname_ = target.hostname();
    } else if (target.has_nea_host()) {
      // TODO(jainrish): Implement NEA host support.
      LOG(DFATAL) << "NEA host support is not implemented yet";
      hostname_ = "";
    }
    port_ = static_cast<uint16_t>(target.port());
    endpoint_ = absl::StrCat(SchemeToString(scheme_), hostname_);
    if (port_ != 0) {
      absl::StrAppend(&endpoint_, ":", port_);
    }
  }

  static std::vector<Host> CreateTestHosts(std::string hostname) {
    return {Host(milotic_grpc_proxy::TargetScheme::TARGET_SCHEME_UNSPECIFIED,
                 std::move(hostname), 0)};
  }

  absl::string_view GetEndpoint() const { return endpoint_; }
  absl::string_view GetHostname() const { return hostname_; }
  uint16_t GetPort() const { return port_; }
  std::string GetHostHeader() const {
    if (port_ == 0) {
      return hostname_;
    }
    return absl::StrCat(hostname_, ":", port_);
  }

 private:
  milotic_grpc_proxy::TargetScheme scheme_ =
      milotic_grpc_proxy::TargetScheme::TARGET_SCHEME_UNSPECIFIED;
  std::string hostname_;
  uint16_t port_ = 0;
  std::string endpoint_;
};

class ProxyRequest {
 public:
  ProxyRequest() = default;
  ProxyRequest(std::vector<Host> hosts, std::string path, std::string body = "",
               HttpHeaders headers = {}, std::string unix_socket_path = "")
      : hosts_(std::move(hosts)),
        path_(std::move(path)),
        unix_socket_path_(std::move(unix_socket_path)),
        body_(std::move(body)),
        headers_(std::move(headers)) {}

  void SetPath(absl::string_view path) { path_ = path; }
  absl::string_view GetPath() const { return path_; }

  void AddHeader(absl::string_view key, absl::string_view value) {
    headers_[key] = value;
  }

  void SetBody(std::string body) { body_ = std::move(body); }
  absl::string_view GetBody() const { return body_; }
  nlohmann::json GetBodyJson() const {
    return nlohmann::json::parse(body_, nullptr, /*allow_exceptions=*/false);
  }

  const HttpHeaders& GetHeaders() const { return headers_; }
  absl::string_view GetHeader(absl::string_view key) const {
    auto it = headers_.find(key);
    if (it == headers_.end()) {
      return "";
    }
    return it->second;
  }

  std::string GetHostUri(int host_index) const {
    if (host_index < 0 || static_cast<size_t>(host_index) >= hosts_.size()) {
      LOG(ERROR) << "Host index is out of range: " << host_index;
      return "";
    }
    return absl::StrCat(hosts_[host_index].GetEndpoint(), path_);
  }
  Host GetHost(int host_index = 0) {
    if (host_index < 0 || static_cast<size_t>(host_index) >= hosts_.size()) {
      LOG(ERROR) << "Host index is out of range" << host_index;
      return Host();
    }
    return hosts_[host_index];
  }

  absl::string_view GetUnixSocketPath() { return unix_socket_path_; }

  std::unique_ptr<ecclesia::HttpClient::HttpRequest> ToHttpClientRequest(
      int host_index = 0) {
    if (host_index < 0 || static_cast<size_t>(host_index) >= hosts_.size()) {
      LOG(ERROR) << "Host index is out of range" << host_index;
      return nullptr;
    }
    auto request = std::make_unique<ecclesia::HttpClient::HttpRequest>();
    request->headers = headers_;
    request->body = body_;
    request->uri = absl::StrCat(hosts_[host_index].GetEndpoint(), path_);
    return request;
  }

 private:
  // The list of hosts to try in order.
  std::vector<Host> hosts_;
  std::string path_;
  std::string unix_socket_path_;
  std::string body_;
  HttpHeaders headers_;
};

class ProxyResponse {
 public:
  explicit ProxyResponse(int code = 0, std::string body = "",
                         HttpHeaders headers = {})
      : code_(code), body_(std::move(body)), headers_(std::move(headers)) {}

  explicit ProxyResponse(ecclesia::HttpClient::HttpResponse response)
      : code_(response.code),
        body_(std::move(response.body)),
        headers_(std::move(response.headers)) {}

  nlohmann::json GetBodyJson() const {
    return nlohmann::json::parse(body_, nullptr, /*allow_exceptions=*/false);
  }

  void SetCode(int code) { code_ = code; }
  int GetCode() const { return code_; }

  const HttpHeaders& GetHeaders() const { return headers_; }
  void AddHeader(absl::string_view key, absl::string_view value) {
    headers_[key] = value;
  }
  void RemoveHeader(absl::string_view key) { headers_.erase(key); }

  void SetBody(absl::string_view body) { body_ = body; }
  absl::string_view GetBody() const { return body_; }

 private:
  int code_;
  std::string body_;
  HttpHeaders headers_;
};

}  // namespace milotic

#endif  // THIRD_PARTY_MILOTIC_INTERNAL_CC_PROXY_REQUEST_RESPONSE_H_