remote_handle.hpp

/*
 * Copyright (c) 2024-2025, NVIDIA CORPORATION.
 *
 * 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
 *
 * https://apache.ac.cn/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.
 */
#pragma once

#include <cassert>
#include <cstddef>
#include <cstring>
#include <iostream>
#include <memory>
#include <optional>
#include <regex>
#include <sstream>
#include <stdexcept>
#include <string>

#include <kvikio/defaults.hpp>
#include <kvikio/error.hpp>
#include <kvikio/parallel_operation.hpp>
#include <kvikio/posix_io.hpp>
#include <kvikio/utils.hpp>

namespace kvikio {

class CurlHandle; // Prototype

class RemoteEndpoint {
 public
 virtual void setopt(CurlHandle& curl) = 0;

 virtual std::string str() const = 0;

 virtual ~RemoteEndpoint() = default;
};

class HttpEndpoint : public RemoteEndpoint {
 private
 std::string _url;

 public
 HttpEndpoint(std::string url);
 void setopt(CurlHandle& curl) override;
 std::string str() const override;
 ~HttpEndpoint() override = default;
};

class S3Endpoint : public RemoteEndpoint {
 private
 std::string _url;
 std::string _aws_sigv4;
 std::string _aws_userpwd;

 static std::string unwrap_or_default(std::optional<std::string> aws_arg,
 std::string const& env_var,
 std::string const& err_msg = "");

 public
 static std::string url_from_bucket_and_object(std::string const& bucket_name,
 std::string const& object_name,
 std::optional<std::string> const& aws_region,
 std::optional<std::string> aws_endpoint_url);

 [[nodiscard]] static std::pair<std::string, std::string> parse_s3_url(std::string const& s3_url);

 S3Endpoint(std::string url,
 std::optional<std::string> aws_region = std::nullopt,
 std::optional<std::string> aws_access_key = std::nullopt,
 std::optional<std::string> aws_secret_access_key = std::nullopt);

 S3Endpoint(std::string const& bucket_name,
 std::string const& object_name,
 std::optional<std::string> aws_region = std::nullopt,
 std::optional<std::string> aws_access_key = std::nullopt,
 std::optional<std::string> aws_secret_access_key = std::nullopt,
 std::optional<std::string> aws_endpoint_url = std::nullopt);

 void setopt(CurlHandle& curl) override;
 std::string str() const override;
 ~S3Endpoint() override = default;
};

class RemoteHandle {
 private
 std::unique_ptr<RemoteEndpoint> _endpoint;
 std::size_t _nbytes;

 public
 RemoteHandle(std::unique_ptr<RemoteEndpoint> endpoint, std::size_t nbytes);

 RemoteHandle(std::unique_ptr<RemoteEndpoint> endpoint);

 // A remote handle is moveable but not copyable.
 RemoteHandle(RemoteHandle&& o) = default;
 RemoteHandle& operator=(RemoteHandle&& o) = default;
 RemoteHandle(RemoteHandle const&) = delete;
 RemoteHandle& operator=(RemoteHandle const&) = delete;

 [[nodiscard]] std::size_t nbytes() const noexcept;

 [[nodiscard]] RemoteEndpoint const& endpoint() const noexcept;

 std::size_t read(void* buf, std::size_t size, std::size_t file_offset = 0);

 std::future<std::size_t> pread(void* buf,
 std::size_t size,
 std::size_t file_offset = 0,
 std::size_t task_size = defaults::task_size());
};

} // namespace kvikio