<commit_before>#include "ostree.h"
#include <json/json.h>
#include <stdio.h>
#include <unistd.h>
#include <boost/algorithm/string.hpp>
#include <boost/algorithm/string/classification.hpp>
#include <boost/algorithm/string/split.hpp>
#include <boost/filesystem.hpp>
#include <fstream>
#include "logger.h"
#include "utils.h"

#include <gio/gio.h>

boost::shared_ptr<OstreeSysroot> Ostree::LoadSysroot(const std::string &path) {
  OstreeSysroot *sysroot = NULL;

  if (path.size()) {
    GFile *fl = g_file_new_for_path(path.c_str());
    sysroot = ostree_sysroot_new(fl);
    g_object_unref(fl);
  } else {
    sysroot = ostree_sysroot_new_default();
  }
  GError *error = NULL;
  if (!ostree_sysroot_load(sysroot, NULL, &error)) {
    if (error != NULL) {
      g_error_free(error);
    }
    g_object_unref(sysroot);
    throw std::runtime_error("could not load sysroot");
  }
  return boost::shared_ptr<OstreeSysroot>(sysroot, g_object_unref);
}

boost::shared_ptr<OstreeDeployment> Ostree::getStagedDeployment(const std::string &path) {
  boost::shared_ptr<OstreeSysroot> sysroot = Ostree::LoadSysroot(path);
  GPtrArray *deployments = NULL;
  OstreeDeployment *res = NULL;

  deployments = ostree_sysroot_get_deployments(sysroot.get());

  if (deployments->len == 0) {
    res = NULL;
  } else {
    OstreeDeployment *d = static_cast<OstreeDeployment *>(deployments->pdata[0]);
    res = static_cast<OstreeDeployment *>(g_object_ref(d));
  }

  g_ptr_array_unref(deployments);
  return boost::shared_ptr<OstreeDeployment>(res, g_object_unref);
}

bool Ostree::addRemote(OstreeRepo *repo, const std::string &remote, const std::string &url,
                       const data::PackageManagerCredentials &cred) {
  GCancellable *cancellable = NULL;
  GError *error = NULL;
  GVariantBuilder b;
  GVariant *options;

  g_variant_builder_init(&b, G_VARIANT_TYPE("a{sv}"));
  g_variant_builder_add(&b, "{s@v}", "gpg-verify", g_variant_new_variant(g_variant_new_boolean(FALSE)));

  if (cred.cert_file.size() && cred.pkey_file.size() && cred.ca_file.size()) {
    g_variant_builder_add(&b, "{s@v}", "tls-client-cert-path",
                          g_variant_new_variant(g_variant_new_string(cred.cert_file.c_str())));
    g_variant_builder_add(&b, "{s@v}", "tls-client-key-path",
                          g_variant_new_variant(g_variant_new_string(cred.pkey_file.c_str())));
    g_variant_builder_add(&b, "{s@v}", "tls-ca-path",
                          g_variant_new_variant(g_variant_new_string(cred.ca_file.c_str())));
  }
  options = g_variant_builder_end(&b);

  if (!ostree_repo_remote_change(repo, NULL, OSTREE_REPO_REMOTE_CHANGE_DELETE_IF_EXISTS, remote.c_str(), url.c_str(),
                                 options, cancellable, &error)) {
    LOGGER_LOG(LVL_error, "Error of adding remote: " << error->message);
    g_error_free(error);
    return false;
  }
  if (!ostree_repo_remote_change(repo, NULL, OSTREE_REPO_REMOTE_CHANGE_ADD_IF_NOT_EXISTS, remote.c_str(), url.c_str(),
                                 options, cancellable, &error)) {
    LOGGER_LOG(LVL_error, "Error of adding remote: " << error->message);
    g_error_free(error);
    return false;
  }
  return true;
}

#include "ostree-1/ostree.h"

OstreePackage::OstreePackage(const std::string &ecu_serial_in, const std::string &ref_name_in,
                             const std::string &desc_in, const std::string &treehub_in)
    : ecu_serial(ecu_serial_in), ref_name(ref_name_in), description(desc_in), pull_uri(treehub_in) {
  std::size_t pos = ref_name.find_last_of("-");
  branch_name = ref_name.substr(0, pos);
  refhash = ref_name.substr(pos + 1, std::string::npos);
  if (branch_name.empty() || refhash.empty()) throw std::runtime_error("malformed OSTree target name: " + ref_name);
}

data::InstallOutcome OstreePackage::install(const data::PackageManagerCredentials &cred, OstreeConfig config) const {
  const char remote[] = "aktualizr-remote";
  const char *const refs[] = {branch_name.c_str()};
  const char *const commit_ids[] = {refhash.c_str()};
  const char *opt_osname = NULL;
  OstreeRepo *repo = NULL;
  GCancellable *cancellable = NULL;
  GError *error = NULL;
  char *revision;
  GVariantBuilder builder;
  GVariant *options;

  if (config.os.size()) {
    opt_osname = config.os.c_str();
  }
  boost::shared_ptr<OstreeSysroot> sysroot = Ostree::LoadSysroot(config.sysroot);
  if (!ostree_sysroot_get_repo(sysroot.get(), &repo, cancellable, &error)) {
    LOGGER_LOG(LVL_error, "could not get repo");
    g_error_free(error);
    return data::InstallOutcome(data::INSTALL_FAILED, "could not get repo");
  }

  if (!Ostree::addRemote(repo, remote, pull_uri, cred)) {
    return data::InstallOutcome(data::INSTALL_FAILED, "Error of adding remote");
  }

  g_variant_builder_init(&builder, G_VARIANT_TYPE("a{sv}"));
  g_variant_builder_add(&builder, "{s@v}", "flags", g_variant_new_variant(g_variant_new_int32(0)));

  g_variant_builder_add(&builder, "{s@v}", "refs", g_variant_new_variant(g_variant_new_strv(refs, 1)));
  g_variant_builder_add(&builder, "{s@v}", "override-commit-ids",
                        g_variant_new_variant(g_variant_new_strv(commit_ids, 1)));

  if (cred.access_token.size()) {
    GVariantBuilder hdr_builder;
    std::string av("Bearer ");
    av += cred.access_token;
    g_variant_builder_init(&hdr_builder, G_VARIANT_TYPE("a(ss)"));
    g_variant_builder_add(&hdr_builder, "(ss)", "Authorization", av.c_str());
    g_variant_builder_add(&builder, "{s@v}", "http-headers",
                          g_variant_new_variant(g_variant_builder_end(&hdr_builder)));
  }
  options = g_variant_ref_sink(g_variant_builder_end(&builder));

  if (!ostree_repo_pull_with_options(repo, remote, options, NULL, cancellable, &error)) {
    LOGGER_LOG(LVL_error, "Error of pulling image: " << error->message);
    data::InstallOutcome install_outcome(data::INSTALL_FAILED, error->message);
    g_error_free(error);
    return install_outcome;
  }

  GKeyFile *origin = ostree_sysroot_origin_new_from_refspec(sysroot.get(), branch_name.c_str());
  if (!ostree_repo_resolve_rev(repo, refhash.c_str(), FALSE, &revision, &error)) {
    LOGGER_LOG(LVL_error, error->message);
    data::InstallOutcome install_outcome(data::INSTALL_FAILED, error->message);
    g_error_free(error);
    return install_outcome;
  }

  OstreeDeployment *merge_deployment = ostree_sysroot_get_merge_deployment(sysroot.get(), opt_osname);
  if (merge_deployment == NULL) {
    LOGGER_LOG(LVL_error, "No merge deployment");
    return data::InstallOutcome(data::INSTALL_FAILED, "No merge deployment");
  }

  if (!ostree_sysroot_prepare_cleanup(sysroot.get(), cancellable, &error)) {
    LOGGER_LOG(LVL_error, error->message);
    data::InstallOutcome install_outcome(data::INSTALL_FAILED, error->message);
    g_error_free(error);
    return install_outcome;
  }

  std::ifstream args_stream("/proc/cmdline");
  std::string args_content((std::istreambuf_iterator<char>(args_stream)), std::istreambuf_iterator<char>());
  std::vector<std::string> args_vector;
  boost::split(args_vector, args_content, boost::is_any_of(" "));

  std::vector<const char *> kargs_strv_vector;
  kargs_strv_vector.reserve(args_vector.size() + 1);

  for (std::vector<std::string>::iterator it = args_vector.begin(); it != args_vector.end(); ++it) {
    kargs_strv_vector.push_back((*it).c_str());
  }
  kargs_strv_vector[args_vector.size()] = 0;
  GStrv kargs_strv = const_cast<char **>(&kargs_strv_vector[0]);

  OstreeDeployment *new_deployment = NULL;
  if (!ostree_sysroot_deploy_tree(sysroot.get(), opt_osname, revision, origin, merge_deployment, kargs_strv,
                                  &new_deployment, cancellable, &error)) {
    LOGGER_LOG(LVL_error, "ostree_sysroot_deploy_tree: " << error->message);
    data::InstallOutcome install_outcome(data::INSTALL_FAILED, error->message);
    g_error_free(error);
    return install_outcome;
  }

  if (!ostree_sysroot_simple_write_deployment(sysroot.get(), NULL, new_deployment, merge_deployment,
                                              OSTREE_SYSROOT_SIMPLE_WRITE_DEPLOYMENT_FLAGS_NONE, cancellable, &error)) {
    LOGGER_LOG(LVL_error, "ostree_sysroot_simple_write_deployment:" << error->message);
    data::InstallOutcome install_outcome(data::INSTALL_FAILED, error->message);
    g_error_free(error);
    return install_outcome;
  }
  LOGGER_LOG(LVL_info, "Performing sync()");
  sync();
  return data::InstallOutcome(data::OK, "Installation successful");
}

OstreeBranch OstreeBranch::getCurrent(const std::string &ecu_serial, const std::string &ostree_sysroot) {
  boost::shared_ptr<OstreeDeployment> staged_deployment = Ostree::getStagedDeployment(ostree_sysroot);

  GKeyFile *origin = ostree_deployment_get_origin(staged_deployment.get());
  const char *ref = ostree_deployment_get_csum(staged_deployment.get());
  char *origin_refspec = g_key_file_get_string(origin, "origin", "refspec", NULL);
  OstreePackage package(ecu_serial, std::string(origin_refspec) + "-" + ref, origin_refspec, "");
  g_free(origin_refspec);
  return OstreeBranch(true, ostree_deployment_get_osname(staged_deployment.get()), package);
}

OstreePackage OstreePackage::fromJson(const Json::Value &json) {
  return OstreePackage(json["ecu_serial"].asString(), json["ref_name"].asString(), json["description"].asString(),
                       json["pull_uri"].asString());
}

Json::Value OstreePackage::toEcuVersion(const Json::Value &custom) const {
  Json::Value installed_image;
  installed_image["filepath"] = ref_name;
  installed_image["fileinfo"]["length"] = 0;
  installed_image["fileinfo"]["hashes"]["sha256"] = refhash;

  Json::Value value;
  value["attacks_detected"] = "";
  value["installed_image"] = installed_image;
  value["ecu_serial"] = ecu_serial;
  value["previous_timeserver_time"] = "1970-01-01T00:00:00Z";
  value["timeserver_time"] = "1970-01-01T00:00:00Z";
  if (custom != Json::nullValue) {
    value["custom"] = custom;
  }
  return value;
}

OstreePackage OstreePackage::getEcu(const std::string &ecu_serial, const std::string &ostree_sysroot) {
  return OstreeBranch::getCurrent(ecu_serial, ostree_sysroot).package;
}

Json::Value Ostree::getInstalledPackages(const std::string &file_path) {
  std::string packages_str = Utils::readFile(file_path);
  std::vector<std::string> package_lines;
  boost::split(package_lines, packages_str, boost::is_any_of("\n"));
  Json::Value packages(Json::arrayValue);
  for (std::vector<std::string>::iterator it = package_lines.begin(); it != package_lines.end(); ++it) {
    if (it->empty()) {
      continue;
    }
    size_t pos = it->find(" ");
    if (pos == std::string::npos) {
      throw std::runtime_error("Wrong packages file format");
    }
    Json::Value package;
    package["name"] = it->substr(0, pos);
    package["version"] = it->substr(pos + 1);
    packages.append(package);
  }
  return packages;
}
<commit_msg>Use kernel arguments from the bootconfig, not the whole /proc/cmdline<commit_after>#include "ostree.h"
#include <json/json.h>
#include <stdio.h>
#include <unistd.h>
#include <boost/algorithm/string.hpp>
#include <boost/algorithm/string/classification.hpp>
#include <boost/algorithm/string/split.hpp>
#include <boost/filesystem.hpp>
#include <fstream>
#include "logger.h"
#include "utils.h"

#include <gio/gio.h>

