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
 *
 *     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.
 */
#pragma once
 
#include <cassert>
#include <cstddef>
#include <cstring>
#include <memory>
#include <optional>
#include <string>
 
#include <kvikio/defaults.hpp>
#include <kvikio/error.hpp>
#include <kvikio/parallel_operation.hpp>
#include <kvikio/posix_io.hpp>
#include <kvikio/utils.hpp>
 
struct curl_slist;
 
namespace kvikio {
 
class CurlHandle;  // Prototype
 
enum class RemoteEndpointType : uint8_t {
  AUTO,  
  S3,    
  S3_PRESIGNED_URL,  
  WEBHDFS,  
  HTTP,  
};
 
class RemoteEndpoint {
 protected:
  RemoteEndpointType _remote_endpoint_type{RemoteEndpointType::AUTO};
  RemoteEndpoint(RemoteEndpointType remote_endpoint_type);
 
 public:
  virtual ~RemoteEndpoint() = default;
 
  virtual void setopt(CurlHandle& curl) = 0;
 
  virtual std::string str() const = 0;
 
  virtual std::size_t get_file_size() = 0;
 
  virtual void setup_range_request(CurlHandle& curl, std::size_t file_offset, std::size_t size) = 0;
 
  [[nodiscard]] RemoteEndpointType remote_endpoint_type() const noexcept;
};
 
class HttpEndpoint : public RemoteEndpoint {
 private:
  std::string _url;
 
 public:
  HttpEndpoint(std::string url);
 
  ~HttpEndpoint() override = default;
  void setopt(CurlHandle& curl) override;
  std::string str() const override;
  std::size_t get_file_size() override;
  void setup_range_request(CurlHandle& curl, std::size_t file_offset, std::size_t size) override;
 
  static bool is_url_valid(std::string const& url) noexcept;
};
 
class S3Endpoint : public RemoteEndpoint {
 private:
  std::string _url;
  std::string _aws_sigv4;
  std::string _aws_userpwd;
  curl_slist* _curl_header_list{};
 
  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 bucket_name,
                                                std::string object_name,
                                                std::optional<std::string> 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,
             std::optional<std::string> aws_session_token     = std::nullopt);
 
  S3Endpoint(std::pair<std::string, std::string> bucket_and_object_names,
             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,
             std::optional<std::string> aws_session_token     = std::nullopt);
 
  ~S3Endpoint() override;
  void setopt(CurlHandle& curl) override;
  std::string str() const override;
  std::size_t get_file_size() override;
  void setup_range_request(CurlHandle& curl, std::size_t file_offset, std::size_t size) override;
 
  static bool is_url_valid(std::string const& url) noexcept;
};
 
class S3EndpointWithPresignedUrl : public RemoteEndpoint {
 private:
  std::string _url;
 
 public:
  explicit S3EndpointWithPresignedUrl(std::string presigned_url);
 
  ~S3EndpointWithPresignedUrl() override = default;
  void setopt(CurlHandle& curl) override;
  std::string str() const override;
  std::size_t get_file_size() override;
  void setup_range_request(CurlHandle& curl, std::size_t file_offset, std::size_t size) override;
 
  static bool is_url_valid(std::string const& url) noexcept;
};
 
class RemoteHandle {
 private:
  std::unique_ptr<RemoteEndpoint> _endpoint;
  std::size_t _nbytes;
 
 public:
  static RemoteHandle open(std::string url,
                           RemoteEndpointType remote_endpoint_type = RemoteEndpointType::AUTO,
                           std::optional<std::vector<RemoteEndpointType>> allow_list = std::nullopt,
                           std::optional<std::size_t> nbytes = std::nullopt);
 
  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]] RemoteEndpointType remote_endpoint_type() const noexcept;
 
  [[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