Implemented OCI mirrors (#3246)

This commit is contained in:
Johan Mabille 2024-03-26 06:20:49 +01:00 committed by GitHub
parent 901f29d9b3
commit bf378ddfb4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 329 additions and 1 deletions

View File

@ -32,7 +32,8 @@ namespace mamba
satisfiablitity_error,
user_interrupted,
incorrect_usage,
invalid_spec
invalid_spec,
download_content
};
class mamba_error : public std::runtime_error

View File

@ -41,7 +41,10 @@ namespace mamba::download
std::string url;
header_list headers;
std::string username = {};
std::string password = {};
MirrorRequest(std::string_view name, std::string_view url, header_list headers = {});
MirrorRequest(const RequestBase& base, std::string_view url, header_list headers = {});
~MirrorRequest() = default;

View File

@ -280,6 +280,16 @@ namespace mamba::download
context.remote_fetch_params.ssl_verify
);
if (!p_request->username.empty())
{
p_handle->set_opt(CURLOPT_USERNAME, p_request->username);
}
if (!p_request->password.empty())
{
p_handle->set_opt(CURLOPT_PASSWORD, p_request->password);
}
p_handle->set_opt(CURLOPT_NOBODY, p_request->check_only);
p_handle->set_opt(CURLOPT_HEADERFUNCTION, &DownloadAttempt::Impl::curl_header_callback);

View File

@ -37,6 +37,12 @@ namespace mamba::download
/*****************
* MirrorRequest *
*****************/
MirrorRequest::MirrorRequest(std::string_view lname, std::string_view lurl, header_list lheaders)
: RequestBase(lname, std::nullopt, false, false)
, url(lurl)
, headers(std::move(lheaders))
{
}
MirrorRequest::MirrorRequest(const RequestBase& base, std::string_view lurl, header_list lheaders)
: RequestBase(base)

View File