boost::shared_ptr<OstreeSysroot> Ostree::LoadSysroot(const std::string &path) {
  OstreeSysroot *sysroot = NULL;

  if (path.size()) {
    GFile *fl = g_file_new_for_path(path.c_str());
    sysroot = ostree_sysroot_new(fl);
    g_object_unref(fl);
  } else {
    sysroot = ostree_sysroot_new_default();
  }
  GError *error = NULL;
  if (!ostree_sysroot_load(sysroot, NULL, &error)) {
    if (error != NULL) {
      g_error_free(error);
    }
    g_object_unref(sysroot);
    throw std::runtime_error("could not load sysroot");
  }
  return boost::shared_ptr<OstreeSysroot>(sysroot, g_object_unref);
}

boost::shared_ptr<OstreeDeployment> Ostree::getStagedDeployment(const std::string &path) {
  boost::shared_ptr<OstreeSysroot> sysroot = Ostree::LoadSysroot(path);
  GPtrArray *deployments = NULL;
  OstreeDeployment *res = NULL;

  deployments = ostree_sysroot_get_deployments(sysroot.get());

  if (deployments->len == 0) {
    res = NULL;
  } else {
    OstreeDeployment *d = static_cast<OstreeDeployment *>(deployments->pdata[0]);
    res = static_cast<OstreeDeployment *>(g_object_ref(d));
  }

  g_ptr_array_unref(deployments);
  return boost::shared_ptr<OstreeDeployment>(res, g_object_unref);
}

bool Ostree::addRemote(OstreeRepo *repo, const std::string &remote, const std::string &url,
                       const data::PackageManagerCredentials &cred) {
  GCancellable *cancellable = NULL;
  GError *error = NULL;
  GVariantBuilder b;
  GVariant *options;

  g_variant_builder_init(&b, G_VARIANT_TYPE("a{sv}"));
  g_variant_builder_add(&b, "{s@v}", "gpg-verify", g_variant_new_variant(g_variant_new_boolean(FALSE)));

  if (cred.cert_file.size() && cred.pkey_file.size() && cred.ca_file.size()) {
    g_variant_builder_add(&b, "{s@v}", "tls-client-cert-path",
                          g_variant_new_variant(g_variant_new_string(cred.cert_file.c_str())));
    g_variant_builder_add(&b, "{s@v}", "tls-client-key-path",
                          g_variant_new_variant(g_variant_new_string(cred.pkey_file.c_str())));
    g_variant_builder_add(&b, "{s@v}", "tls-ca-path",
                          g_variant_new_variant(g_variant_new_string(cred.ca_file.c_str())));
  }
  options = g_variant_builder_end(&b);

  if (!ostree_repo_remote_change(repo, NULL, OSTREE_REPO_REMOTE_CHANGE_DELETE_IF_EXISTS, remote.c_str(), url.c_str(),
                                 options, cancellable, &error)) {
    LOGGER_LOG(LVL_error, "Error of adding remote: " << error->message);
    g_error_free(error);
    return false;
  }
  if (!ostree_repo_remote_change(repo, NULL, OSTREE_REPO_REMOTE_CHANGE_ADD_IF_NOT_EXISTS, remote.c_str(), url.c_str(),
                                 options, cancellable, &error)) {
    LOGGER_LOG(LVL_error, "Error of adding remote: " << error->message);
    g_error_free(error);
    return false;
  }
  return true;
}

#include "ostree-1/ostree.h"

OstreePackage::OstreePackage(const std::string &ecu_serial_in, const std::string &ref_name_in,
                             const std::string &desc_in, const std::string &treehub_in)
    : ecu_serial(ecu_serial_in), ref_name(ref_name_in), description(desc_in), pull_uri(treehub_in) {
  std::size_t pos = ref_name.find_last_of("-");
  branch_name = ref_name.substr(0, pos);
  refhash = ref_name.substr(pos + 1, std::string::npos);
  if (branch_name.empty() || refhash.empty()) throw std::runtime_error("malformed OSTree target name: " + ref_name);
}

data::InstallOutcome OstreePackage::install(const data::PackageManagerCredentials &cred, OstreeConfig config) const {
  const char remote[] = "aktualizr-remote";
  const char *const refs[] = {branch_name.c_str()};
  const char *const commit_ids[] = {refhash.c_str()};
  const char *opt_osname = NULL;
  OstreeRepo *repo = NULL;
  GCancellable *cancellable = NULL;
  GError *error = NULL;
  char *revision;
  GVariantBuilder builder;
  GVariant *options;

  if (config.os.size()) {
    opt_osname = config.os.c_str();
  }
  boost::shared_ptr<OstreeSysroot> sysroot = Ostree::LoadSysroot(config.sysroot);
  if (!ostree_sysroot_get_repo(sysroot.get(), &repo, cancellable, &error)) {
    LOGGER_LOG(LVL_error, "could not get repo");
    g_error_free(error);
    return data::InstallOutcome(data::INSTALL_FAILED, "could not get repo");
  }

  if (!Ostree::addRemote(repo, remote, pull_uri, cred)) {
    return data::InstallOutcome(data::INSTALL_FAILED, "Error of adding remote");
  }

  g_variant_builder_init(&builder, G_VARIANT_TYPE("a{sv}"));
  g_variant_builder_add(&builder, "{s@v}", "flags", g_variant_new_variant(g_variant_new_int32(0)));

  g_variant_builder_add(&builder, "{s@v}", "refs", g_variant_new_variant(g_variant_new_strv(refs, 1)));
  g_variant_builder_add(&builder, "{s@v}", "override-commit-ids",
                        g_variant_new_variant(g_variant_new_strv(commit_ids, 1)));

  if (cred.access_token.size()) {
    GVariantBuilder hdr_builder;
    std::string av("Bearer ");
    av += cred.access_token;
    g_variant_builder_init(&hdr_builder, G_VARIANT_TYPE("a(ss)"));
    g_variant_builder_add(&hdr_builder, "(ss)", "Authorization", av.c_str());
    g_variant_builder_add(&builder, "{s@v}", "http-headers",
                          g_variant_new_variant(g_variant_builder_end(&hdr_builder)));
  }
  options = g_variant_ref_sink(g_variant_builder_end(&builder));

  if (!ostree_repo_pull_with_options(repo, remote, options, NULL, cancellable, &error)) {
    LOGGER_LOG(LVL_error, "Error of pulling image: " << error->message);
    data::InstallOutcome install_outcome(data::INSTALL_FAILED, error->message);
    g_error_free(error);
    return install_outcome;
  }

  GKeyFile *origin = ostree_sysroot_origin_new_from_refspec(sysroot.get(), branch_name.c_str());
  if (!ostree_repo_resolve_rev(repo, refhash.c_str(), FALSE, &revision, &error)) {
    LOGGER_LOG(LVL_error, error->message);
    data::InstallOutcome install_outcome(data::INSTALL_FAILED, error->message);
    g_error_free(error);
    return install_outcome;
  }

  OstreeDeployment *merge_deployment = ostree_sysroot_get_merge_deployment(sysroot.get(), opt_osname);
  if (merge_deployment == NULL) {
    LOGGER_LOG(LVL_error, "No merge deployment");
    return data::InstallOutcome(data::INSTALL_FAILED, "No merge deployment");
  }

  if (!ostree_sysroot_prepare_cleanup(sysroot.get(), cancellable, &error)) {
    LOGGER_LOG(LVL_error, error->message);
    data::InstallOutcome install_outcome(data::INSTALL_FAILED, error->message);
    g_error_free(error);
    return install_outcome;
  }

  std::string args_content =
      std::string(ostree_bootconfig_parser_get(ostree_deployment_get_bootconfig(merge_deployment), "options"));
  std::vector<std::string> args_vector;
  boost::split(args_vector, args_content, boost::is_any_of(" "));

  std::vector<const char *> kargs_strv_vector;
  kargs_strv_vector.reserve(args_vector.size() + 1);

  for (std::vector<std::string>::iterator it = args_vector.begin(); it != args_vector.end(); ++it) {
    kargs_strv_vector.push_back((*it).c_str());
  }
  kargs_strv_vector[args_vector.size()] = 0;
  GStrv kargs_strv = const_cast<char **>(&kargs_strv_vector[0]);

  OstreeDeployment *new_deployment = NULL;
  if (!ostree_sysroot_deploy_tree(sysroot.get(), opt_osname, revision, origin, merge_deployment, kargs_strv,
                                  &new_deployment, cancellable, &error)) {
    LOGGER_LOG(LVL_error, "ostree_sysroot_deploy_tree: " << error->message);
    data::InstallOutcome install_outcome(data::INSTALL_FAILED, error->message);
    g_error_free(error);
    return install_outcome;
  }

  if (!ostree_sysroot_simple_write_deployment(sysroot.get(), NULL, new_deployment, merge_deployment,
                                              OSTREE_SYSROOT_SIMPLE_WRITE_DEPLOYMENT_FLAGS_NONE, cancellable, &error)) {
    LOGGER_LOG(LVL_error, "ostree_sysroot_simple_write_deployment:" << error->message);
    data::InstallOutcome install_outcome(data::INSTALL_FAILED, error->message);
    g_error_free(error);
    return install_outcome;
  }
  LOGGER_LOG(LVL_info, "Performing sync()");
  sync();
  return data::InstallOutcome(data::OK, "Installation successful");
}

OstreeBranch OstreeBranch::getCurrent(const std::string &ecu_serial, const std::string &ostree_sysroot) {
  boost::shared_ptr<OstreeDeployment> staged_deployment = Ostree::getStagedDeployment(ostree_sysroot);

  GKeyFile *origin = ostree_deployment_get_origin(staged_deployment.get());
  const char *ref = ostree_deployment_get_csum(staged_deployment.get());
  char *origin_refspec = g_key_file_get_string(origin, "origin", "refspec", NULL);
  OstreePackage package(ecu_serial, std::string(origin_refspec) + "-" + ref, origin_refspec, "");
  g_free(origin_refspec);
  return OstreeBranch(true, ostree_deployment_get_osname(staged_deployment.get()), package);
}

OstreePackage OstreePackage::fromJson(const Json::Value &json) {
  return OstreePackage(json["ecu_serial"].asString(), json["ref_name"].asString(), json["description"].asString(),
                       json["pull_uri"].asString());
}

Json::Value OstreePackage::toEcuVersion(const Json::Value &custom) const {
  Json::Value installed_image;
  installed_image["filepath"] = ref_name;
  installed_image["fileinfo"]["length"] = 0;
  installed_image["fileinfo"]["hashes"]["sha256"] = refhash;

  Json::Value value;
  value["attacks_detected"] = "";
  value["installed_image"] = installed_image;
  value["ecu_serial"] = ecu_serial;
  value["previous_timeserver_time"] = "1970-01-01T00:00:00Z";
  value["timeserver_time"] = "1970-01-01T00:00:00Z";
  if (custom != Json::nullValue) {
    value["custom"] = custom;
  }
  return value;
}

OstreePackage OstreePackage::getEcu(const std::string &ecu_serial, const std::string &ostree_sysroot) {
  return OstreeBranch::getCurrent(ecu_serial, ostree_sysroot).package;
}

Json::Value Ostree::getInstalledPackages(const std::string &file_path) {
  std::string packages_str = Utils::readFile(file_path);
  std::vector<std::string> package_lines;
  boost::split(package_lines, packages_str, boost::is_any_of("\n"));
  Json::Value packages(Json::arrayValue);
  for (std::vector<std::string>::iterator it = package_lines.begin(); it != package_lines.end(); ++it) {
    if (it->empty()) {
      continue;
    }
    size_t pos = it->find(" ");
    if (pos == std::string::npos) {
      throw std::runtime_error("Wrong packages file format");
    }
    Json::Value package;
    package["name"] = it->substr(0, pos);
    package["version"] = it->substr(pos + 1);
    packages.append(package);
  }
  return packages;
}
<|endoftext|>


<commit_before>//=======================================================================
// Copyright 2014-2015 David Simmons-Duffin.
// Distributed under the MIT License.
// (See accompanying file LICENSE or copy at
//  http://opensource.org/licenses/MIT)
//=======================================================================

// Currently, SDPB uses tinyxml2 to parse an input file into an XML
// tree, which is stored entirely in memory.  This tree is then
// transformed into the appropriate data structures using the parse
// functions below.  In the future, this could be made more memory
// efficient by avoiding building the XML tree in memory.

// See the manual for a description of the correct XML input format.

#include <vector>
#include "types.h"
#include "parse.h"
#include "Polynomial.h"
#include "SDP.h"
#include "tinyxml2/tinyxml2.h"

using std::vector;
using tinyxml2::XMLDocument;
using tinyxml2::XMLElement;
using boost::filesystem::path;

// parse a bunch of adjacent elements with tag `name' into a vector<T>
// using the function `parse' for each element.
template <class T>
vector<T> parseMany(const char *name, T(*parse)(XMLElement *),
                    XMLElement *elt) {
  XMLElement *e;
  vector<T> v;
  for (e = elt->FirstChildElement(name);
       e != NULL;
       e = e->NextSiblingElement(name)) {
    v.push_back(parse(e));
  }
  return v;
}

Real parseReal(XMLElement *xml) {
  return Real(xml->GetText());
}

int parseInt(XMLElement *xml) {
  return atoi(xml->GetText());
}

Vector parseVector(XMLElement *xml) {
  return parseMany("elt", parseReal, xml);
}

Matrix parseMatrix(XMLElement *xml) {
  Matrix m;
  m.rows     = parseInt(xml->FirstChildElement("rows"));
  m.cols     = parseInt(xml->FirstChildElement("cols"));
  m.elements = parseVector(xml->FirstChildElement("elements"));
  return m;
}

Polynomial parsePolynomial(XMLElement *xml) {
  Polynomial p;
  p.coefficients = parseMany("coeff", parseReal, xml);
  return p;
}

vector<Polynomial> parsePolynomialVector(XMLElement *xml) {
  return parseMany("polynomial", parsePolynomial, xml);
}

PolynomialVectorMatrix parsePolynomialVectorMatrix(XMLElement *xml) {
  PolynomialVectorMatrix m;
  m.rows           = parseInt(xml->FirstChildElement("rows"));
  m.cols           = parseInt(xml->FirstChildElement("cols"));
  m.elements       = parseMany("polynomialVector", parsePolynomialVector,
                               xml->FirstChildElement("elements"));
  m.samplePoints   = parseVector(xml->FirstChildElement("samplePoints"));
  m.sampleScalings = parseVector(xml->FirstChildElement("sampleScalings"));
  m.bilinearBasis  = parsePolynomialVector(xml->FirstChildElement("bilinearBasis"));
  return m;
}

SDP parseBootstrapSDP(XMLElement *xml) {
  return bootstrapSDP(parseVector(xml->FirstChildElement("objective")),
                      parseMany("polynomialVectorMatrix",
                                parsePolynomialVectorMatrix,
                                xml->FirstChildElement("polynomialVectorMatrices")));
}

SDP readBootstrapSDP(const path sdpFile) {
  XMLDocument doc;
#ifndef BOOST_WINDOWS_API
  doc.LoadFile(sdpFile.c_str());
#else
  doc.LoadFile(sdpFile.string().c_str());
#endif
return parseBootstrapSDP(doc.FirstChildElement("sdp"));
}
<commit_msg>cosmetic changes<commit_after>//=======================================================================
// Copyright 2014-2015 David Simmons-Duffin.
// Distributed under the MIT License.
// (See accompanying file LICENSE or copy at
//  http://opensource.org/licenses/MIT)
//=======================================================================

// Currently, SDPB uses tinyxml2 to parse an input file into an XML
// tree, which is stored entirely in memory.  This tree is then
// transformed into the appropriate data structures using the parse
// functions below.  In the future, this could be made more memory
// efficient by avoiding building the XML tree in memory.

// See the manual for a description of the correct XML input format.

#include <vector>
#include "types.h"
#include "parse.h"
#include "Polynomial.h"
#include "SDP.h"
#include "tinyxml2/tinyxml2.h"

using std::vector;
using tinyxml2::XMLDocument;
using tinyxml2::XMLElement;
using boost::filesystem::path;

// parse a bunch of adjacent elements with tag `name' into a vector<T>
// using the function `parse' for each element.
template <class T>
vector<T> parseMany(const char *name, T(*parse)(XMLElement *),
                    XMLElement *elt) {
  XMLElement *e;
  vector<T> v;
  for (e = elt->FirstChildElement(name);
       e != NULL;
       e = e->NextSiblingElement(name)) {
    v.push_back(parse(e));
  }
  return v;
}

Real parseReal(XMLElement *xml) {
  return Real(xml->GetText());
}

int parseInt(XMLElement *xml) {
  return atoi(xml->GetText());
}

Vector parseVector(XMLElement *xml) {
  return parseMany("elt", parseReal, xml);
}

Matrix parseMatrix(XMLElement *xml) {
  Matrix m;
  m.rows     = parseInt(xml->FirstChildElement("rows"));
  m.cols     = parseInt(xml->FirstChildElement("cols"));
  m.elements = parseVector(xml->FirstChildElement("elements"));
  return m;
}

Polynomial parsePolynomial(XMLElement *xml) {
  Polynomial p;
  p.coefficients = parseMany("coeff", parseReal, xml);
  return p;
}

vector<Polynomial> parsePolynomialVector(XMLElement *xml) {
  return parseMany("polynomial", parsePolynomial, xml);
}

PolynomialVectorMatrix parsePolynomialVectorMatrix(XMLElement *xml) {
  PolynomialVectorMatrix m;
  m.rows           = parseInt(xml->FirstChildElement("rows"));
  m.cols           = parseInt(xml->FirstChildElement("cols"));
  m.elements       = parseMany("polynomialVector", parsePolynomialVector,
                               xml->FirstChildElement("elements"));
  m.samplePoints   = parseVector(xml->FirstChildElement("samplePoints"));
  m.sampleScalings = parseVector(xml->FirstChildElement("sampleScalings"));
  m.bilinearBasis  = parsePolynomialVector(xml->FirstChildElement("bilinearBasis"));
  return m;
}

SDP parseBootstrapSDP(XMLElement *xml) {
  return bootstrapSDP(parseVector(xml->FirstChildElement("objective")),
                      parseMany("polynomialVectorMatrix",
                                parsePolynomialVectorMatrix,
                                xml->FirstChildElement("polynomialVectorMatrices")));
}

SDP readBootstrapSDP(const path sdpFile) {
  XMLDocument doc;

#ifndef BOOST_WINDOWS_API
  doc.LoadFile(sdpFile.c_str());
#else
  doc.LoadFile(sdpFile.string().c_str());
#endif

  return parseBootstrapSDP(doc.FirstChildElement("sdp"));
}
<|endoftext|>


<commit_before>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/c/experimental/filesystem/plugins/s3/s3_filesystem.h"

#include <aws/core/config/AWSProfileConfigLoader.h>
#include <stdlib.h>
#include <string.h>

#include "absl/strings/ascii.h"
#include "absl/strings/numbers.h"
#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
#include "tensorflow/c/experimental/filesystem/plugins/s3/aws_crypto.h"
#include "tensorflow/c/tf_status.h"

// Implementation of a filesystem for S3 environments.
// This filesystem will support `s3://` URI schemes.
constexpr char kS3ClientAllocationTag[] = "S3ClientAllocation";
constexpr int64_t kS3TimeoutMsec = 300000;  // 5 min

constexpr char kExecutorTag[] = "TransferManagerExecutorAllocation";
constexpr int kExecutorPoolSize = 25;

constexpr uint64_t kS3MultiPartUploadChunkSize = 50 * 1024 * 1024;    // 50 MB
constexpr uint64_t kS3MultiPartDownloadChunkSize = 50 * 1024 * 1024;  // 50 MB

static void* plugin_memory_allocate(size_t size) { return calloc(1, size); }
static void plugin_memory_free(void* ptr) { free(ptr); }

static void ParseS3Path(const Aws::String& fname, bool object_empty_ok,
                        Aws::String* bucket, Aws::String* object,
                        TF_Status* status) {
  size_t scheme_end = fname.find("://") + 2;
  if (fname.substr(0, scheme_end + 1) != "s3://") {
    TF_SetStatus(status, TF_INVALID_ARGUMENT,
                 "S3 path doesn't start with 's3://'.");
    return;
  }

  size_t bucket_end = fname.find("/", scheme_end + 1);
  if (bucket_end == std::string::npos) {
    TF_SetStatus(status, TF_INVALID_ARGUMENT,
                 "S3 path doesn't contain a bucket name.");
    return;
  }

  *bucket = fname.substr(scheme_end + 1, bucket_end - scheme_end - 1);
  *object = fname.substr(bucket_end + 1);

  if (object->empty() && !object_empty_ok) {
    TF_SetStatus(status, TF_INVALID_ARGUMENT,
                 "S3 path doesn't contain an object name.");
  }
}

static Aws::Client::ClientConfiguration& GetDefaultClientConfig() {
  ABSL_CONST_INIT static absl::Mutex cfg_lock(absl::kConstInit);
  static bool init(false);
  static Aws::Client::ClientConfiguration cfg;

  absl::MutexLock l(&cfg_lock);

  if (!init) {
    const char* endpoint = getenv("S3_ENDPOINT");
    if (endpoint) cfg.endpointOverride = Aws::String(endpoint);
    const char* region = getenv("AWS_REGION");
    // TODO (yongtang): `S3_REGION` should be deprecated after 2.0.
    if (!region) region = getenv("S3_REGION");
    if (region) {
      cfg.region = Aws::String(region);
    } else {
      // Load config file (e.g., ~/.aws/config) only if AWS_SDK_LOAD_CONFIG
      // is set with a truthy value.
      const char* load_config_env = getenv("AWS_SDK_LOAD_CONFIG");
      std::string load_config =
          load_config_env ? absl::AsciiStrToLower(load_config_env) : "";
      if (load_config == "true" || load_config == "1") {
        Aws::String config_file;
        // If AWS_CONFIG_FILE is set then use it, otherwise use ~/.aws/config.
        const char* config_file_env = getenv("AWS_CONFIG_FILE");
        if (config_file_env) {
          config_file = config_file_env;
        } else {
          const char* home_env = getenv("HOME");
          if (home_env) {
            config_file = home_env;
            config_file += "/.aws/config";
          }
        }
        Aws::Config::AWSConfigFileProfileConfigLoader loader(config_file);
        loader.Load();
        auto profiles = loader.GetProfiles();
        if (!profiles["default"].GetRegion().empty())
          cfg.region = profiles["default"].GetRegion();
      }
    }
    const char* use_https = getenv("S3_USE_HTTPS");
    if (use_https) {
      if (use_https[0] == '0')
        cfg.scheme = Aws::Http::Scheme::HTTP;
      else
        cfg.scheme = Aws::Http::Scheme::HTTPS;
    }
    const char* verify_ssl = getenv("S3_VERIFY_SSL");
    if (verify_ssl) {
      if (verify_ssl[0] == '0')
        cfg.verifySSL = false;
      else
        cfg.verifySSL = true;
    }
    // if these timeouts are low, you may see an error when
    // uploading/downloading large files: Unable to connect to endpoint
    int64_t timeout;
    cfg.connectTimeoutMs =
        absl::SimpleAtoi(getenv("S3_CONNECT_TIMEOUT_MSEC"), &timeout)
            ? timeout
            : kS3TimeoutMsec;
    cfg.requestTimeoutMs =
        absl::SimpleAtoi(getenv("S3_REQUEST_TIMEOUT_MSEC"), &timeout)
            ? timeout
            : kS3TimeoutMsec;
    const char* ca_file = getenv("S3_CA_FILE");
    if (ca_file) cfg.caFile = Aws::String(ca_file);
    const char* ca_path = getenv("S3_CA_PATH");
    if (ca_path) cfg.caPath = Aws::String(ca_path);
    init = true;
  }
  return cfg;
};

static void GetS3Client(tf_s3_filesystem::S3File* s3_file) {
  absl::MutexLock l(&s3_file->initialization_lock);

  if (s3_file->s3_client.get() == nullptr) {
    Aws::SDKOptions options;
    options.cryptoOptions.sha256Factory_create_fn = []() {
      return Aws::MakeShared<tf_s3_filesystem::AWSSHA256Factory>(
          tf_s3_filesystem::AWSCryptoAllocationTag);
    };
    options.cryptoOptions.sha256HMACFactory_create_fn = []() {
      return Aws::MakeShared<tf_s3_filesystem::AWSSHA256HmacFactory>(
          tf_s3_filesystem::AWSCryptoAllocationTag);
    };
    options.cryptoOptions.secureRandomFactory_create_fn = []() {
      return Aws::MakeShared<tf_s3_filesystem::AWSSecureRandomFactory>(
          tf_s3_filesystem::AWSCryptoAllocationTag);
    };
    Aws::InitAPI(options);

    // The creation of S3Client disables virtual addressing:
    //   S3Client(clientConfiguration, signPayloads, useVirtualAddressing =
    //   true)
    // The purpose is to address the issue encountered when there is an `.`
    // in the bucket name. Due to TLS hostname validation or DNS rules,
    // the bucket may not be resolved. Disabling of virtual addressing
    // should address the issue. See GitHub issue 16397 for details.
    s3_file->s3_client = Aws::MakeShared<Aws::S3::S3Client>(
        kS3ClientAllocationTag, GetDefaultClientConfig(),
        Aws::Client::AWSAuthV4Signer::PayloadSigningPolicy::Never, false);
  }
}

static void GetExecutor(tf_s3_filesystem::S3File* s3_file) {
  absl::MutexLock l(&s3_file->initialization_lock);

  if (s3_file->executor.get() == nullptr) {
    s3_file->executor =
        Aws::MakeShared<Aws::Utils::Threading::PooledThreadExecutor>(
            kExecutorTag, kExecutorPoolSize);
  }
}

static void ShutdownClient(Aws::S3::S3Client* s3_client) {
  if (s3_client != nullptr) {
    delete s3_client;
    Aws::SDKOptions options;
    Aws::ShutdownAPI(options);
  }
}