@ -4,11 +4,18 @@
//
// The full license is in the file LICENSE, distributed with this software.
#include <spdlog/spdlog.h>
#include "mamba/core/output.hpp"
#include "mamba/util/string.hpp"
#include "mamba/util/url_manip.hpp"
#include "nlohmann/json.hpp"
#include "mirror_impl.hpp"
namespace nl = nlohmann;
namespace mamba::download
{
/************************************
@ -58,6 +65,251 @@ namespace mamba::download
{ return MirrorRequest(dl_request, util::url_concat(url, dl_request.url_path)); } };
}
/****************************
* OCIMirror implementation *
****************************/
namespace
{
std::pair<std::string, std::string> split_path_tag(const std::string& path)
{
// for OCI, if we have a filename like "xtensor-0.23.10-h2acdbc0_0.tar.bz2"
// we want to split it to `xtensor:0.23.10-h2acdbc0-0`
if (util::ends_with(path, ".json"))
{
return { path, "latest" };
}
std::pair<std::string, std::string> result;
auto parts = util::rsplit(path, "-", 2);
if (parts.size() < 2)
{
LOG_ERROR << "Could not split filename into enough parts";
throw std::runtime_error("Could not split filename into enough parts");
}
result.first = parts[0];
std::string tag;
if (parts.size() > 2)
{
std::string last_part = parts[2].substr(0, parts[2].find_first_of("."));
tag = fmt::format("{}-{}", parts[1], last_part);
}
else
{
tag = parts[1];
}
util::replace_all(tag, "_", "-");
result.second = tag;
LOG_INFO << "Splitting " << path << " to name: " << result.first
<< " tag: " << result.second;
return result;
}
nl::json parse_json_nothrow(const std::string& value)
{
try
{
auto j = nl::json::parse(value);
return j;
}
catch (const nlohmann::detail::parse_error& e)
{
spdlog::error("Could not parse JSON\n{}", value);
spdlog::error("Error message: {}", e.what());
return nl::json::object();
}
}
}
OCIMirror::OCIMirror(
std::string url,
std::string repo_prefix,
std::string scope,
std::string username,
std::string password
)
: Mirror(OCIMirror::make_id(url))
, m_url(std::move(url))
, m_repo_prefix(std::move(repo_prefix))
, m_scope(std::move(scope))
, m_username(std::move(username))
, m_password(std::move(password))
, m_path_map()
{
}
MirrorID OCIMirror::make_id(std::string url)
{
return MirrorID(std::move(url));
}
auto OCIMirror::get_request_generators_impl(const std::string& url_path) const
-> request_generator_list
{
// NB: This method can be executed by many threads in parallel. Therefore,
// data should not be captured in lambda used for building the request, as
// inserting a new AuthenticationData object may relocate preexisting ones.
auto [split_path, split_tag] = split_path_tag(url_path);
auto* data = get_authentication_data(split_path);
if (!data)
{
m_path_map[split_path].reset(new AuthenticationData);
data = m_path_map[split_path].get();
}
request_generator_list req_gen;
if (data->token.empty() && need_authentication())
{
req_gen.push_back([this, split_path](const Request& dl_request, const Content*)
{ return build_authentication_request(dl_request, split_path); });
}
if (data->sha256sum.empty())
{
req_gen.push_back([this, split_path, split_tag](const Request& dl_request, const Content*)
{ return build_manifest_request(dl_request, split_path, split_tag); });
}
req_gen.push_back([this, split_path](const Request& dl_request, const Content*)
{ return build_blob_request(dl_request, split_path); });
return req_gen;
}
MirrorRequest OCIMirror::build_authentication_request(
const Request& initial_request,
const std::string& split_path
) const
{
AuthenticationData* data = get_authentication_data(split_path);
std::string auth_url = get_authentication_url(split_path);
MirrorRequest req(initial_request.name, auth_url);
req.username = m_username;
req.password = m_password;
req.on_success = [data](const Success& success) -> expected_t<void>
{
const Buffer& buf = std::get<Buffer>(success.content);
auto j = parse_json_nothrow(buf.value);
if (j.contains("token"))
{
data->token = j["token"].get<std::string>();
return expected_t<void>();
}
else
{
return make_unexpected(
"Could not retrieve authentication token",
mamba_error_code::download_content
);
}
};
return req;
}
MirrorRequest OCIMirror::build_manifest_request(
const Request& initial_request,
const std::string& split_path,
const std::string& split_tag
) const
{
AuthenticationData* data = get_authentication_data(split_path);
std::string manifest_url = get_manifest_url(split_path, split_tag);
std::vector<std::string> headers = { get_authentication_header(data->token),
"Accept: application/vnd.oci.image.manifest.v1+json" };
MirrorRequest req(initial_request.name, manifest_url, std::move(headers));
req.on_success = [data](const Success& success) -> expected_t<void>
{
const Buffer& buf = std::get<Buffer>(success.content);
auto j = parse_json_nothrow(buf.value);
if (j.contains("layers"))
{
std::string digest = j["layers"][0]["digest"];
assert(util::starts_with(digest, "sha256:"));
data->sha256sum = digest.substr(sizeof("sha256:") - 1);
return expected_t<void>();
}
else
{
return make_unexpected("Could not retrieve sha256", mamba_error_code::download_content);
}
};
return req;
}
MirrorRequest
OCIMirror::build_blob_request(const Request& initial_request, const std::string& split_path) const
{
const AuthenticationData* data = get_authentication_data(split_path);
std::string url = get_blob_url(split_path, data->sha256sum);
return MirrorRequest(initial_request, url);
}
bool OCIMirror::need_authentication() const
{
return !m_username.empty() && !m_password.empty();
}
std::string OCIMirror::get_repo(const std::string& repo) const
{
if (!m_repo_prefix.empty())
{
return fmt::format("{}/{}", m_repo_prefix, repo);
}
else
{
return repo;
}
}
std::string OCIMirror::get_authentication_url(const std::string& repo) const
{
return fmt::format("{}/token?scope=repository:{}:{}", m_url, get_repo(repo), m_scope);
}
std::string OCIMirror::get_authentication_header(const std::string& token) const
{
if (!need_authentication() || token.empty())
{
return {};
}
else
{
return fmt::format("Authorization: Bearer {}", token);
}
}
std::string OCIMirror::get_manifest_url(const std::string& repo, const std::string& reference) const
{
return fmt::format("{}/v2/{}/manifests/{}", m_url, get_repo(repo), reference);
}
std::string OCIMirror::get_blob_url(const std::string& repo, const std::string& sha256sum) const
{
// Should be this format:
// https://ghcr.io/v2/wolfv/artifact/blobs/sha256:c5be3ea75353851e1fcf3a298af3b6cfd2af3d7ff018ce52657b6dbd8f986aa4
return fmt::format("{}/v2/{}/blobs/sha256:{}", m_url, get_repo(repo), sha256sum);
}
auto OCIMirror::get_authentication_data(const std::string& split_path) const
-> AuthenticationData*
{
auto it = m_path_map.find(split_path);
if (it != m_path_map.end())
{
return it->second.get();
}
return nullptr;
}
/******************************
* make_mirror implementation *
******************************/

View File

@ -7,6 +7,8 @@
#ifndef MAMBA_DL_MIRROR_IMPL_HPP
#define MAMBA_DL_MIRROR_IMPL_HPP
#include <unordered_map>
#include "mamba/download/mirror.hpp"
namespace mamba::download
@ -43,6 +45,60 @@ namespace mamba::download
std::string m_url;
};
class OCIMirror : public Mirror
{
public:
explicit OCIMirror(
std::string url,
std::string repo_prefix,
std::string scope,
std::string username = {},
std::string password = {}
);
static MirrorID make_id(std::string url);
private:
struct AuthenticationData
{
std::string sha256sum = {}; // TODO what about other checksums types? i.e md5
std::string token = {};
};
using request_generator_list = Mirror::request_generator_list;
request_generator_list
get_request_generators_impl(const std::string& url_path) const override;
MirrorRequest
build_authentication_request(const Request& initial_request, const std::string& split_path) const;
MirrorRequest build_manifest_request(
const Request& initial_request,
const std::string& split_path,
const std::string& split_tag
) const;
MirrorRequest
build_blob_request(const Request& initial_request, const std::string& split_path) const;
bool need_authentication() const;
std::string get_repo(const std::string& repo) const;
std::string get_authentication_url(const std::string& repo) const;
std::string get_authentication_header(const std::string& token) const;
std::string get_manifest_url(const std::string& repo, const std::string& reference) const;
std::string get_blob_url(const std::string& repo, const std::string& sha256sum) const;
AuthenticationData* get_authentication_data(const std::string& split_path) const;
std::string m_url;
std::string m_repo_prefix;
std::string m_scope;
std::string m_username;
std::string m_password;
mutable std::unordered_map<std::string, std::unique_ptr<AuthenticationData>> m_path_map;
};
}
#endif