// SECTION 1. Implementation for `TF_RandomAccessFile`
// ----------------------------------------------------------------------------
namespace tf_random_access_file {

// TODO(vnvo2409): Implement later

}  // namespace tf_random_access_file

// SECTION 2. Implementation for `TF_WritableFile`
// ----------------------------------------------------------------------------
namespace tf_writable_file {

// TODO(vnvo2409): Implement later

}  // namespace tf_writable_file

// SECTION 3. Implementation for `TF_ReadOnlyMemoryRegion`
// ----------------------------------------------------------------------------
namespace tf_read_only_memory_region {

// TODO(vnvo2409): Implement later

}  // namespace tf_read_only_memory_region

// SECTION 4. Implementation for `TF_Filesystem`, the actual filesystem
// ----------------------------------------------------------------------------
namespace tf_s3_filesystem {
S3File::S3File()
    : s3_client(nullptr, ShutdownClient),
      executor(nullptr),
      transfer_managers(),
      multi_part_chunk_sizes(),
      use_multi_part_download(true),
      initialization_lock() {
  uint64_t temp_value;
  multi_part_chunk_sizes[Aws::Transfer::TransferDirection::UPLOAD] =
      absl::SimpleAtoi(getenv("S3_MULTI_PART_UPLOAD_CHUNK_SIZE"), &temp_value)
          ? temp_value
          : kS3MultiPartUploadChunkSize;
  multi_part_chunk_sizes[Aws::Transfer::TransferDirection::DOWNLOAD] =
      absl::SimpleAtoi(getenv("S3_MULTI_PART_DOWNLOAD_CHUNK_SIZE"), &temp_value)
          ? temp_value
          : kS3MultiPartDownloadChunkSize;
  use_multi_part_download =
      absl::SimpleAtoi(getenv("S3_DISABLE_MULTI_PART_DOWNLOAD"), &temp_value)
          ? (temp_value != 1)
          : use_multi_part_download;
  transfer_managers.emplace(Aws::Transfer::TransferDirection::UPLOAD, nullptr);
  transfer_managers.emplace(Aws::Transfer::TransferDirection::DOWNLOAD,
                            nullptr);
}
void Init(TF_Filesystem* filesystem, TF_Status* status) {
  filesystem->plugin_filesystem = new S3File();
  TF_SetStatus(status, TF_OK, "");
}

void Cleanup(TF_Filesystem* filesystem) {
  auto s3_file = static_cast<S3File*>(filesystem->plugin_filesystem);
  delete s3_file;
}

// TODO(vnvo2409): Implement later

}  // namespace tf_s3_filesystem

static void ProvideFilesystemSupportFor(TF_FilesystemPluginOps* ops,
                                        const char* uri) {
  TF_SetFilesystemVersionMetadata(ops);
  ops->scheme = strdup(uri);
}

void TF_InitPlugin(TF_FilesystemPluginInfo* info) {
  info->plugin_memory_allocate = plugin_memory_allocate;
  info->plugin_memory_free = plugin_memory_free;
  info->num_schemes = 1;
  info->ops = static_cast<TF_FilesystemPluginOps*>(
      plugin_memory_allocate(info->num_schemes * sizeof(info->ops[0])));
  ProvideFilesystemSupportFor(&info->ops[0], "s3");
}
<commit_msg>Add get transfer manager<commit_after>/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/c/experimental/filesystem/plugins/s3/s3_filesystem.h"

#include <aws/core/config/AWSProfileConfigLoader.h>
#include <stdlib.h>
#include <string.h>

#include "absl/strings/ascii.h"
#include "absl/strings/numbers.h"
#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
#include "tensorflow/c/experimental/filesystem/plugins/s3/aws_crypto.h"
#include "tensorflow/c/tf_status.h"

// Implementation of a filesystem for S3 environments.
// This filesystem will support `s3://` URI schemes.
constexpr char kS3ClientAllocationTag[] = "S3ClientAllocation";
constexpr int64_t kS3TimeoutMsec = 300000;  // 5 min

constexpr char kExecutorTag[] = "TransferManagerExecutorAllocation";
constexpr int kExecutorPoolSize = 25;

constexpr uint64_t kS3MultiPartUploadChunkSize = 50 * 1024 * 1024;    // 50 MB
constexpr uint64_t kS3MultiPartDownloadChunkSize = 50 * 1024 * 1024;  // 50 MB

static void* plugin_memory_allocate(size_t size) { return calloc(1, size); }
static void plugin_memory_free(void* ptr) { free(ptr); }

static void ParseS3Path(const Aws::String& fname, bool object_empty_ok,
                        Aws::String* bucket, Aws::String* object,
                        TF_Status* status) {
  size_t scheme_end = fname.find("://") + 2;
  if (fname.substr(0, scheme_end + 1) != "s3://") {
    TF_SetStatus(status, TF_INVALID_ARGUMENT,
                 "S3 path doesn't start with 's3://'.");
    return;
  }

  size_t bucket_end = fname.find("/", scheme_end + 1);
  if (bucket_end == std::string::npos) {
    TF_SetStatus(status, TF_INVALID_ARGUMENT,
                 "S3 path doesn't contain a bucket name.");
    return;
  }

  *bucket = fname.substr(scheme_end + 1, bucket_end - scheme_end - 1);
  *object = fname.substr(bucket_end + 1);

  if (object->empty() && !object_empty_ok) {
    TF_SetStatus(status, TF_INVALID_ARGUMENT,
                 "S3 path doesn't contain an object name.");
  }
}

static Aws::Client::ClientConfiguration& GetDefaultClientConfig() {
  ABSL_CONST_INIT static absl::Mutex cfg_lock(absl::kConstInit);
  static bool init(false);
  static Aws::Client::ClientConfiguration cfg;

  absl::MutexLock l(&cfg_lock);

  if (!init) {
    const char* endpoint = getenv("S3_ENDPOINT");
    if (endpoint) cfg.endpointOverride = Aws::String(endpoint);
    const char* region = getenv("AWS_REGION");
    // TODO (yongtang): `S3_REGION` should be deprecated after 2.0.
    if (!region) region = getenv("S3_REGION");
    if (region) {
      cfg.region = Aws::String(region);
    } else {
      // Load config file (e.g., ~/.aws/config) only if AWS_SDK_LOAD_CONFIG
      // is set with a truthy value.
      const char* load_config_env = getenv("AWS_SDK_LOAD_CONFIG");
      std::string load_config =
          load_config_env ? absl::AsciiStrToLower(load_config_env) : "";
      if (load_config == "true" || load_config == "1") {
        Aws::String config_file;
        // If AWS_CONFIG_FILE is set then use it, otherwise use ~/.aws/config.
        const char* config_file_env = getenv("AWS_CONFIG_FILE");
        if (config_file_env) {
          config_file = config_file_env;
        } else {
          const char* home_env = getenv("HOME");
          if (home_env) {
            config_file = home_env;
            config_file += "/.aws/config";
          }
        }
        Aws::Config::AWSConfigFileProfileConfigLoader loader(config_file);
        loader.Load();
        auto profiles = loader.GetProfiles();
        if (!profiles["default"].GetRegion().empty())
          cfg.region = profiles["default"].GetRegion();
      }
    }
    const char* use_https = getenv("S3_USE_HTTPS");
    if (use_https) {
      if (use_https[0] == '0')
        cfg.scheme = Aws::Http::Scheme::HTTP;
      else
        cfg.scheme = Aws::Http::Scheme::HTTPS;
    }
    const char* verify_ssl = getenv("S3_VERIFY_SSL");
    if (verify_ssl) {
      if (verify_ssl[0] == '0')
        cfg.verifySSL = false;
      else
        cfg.verifySSL = true;
    }
    // if these timeouts are low, you may see an error when
    // uploading/downloading large files: Unable to connect to endpoint
    int64_t timeout;
    cfg.connectTimeoutMs =
        absl::SimpleAtoi(getenv("S3_CONNECT_TIMEOUT_MSEC"), &timeout)
            ? timeout
            : kS3TimeoutMsec;
    cfg.requestTimeoutMs =
        absl::SimpleAtoi(getenv("S3_REQUEST_TIMEOUT_MSEC"), &timeout)
            ? timeout
            : kS3TimeoutMsec;
    const char* ca_file = getenv("S3_CA_FILE");
    if (ca_file) cfg.caFile = Aws::String(ca_file);
    const char* ca_path = getenv("S3_CA_PATH");
    if (ca_path) cfg.caPath = Aws::String(ca_path);
    init = true;
  }
  return cfg;
};

static void GetS3Client(tf_s3_filesystem::S3File* s3_file) {
  absl::MutexLock l(&s3_file->initialization_lock);

  if (s3_file->s3_client.get() == nullptr) {
    Aws::SDKOptions options;
    options.cryptoOptions.sha256Factory_create_fn = []() {
      return Aws::MakeShared<tf_s3_filesystem::AWSSHA256Factory>(
          tf_s3_filesystem::AWSCryptoAllocationTag);
    };
    options.cryptoOptions.sha256HMACFactory_create_fn = []() {
      return Aws::MakeShared<tf_s3_filesystem::AWSSHA256HmacFactory>(
          tf_s3_filesystem::AWSCryptoAllocationTag);
    };
    options.cryptoOptions.secureRandomFactory_create_fn = []() {
      return Aws::MakeShared<tf_s3_filesystem::AWSSecureRandomFactory>(
          tf_s3_filesystem::AWSCryptoAllocationTag);
    };
    Aws::InitAPI(options);

    // The creation of S3Client disables virtual addressing:
    //   S3Client(clientConfiguration, signPayloads, useVirtualAddressing =
    //   true)
    // The purpose is to address the issue encountered when there is an `.`
    // in the bucket name. Due to TLS hostname validation or DNS rules,
    // the bucket may not be resolved. Disabling of virtual addressing
    // should address the issue. See GitHub issue 16397 for details.
    s3_file->s3_client = Aws::MakeShared<Aws::S3::S3Client>(
        kS3ClientAllocationTag, GetDefaultClientConfig(),
        Aws::Client::AWSAuthV4Signer::PayloadSigningPolicy::Never, false);
  }
}

static void GetExecutor(tf_s3_filesystem::S3File* s3_file) {
  absl::MutexLock l(&s3_file->initialization_lock);

  if (s3_file->executor.get() == nullptr) {
    s3_file->executor =
        Aws::MakeShared<Aws::Utils::Threading::PooledThreadExecutor>(
            kExecutorTag, kExecutorPoolSize);
  }
}

static void GetTransferManager(
    const Aws::Transfer::TransferDirection& direction,
    tf_s3_filesystem::S3File* s3_file) {
  absl::MutexLock l(&s3_file->initialization_lock);

  if (s3_file->transfer_managers[direction].get() == nullptr) {
    GetS3Client(s3_file);
    GetExecutor(s3_file);
    Aws::Transfer::TransferManagerConfiguration config(s3_file->executor.get());
    config.s3Client = s3_file->s3_client;
    config.bufferSize = s3_file->multi_part_chunk_sizes[direction];
    // must be larger than pool size * multi part chunk size
    config.transferBufferMaxHeapSize =
        (kExecutorPoolSize + 1) * s3_file->multi_part_chunk_sizes[direction];
    s3_file->transfer_managers[direction] =
        Aws::Transfer::TransferManager::Create(config);
  }
}

static void ShutdownClient(Aws::S3::S3Client* s3_client) {
  if (s3_client != nullptr) {
    delete s3_client;
    Aws::SDKOptions options;
    Aws::ShutdownAPI(options);
  }
}

// SECTION 1. Implementation for `TF_RandomAccessFile`
// ----------------------------------------------------------------------------
namespace tf_random_access_file {

// TODO(vnvo2409): Implement later

}  // namespace tf_random_access_file

// SECTION 2. Implementation for `TF_WritableFile`
// ----------------------------------------------------------------------------
namespace tf_writable_file {

// TODO(vnvo2409): Implement later

}  // namespace tf_writable_file

// SECTION 3. Implementation for `TF_ReadOnlyMemoryRegion`
// ----------------------------------------------------------------------------
namespace tf_read_only_memory_region {

// TODO(vnvo2409): Implement later

}  // namespace tf_read_only_memory_region

// SECTION 4. Implementation for `TF_Filesystem`, the actual filesystem
// ----------------------------------------------------------------------------
namespace tf_s3_filesystem {
S3File::S3File()
    : s3_client(nullptr, ShutdownClient),
      executor(nullptr),
      transfer_managers(),
      multi_part_chunk_sizes(),
      use_multi_part_download(true),
      initialization_lock() {
  uint64_t temp_value;
  multi_part_chunk_sizes[Aws::Transfer::TransferDirection::UPLOAD] =
      absl::SimpleAtoi(getenv("S3_MULTI_PART_UPLOAD_CHUNK_SIZE"), &temp_value)
          ? temp_value
          : kS3MultiPartUploadChunkSize;
  multi_part_chunk_sizes[Aws::Transfer::TransferDirection::DOWNLOAD] =
      absl::SimpleAtoi(getenv("S3_MULTI_PART_DOWNLOAD_CHUNK_SIZE"), &temp_value)
          ? temp_value
          : kS3MultiPartDownloadChunkSize;
  use_multi_part_download =
      absl::SimpleAtoi(getenv("S3_DISABLE_MULTI_PART_DOWNLOAD"), &temp_value)
          ? (temp_value != 1)
          : use_multi_part_download;
  transfer_managers.emplace(Aws::Transfer::TransferDirection::UPLOAD, nullptr);
  transfer_managers.emplace(Aws::Transfer::TransferDirection::DOWNLOAD,
                            nullptr);
}
void Init(TF_Filesystem* filesystem, TF_Status* status) {
  filesystem->plugin_filesystem = new S3File();
  TF_SetStatus(status, TF_OK, "");
}

void Cleanup(TF_Filesystem* filesystem) {
  auto s3_file = static_cast<S3File*>(filesystem->plugin_filesystem);
  delete s3_file;
}

// TODO(vnvo2409): Implement later

}  // namespace tf_s3_filesystem

static void ProvideFilesystemSupportFor(TF_FilesystemPluginOps* ops,
                                        const char* uri) {
  TF_SetFilesystemVersionMetadata(ops);
  ops->scheme = strdup(uri);
}

void TF_InitPlugin(TF_FilesystemPluginInfo* info) {
  info->plugin_memory_allocate = plugin_memory_allocate;
  info->plugin_memory_free = plugin_memory_free;
  info->num_schemes = 1;
  info->ops = static_cast<TF_FilesystemPluginOps*>(
      plugin_memory_allocate(info->num_schemes * sizeof(info->ops[0])));
  ProvideFilesystemSupportFor(&info->ops[0], "s3");
}
<|endoftext|>


<commit_before>#include <sstream>
#include <vector>

#include "extensions/filters/network/redis_proxy/command_splitter_impl.h"

#include "test/integration/integration.h"

using testing::Return;

namespace Envoy {
namespace {

// This is a basic redis_proxy configuration with a single host
// in the cluster. The load balancing policy must be set
// to random for proper test operation.

const std::string CONFIG = R"EOF(
admin:
  access_log_path: /dev/null
  address:
    socket_address:
      address: 127.0.0.1
      port_value: 0
static_resources:
  listeners:
    name: listener_0
    address:
      socket_address:
        address: 127.0.0.1
        port_value: 0
    filter_chains:
      filters:
        name: envoy.redis_proxy
        config:
          stat_prefix: redis_stats
          cluster: cluster_0
          settings:
            op_timeout: 5s
  clusters:
    - name: cluster_0
      lb_policy: CLUSTER_PROVIDED
      hosts:
      - socket_address:
          address: 127.0.0.1
          port_value: 0
      cluster_type:
        name: envoy.clusters.redis
        typed_config:
          "@type": type.googleapis.com/google.protobuf.Struct
          value:
            cluster_refresh_rate: 1s
            cluster_refresh_timeout: 4s
)EOF";

// This is the basic redis_proxy configuration with an upstream
// authentication password specified.

const std::string CONFIG_WITH_AUTH = CONFIG + R"EOF(
      extension_protocol_options:
        envoy.redis_proxy: { auth_password: { inline_string: somepassword }}
)EOF";

// This function encodes commands as an array of bulkstrings as transmitted by Redis clients to
// Redis servers, according to the Redis protocol.
std::string makeBulkStringArray(std::vector<std::string>&& command_strings) {
  std::stringstream result;

  result << "*" << command_strings.size() << "\r\n";
  for (uint64_t i = 0; i < command_strings.size(); i++) {
    result << "$" << command_strings[i].size() << "\r\n";
    result << command_strings[i] << "\r\n";
  }

  return result.str();
}

class RedisClusterIntegrationTest : public testing::TestWithParam<Network::Address::IpVersion>,
                                    public BaseIntegrationTest {
public:
  RedisClusterIntegrationTest(const std::string& config = CONFIG, int num_upstreams = 2)
      : BaseIntegrationTest(GetParam(), config), num_upstreams_(num_upstreams),
        version_(GetParam()) {}

  void TearDown() override {
    test_server_.reset();
    fake_upstreams_.clear();
  }

  void initialize() override {
    setUpstreamCount(num_upstreams_);
    setDeterministic();
    config_helper_.renameListener("redis_proxy");

    // Change the port for each of the discovery host in cluster_0.
    config_helper_.addConfigModifier([this](envoy::config::bootstrap::v2::Bootstrap& bootstrap) {
      uint32_t upstream_idx = 0;
      auto* cluster_0 = bootstrap.mutable_static_resources()->mutable_clusters(0);

      for (int j = 0; j < cluster_0->hosts_size(); ++j) {
        if (cluster_0->mutable_hosts(j)->has_socket_address()) {
          auto* host_socket_addr = cluster_0->mutable_hosts(j)->mutable_socket_address();
          RELEASE_ASSERT(fake_upstreams_.size() > upstream_idx, "");
          host_socket_addr->set_address(
              fake_upstreams_[upstream_idx]->localAddress()->ip()->addressAsString());
          host_socket_addr->set_port_value(
              fake_upstreams_[upstream_idx++]->localAddress()->ip()->port());
        }
      }
    });

    BaseIntegrationTest::initialize();

    mock_rng_ = dynamic_cast<Runtime::MockRandomGenerator*>(&test_server_->server().random());
    // Abort now if we cannot downcast the server's random number generator pointer.
    ASSERT_TRUE(mock_rng_ != nullptr);
    // Ensure that fake_upstreams_[0] is the load balancer's host of choice by default.
    ON_CALL(*mock_rng_, random()).WillByDefault(Return(random_index_));
  }

protected:
  /**
   * A single step of a larger test involving a fake Redis client and a specific Redis server.
   * @param upstream a handle to the server that will respond to the request.
   * @param request supplies Redis client data to transmit to the Redis server.
   * @param response supplies Redis server data to transmit to the client.
   * @param redis_client a handle to the fake redis client that sends the request.
   * @param fake_upstream_connection supplies a handle to connection from the proxy to the fake
   * server.
   * @param auth_password supplies the fake upstream's server password, if not an empty string.
   */
  void roundtripToUpstreamStep(FakeUpstreamPtr& upstream, const std::string& request,
                               const std::string& response, IntegrationTcpClientPtr& redis_client,
                               FakeRawConnectionPtr& fake_upstream_connection,
                               const std::string& auth_password) {
    std::string proxy_to_server;
    bool expect_auth_command = false;
    std::string ok = "+OK\r\n";

    redis_client->clearData();
    redis_client->write(request);

    if (fake_upstream_connection.get() == nullptr) {
      expect_auth_command = (!auth_password.empty());
      EXPECT_TRUE(upstream->waitForRawConnection(fake_upstream_connection));
    }

    if (expect_auth_command) {
      std::string auth_command = makeBulkStringArray({"auth", auth_password});
      EXPECT_TRUE(fake_upstream_connection->waitForData(auth_command.size() + request.size(),
                                                        &proxy_to_server));
      // The original request should be the same as the data received by the server.
      EXPECT_EQ(auth_command + request, proxy_to_server);
      // Send back an OK for the auth command.
      EXPECT_TRUE(fake_upstream_connection->write(ok));

    } else {
      EXPECT_TRUE(fake_upstream_connection->waitForData(request.size(), &proxy_to_server));
      // The original request should be the same as the data received by the server.
      EXPECT_EQ(request, proxy_to_server);
    }

    EXPECT_TRUE(fake_upstream_connection->write(response));
    redis_client->waitForData(response);
    // The original response should be received by the fake Redis client.
    EXPECT_EQ(response, redis_client->data());
  }

  /**
   * Simple bi-directional test between a fake Redis client and Redis server.
   * @param request supplies Redis client data to transmit to the Redis server.
   * @param response supplies Redis server data to transmit to the client.
   */
  void simpleRequestAndResponse(const int stream_index, const std::string& request,
                                const std::string& response) {
    IntegrationTcpClientPtr redis_client = makeTcpConnection(lookupPort("redis_proxy"));
    FakeRawConnectionPtr fake_upstream_connection;

    roundtripToUpstreamStep(fake_upstreams_[stream_index], request, response, redis_client,
                            fake_upstream_connection, "");

    redis_client->close();
    EXPECT_TRUE(fake_upstream_connection->close());
  }

  void expectCallClusterSlot(int stream_index, std::string& response,
                             const std::string& auth_password = "") {
    std::string cluster_slot_request = makeBulkStringArray({"CLUSTER", "SLOTS"});

    fake_upstreams_[stream_index]->set_allow_unexpected_disconnects(true);

    std::string proxied_cluster_slot_request;

    FakeRawConnectionPtr fake_upstream_connection_;
    EXPECT_TRUE(fake_upstreams_[stream_index]->waitForRawConnection(fake_upstream_connection_));
    if (auth_password.empty()) {
      EXPECT_TRUE(fake_upstream_connection_->waitForData(cluster_slot_request.size(),
                                                         &proxied_cluster_slot_request));
      EXPECT_EQ(cluster_slot_request, proxied_cluster_slot_request);
    } else {
      std::string auth_request = makeBulkStringArray({"auth", auth_password});
      std::string ok = "+OK\r\n";

      EXPECT_TRUE(fake_upstream_connection_->waitForData(
          auth_request.size() + cluster_slot_request.size(), &proxied_cluster_slot_request));
      EXPECT_EQ(auth_request + cluster_slot_request, proxied_cluster_slot_request);
      EXPECT_TRUE(fake_upstream_connection_->write(ok));
    }

    EXPECT_TRUE(fake_upstream_connection_->write(response));
    EXPECT_TRUE(fake_upstream_connection_->close());
  }

  /**
   * Simple response for a single slot redis cluster with a master and slave.
   * @param master the ip of the master node.
   * @param slave the ip of the slave node.
   * @return The cluster slot response.
   */
  std::string singleSlotMasterSlave(const Network::Address::Ip* master,
                                    const Network::Address::Ip* slave) {
    int64_t start_slot = 0;
    int64_t end_slot = 16383;

    std::stringstream resp;
    resp << "*1\r\n"
         << "*4\r\n"
         << ":" << start_slot << "\r\n"
         << ":" << end_slot << "\r\n"
         << makeIp(master->addressAsString(), master->port())
         << makeIp(slave->addressAsString(), slave->port());

    return resp.str();
  }

  /**
   * Simple response for 2 slot redis cluster with 2 nodes.
   * @param slot1 the ip of the master node of slot1.
   * @param slot2 the ip of the master node of slot2.
   * @return The cluster slot response.
   */
  std::string twoSlots(const Network::Address::Ip* slot1, const Network::Address::Ip* slot2,
                       int64_t start_slot1 = 0, int64_t end_slot1 = 10000,
                       int64_t start_slot2 = 10000, int64_t end_slot2 = 16383) {
    std::stringstream resp;
    resp << "*2\r\n"
         << "*3\r\n"
         << ":" << start_slot1 << "\r\n"
         << ":" << end_slot1 << "\r\n"
         << makeIp(slot1->addressAsString(), slot1->port()) << "*3\r\n"
         << ":" << start_slot2 << "\r\n"
         << ":" << end_slot2 << "\r\n"
         << makeIp(slot2->addressAsString(), slot2->port());
    return resp.str();
  }

  std::string makeIp(const std::string& address, uint32_t port) {
    return fmt::format("*2\r\n${0}\r\n{1}\r\n:{2}\r\n", address.size(), address, port);
  }

  Runtime::MockRandomGenerator* mock_rng_{};
  const int num_upstreams_;
  const Network::Address::IpVersion version_;
  int random_index_;
};

class RedisClusterWithAuthIntegrationTest : public RedisClusterIntegrationTest {
public:
  RedisClusterWithAuthIntegrationTest(const std::string& config = CONFIG_WITH_AUTH,
                                      int num_upstreams = 2)
      : RedisClusterIntegrationTest(config, num_upstreams) {}
};

INSTANTIATE_TEST_SUITE_P(IpVersions, RedisClusterIntegrationTest,
                         testing::ValuesIn(TestEnvironment::getIpVersionsForTest()),
                         TestUtility::ipTestParamsToString);

INSTANTIATE_TEST_SUITE_P(IpVersions, RedisClusterWithAuthIntegrationTest,
                         testing::ValuesIn(TestEnvironment::getIpVersionsForTest()),
                         TestUtility::ipTestParamsToString);

// This test sends a simple "get foo" command from a fake
// downstream client through the proxy to a fake upstream
// Redis cluster with a single slot with master and slave.
// The fake server sends a valid response back to the client.
// The request and response should make it through the envoy
// proxy server code unchanged.

TEST_P(RedisClusterIntegrationTest, SingleSlotMasterSlave) {
  random_index_ = 0;

  on_server_init_function_ = [this]() {
    std::string cluster_slot_response = singleSlotMasterSlave(
        fake_upstreams_[0]->localAddress()->ip(), fake_upstreams_[1]->localAddress()->ip());
    expectCallClusterSlot(random_index_, cluster_slot_response);
  };

  initialize();

  // foo hashes to slot 12182 which is in upstream 0
  simpleRequestAndResponse(0, makeBulkStringArray({"get", "foo"}), "$3\r\nbar\r\n");
}

// This test sends a simple "get foo" command from a fake
// downstream client through the proxy to a fake upstream
// Redis cluster with 2 slots. The fake server sends a valid response
// back to the client. The request and response should
// make it through the envoy proxy server code unchanged.

TEST_P(RedisClusterIntegrationTest, TwoSlot) {
  random_index_ = 0;

  on_server_init_function_ = [this]() {
    std::string cluster_slot_response = twoSlots(fake_upstreams_[0]->localAddress()->ip(),
                                                 fake_upstreams_[1]->localAddress()->ip());
    expectCallClusterSlot(random_index_, cluster_slot_response);
  };

  initialize();

  // foobar hashes to slot 12325 which is in upstream 1
  simpleRequestAndResponse(1, makeBulkStringArray({"get", "foobar"}), "$3\r\nbar\r\n");
  // bar hashes to slot 5061 which is in upstream 0
  simpleRequestAndResponse(0, makeBulkStringArray({"get", "bar"}), "$3\r\nbar\r\n");
  // foo hashes to slot 12182 which is in upstream 1
  simpleRequestAndResponse(1, makeBulkStringArray({"get", "foo"}), "$3\r\nbar\r\n");
}

// This test sends a simple "get foo" command from a fake
// downstream client through the proxy to a fake upstream
// Redis cluster with a single slot with master and slave.
// The fake server sends a valid response back to the client.
// The request and response should make it through the envoy
// proxy server code unchanged.
//
// In this scenario, the fake server will receive 2 auth commands:
// one as part of a topology discovery connection (before sending a
// "cluster slots" command), and one to authenticate the connection
// that carries the "get foo" request.

TEST_P(RedisClusterWithAuthIntegrationTest, SingleSlotMasterSlave) {
  random_index_ = 0;

  on_server_init_function_ = [this]() {
    std::string cluster_slot_response = singleSlotMasterSlave(
        fake_upstreams_[0]->localAddress()->ip(), fake_upstreams_[1]->localAddress()->ip());
    expectCallClusterSlot(0, cluster_slot_response, "somepassword");
  };

  initialize();

  IntegrationTcpClientPtr redis_client = makeTcpConnection(lookupPort("redis_proxy"));
  FakeRawConnectionPtr fake_upstream_connection;

  roundtripToUpstreamStep(fake_upstreams_[random_index_], makeBulkStringArray({"get", "foo"}),
                          "$3\r\nbar\r\n", redis_client, fake_upstream_connection, "somepassword");

  redis_client->close();
  EXPECT_TRUE(fake_upstream_connection->close());
}

} // namespace
} // namespace Envoy
<commit_msg>fix static initialization fiasco (#7684)<commit_after>#include <sstream>
#include <vector>

#include "common/common/macros.h"

#include "extensions/filters/network/redis_proxy/command_splitter_impl.h"

#include "test/integration/integration.h"

using testing::Return;

namespace Envoy {
namespace {

// This is a basic redis_proxy configuration with a single host
// in the cluster. The load balancing policy must be set
// to random for proper test operation.

const std::string& testConfig() {
  CONSTRUCT_ON_FIRST_USE(std::string, R"EOF(
admin:
  access_log_path: /dev/null
  address:
    socket_address:
      address: 127.0.0.1
      port_value: 0
static_resources:
  listeners:
    name: listener_0
    address:
      socket_address:
        address: 127.0.0.1
        port_value: 0
    filter_chains:
      filters:
        name: envoy.redis_proxy
        config:
          stat_prefix: redis_stats
          cluster: cluster_0
          settings:
            op_timeout: 5s
  clusters:
    - name: cluster_0
      lb_policy: CLUSTER_PROVIDED
      hosts:
      - socket_address:
          address: 127.0.0.1
          port_value: 0
      cluster_type:
        name: envoy.clusters.redis
        typed_config:
          "@type": type.googleapis.com/google.protobuf.Struct
          value:
            cluster_refresh_rate: 1s
            cluster_refresh_timeout: 4s
)EOF");
}

// This is the basic redis_proxy configuration with an upstream
// authentication password specified.

const std::string& testConfigWithAuth() {
  CONSTRUCT_ON_FIRST_USE(std::string, testConfig() + R"EOF(
      extension_protocol_options:
        envoy.redis_proxy: { auth_password: { inline_string: somepassword }}
)EOF");
}

// This function encodes commands as an array of bulkstrings as transmitted by Redis clients to
// Redis servers, according to the Redis protocol.
std::string makeBulkStringArray(std::vector<std::string>&& command_strings) {
  std::stringstream result;

  result << "*" << command_strings.size() << "\r\n";
  for (uint64_t i = 0; i < command_strings.size(); i++) {
    result << "$" << command_strings[i].size() << "\r\n";
    result << command_strings[i] << "\r\n";
  }

  return result.str();
}

class RedisClusterIntegrationTest : public testing::TestWithParam<Network::Address::IpVersion>,
                                    public BaseIntegrationTest {
public:
  RedisClusterIntegrationTest(const std::string& config = testConfig(), int num_upstreams = 2)
      : BaseIntegrationTest(GetParam(), config), num_upstreams_(num_upstreams),
        version_(GetParam()) {}

  void TearDown() override {
    test_server_.reset();
    fake_upstreams_.clear();
  }

  void initialize() override {
    setUpstreamCount(num_upstreams_);
    setDeterministic();
    config_helper_.renameListener("redis_proxy");

    // Change the port for each of the discovery host in cluster_0.
    config_helper_.addConfigModifier([this](envoy::config::bootstrap::v2::Bootstrap& bootstrap) {
      uint32_t upstream_idx = 0;
      auto* cluster_0 = bootstrap.mutable_static_resources()->mutable_clusters(0);

      for (int j = 0; j < cluster_0->hosts_size(); ++j) {
        if (cluster_0->mutable_hosts(j)->has_socket_address()) {
          auto* host_socket_addr = cluster_0->mutable_hosts(j)->mutable_socket_address();
          RELEASE_ASSERT(fake_upstreams_.size() > upstream_idx, "");
          host_socket_addr->set_address(
              fake_upstreams_[upstream_idx]->localAddress()->ip()->addressAsString());
          host_socket_addr->set_port_value(
              fake_upstreams_[upstream_idx++]->localAddress()->ip()->port());
        }
      }
    });

    BaseIntegrationTest::initialize();

    mock_rng_ = dynamic_cast<Runtime::MockRandomGenerator*>(&test_server_->server().random());
    // Abort now if we cannot downcast the server's random number generator pointer.
    ASSERT_TRUE(mock_rng_ != nullptr);
    // Ensure that fake_upstreams_[0] is the load balancer's host of choice by default.
    ON_CALL(*mock_rng_, random()).WillByDefault(Return(random_index_));
  }

protected:
  /**
   * A single step of a larger test involving a fake Redis client and a specific Redis server.
   * @param upstream a handle to the server that will respond to the request.
   * @param request supplies Redis client data to transmit to the Redis server.
   * @param response supplies Redis server data to transmit to the client.
   * @param redis_client a handle to the fake redis client that sends the request.
   * @param fake_upstream_connection supplies a handle to connection from the proxy to the fake
   * server.
   * @param auth_password supplies the fake upstream's server password, if not an empty string.
   */
  void roundtripToUpstreamStep(FakeUpstreamPtr& upstream, const std::string& request,
                               const std::string& response, IntegrationTcpClientPtr& redis_client,
                               FakeRawConnectionPtr& fake_upstream_connection,
                               const std::string& auth_password) {
    std::string proxy_to_server;
    bool expect_auth_command = false;
    std::string ok = "+OK\r\n";

    redis_client->clearData();
    redis_client->write(request);

    if (fake_upstream_connection.get() == nullptr) {
      expect_auth_command = (!auth_password.empty());
      EXPECT_TRUE(upstream->waitForRawConnection(fake_upstream_connection));
    }

    if (expect_auth_command) {
      std::string auth_command = makeBulkStringArray({"auth", auth_password});
      EXPECT_TRUE(fake_upstream_connection->waitForData(auth_command.size() + request.size(),
                                                        &proxy_to_server));
      // The original request should be the same as the data received by the server.
      EXPECT_EQ(auth_command + request, proxy_to_server);
      // Send back an OK for the auth command.
      EXPECT_TRUE(fake_upstream_connection->write(ok));

    } else {
      EXPECT_TRUE(fake_upstream_connection->waitForData(request.size(), &proxy_to_server));
      // The original request should be the same as the data received by the server.
      EXPECT_EQ(request, proxy_to_server);
    }

    EXPECT_TRUE(fake_upstream_connection->write(response));
    redis_client->waitForData(response);
    // The original response should be received by the fake Redis client.
    EXPECT_EQ(response, redis_client->data());
  }

  /**
   * Simple bi-directional test between a fake Redis client and Redis server.
   * @param request supplies Redis client data to transmit to the Redis server.
   * @param response supplies Redis server data to transmit to the client.
   */
  void simpleRequestAndResponse(const int stream_index, const std::string& request,
                                const std::string& response) {
    IntegrationTcpClientPtr redis_client = makeTcpConnection(lookupPort("redis_proxy"));
    FakeRawConnectionPtr fake_upstream_connection;

    roundtripToUpstreamStep(fake_upstreams_[stream_index], request, response, redis_client,
                            fake_upstream_connection, "");

    redis_client->close();
    EXPECT_TRUE(fake_upstream_connection->close());
  }

  void expectCallClusterSlot(int stream_index, std::string& response,
                             const std::string& auth_password = "") {
    std::string cluster_slot_request = makeBulkStringArray({"CLUSTER", "SLOTS"});

    fake_upstreams_[stream_index]->set_allow_unexpected_disconnects(true);

    std::string proxied_cluster_slot_request;

    FakeRawConnectionPtr fake_upstream_connection_;
    EXPECT_TRUE(fake_upstreams_[stream_index]->waitForRawConnection(fake_upstream_connection_));
    if (auth_password.empty()) {
      EXPECT_TRUE(fake_upstream_connection_->waitForData(cluster_slot_request.size(),
                                                         &proxied_cluster_slot_request));
      EXPECT_EQ(cluster_slot_request, proxied_cluster_slot_request);
    } else {
      std::string auth_request = makeBulkStringArray({"auth", auth_password});
      std::string ok = "+OK\r\n";

      EXPECT_TRUE(fake_upstream_connection_->waitForData(
          auth_request.size() + cluster_slot_request.size(), &proxied_cluster_slot_request));
      EXPECT_EQ(auth_request + cluster_slot_request, proxied_cluster_slot_request);
      EXPECT_TRUE(fake_upstream_connection_->write(ok));
    }

    EXPECT_TRUE(fake_upstream_connection_->write(response));
    EXPECT_TRUE(fake_upstream_connection_->close());
  }

  /**
   * Simple response for a single slot redis cluster with a master and slave.
   * @param master the ip of the master node.
   * @param slave the ip of the slave node.
   * @return The cluster slot response.
   */
  std::string singleSlotMasterSlave(const Network::Address::Ip* master,
                                    const Network::Address::Ip* slave) {
    int64_t start_slot = 0;
    int64_t end_slot = 16383;

    std::stringstream resp;
    resp << "*1\r\n"
         << "*4\r\n"
         << ":" << start_slot << "\r\n"
         << ":" << end_slot << "\r\n"
         << makeIp(master->addressAsString(), master->port())
         << makeIp(slave->addressAsString(), slave->port());

    return resp.str();
  }

  /**
   * Simple response for 2 slot redis cluster with 2 nodes.
   * @param slot1 the ip of the master node of slot1.
   * @param slot2 the ip of the master node of slot2.
   * @return The cluster slot response.
   */
  std::string twoSlots(const Network::Address::Ip* slot1, const Network::Address::Ip* slot2,
                       int64_t start_slot1 = 0, int64_t end_slot1 = 10000,
                       int64_t start_slot2 = 10000, int64_t end_slot2 = 16383) {
    std::stringstream resp;
    resp << "*2\r\n"
         << "*3\r\n"
         << ":" << start_slot1 << "\r\n"
         << ":" << end_slot1 << "\r\n"
         << makeIp(slot1->addressAsString(), slot1->port()) << "*3\r\n"
         << ":" << start_slot2 << "\r\n"
         << ":" << end_slot2 << "\r\n"
         << makeIp(slot2->addressAsString(), slot2->port());
    return resp.str();
  }

  std::string makeIp(const std::string& address, uint32_t port) {
    return fmt::format("*2\r\n${0}\r\n{1}\r\n:{2}\r\n", address.size(), address, port);
  }

  Runtime::MockRandomGenerator* mock_rng_{};
  const int num_upstreams_;
  const Network::Address::IpVersion version_;
  int random_index_;
};

class RedisClusterWithAuthIntegrationTest : public RedisClusterIntegrationTest {
public:
  RedisClusterWithAuthIntegrationTest(const std::string& config = testConfigWithAuth(),
                                      int num_upstreams = 2)
      : RedisClusterIntegrationTest(config, num_upstreams) {}
};

INSTANTIATE_TEST_SUITE_P(IpVersions, RedisClusterIntegrationTest,
                         testing::ValuesIn(TestEnvironment::getIpVersionsForTest()),
                         TestUtility::ipTestParamsToString);

INSTANTIATE_TEST_SUITE_P(IpVersions, RedisClusterWithAuthIntegrationTest,
                         testing::ValuesIn(TestEnvironment::getIpVersionsForTest()),
                         TestUtility::ipTestParamsToString);

// This test sends a simple "get foo" command from a fake
// downstream client through the proxy to a fake upstream
// Redis cluster with a single slot with master and slave.
// The fake server sends a valid response back to the client.
// The request and response should make it through the envoy
// proxy server code unchanged.

TEST_P(RedisClusterIntegrationTest, SingleSlotMasterSlave) {
  random_index_ = 0;

  on_server_init_function_ = [this]() {
    std::string cluster_slot_response = singleSlotMasterSlave(
        fake_upstreams_[0]->localAddress()->ip(), fake_upstreams_[1]->localAddress()->ip());
    expectCallClusterSlot(random_index_, cluster_slot_response);
  };

  initialize();

  // foo hashes to slot 12182 which is in upstream 0
  simpleRequestAndResponse(0, makeBulkStringArray({"get", "foo"}), "$3\r\nbar\r\n");
}

// This test sends a simple "get foo" command from a fake
// downstream client through the proxy to a fake upstream
// Redis cluster with 2 slots. The fake server sends a valid response
// back to the client. The request and response should
// make it through the envoy proxy server code unchanged.

TEST_P(RedisClusterIntegrationTest, TwoSlot) {
  random_index_ = 0;

  on_server_init_function_ = [this]() {
    std::string cluster_slot_response = twoSlots(fake_upstreams_[0]->localAddress()->ip(),
                                                 fake_upstreams_[1]->localAddress()->ip());
    expectCallClusterSlot(random_index_, cluster_slot_response);
  };

  initialize();

  // foobar hashes to slot 12325 which is in upstream 1
  simpleRequestAndResponse(1, makeBulkStringArray({"get", "foobar"}), "$3\r\nbar\r\n");
  // bar hashes to slot 5061 which is in upstream 0
  simpleRequestAndResponse(0, makeBulkStringArray({"get", "bar"}), "$3\r\nbar\r\n");
  // foo hashes to slot 12182 which is in upstream 1
  simpleRequestAndResponse(1, makeBulkStringArray({"get", "foo"}), "$3\r\nbar\r\n");
}

// This test sends a simple "get foo" command from a fake
// downstream client through the proxy to a fake upstream
// Redis cluster with a single slot with master and slave.
// The fake server sends a valid response back to the client.
// The request and response should make it through the envoy
// proxy server code unchanged.
//
// In this scenario, the fake server will receive 2 auth commands:
// one as part of a topology discovery connection (before sending a
// "cluster slots" command), and one to authenticate the connection
// that carries the "get foo" request.

TEST_P(RedisClusterWithAuthIntegrationTest, SingleSlotMasterSlave) {
  random_index_ = 0;

  on_server_init_function_ = [this]() {
    std::string cluster_slot_response = singleSlotMasterSlave(
        fake_upstreams_[0]->localAddress()->ip(), fake_upstreams_[1]->localAddress()->ip());
    expectCallClusterSlot(0, cluster_slot_response, "somepassword");
  };

  initialize();

  IntegrationTcpClientPtr redis_client = makeTcpConnection(lookupPort("redis_proxy"));
  FakeRawConnectionPtr fake_upstream_connection;

  roundtripToUpstreamStep(fake_upstreams_[random_index_], makeBulkStringArray({"get", "foo"}),
                          "$3\r\nbar\r\n", redis_client, fake_upstream_connection, "somepassword");

  redis_client->close();
  EXPECT_TRUE(fake_upstream_connection->close());
}

} // namespace
} // namespace Envoy
<|endoftext|>


<commit_before>/*=========================================================================

  Program:   Monteverdi2
  Language:  C++

  Copyright (c) Centre National d'Etudes Spatiales. All rights reserved.
  See Copyright.txt for details.

  Monteverdi2 is distributed under the CeCILL licence version 2. See
  Licence_CeCILL_V2-en.txt or
  http://www.cecill.info/licences/Licence_CeCILL_V2-en.txt for more details.

  This software is distributed WITHOUT ANY WARRANTY; without even
  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
  PURPOSE.  See the above copyright notices for more information.

=========================================================================*/

//
// Qt includes (sorted by alphabetic order)
//// Must be included before system/custom includes.

//
// System includes (sorted by alphabetic order)

//
// ITK includes (sorted by alphabetic order)
#include "itkImageRegionConstIteratorWithIndex.h"
#include "itkImageRegionConstIterator.h"
//
// OTB includes (sorted by alphabetic order)

//
// Monteverdi includes (sorted by alphabetic order)
#include "mvdVectorImageModel.h"

namespace mvd
{
/*
  TRANSLATOR mvd::VectorImageModel

  Necessary for lupdate to be aware of C++ namespaces.

  Context comment for translator.
*/

/*******************************************************************************/
VectorImageModel
::VectorImageModel( QObject* parent ) :
  AbstractImageModel( parent ),
  m_RasterizedBuffer( NULL ),
  m_PreviousRegion(ImageRegionType())
{
}

/*******************************************************************************/
VectorImageModel
::~VectorImageModel()
{
  this->ClearBuffer();
}

/*******************************************************************************/
void
VectorImageModel
::loadFile( const QString& filename )
{
  DefaultImageFileReaderType::Pointer imageFileReader(
    DefaultImageFileReaderType::New()
  );

  imageFileReader->SetFileName( filename.toLatin1().data() );
  imageFileReader->UpdateOutputInformation();

  m_ImageFileReader = imageFileReader;

  // initialize the channel list for the rendering needs following the
  // input image
  // TODO : See if if needs to be moved somewhere else
  // TODO : use the default display
  if (m_ImageFileReader->GetOutput()->GetNumberOfComponentsPerPixel()  < 3)
    {
    m_Channels.resize(1);
    m_Channels[0]  = 0;
    }
  else
    {
    m_Channels.resize(3);
    m_Channels[0]  = 0;
    m_Channels[1]  = 1;
    m_Channels[2]  = 2;
    }
}

void
VectorImageModel
::ClearBuffer()
{
  // Delete previous buffer if needed
  if (m_RasterizedBuffer != NULL)
    {
    delete[] m_RasterizedBuffer;
    m_RasterizedBuffer = NULL;
    }
}

unsigned char *
VectorImageModel
::RasterizeRegion( const ImageRegionType& region)
{
  m_Region = region;
  // Don't do anything if the region did not changed
  // moveMoveEvent (Drag region).
  if ( m_PreviousRegion != region )
    {
    // check the current and the previous region have some pixels in
    // common
    ImageRegionType  tempPreviousRegionRegion = m_PreviousRegion;
    bool res =  tempPreviousRegionRegion.Crop(region);

    // if the first time, image must be read
    if ( m_PreviousRegion != ImageRegionType() && res)
      {
      // Compute loaeded region, and region not loaded yet within the
      // new requested region
      this->ComputeRegionsToLoad(m_Region);

      // Copy the previous buffer into a temporary buf to get the
      // previously loaded data
      unsigned char * previousRasterizedBuffer =
        new unsigned char[3 * m_PreviousRegion.GetNumberOfPixels()];

      std::memcpy(previousRasterizedBuffer, m_RasterizedBuffer, 3 * m_PreviousRegion.GetNumberOfPixels());

      // Clear the previous buffer
      this->ClearBuffer();

      // Allocate new memory
      unsigned int nbPixels = 3 * region.GetNumberOfPixels();
      m_RasterizedBuffer = new unsigned char[nbPixels];

      // Copy the already loaded pixels into the m_RasterizedBuffer
      unsigned int previousNbPixels = m_PreviousRegion.GetNumberOfPixels();

      for (unsigned int idx = 0; idx < previousNbPixels; idx++)
        {
        // compose the image index from the buffer index
        IndexType imageIndex;
        imageIndex = ComputeImageIndexFromFlippedBuffer(idx, m_PreviousRegion);

        if (m_AlreadyLoadedRegion.IsInside(imageIndex))
          {
          // Get the buffer index relative to the imageIndex in the
          // new requested region
          unsigned int newBufferindex = 0;
          newBufferindex = ComputeXAxisFlippedBufferIndex(imageIndex, m_Region);

          // Copy the already loaded values into the new buffer
          m_RasterizedBuffer[newBufferindex]     = previousRasterizedBuffer[3*idx];
          m_RasterizedBuffer[newBufferindex + 1] = previousRasterizedBuffer[3*idx + 1];
          m_RasterizedBuffer[newBufferindex + 2] = previousRasterizedBuffer[3*idx + 2];
          }
        }

      // Get the image pixels within not loaded region and add them to
      // the buffer
      for (unsigned int idx = 0; idx < m_RegionsToLoadVector.size() ; idx++)
        {
        this->DumpImagePixelsWithinRegionIntoBuffer(m_RegionsToLoadVector[idx]);
        }

      // free the previous buffer memory (copied one)
      if (previousRasterizedBuffer != NULL)
        {
        delete[] previousRasterizedBuffer;
        previousRasterizedBuffer = NULL;
        }
      }
    else
      {
      // Delete previous buffer if needed
      this->ClearBuffer();

      // Allocate new memory
      m_RasterizedBuffer = new unsigned char[3 * region.GetNumberOfPixels()];

      // rasterize the region
      this->DumpImagePixelsWithinRegionIntoBuffer(region);
      }
    }

  // Store the region
  m_PreviousRegion = region;

  // if ok return the  buffer
  return m_RasterizedBuffer;
}

void
VectorImageModel
::DumpImagePixelsWithinRegionIntoBuffer(const ImageRegionType& region)
{
  // Before doing anything, check if region is inside the buffered
  // region of image
  unsigned int currentIndex = 0;

  // TODO : add some checking
  const DefaultImageType* image =  this->GetOutput(currentIndex);

  // some checking
  if (!image->GetBufferedRegion().IsInside(region))
    {
    //itkExceptionMacro(<< "Region to read is oustside of the buffered region.");
    }

  // Extract the region of interest in the image
  m_ExtractFilter = ExtractFilterType::New();
  m_ExtractFilter->SetInput(image);
  m_ExtractFilter->SetExtractionRegion(region);

  // Use the rendering filter to get
  m_RenderingFilter = RenderingFilterType::New();
  m_RenderingFilter->SetInput(m_ExtractFilter->GetOutput());
  m_RenderingFilter->GetRenderingFunction()->SetAutoMinMax(false);
  m_RenderingFilter->GetRenderingFunction()->SetChannelList(m_Channels);
  m_RenderingFilter->GetOutput()->SetRequestedRegion(region);
  m_RenderingFilter->Update();

  // Declare the iterator
  itk::ImageRegionConstIteratorWithIndex< RenderingFilterType::OutputImageType >
    it(m_RenderingFilter->GetOutput(), region);

  // Go to begin
  it.GoToBegin();

  while (!it.IsAtEnd())
    {
    // Fill the buffer
    unsigned int index = 0;
    index = ComputeXAxisFlippedBufferIndex(it.GetIndex(), m_Region);

    // Fill the buffer
    m_RasterizedBuffer[index]  = it.Get()[0];
    m_RasterizedBuffer[index + 1] = it.Get()[1];
    m_RasterizedBuffer[index + 2] = it.Get()[2];
    ++it;
    }
}

void
VectorImageModel
::ComputeRegionsToLoad(const ImageRegionType& region)
{
  // Initialize the region and clear vector
  m_AlreadyLoadedRegion = m_PreviousRegion;
  m_RegionsToLoadVector.clear();

  // 4 regions to compute
  ImageRegionType upperRegion;
  ImageRegionType lowerRegion;
  ImageRegionType leftRegion;
  ImageRegionType rightRegion;

  // local variables
  IndexType index;
  ImageRegionType::SizeType size;

  // Compute the already loaded region as a simple Crop
  bool res =  m_AlreadyLoadedRegion.Crop(region);

  // ------ upper region
  upperRegion.SetIndex(region.GetIndex());
  size[0] = region.GetSize()[0];
  size[1] = vcl_abs( region.GetIndex()[1] - m_AlreadyLoadedRegion.GetIndex()[1]);
  upperRegion.SetSize(size);

  // ------ lower region
  index[0] = region.GetIndex()[0];
  index[1] = m_AlreadyLoadedRegion.GetSize()[1] + m_AlreadyLoadedRegion.GetIndex()[1];
  lowerRegion.SetIndex(index);

  size[0] = region.GetSize()[0];
  size[1] = vcl_abs( region.GetIndex()[1] + region.GetSize()[1]
                     - m_AlreadyLoadedRegion.GetIndex()[1] - m_AlreadyLoadedRegion.GetSize()[1] );
  lowerRegion.SetSize(size);

  // ------- left region
  index[0] = region.GetIndex()[0];
  index[1] = m_AlreadyLoadedRegion.GetIndex()[1];
  leftRegion.SetIndex(index);

  size[0] = vcl_abs(region.GetIndex(0) - m_AlreadyLoadedRegion.GetIndex()[0]);
  size[1] = m_AlreadyLoadedRegion.GetSize()[1];
  leftRegion.SetSize(size);

  // -------- right region
  index[0] = m_AlreadyLoadedRegion.GetIndex()[0] + m_AlreadyLoadedRegion.GetSize()[0];
  index[1] = m_AlreadyLoadedRegion.GetIndex()[1];
  rightRegion.SetIndex(index);

  size[0] = vcl_abs(region.GetSize()[0] + region.GetIndex()[0]
                    - m_AlreadyLoadedRegion.GetIndex()[0] - m_AlreadyLoadedRegion.GetSize()[0]);
  size[1] = m_AlreadyLoadedRegion.GetSize()[1];
  rightRegion.SetSize(size);

  // add the upper region if any pixel
  if ( upperRegion .GetNumberOfPixels() > 0 )
    m_RegionsToLoadVector.push_back(upperRegion);

  // add the lower region if any pixel
  if ( lowerRegion.GetNumberOfPixels() > 0 )
    m_RegionsToLoadVector.push_back(lowerRegion);

  // add the left region if any pixel
  if ( leftRegion.GetNumberOfPixels() > 0 )
    m_RegionsToLoadVector.push_back(leftRegion);

  // add the right region if any pixel
  if ( rightRegion.GetNumberOfPixels() > 0 )
    m_RegionsToLoadVector.push_back(rightRegion);
}

/*******************************************************************************/
/* SLOTS                                                                       */
/*******************************************************************************/

/*******************************************************************************/

} // end namespace 'mvd'
<commit_msg>ENH: unnecessary include<commit_after>/*=========================================================================

  Program:   Monteverdi2
  Language:  C++

  Copyright (c) Centre National d'Etudes Spatiales. All rights reserved.
  See Copyright.txt for details.

  Monteverdi2 is distributed under the CeCILL licence version 2. See
  Licence_CeCILL_V2-en.txt or
  http://www.cecill.info/licences/Licence_CeCILL_V2-en.txt for more details.

  This software is distributed WITHOUT ANY WARRANTY; without even
  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
  PURPOSE.  See the above copyright notices for more information.

=========================================================================*/

//
// Qt includes (sorted by alphabetic order)
//// Must be included before system/custom includes.

//
// System includes (sorted by alphabetic order)

//
// ITK includes (sorted by alphabetic order)
#include "itkImageRegionConstIteratorWithIndex.h"

//
// OTB includes (sorted by alphabetic order)

//
// Monteverdi includes (sorted by alphabetic order)
#include "mvdVectorImageModel.h"

namespace mvd
{
/*
  TRANSLATOR mvd::VectorImageModel

  Necessary for lupdate to be aware of C++ namespaces.

  Context comment for translator.
*/

/*******************************************************************************/
VectorImageModel
::VectorImageModel( QObject* parent ) :
  AbstractImageModel( parent ),
  m_RasterizedBuffer( NULL ),
  m_PreviousRegion(ImageRegionType())
{
}

/*******************************************************************************/
VectorImageModel
::~VectorImageModel()
{
  this->ClearBuffer();
}

/*******************************************************************************/
void
VectorImageModel
::loadFile( const QString& filename )
{
  DefaultImageFileReaderType::Pointer imageFileReader(
    DefaultImageFileReaderType::New()
  );

  imageFileReader->SetFileName( filename.toLatin1().data() );
  imageFileReader->UpdateOutputInformation();

  m_ImageFileReader = imageFileReader;

  // initialize the channel list for the rendering needs following the
  // input image
  // TODO : See if if needs to be moved somewhere else
  // TODO : use the default display
  if (m_ImageFileReader->GetOutput()->GetNumberOfComponentsPerPixel()  < 3)
    {
    m_Channels.resize(1);
    m_Channels[0]  = 0;
    }
  else
    {
    m_Channels.resize(3);
    m_Channels[0]  = 0;
    m_Channels[1]  = 1;
    m_Channels[2]  = 2;
    }
}

void
VectorImageModel
::ClearBuffer()
{
  // Delete previous buffer if needed
  if (m_RasterizedBuffer != NULL)
    {
    delete[] m_RasterizedBuffer;
    m_RasterizedBuffer = NULL;
    }
}

unsigned char *
VectorImageModel
::RasterizeRegion( const ImageRegionType& region)
{
  m_Region = region;
  // Don't do anything if the region did not changed
  // moveMoveEvent (Drag region).
  if ( m_PreviousRegion != region )
    {
    // check the current and the previous region have some pixels in
    // common
    ImageRegionType  tempPreviousRegionRegion = m_PreviousRegion;
    bool res =  tempPreviousRegionRegion.Crop(region);

    // if the first time, image must be read
    if ( m_PreviousRegion != ImageRegionType() && res)
      {
      // Compute loaeded region, and region not loaded yet within the
      // new requested region
      this->ComputeRegionsToLoad(m_Region);

      // Copy the previous buffer into a temporary buf to get the
      // previously loaded data
      unsigned char * previousRasterizedBuffer =
        new unsigned char[3 * m_PreviousRegion.GetNumberOfPixels()];

      std::memcpy(previousRasterizedBuffer, m_RasterizedBuffer, 3 * m_PreviousRegion.GetNumberOfPixels());

      // Clear the previous buffer
      this->ClearBuffer();

      // Allocate new memory
      unsigned int nbPixels = 3 * region.GetNumberOfPixels();
      m_RasterizedBuffer = new unsigned char[nbPixels];

      // Copy the already loaded pixels into the m_RasterizedBuffer
      unsigned int previousNbPixels = m_PreviousRegion.GetNumberOfPixels();

      for (unsigned int idx = 0; idx < previousNbPixels; idx++)
        {
        // compose the image index from the buffer index
        IndexType imageIndex;
        imageIndex = ComputeImageIndexFromFlippedBuffer(idx, m_PreviousRegion);

        if (m_AlreadyLoadedRegion.IsInside(imageIndex))
          {
          // Get the buffer index relative to the imageIndex in the
          // new requested region
          unsigned int newBufferindex = 0;
          newBufferindex = ComputeXAxisFlippedBufferIndex(imageIndex, m_Region);

          // Copy the already loaded values into the new buffer
          m_RasterizedBuffer[newBufferindex]     = previousRasterizedBuffer[3*idx];
          m_RasterizedBuffer[newBufferindex + 1] = previousRasterizedBuffer[3*idx + 1];
          m_RasterizedBuffer[newBufferindex + 2] = previousRasterizedBuffer[3*idx + 2];
          }
        }

      // Get the image pixels within not loaded region and add them to
      // the buffer
      for (unsigned int idx = 0; idx < m_RegionsToLoadVector.size() ; idx++)
        {
        this->DumpImagePixelsWithinRegionIntoBuffer(m_RegionsToLoadVector[idx]);
        }

      // free the previous buffer memory (copied one)
      if (previousRasterizedBuffer != NULL)
        {
        delete[] previousRasterizedBuffer;
        previousRasterizedBuffer = NULL;
        }
      }
    else
      {
      // Delete previous buffer if needed
      this->ClearBuffer();

      // Allocate new memory
      m_RasterizedBuffer = new unsigned char[3 * region.GetNumberOfPixels()];

      // rasterize the region
      this->DumpImagePixelsWithinRegionIntoBuffer(region);
      }
    }

  // Store the region
  m_PreviousRegion = region;

  // if ok return the  buffer
  return m_RasterizedBuffer;
}

void
VectorImageModel
::DumpImagePixelsWithinRegionIntoBuffer(const ImageRegionType& region)
{
  // Before doing anything, check if region is inside the buffered
  // region of image
  unsigned int currentIndex = 0;

  // TODO : add some checking
  const DefaultImageType* image =  this->GetOutput(currentIndex);

  // some checking
  if (!image->GetBufferedRegion().IsInside(region))
    {
    //itkExceptionMacro(<< "Region to read is oustside of the buffered region.");
    }

  // Extract the region of interest in the image
  m_ExtractFilter = ExtractFilterType::New();
  m_ExtractFilter->SetInput(image);
  m_ExtractFilter->SetExtractionRegion(region);

  // Use the rendering filter to get
  m_RenderingFilter = RenderingFilterType::New();
  m_RenderingFilter->SetInput(m_ExtractFilter->GetOutput());
  m_RenderingFilter->GetRenderingFunction()->SetAutoMinMax(false);
  m_RenderingFilter->GetRenderingFunction()->SetChannelList(m_Channels);
  m_RenderingFilter->GetOutput()->SetRequestedRegion(region);
  m_RenderingFilter->Update();

  // Declare the iterator
  itk::ImageRegionConstIteratorWithIndex< RenderingFilterType::OutputImageType >
    it(m_RenderingFilter->GetOutput(), region);

  // Go to begin
  it.GoToBegin();

  while (!it.IsAtEnd())
    {
    // Fill the buffer
    unsigned int index = 0;
    index = ComputeXAxisFlippedBufferIndex(it.GetIndex(), m_Region);

    // Fill the buffer
    m_RasterizedBuffer[index]  = it.Get()[0];
    m_RasterizedBuffer[index + 1] = it.Get()[1];
    m_RasterizedBuffer[index + 2] = it.Get()[2];
    ++it;
    }
}

void
VectorImageModel
::ComputeRegionsToLoad(const ImageRegionType& region)
{
  // Initialize the region and clear vector
  m_AlreadyLoadedRegion = m_PreviousRegion;
  m_RegionsToLoadVector.clear();

  // 4 regions to compute
  ImageRegionType upperRegion;
  ImageRegionType lowerRegion;
  ImageRegionType leftRegion;
  ImageRegionType rightRegion;

  // local variables
  IndexType index;
  ImageRegionType::SizeType size;

  // Compute the already loaded region as a simple Crop
  bool res =  m_AlreadyLoadedRegion.Crop(region);

  // ------ upper region
  upperRegion.SetIndex(region.GetIndex());
  size[0] = region.GetSize()[0];
  size[1] = vcl_abs( region.GetIndex()[1] - m_AlreadyLoadedRegion.GetIndex()[1]);
  upperRegion.SetSize(size);

  // ------ lower region
  index[0] = region.GetIndex()[0];
  index[1] = m_AlreadyLoadedRegion.GetSize()[1] + m_AlreadyLoadedRegion.GetIndex()[1];
  lowerRegion.SetIndex(index);

  size[0] = region.GetSize()[0];
  size[1] = vcl_abs( region.GetIndex()[1] + region.GetSize()[1]
                     - m_AlreadyLoadedRegion.GetIndex()[1] - m_AlreadyLoadedRegion.GetSize()[1] );
  lowerRegion.SetSize(size);

  // ------- left region
  index[0] = region.GetIndex()[0];
  index[1] = m_AlreadyLoadedRegion.GetIndex()[1];
  leftRegion.SetIndex(index);

  size[0] = vcl_abs(region.GetIndex(0) - m_AlreadyLoadedRegion.GetIndex()[0]);
  size[1] = m_AlreadyLoadedRegion.GetSize()[1];
  leftRegion.SetSize(size);

  // -------- right region
  index[0] = m_AlreadyLoadedRegion.GetIndex()[0] + m_AlreadyLoadedRegion.GetSize()[0];
  index[1] = m_AlreadyLoadedRegion.GetIndex()[1];
  rightRegion.SetIndex(index);

  size[0] = vcl_abs(region.GetSize()[0] + region.GetIndex()[0]
                    - m_AlreadyLoadedRegion.GetIndex()[0] - m_AlreadyLoadedRegion.GetSize()[0]);
  size[1] = m_AlreadyLoadedRegion.GetSize()[1];
  rightRegion.SetSize(size);

  // add the upper region if any pixel
  if ( upperRegion .GetNumberOfPixels() > 0 )
    m_RegionsToLoadVector.push_back(upperRegion);

  // add the lower region if any pixel
  if ( lowerRegion.GetNumberOfPixels() > 0 )
    m_RegionsToLoadVector.push_back(lowerRegion);

  // add the left region if any pixel
  if ( leftRegion.GetNumberOfPixels() > 0 )
    m_RegionsToLoadVector.push_back(leftRegion);

  // add the right region if any pixel
  if ( rightRegion.GetNumberOfPixels() > 0 )
    m_RegionsToLoadVector.push_back(rightRegion);
}

/*******************************************************************************/
/* SLOTS                                                                       */
/*******************************************************************************/

/*******************************************************************************/

} // end namespace 'mvd'
<|endoftext|>