diff --git a/cmake_modules/IcebergThirdpartyToolchain.cmake b/cmake_modules/IcebergThirdpartyToolchain.cmake index fef63f651..8e9cc8b0c 100644 --- a/cmake_modules/IcebergThirdpartyToolchain.cmake +++ b/cmake_modules/IcebergThirdpartyToolchain.cmake @@ -937,6 +937,14 @@ endif() if(ICEBERG_BUILD_BUNDLE) resolve_arrow_dependency() + # An installed Arrow does not export its bundled AWS SDK, which the S3 FileIO + # uses directly. Use the system SDK, which must match Arrow's. + if(ICEBERG_AWSSDK_BUNDLED AND NOT ARROW_VENDORED) + message(STATUS "Using the system AWS SDK, since Arrow is provided") + set(ICEBERG_AWSSDK_BUNDLED FALSE) + set(ICEBERG_AWSSDK_COMPONENTS core) + resolve_aws_sdk_dependency() + endif() resolve_avro_dependency() resolve_zstd_dependency() endif() diff --git a/mkdocs/docs/file-io.md b/mkdocs/docs/file-io.md index 63ce34de9..62f998b1f 100644 --- a/mkdocs/docs/file-io.md +++ b/mkdocs/docs/file-io.md @@ -63,6 +63,7 @@ each file location's scheme. | `s3.access-key-id` | `admin` | Static access key ID; must be set together with the secret key | | `s3.secret-access-key` | `password` | Static secret access key | | `s3.session-token` | `AQoDYXdzEJr...` | Session token, for temporary credentials. Ignored unless both static keys are set | +| `s3.session-token-expires-at-ms` | `1767225600000` | When the session token expires, in epoch milliseconds. Vended credentials are refreshed from the catalog five minutes before, including for files already open; once they have expired and cannot be refreshed, new operations under their prefix fail | | `client.region` | `us-east-1` | Region to sign requests for | | `s3.endpoint` | `https://127.0.0.1:9000` | Endpoint to use instead of the AWS one. When absent, the `AWS_ENDPOINT_URL_S3` / `AWS_ENDPOINT_URL` environment variables are consulted | | `s3.path-style-access` | `true` | Address buckets as a path (`endpoint/bucket`) instead of a virtual host (`bucket.endpoint`). Only takes effect together with a custom endpoint | diff --git a/src/iceberg/CMakeLists.txt b/src/iceberg/CMakeLists.txt index 717f1da4b..ebd9a6723 100644 --- a/src/iceberg/CMakeLists.txt +++ b/src/iceberg/CMakeLists.txt @@ -371,6 +371,16 @@ if(ICEBERG_BUILD_BUNDLE) endforeach() endif() + # The S3 FileIO signs with its own AWS credentials provider. + if(ICEBERG_S3) + list(APPEND ICEBERG_BUNDLE_STATIC_BUILD_INTERFACE_LIBS aws-cpp-sdk-core) + list(APPEND ICEBERG_BUNDLE_SHARED_BUILD_INTERFACE_LIBS aws-cpp-sdk-core) + if(NOT ICEBERG_AWSSDK_BUNDLED) + list(APPEND ICEBERG_BUNDLE_STATIC_INSTALL_INTERFACE_LIBS aws-cpp-sdk-core) + list(APPEND ICEBERG_BUNDLE_SHARED_INSTALL_INTERFACE_LIBS aws-cpp-sdk-core) + endif() + endif() + add_iceberg_lib(iceberg_bundle SOURCES ${ICEBERG_BUNDLE_SOURCES} diff --git a/src/iceberg/arrow/s3/arrow_s3_file_io.cc b/src/iceberg/arrow/s3/arrow_s3_file_io.cc index 24b89bdd6..a38e86421 100644 --- a/src/iceberg/arrow/s3/arrow_s3_file_io.cc +++ b/src/iceberg/arrow/s3/arrow_s3_file_io.cc @@ -19,6 +19,7 @@ #include #include +#include #include #include #include @@ -34,12 +35,15 @@ #include #include #include +#include #include #include #include #if ICEBERG_S3_ENABLED # include +# include +# include #endif #include "iceberg/arrow/arrow_io_internal.h" @@ -52,6 +56,7 @@ #include "iceberg/util/macros.h" #include "iceberg/util/property_util.h" #include "iceberg/util/string_util.h" +#include "iceberg/util/timepoint.h" namespace iceberg::arrow { @@ -295,33 +300,437 @@ Result<::arrow::fs::S3Options> ConfigureS3Options( namespace { +// Build an Arrow S3 filesystem. If a provider is supplied, the AWS signer gets +// access keys from it instead of using the keys in `properties`. Result> BuildArrowS3FileSystem( - const std::unordered_map& properties) { + const std::unordered_map& properties, + std::shared_ptr credentials_provider = nullptr) { ICEBERG_RETURN_UNEXPECTED(EnsureS3Initialized()); ICEBERG_ASSIGN_OR_RAISE(auto options, ConfigureS3Options(properties)); + if (credentials_provider != nullptr) { + options.credentials_provider = std::move(credentials_provider); + options.credentials_kind = ::arrow::fs::S3CredentialsKind::Explicit; + } ICEBERG_ARROW_ASSIGN_OR_RETURN(auto fs, ::arrow::fs::S3FileSystem::Make(options)); return std::shared_ptr<::arrow::fs::FileSystem>(std::move(fs)); } -// Rewrites any alias of `s3://` (any case — routing is case-insensitive) to -// exactly that, so locations and credential prefixes compare equal. An alias -// missing from kS3Schemes would silently stop matching its credential. -std::string CanonicalizeS3Scheme(std::string_view location) { +// Make the URI scheme lower-case for matching. Keep different schemes, such as +// `s3` and `oss`, separate because they can use different endpoints and keys. +std::string NormalizeS3Location(std::string_view location) { const auto separator = location.find("://"); if (separator != std::string_view::npos && IsS3Scheme(location.substr(0, separator))) { - return std::string("s3://").append(location.substr(separator + 3)); + // The scheme matched an ASCII alias, so only ASCII letters change. + return StringUtils::ToLower(location.substr(0, separator)) + .append(location.substr(separator)); } return std::string(location); } +// Start refreshing five minutes before the earliest token expires. +constexpr auto kRefreshLeadTime = std::chrono::minutes(5); + +// Wait this long before retrying a failed refresh. +constexpr std::chrono::milliseconds kRefreshRetryBackoff = std::chrono::seconds(30); + +// Never use a retry delay shorter than this value. +constexpr std::chrono::milliseconds kMinRefreshRetryBackoff = std::chrono::seconds(1); + +// An expired request waits this long for another request's refresh to finish. +constexpr auto kExpiredCredentialWait = std::chrono::seconds(10); + +// Return true when the config has both parts of a static key pair. +bool HasKeys(const std::unordered_map& config) { + return FindProperty(config, S3Properties::kAccessKeyId) != nullptr && + FindProperty(config, S3Properties::kSecretAccessKey) != nullptr; +} + +// Convert a complete static key pair to the AWS SDK credential object. +Aws::Auth::AWSCredentials AwsCredentialsOf( + const std::unordered_map& config) { + Aws::Auth::AWSCredentials credentials; + credentials.SetAWSAccessKeyId( + FindProperty(config, S3Properties::kAccessKeyId)->c_str()); + credentials.SetAWSSecretKey( + FindProperty(config, S3Properties::kSecretAccessKey)->c_str()); + if (const auto* session_token = FindProperty(config, S3Properties::kSessionToken); + session_token != nullptr) { + credentials.SetSessionToken(session_token->c_str()); + } + return credentials; +} + +struct ResolvedCredential { + std::string prefix; + std::unordered_map properties; + // These settings identify the client configuration independently of its + // short-lived credentials. They protect a client from receiving keys meant + // for a different endpoint, region, or addressing mode. + std::unordered_map client_settings; + std::optional credentials; +}; + +// Merge FileIO defaults with one vended credential and derive all values used +// by the client and refresh store in one pass. +ResolvedCredential ResolveCredential( + const std::unordered_map& default_properties, + const StorageCredential& credential) { + ResolvedCredential resolved; + resolved.prefix = NormalizeS3Location(credential.prefix); + resolved.properties = default_properties; + for (const auto& [key, value] : credential.config) { + resolved.properties[key] = value; + } + resolved.client_settings = resolved.properties; + for (const auto key : + {S3Properties::kAccessKeyId, S3Properties::kSecretAccessKey, + S3Properties::kSessionToken, S3Properties::kSessionTokenExpiresAtMs}) { + resolved.client_settings.erase(std::string(key)); + } + if (HasKeys(resolved.properties)) { + resolved.credentials = AwsCredentialsOf(resolved.properties); + } + return resolved; +} + +// Read the expiry of the token used by these properties. No session token +// means there is no expiry to track; a bad token expiry is an error. +Result> ParseCredentialExpiry( + const std::unordered_map& properties, + std::string_view prefix) { + // Only a non-empty session token paired with access keys needs an expiry. + const auto* token = FindProperty(properties, S3Properties::kSessionToken); + if (!HasKeys(properties) || token == nullptr || token->empty()) { + return std::nullopt; + } + const auto* value = FindProperty(properties, S3Properties::kSessionTokenExpiresAtMs); + if (value == nullptr) { + return InvalidArgument(R"(Credential "{}" has a session token but no "{}")", prefix, + S3Properties::kSessionTokenExpiresAtMs); + } + auto millis = StringUtils::ParseNumber(*value); + if (!millis.has_value()) { + return InvalidArgument( + R"(Credential "{}" has a session token but an unparseable "{}" value "{}")", + prefix, S3Properties::kSessionTokenExpiresAtMs, *value); + } + return TimePointMsFromUnixMs(*millis); +} + +// Keep the first credential for each prefix. This is also the order used when +// a duplicate prefix is returned by the catalog. +std::vector S3CredentialsOf( + const std::vector& credentials) { + std::vector s3_credentials; + std::unordered_set prefixes; + for (const auto& credential : credentials) { + if (IsS3CredentialPrefix(credential.prefix) && + prefixes.insert(NormalizeS3Location(credential.prefix)).second) { + s3_credentials.push_back(&credential); + } + } + return s3_credentials; +} + +struct CredentialSlot { + std::string prefix; + std::unordered_map settings; + std::optional credentials; + std::optional expires_at; + bool expiry_unknown = false; + + bool ExpiresBy(TimePointMs time) const { + return expires_at.has_value() && time >= *expires_at; + } +}; + +class CredentialStore { + public: + explicit CredentialStore( + std::unordered_map default_properties) + : default_properties_(std::move(default_properties)) {} + + ResolvedCredential Resolve(const StorageCredential& credential) const { + return ResolveCredential(default_properties_, credential); + } + + /// \brief Make `credentials` current, and `provider` too when one is given. + void Install(std::vector credentials, + std::optional> provider) { + std::unique_lock lock(mutex_); + if (provider.has_value()) { + provider_ = std::move(*provider); + } + InstallLocked(std::move(credentials)); + } + + std::shared_ptr SlotFor( + std::string prefix, std::unordered_map settings, + std::optional initial_credentials) { + std::unique_lock lock(mutex_); + return SlotForLocked(std::move(prefix), std::move(settings), + std::move(initial_credentials)); + } + + std::vector credentials() const { + std::shared_lock lock(mutex_); + return storage_credentials_; + } + + bool IsExpired(const std::string& prefix) const { + std::shared_lock lock(mutex_); + const auto current = current_slots_.find(prefix); + return current != current_slots_.end() && + current->second->ExpiresBy(CurrentTimePointMs()); + } + + void MaybeRefreshCredentials() { + // Only one caller loads from the catalog. Other callers keep using the + // current keys while they are still valid. + { + std::shared_lock lock(mutex_); + if (!provider_ || !RefreshDue()) { + return; + } + } + + std::unique_lock refresh_lock(refresh_mutex_, std::defer_lock); + if (!refresh_lock.try_lock()) { + { + std::shared_lock lock(mutex_); + if (!Expired()) { + return; + } + } + if (!refresh_lock.try_lock_for(kExpiredCredentialWait)) { + return; + } + } + + std::shared_ptr provider; + uint64_t generation = 0; + { + std::shared_lock lock(mutex_); + if (!provider_ || !RefreshDue()) { + return; + } + provider = provider_; + generation = credential_generation_; + } + + auto refreshed = provider->Load(); + + std::unique_lock lock(mutex_); + if (credential_generation_ != generation) { + return; // An install replaced the credentials this refresh started from. + } + const Status status = refreshed.has_value() + ? ApplyRefreshLocked(*refreshed) + : Status(::iceberg::unexpected(refreshed.error())); + if (status.has_value() && !RefreshDue()) { + return; + } + const auto delay = BackoffDelay(); + retry_refresh_at_ = std::chrono::steady_clock::now() + delay; + if (!status.has_value()) { + ICEBERG_LOG_WARN( + "Failed to refresh vended storage credentials ({}); keeping the current " + "ones and retrying in {}ms", + status.error().message, delay.count()); + } + } + + std::optional CredentialsFor(const CredentialSlot& slot) { + MaybeRefreshCredentials(); + std::shared_lock lock(mutex_); + return slot.credentials; + } + + private: + void InstallLocked(std::vector credentials) { + std::unordered_map> current_slots; + for (const auto* credential : S3CredentialsOf(credentials)) { + auto resolved = Resolve(*credential); + auto parsed_expiry = ParseCredentialExpiry(resolved.properties, credential->prefix); + if (!parsed_expiry.has_value()) { + ICEBERG_LOG_WARN("{}", parsed_expiry.error().message); + } + auto slot = SlotForLocked(std::move(resolved.prefix), + std::move(resolved.client_settings), std::nullopt); + // Clients from an earlier install may share this slot; they move to these + // keys too. + slot->credentials = std::move(resolved.credentials); + slot->expires_at = parsed_expiry.value_or(std::nullopt); + slot->expiry_unknown = !parsed_expiry.has_value(); + current_slots.emplace(slot->prefix, std::move(slot)); + } + retry_refresh_at_ = {}; + ++credential_generation_; + current_slots_ = std::move(current_slots); + storage_credentials_ = std::move(credentials); + } + + Status ApplyRefreshLocked(const std::vector& credentials) { + // Check every credential, skipped ones included, before changing any slot. + for (const auto& credential : credentials) { + ICEBERG_RETURN_UNEXPECTED(credential.Validate()); + } + bool updated = false; + for (const auto* credential : S3CredentialsOf(credentials)) { + auto resolved = Resolve(*credential); + const auto& prefix = resolved.prefix; + const auto current = current_slots_.find(prefix); + if (current == current_slots_.end()) { + continue; + } + const auto& settings = resolved.client_settings; + const bool has_keys = resolved.credentials.has_value(); + auto& slot = *current->second; + if (slot.settings != settings || slot.credentials.has_value() != has_keys) { + ICEBERG_LOG_WARN( + "Ignoring refreshed credentials for '{}' because their client settings " + "changed", + prefix); + continue; + } + auto parsed_expiry = ParseCredentialExpiry(resolved.properties, credential->prefix); + if (!parsed_expiry.has_value()) { + ICEBERG_LOG_WARN("Ignoring refreshed credentials for '{}': {}", prefix, + parsed_expiry.error().message); + continue; + } + if (parsed_expiry->has_value() && **parsed_expiry <= CurrentTimePointMs()) { + ICEBERG_LOG_WARN("Ignoring already expired refreshed credentials for '{}'", + prefix); + continue; + } + slot.credentials = std::move(resolved.credentials); + slot.expires_at = *parsed_expiry; + slot.expiry_unknown = false; + for (auto& installed : storage_credentials_) { + if (NormalizeS3Location(installed.prefix) == prefix) { + installed = *credential; + break; + } + } + updated = true; + } + if (!updated && !current_slots_.empty()) { + return InvalidArgument("Refreshed credentials did not update any active prefix"); + } + ++credential_generation_; + return {}; + } + + std::shared_ptr SlotForLocked( + std::string prefix, std::unordered_map settings, + std::optional initial_credentials) { + for (const auto& slot : known_slots_) { + if (slot->prefix == prefix && slot->settings == settings) { + return slot; + } + } + auto slot = std::make_shared( + CredentialSlot{.prefix = std::move(prefix), + .settings = std::move(settings), + .credentials = std::move(initial_credentials)}); + known_slots_.push_back(slot); + return slot; + } + + bool RefreshDue() const { + if (std::chrono::steady_clock::now() < retry_refresh_at_) { + return false; + } + const auto refresh_by = CurrentTimePointMs() + kRefreshLeadTime; + return std::ranges::any_of(current_slots_, [refresh_by](const auto& entry) { + return entry.second->expiry_unknown || entry.second->ExpiresBy(refresh_by); + }); + } + + bool Expired() const { + const auto now = CurrentTimePointMs(); + return std::ranges::any_of(current_slots_, [now](const auto& entry) { + return entry.second->ExpiresBy(now); + }); + } + + // Retry after kRefreshRetryBackoff, or sooner if a token expires first, but + // never sooner than kMinRefreshRetryBackoff. + std::chrono::milliseconds BackoffDelay() const { + std::optional earliest; + for (const auto& [_, slot] : current_slots_) { + if (slot->expires_at.has_value() && + (!earliest.has_value() || *slot->expires_at < *earliest)) { + earliest = slot->expires_at; + } + } + const auto now = CurrentTimePointMs(); + if (earliest.has_value() && *earliest > now) { + return std::clamp( + *earliest - now, kMinRefreshRetryBackoff, kRefreshRetryBackoff); + } + return kRefreshRetryBackoff; + } + + const std::unordered_map default_properties_; + mutable std::shared_mutex mutex_; + std::vector storage_credentials_; + std::vector> known_slots_; + std::unordered_map> current_slots_; + std::shared_ptr provider_; + std::chrono::steady_clock::time_point retry_refresh_at_; + uint64_t credential_generation_ = 0; + std::timed_mutex refresh_mutex_; +}; + +// Supplies one S3 client's keys from its slot in the shared store, which lets +// clients and handles outlive the FileIO. The last keys seen stay in use while +// the slot has none, as after the prefix is reinstalled without keys. +class VendedCredentialsProvider final : public Aws::Auth::AWSCredentialsProvider { + public: + VendedCredentialsProvider(std::shared_ptr store, + std::shared_ptr> installed, + std::shared_ptr slot, + Aws::Auth::AWSCredentials credentials) + : store_(std::move(store)), + installed_(std::move(installed)), + slot_(std::move(slot)), + credentials_(std::move(credentials)) {} + + Aws::Auth::AWSCredentials GetAWSCredentials() override { + // The SDK also asks while building the client. Until its install, the slot + // may hold the keys being replaced, and asking could start their refresh. + if (!installed_->load()) { + std::lock_guard lock(mutex_); + return credentials_; + } + auto current = store_->CredentialsFor(*slot_); + std::lock_guard lock(mutex_); + if (current.has_value()) { + credentials_ = std::move(current).value(); + } + return credentials_; + } + + private: + const std::shared_ptr store_; + // Set once the install that built this client has made its keys current. + const std::shared_ptr> installed_; + const std::shared_ptr slot_; + std::mutex mutex_; + Aws::Auth::AWSCredentials credentials_; +}; + class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials { public: ArrowS3FileIO(std::shared_ptr<::arrow::fs::FileSystem> arrow_fs, std::unordered_map default_properties, size_t delete_parallelism) : default_file_io_(std::make_shared(std::move(arrow_fs))), - default_properties_(std::move(default_properties)), - delete_parallelism_(delete_parallelism) {} + delete_parallelism_(delete_parallelism), + credential_store_( + std::make_shared(std::move(default_properties))) {} Result> NewInputFile(std::string file_location) override; @@ -338,87 +747,100 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials { const std::vector& storage_credentials) override; std::vector credentials() const override { - std::shared_lock lock(mutex_); - return storage_credentials_; + return credential_store_->credentials(); } + Status InitializeStorageCredentials( + const std::vector& storage_credentials, + std::shared_ptr provider) override; + SupportsStorageCredentials* AsSupportsStorageCredentials() override { return this; } private: - /// \brief Delegate serving `location`, pinned by the caller against a - /// concurrent credential install. - std::shared_ptr FileIOForPath(std::string_view location); - using DelegatesByPrefix = std::vector>>; - /// \brief Longest-prefix match against one consistent view of the delegates. - static std::shared_ptr MatchDelegate( - const std::shared_ptr& fallback, - const DelegatesByPrefix& by_prefix, std::string_view location); + /// \brief Delegate with the longest prefix of `location`, pinned by the caller + /// against a concurrent credential install. Fails with TokenExpired once that + /// prefix's keys have expired. + Result> FileIOForPath(std::string_view location); + + /// \brief Build delegates for `storage_credentials` and make them current, + /// replacing the refresh provider too when one is given. + Status InstallStorageCredentials( + const std::vector& storage_credentials, + std::optional> provider); - /// \brief Build a delegate for each credential this FileIO can serve. + /// \brief Build a delegate for each credential this FileIO can serve. Their + /// clients take keys from the store once `installed` is set. /// /// Runs without holding `mutex_`: building an S3 client can block (without /// static keys the AWS SDK may wait on the EC2 metadata service), which would /// stall every concurrent operation. Reads no mutable member state. Result BuildDelegates( - const std::vector& storage_credentials) const; - - /// \brief Swap in credentials and delegates, handing back the retired ones. - /// - /// Callers must hold `mutex_` exclusively and let the returned generation - /// destruct only after releasing it: tearing down an S3 client can block on - /// in-flight requests, which would stall every operation. - void InstallCredentials(std::vector& storage_credentials, - DelegatesByPrefix& delegates); + const std::vector& storage_credentials, + const std::shared_ptr>& installed) const; std::shared_ptr default_file_io_; - std::unordered_map default_properties_; const size_t delete_parallelism_; + const std::shared_ptr credential_store_; // Guards everything below; shared because reads happen per file operation. mutable std::shared_mutex mutex_; - std::vector storage_credentials_; DelegatesByPrefix file_io_by_prefix_; }; +Status ArrowS3FileIO::InitializeStorageCredentials( + const std::vector& storage_credentials, + std::shared_ptr provider) { + return InstallStorageCredentials(storage_credentials, std::move(provider)); +} + Status ArrowS3FileIO::SetStorageCredentials( const std::vector& storage_credentials) { - ICEBERG_ASSIGN_OR_RAISE(auto delegates, BuildDelegates(storage_credentials)); + return InstallStorageCredentials(storage_credentials, std::nullopt); +} + +Status ArrowS3FileIO::InstallStorageCredentials( + const std::vector& storage_credentials, + std::optional> provider) { + auto installed = std::make_shared>(false); + ICEBERG_ASSIGN_OR_RAISE(auto delegates, BuildDelegates(storage_credentials, installed)); auto credentials = storage_credentials; { std::unique_lock lock(mutex_); - InstallCredentials(credentials, delegates); + file_io_by_prefix_.swap(delegates); + credential_store_->Install(std::move(credentials), std::move(provider)); + installed->store(true); } - // `credentials` and `delegates` now hold the retired generation and destruct - // here, outside the lock. + // `delegates` now holds the retired clients and destructs here, outside the + // lock: tearing down an S3 client can block on its in-flight requests. return {}; } Result ArrowS3FileIO::BuildDelegates( - const std::vector& storage_credentials) const { + const std::vector& storage_credentials, + const std::shared_ptr>& installed) const { DelegatesByPrefix delegates; delegates.reserve(storage_credentials.size()); - // TODO(gangwu): Refresh vended credentials via credentials.uri before tokens expire. for (const auto& credential : storage_credentials) { ICEBERG_RETURN_UNEXPECTED(credential.Validate()); - // A server may vend credentials for several storage systems at once; - // non-S3 prefixes are skipped, not rejected (Java S3FileIO filters - // credentials by the "s3" prefix). - if (!IsS3CredentialPrefix(credential.prefix)) { - continue; - } - auto properties = default_properties_; - for (const auto& [key, value] : credential.config) { - properties[key] = value; + } + for (const auto* credential : S3CredentialsOf(storage_credentials)) { + auto resolved = credential_store_->Resolve(*credential); + std::shared_ptr keys; + if (resolved.credentials.has_value()) { + auto slot = credential_store_->SlotFor( + resolved.prefix, std::move(resolved.client_settings), resolved.credentials); + keys = std::make_shared( + credential_store_, installed, std::move(slot), + *std::move(resolved.credentials)); } - ICEBERG_ASSIGN_OR_RAISE(auto fs, BuildArrowS3FileSystem(properties)); - delegates.emplace_back(CanonicalizeS3Scheme(credential.prefix), + ICEBERG_ASSIGN_OR_RAISE(auto fs, + BuildArrowS3FileSystem(resolved.properties, std::move(keys))); + delegates.emplace_back(std::move(resolved.prefix), std::make_shared(std::move(fs))); } if (delegates.empty() && !storage_credentials.empty()) { - // Silent skipping of every vended credential is hard to diagnose: S3 access - // would proceed with the default credentials and fail only at IO time. ICEBERG_LOG_WARN( "None of the {} vended storage credential(s) has an S3-compatible prefix; " "S3 access will use the default credentials", @@ -427,69 +849,63 @@ Result ArrowS3FileIO::BuildDelegates( return delegates; } -void ArrowS3FileIO::InstallCredentials( - std::vector& storage_credentials, DelegatesByPrefix& delegates) { - file_io_by_prefix_.swap(delegates); - storage_credentials_.swap(storage_credentials); -} - -std::shared_ptr ArrowS3FileIO::MatchDelegate( - const std::shared_ptr& fallback, - const DelegatesByPrefix& by_prefix, std::string_view location) { - if (by_prefix.empty()) { - return fallback; - } - const std::string canonical = CanonicalizeS3Scheme(location); - auto best = fallback; - size_t best_len = 0; - for (const auto& [prefix, file_io] : by_prefix) { - if (prefix.size() > best_len && canonical.starts_with(prefix)) { - best = file_io; - best_len = prefix.size(); - } - } - return best; -} - -std::shared_ptr ArrowS3FileIO::FileIOForPath( +Result> ArrowS3FileIO::FileIOForPath( std::string_view location) { + credential_store_->MaybeRefreshCredentials(); std::shared_lock lock(mutex_); - return MatchDelegate(default_file_io_, file_io_by_prefix_, location); + if (file_io_by_prefix_.empty()) { + return default_file_io_; + } + const std::string normalized = NormalizeS3Location(location); + const DelegatesByPrefix::value_type* best = nullptr; + for (const auto& delegate : file_io_by_prefix_) { + const auto& prefix = delegate.first; + if (prefix.size() > (best != nullptr ? best->first.size() : 0) && + normalized.starts_with(prefix)) { + best = &delegate; + } + } + if (best == nullptr) { + return default_file_io_; + } + if (credential_store_->IsExpired(best->first)) { + return TokenExpired( + "Vended credentials for \"{}\" have expired and could not be refreshed", + best->first); + } + return best->second; } Result> ArrowS3FileIO::NewInputFile( std::string file_location) { - return FileIOForPath(file_location)->NewInputFile(std::move(file_location)); + ICEBERG_ASSIGN_OR_RAISE(auto file_io, FileIOForPath(file_location)); + return file_io->NewInputFile(std::move(file_location)); } Result> ArrowS3FileIO::NewInputFile(std::string file_location, size_t length) { - return FileIOForPath(file_location)->NewInputFile(std::move(file_location), length); + ICEBERG_ASSIGN_OR_RAISE(auto file_io, FileIOForPath(file_location)); + return file_io->NewInputFile(std::move(file_location), length); } Result> ArrowS3FileIO::NewOutputFile( std::string file_location) { - return FileIOForPath(file_location)->NewOutputFile(std::move(file_location)); + ICEBERG_ASSIGN_OR_RAISE(auto file_io, FileIOForPath(file_location)); + return file_io->NewOutputFile(std::move(file_location)); } Status ArrowS3FileIO::DeleteFile(const std::string& file_location) { - return FileIOForPath(file_location)->DeleteFile(file_location); + ICEBERG_ASSIGN_OR_RAISE(auto file_io, FileIOForPath(file_location)); + return file_io->DeleteFile(file_location); } Status ArrowS3FileIO::DeleteFiles(const std::vector& file_locations) { - // Keep one delegate snapshot for the whole batch. - std::shared_ptr fallback; - DelegatesByPrefix by_prefix; - { - std::shared_lock lock(mutex_); - fallback = default_file_io_; - by_prefix = file_io_by_prefix_; - } - return BulkDeleteFiles( - file_locations, delete_executor.Get(), delete_parallelism_, - [&](const std::string& location) { - return MatchDelegate(fallback, by_prefix, location)->DeleteFile(location); - }); + // Each file is resolved afresh, as a long batch may outlive its credentials. + return BulkDeleteFiles(file_locations, delete_executor.Get(), delete_parallelism_, + [this](const std::string& location) -> Status { + ICEBERG_ASSIGN_OR_RAISE(auto file_io, FileIOForPath(location)); + return file_io->DeleteFile(location); + }); } } // namespace diff --git a/src/iceberg/arrow/s3/s3_properties.h b/src/iceberg/arrow/s3/s3_properties.h index fe0a9aa3c..79609c878 100644 --- a/src/iceberg/arrow/s3/s3_properties.h +++ b/src/iceberg/arrow/s3/s3_properties.h @@ -42,6 +42,9 @@ struct S3Properties { static constexpr std::string_view kSecretAccessKey = "s3.secret-access-key"; /// AWS session token (for temporary credentials) static constexpr std::string_view kSessionToken = "s3.session-token"; + /// Epoch milliseconds at which a vended session token stops being valid + static constexpr std::string_view kSessionTokenExpiresAtMs = + "s3.session-token-expires-at-ms"; /// AWS region, standard Iceberg client property. static constexpr std::string_view kClientRegion = "client.region"; /// Custom endpoint override (for S3-compatible object stores) @@ -58,9 +61,11 @@ struct S3Properties { static constexpr std::string_view kDeleteNumThreads = "s3.delete.num-threads"; }; -/// \brief URI schemes served by the Arrow S3 FileIO, lower-case. +/// \brief URI schemes served by the Arrow S3 FileIO, lower-case. These schemes +/// share the S3 client implementation but remain distinct when routing +/// credentials. /// -/// Single source of truth: registration, IsS3Scheme and alias canonicalization +/// Single source of truth: registration, IsS3Scheme and alias normalization /// all derive from this list, so a new alias only has to be added here. /// /// `oss` is served because the store is S3-compatible; see the FileIO docs. @@ -88,7 +93,7 @@ inline constexpr bool IsS3Scheme(std::string_view scheme) { /// \brief Return whether a storage credential prefix belongs to S3. /// /// Accepts the bare `s3` prefix (exactly: a case variant would be stored as a -/// key no canonicalized location can match) or a URI prefix such as +/// key no normalized location can match) or a URI prefix such as /// `s3a://bucket`. Bare aliases such as `oss` are deliberately rejected, /// matching Java S3FileIO's credential filter. inline constexpr bool IsS3CredentialPrefix(std::string_view prefix) { diff --git a/src/iceberg/catalog/rest/catalog_properties.h b/src/iceberg/catalog/rest/catalog_properties.h index da5bd550a..a1f205a60 100644 --- a/src/iceberg/catalog/rest/catalog_properties.h +++ b/src/iceberg/catalog/rest/catalog_properties.h @@ -72,6 +72,11 @@ class ICEBERG_REST_EXPORT RestCatalogProperties "rest-scan-planning.poll-num-retries", 10}; /// \brief The prefix for HTTP headers. inline static constexpr std::string_view kHeaderPrefix = "header."; + /// \brief HTTP connection timeout in milliseconds, matching Java REST clients. + inline static Entry kClientConnectionTimeoutMs{ + "rest.client.connection-timeout-ms", 0}; + /// \brief HTTP socket/read timeout in milliseconds, matching Java REST clients. + inline static Entry kClientSocketTimeoutMs{"rest.client.socket-timeout-ms", 0}; /// \brief Create a default RestCatalogProperties instance. static RestCatalogProperties default_properties(); diff --git a/src/iceberg/catalog/rest/http_client.cc b/src/iceberg/catalog/rest/http_client.cc index d48471cde..871a21e67 100644 --- a/src/iceberg/catalog/rest/http_client.cc +++ b/src/iceberg/catalog/rest/http_client.cc @@ -146,9 +146,27 @@ ErrorResponse BuildDefaultErrorResponse(const cpr::Response& response) { }; } +/// \brief Fails a response whose transfer failed, even after a success status. +Status CheckTransfer(const cpr::Response& response) { + if (response.error.code == cpr::ErrorCode::OPERATION_TIMEDOUT) { + return IOError("HTTP request timed out: {}", response.error.message); + } + if (response.error) { + if (IsSuccessful(response.status_code)) { + return IOError("HTTP request failed after status {}: {}", response.status_code, + response.error.message); + } + return IOError("HTTP request failed: {}", response.error.message); + } + return {}; +} + /// \brief Handles failure responses by invoking the provided error handler. +/// +/// A failed transfer is checked first, as its status code may be 0 or a success. Status HandleFailureResponse(const cpr::Response& response, const ErrorHandler& error_handler) { + ICEBERG_RETURN_UNEXPECTED(CheckTransfer(response)); if (IsSuccessful(response.status_code)) { return {}; } @@ -161,8 +179,12 @@ Status HandleFailureResponse(const cpr::Response& response, } // namespace -HttpClient::HttpClient(std::unordered_map default_headers) +HttpClient::HttpClient(std::unordered_map default_headers, + std::chrono::milliseconds connection_timeout, + std::chrono::milliseconds socket_timeout) : default_headers_{std::move(default_headers)}, + connection_timeout_{connection_timeout}, + socket_timeout_{socket_timeout}, connection_pool_{std::make_unique()} { // Set default Content-Type for all requests (including GET/HEAD/DELETE). // Many systems require that content type is set regardless and will fail, @@ -181,8 +203,9 @@ Result HttpClient::Get( ICEBERG_ASSIGN_OR_RAISE(auto authenticated, AuthenticateRequest(session, HttpMethod::kGet, std::move(url), MergeHeaders(default_headers_, headers))); - cpr::Response response = cpr::Get(cpr::Url{authenticated.url}, - ToCprHeader(authenticated), *connection_pool_); + cpr::Response response = + cpr::Get(cpr::Url{authenticated.url}, ToCprHeader(authenticated), *connection_pool_, + cpr::ConnectTimeout{connection_timeout_}, cpr::Timeout{socket_timeout_}); ICEBERG_RETURN_UNEXPECTED(HandleFailureResponse(response, error_handler)); HttpResponse http_response; @@ -200,7 +223,8 @@ Result HttpClient::Post( MergeHeaders(default_headers_, headers), body)); cpr::Response response = cpr::Post(cpr::Url{authenticated.url}, cpr::Body{authenticated.body}, - ToCprHeader(authenticated), *connection_pool_); + ToCprHeader(authenticated), *connection_pool_, + cpr::ConnectTimeout{connection_timeout_}, cpr::Timeout{socket_timeout_}); ICEBERG_RETURN_UNEXPECTED(HandleFailureResponse(response, error_handler)); HttpResponse http_response; @@ -230,7 +254,8 @@ Result HttpClient::PostForm( std::move(encoded_body))); cpr::Response response = cpr::Post(cpr::Url{authenticated.url}, cpr::Body{authenticated.body}, - ToCprHeader(authenticated), *connection_pool_); + ToCprHeader(authenticated), *connection_pool_, + cpr::ConnectTimeout{connection_timeout_}, cpr::Timeout{socket_timeout_}); ICEBERG_RETURN_UNEXPECTED(HandleFailureResponse(response, error_handler)); HttpResponse http_response; @@ -244,8 +269,9 @@ Result HttpClient::Head( ICEBERG_ASSIGN_OR_RAISE(auto authenticated, AuthenticateRequest(session, HttpMethod::kHead, path, MergeHeaders(default_headers_, headers))); - cpr::Response response = cpr::Head(cpr::Url{authenticated.url}, - ToCprHeader(authenticated), *connection_pool_); + cpr::Response response = cpr::Head( + cpr::Url{authenticated.url}, ToCprHeader(authenticated), *connection_pool_, + cpr::ConnectTimeout{connection_timeout_}, cpr::Timeout{socket_timeout_}); ICEBERG_RETURN_UNEXPECTED(HandleFailureResponse(response, error_handler)); HttpResponse http_response; @@ -262,8 +288,9 @@ Result HttpClient::Delete( auto authenticated, AuthenticateRequest(session, HttpMethod::kDelete, std::move(url), MergeHeaders(default_headers_, headers))); - cpr::Response response = cpr::Delete(cpr::Url{authenticated.url}, - ToCprHeader(authenticated), *connection_pool_); + cpr::Response response = cpr::Delete( + cpr::Url{authenticated.url}, ToCprHeader(authenticated), *connection_pool_, + cpr::ConnectTimeout{connection_timeout_}, cpr::Timeout{socket_timeout_}); ICEBERG_RETURN_UNEXPECTED(HandleFailureResponse(response, error_handler)); HttpResponse http_response; diff --git a/src/iceberg/catalog/rest/http_client.h b/src/iceberg/catalog/rest/http_client.h index 9400202dc..a18506404 100644 --- a/src/iceberg/catalog/rest/http_client.h +++ b/src/iceberg/catalog/rest/http_client.h @@ -19,6 +19,7 @@ #pragma once +#include #include #include #include @@ -71,7 +72,10 @@ class ICEBERG_REST_EXPORT HttpResponse { /// \brief HTTP client for making requests to Iceberg REST Catalog API. class ICEBERG_REST_EXPORT HttpClient { public: - explicit HttpClient(std::unordered_map default_headers = {}); + /// \param connection_timeout,socket_timeout Zero leaves requests unbounded. + explicit HttpClient(std::unordered_map default_headers = {}, + std::chrono::milliseconds connection_timeout = {}, + std::chrono::milliseconds socket_timeout = {}); virtual ~HttpClient(); HttpClient(const HttpClient&) = delete; @@ -112,6 +116,8 @@ class ICEBERG_REST_EXPORT HttpClient { private: std::unordered_map default_headers_; + std::chrono::milliseconds connection_timeout_; + std::chrono::milliseconds socket_timeout_; std::unique_ptr connection_pool_; }; diff --git a/src/iceberg/catalog/rest/rest_catalog.cc b/src/iceberg/catalog/rest/rest_catalog.cc index c45067310..a945e25d4 100644 --- a/src/iceberg/catalog/rest/rest_catalog.cc +++ b/src/iceberg/catalog/rest/rest_catalog.cc @@ -19,6 +19,9 @@ #include "iceberg/catalog/rest/rest_catalog.h" +#include +#include +#include #include #include #include @@ -129,6 +132,18 @@ std::string RestTableName(std::string_view catalog_name, return std::string(catalog_name) + '.' + identifier.ToString(); } +/// \brief Report a write whose response transfer failed as having an unknown +/// outcome, since the server may have applied it. +Result UnknownIfTransferFailed(Result response, + std::string_view operation, + std::string_view state) { + if (!response.has_value() && response.error().kind == ErrorKind::kIOError) { + return CommitStateUnknown("{} response transfer failed; the {} state is unknown: {}", + operation, state, response.error().message); + } + return response; +} + /// \brief Fetch server configuration from the REST catalog server. Result FetchServerConfig(HttpClient& client, const ResourcePaths& paths, const RestCatalogProperties& current_config, @@ -437,7 +452,14 @@ Result> RestCatalog::Make( if (!client_factory) { client_factory = [](const RestCatalogProperties& properties) -> Result> { - return std::make_shared(properties.ExtractHeaders()); + // A value that is not positive leaves requests unbounded. + const auto timeout_of = [&properties](const auto& entry) { + return std::chrono::milliseconds(std::max(0, properties.Get(entry))); + }; + return std::make_shared( + properties.ExtractHeaders(), + timeout_of(RestCatalogProperties::kClientConnectionTimeoutMs), + timeout_of(RestCatalogProperties::kClientSocketTimeoutMs)); }; } ICEBERG_ASSIGN_OR_RAISE(auto uri, config.Uri()); @@ -777,9 +799,11 @@ Result RestCatalog::CreateTableInternal( ICEBERG_ASSIGN_OR_RAISE(auto request_json, ToJson(request)); ICEBERG_ASSIGN_OR_RAISE(auto json_request, ToJsonString(request_json)); - ICEBERG_ASSIGN_OR_RAISE(const auto response, - client_->Post(path, json_request, /*headers=*/{}, - *TableErrorHandler::Instance(), session)); + ICEBERG_ASSIGN_OR_RAISE( + const auto response, + UnknownIfTransferFailed(client_->Post(path, json_request, /*headers=*/{}, + *TableErrorHandler::Instance(), session), + "Create table", "table")); ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body())); ICEBERG_ASSIGN_OR_RAISE(auto load_result, LoadTableResultFromJson(json)); @@ -828,7 +852,9 @@ Result RestCatalog::UpdateTableInternal( } ICEBERG_ASSIGN_OR_RAISE( const auto response, - client_->Post(path, json_request, /*headers=*/{}, *error_handler, session)); + UnknownIfTransferFailed( + client_->Post(path, json_request, /*headers=*/{}, *error_handler, session), + "Commit", "commit")); ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body())); ICEBERG_ASSIGN_OR_RAISE(auto commit_response, CommitTableResponseFromJson(json)); diff --git a/src/iceberg/test/CMakeLists.txt b/src/iceberg/test/CMakeLists.txt index 292132bb6..48191ef37 100644 --- a/src/iceberg/test/CMakeLists.txt +++ b/src/iceberg/test/CMakeLists.txt @@ -328,6 +328,7 @@ if(ICEBERG_BUILD_REST) catalog_properties_test.cc error_handlers_test.cc endpoint_test.cc + http_client_test.cc rest_catalog_unit_test.cc rest_file_io_test.cc rest_json_serde_test.cc diff --git a/src/iceberg/test/arrow_s3_file_io_test.cc b/src/iceberg/test/arrow_s3_file_io_test.cc index 6bb3f4178..9a5eace0d 100644 --- a/src/iceberg/test/arrow_s3_file_io_test.cc +++ b/src/iceberg/test/arrow_s3_file_io_test.cc @@ -22,10 +22,12 @@ #include #include #include +#include #include #include #include #include +#include #include #include #include @@ -55,7 +57,9 @@ #include "iceberg/test/executor.h" #include "iceberg/test/logging_test_helpers.h" #include "iceberg/test/matchers.h" +#include "iceberg/test/mock_io.h" #include "iceberg/util/macros.h" +#include "iceberg/util/timepoint.h" namespace { @@ -157,10 +161,27 @@ class ArrowS3FileIOTest : public ::testing::Test { std::optional base_uri_; }; -bool HasWarning(const CapturingLogger& logger) { +bool HasWarning(const CapturingLogger& logger, std::string_view substring = {}) { const auto records = logger.records(); - return std::ranges::any_of( - records, [](const LogMessage& record) { return record.level == LogLevel::kWarn; }); + return std::ranges::any_of(records, [substring](const LogMessage& record) { + return record.level == LogLevel::kWarn && + record.message.find(substring) != std::string::npos; + }); +} + +std::string ExpiresAtMs(std::chrono::milliseconds valid_for) { + return std::to_string(UnixMsFromTimePointMs(CurrentTimePointMs() + valid_for)); +} + +std::vector ExpiringCredentials(std::chrono::milliseconds valid_for, + std::string_view access_key, + std::string prefix = "s3") { + return {{.prefix = std::move(prefix), + .config = {{std::string(S3Properties::kAccessKeyId), std::string(access_key)}, + {std::string(S3Properties::kSecretAccessKey), "secret"}, + {std::string(S3Properties::kSessionToken), "token"}, + {std::string(S3Properties::kSessionTokenExpiresAtMs), + ExpiresAtMs(valid_for)}}}}; } Status CheckReadWrite(FileIO& io, const std::string& object_uri, @@ -171,6 +192,16 @@ Status CheckReadWrite(FileIO& io, const std::string& object_uri, return io.DeleteFile(object_uri); } +// Reads `length` bytes of `file` through a fresh stream. +Result ReadAll(InputFile& file, size_t length) { + ICEBERG_ASSIGN_OR_RAISE(auto stream, file.Open()); + std::string content(length, '\0'); + ICEBERG_RETURN_UNEXPECTED( + stream->ReadFully(0, std::as_writable_bytes(std::span(content)))); + ICEBERG_RETURN_UNEXPECTED(stream->Close()); + return content; +} + class InlineDeleteExecutor : public Executor { public: Status Submit(ExecutorTask task) override { @@ -510,6 +541,422 @@ TEST_F(ArrowS3FileIOTest, WarnsWhenNoCredentialApplies) { EXPECT_TRUE(HasWarning(*logger)); } +TEST_F(ArrowS3FileIOTest, RefreshesOnlyCredentialsCloseToExpiry) { + struct Case { + std::string_view name; + std::vector initial; + bool refreshes; + }; + const std::vector cases = { + {.name = "close to expiry", + .initial = ExpiringCredentials(std::chrono::minutes(1), "expiring-key"), + .refreshes = true}, + {.name = "not close to expiry", + .initial = ExpiringCredentials(std::chrono::hours(1), "valid-key"), + .refreshes = false}, + {.name = "static keys", + .initial = {{.prefix = "s3", .config = BadS3Credentials()}}, + .refreshes = false}}; + for (const auto& test_case : cases) { + SCOPED_TRACE(test_case.name); + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO({})); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + const auto refreshed = ExpiringCredentials(std::chrono::hours(1), "refreshed-key"); + auto provider = std::make_shared(); + ON_CALL(*provider, Load()).WillByDefault(::testing::Return(refreshed)); + EXPECT_CALL(*provider, Load()).Times(test_case.refreshes ? 1 : 0); + ASSERT_THAT(credentialed->InitializeStorageCredentials(test_case.initial, provider), + IsOk()); + + // The replacement is not close to expiry, so the second does not refresh. + EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(credentialed->credentials(), + test_case.refreshes ? refreshed : test_case.initial); + } +} + +TEST_F(ArrowS3FileIOTest, KeepsProviderWhenCredentialsAreUpdated) { + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO({})); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + const auto refreshed = ExpiringCredentials(std::chrono::hours(2), "refreshed-key"); + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()).WillOnce(::testing::Return(refreshed)); + ASSERT_THAT(credentialed->InitializeStorageCredentials( + ExpiringCredentials(std::chrono::hours(1), "initial-key"), provider), + IsOk()); + ASSERT_THAT(credentialed->SetStorageCredentials( + ExpiringCredentials(std::chrono::minutes(1), "expiring-key")), + IsOk()); + + EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(credentialed->credentials(), refreshed); +} + +TEST_F(ArrowS3FileIOTest, FailedInitializationDoesNotInstallProvider) { + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO({})); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + const auto initial = ExpiringCredentials(std::chrono::minutes(1), "initial-key"); + ASSERT_THAT(credentialed->InitializeStorageCredentials(initial, nullptr), IsOk()); + auto invalid = initial; + invalid.front().prefix.clear(); + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()).Times(0); + EXPECT_THAT(credentialed->InitializeStorageCredentials(invalid, provider), + IsError(ErrorKind::kValidationFailed)); + + EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(credentialed->credentials(), initial); +} + +TEST_F(ArrowS3FileIOTest, RefreshesOnceWhenCredentialsHaveExpired) { + auto result = MakeS3FileIO({}); + ASSERT_THAT(result, IsOk()); + auto* credentialed = result.value()->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + std::mutex mutex; + std::condition_variable cv; + bool refresh_started = false; + bool release_refresh = false; + std::atomic refresh_calls = 0; + + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()) + .WillRepeatedly([&]() -> Result> { + ++refresh_calls; + std::unique_lock lock(mutex); + refresh_started = true; + cv.notify_all(); + cv.wait(lock, [&] { return release_refresh; }); + return ExpiringCredentials(std::chrono::hours(1), "refreshed-key"); + }); + ASSERT_THAT(credentialed->InitializeStorageCredentials( + ExpiringCredentials(-std::chrono::minutes(1), "expired-key"), provider), + IsOk()); + + std::thread winner( + [&] { EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk()); }); + { + std::unique_lock lock(mutex); + cv.wait(lock, [&] { return refresh_started; }); + } + + bool loser_ready = false; + std::thread loser([&] { + { + std::lock_guard lock(mutex); + loser_ready = true; + } + cv.notify_all(); + EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk()); + }); + { + std::unique_lock lock(mutex); + cv.wait(lock, [&] { return loser_ready; }); + release_refresh = true; + } + cv.notify_all(); + winner.join(); + loser.join(); + + EXPECT_EQ(refresh_calls, 1); + EXPECT_THAT(credentialed->credentials(), ::testing::Not(::testing::IsEmpty())); +} + +TEST_F(ArrowS3FileIOTest, RefreshDoesNotUndoCredentialsInstalledWhileItRan) { + auto result = MakeS3FileIO({}); + ASSERT_THAT(result, IsOk()); + auto* credentialed = result.value()->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + std::mutex mutex; + std::condition_variable cv; + bool refresh_started = false; + bool release_refresh = false; + + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()) + .WillRepeatedly([&]() -> Result> { + std::unique_lock lock(mutex); + refresh_started = true; + cv.notify_all(); + cv.wait(lock, [&] { return release_refresh; }); + return ExpiringCredentials(std::chrono::hours(1), "fetched-by-refresh"); + }); + ASSERT_THAT(credentialed->InitializeStorageCredentials( + ExpiringCredentials(std::chrono::minutes(1), "expiring-key"), provider), + IsOk()); + + std::thread operation( + [&] { EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk()); }); + + const auto installed = + ExpiringCredentials(std::chrono::hours(2), "installed-meanwhile"); + { + std::unique_lock lock(mutex); + cv.wait(lock, [&] { return refresh_started; }); + } + auto install_status = credentialed->SetStorageCredentials(installed); + { + std::lock_guard lock(mutex); + release_refresh = true; + } + cv.notify_all(); + operation.join(); + + ASSERT_THAT(install_status, IsOk()); + EXPECT_EQ(credentialed->credentials(), installed); +} + +TEST_F(ArrowS3FileIOTest, RefreshesSessionCredentialsWithoutAUsableExpiry) { + for (std::string_view expiry : {"", "not-a-number"}) { + SCOPED_TRACE(expiry); + auto result = MakeS3FileIO({}); + ASSERT_THAT(result, IsOk()); + auto* credentialed = result.value()->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + int refresh_calls = 0; + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()) + .WillRepeatedly([&]() -> Result> { + ++refresh_calls; + return ExpiringCredentials(std::chrono::hours(1), "refreshed-key"); + }); + + auto logger = std::make_shared(); + ScopedDefaultLogger scoped(logger); + std::unordered_map config = { + {std::string(S3Properties::kAccessKeyId), "access-key"}, + {std::string(S3Properties::kSecretAccessKey), "secret"}, + {std::string(S3Properties::kSessionToken), "token"}}; + if (!expiry.empty()) { + config[std::string(S3Properties::kSessionTokenExpiresAtMs)] = std::string(expiry); + } + ASSERT_THAT(credentialed->InitializeStorageCredentials( + {{.prefix = "s3", .config = std::move(config)}}, provider), + IsOk()); + EXPECT_TRUE(HasWarning(*logger, "session token")); + + EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(refresh_calls, 1); + } +} + +TEST_F(ArrowS3FileIOTest, BacksOffWhenTheReplacementIsAlsoCloseToExpiry) { + auto result = MakeS3FileIO({}); + ASSERT_THAT(result, IsOk()); + auto* credentialed = result.value()->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + int refresh_calls = 0; + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()) + .WillRepeatedly([&]() -> Result> { + ++refresh_calls; + return ExpiringCredentials(std::chrono::minutes(1), "short-lived-key"); + }); + ASSERT_THAT(credentialed->InitializeStorageCredentials( + ExpiringCredentials(std::chrono::minutes(1), "expiring-key"), provider), + IsOk()); + + for (int i = 0; i < 3; ++i) { + EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk()); + } + EXPECT_EQ(refresh_calls, 1); +} + +TEST_F(ArrowS3FileIOTest, KeepsCredentialsWhenTheRefreshFails) { + struct Case { + std::string_view name; + std::optional expires_at_ms; + std::string prefix = "s3"; + std::unordered_map settings; + bool keys = true; + bool load_fails = false; + }; + const auto valid = ExpiresAtMs(std::chrono::hours(1)); + const std::vector cases = { + {.name = "catalog error", .expires_at_ms = valid, .load_fails = true}, + {.name = "expired", .expires_at_ms = ExpiresAtMs(-std::chrono::minutes(1))}, + {.name = "no expiry", .expires_at_ms = std::nullopt}, + {.name = "unparseable expiry", .expires_at_ms = "not-a-number"}, + // A refresh may only rotate keys. + {.name = "prefix changed", .expires_at_ms = valid, .prefix = "s3://other/"}, + {.name = "setting changed", + .expires_at_ms = valid, + .settings = {{std::string(S3Properties::kConnectTimeoutMs), "5000"}}}, + {.name = "no keys", .expires_at_ms = valid, .keys = false}}; + for (const auto& test_case : cases) { + SCOPED_TRACE(test_case.name); + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO({})); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + auto config = test_case.settings; + if (test_case.keys) { + config.insert({{std::string(S3Properties::kAccessKeyId), "unusable-key"}, + {std::string(S3Properties::kSecretAccessKey), "secret"}}); + } + config.emplace(std::string(S3Properties::kSessionToken), "token"); + if (test_case.expires_at_ms.has_value()) { + config[std::string(S3Properties::kSessionTokenExpiresAtMs)] = + *test_case.expires_at_ms; + } + const std::vector unusable = { + {.prefix = test_case.prefix, .config = std::move(config)}}; + auto provider = std::make_shared(); + // Once: a failed refresh backs off. + EXPECT_CALL(*provider, Load()) + .WillOnce([&]() -> Result> { + if (test_case.load_fails) { + return NotFound("catalog unreachable"); + } + return unusable; + }); + const auto expiring = ExpiringCredentials(std::chrono::minutes(1), "expiring-key"); + ASSERT_THAT(credentialed->InitializeStorageCredentials(expiring, provider), IsOk()); + + auto logger = std::make_shared(); + ScopedDefaultLogger scoped(logger); + EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(credentialed->credentials(), expiring); + EXPECT_TRUE(HasWarning(*logger, "Failed to refresh")); + EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsOk()); + } +} + +TEST_F(ArrowS3FileIOTest, RejectsARefreshWithAnInvalidCredential) { + // Usable S3 keys do not make up for an invalid credential beside them, even + // one this FileIO would skip. + const std::vector invalid = { + {.prefix = "", .config = {{"k", "v"}}}, + {.prefix = "gs://bucket", .config = {}}, + {.prefix = "s3://other/", .config = {}}}; + for (const auto& bad : invalid) { + SCOPED_TRACE(bad.prefix); + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO({})); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + auto refreshed = ExpiringCredentials(std::chrono::hours(1), "refreshed-key"); + refreshed.push_back(bad); + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()).WillOnce(::testing::Return(refreshed)); + const auto expiring = ExpiringCredentials(std::chrono::minutes(1), "expiring-key"); + ASSERT_THAT(credentialed->InitializeStorageCredentials(expiring, provider), IsOk()); + + auto logger = std::make_shared(); + ScopedDefaultLogger scoped(logger); + EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(credentialed->credentials(), expiring); + EXPECT_TRUE(HasWarning(*logger, "Failed to refresh")); + } +} + +TEST_F(ArrowS3FileIOTest, AcceptsRefreshedPrefixesInAnotherOrder) { + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO({})); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + auto initial = + ExpiringCredentials(std::chrono::minutes(1), "expiring-key", "s3://first/"); + initial.push_back( + ExpiringCredentials(std::chrono::minutes(1), "expiring-key", "s3://second/") + .front()); + auto refreshed = + ExpiringCredentials(std::chrono::hours(1), "refreshed-key", "s3://second/"); + refreshed.push_back( + ExpiringCredentials(std::chrono::hours(1), "refreshed-key", "s3://first/").front()); + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()).WillOnce(::testing::Return(refreshed)); + ASSERT_THAT(credentialed->InitializeStorageCredentials(initial, provider), IsOk()); + + EXPECT_THAT(io->NewInputFile("s3://first/key"), IsOk()); + EXPECT_THAT(credentialed->credentials(), + ::testing::UnorderedElementsAreArray(refreshed)); +} + +TEST_F(ArrowS3FileIOTest, TracksTheExpiryOfTheTokenItSignsWith) { + const std::unordered_map expired_token = { + {std::string(S3Properties::kSessionToken), "token"}, + {std::string(S3Properties::kSessionTokenExpiresAtMs), "0"}}; + auto keys_and = [](std::unordered_map config) { + config.merge(BadS3Credentials()); + return config; + }; + struct Case { + std::string_view name; + std::unordered_map properties; + std::vector vended; + bool expired; + }; + const std::vector cases = { + // Without an access key pair the AWS default chain signs instead. + {.name = "token without keys", + .vended = {{.prefix = "s3://bucket/", .config = expired_token}}, + .expired = false}, + {.name = "expiry without a token", + .vended = {{.prefix = "s3://bucket/", + .config = keys_and( + {{std::string(S3Properties::kSessionTokenExpiresAtMs), "0"}})}}, + .expired = false}, + // Inherited from the FileIO's properties, it signs with the vended keys. + {.name = "inherited token", + .properties = expired_token, + .vended = {{.prefix = "s3://bucket/", .config = BadS3Credentials()}}, + .expired = true}, + // Only the first credential for a prefix is used. + {.name = "repeated prefix", + .vended = {{.prefix = "s3://bucket/", .config = BadS3Credentials()}, + {.prefix = "s3://bucket/", .config = keys_and(expired_token)}}, + .expired = false}}; + for (const auto& test_case : cases) { + SCOPED_TRACE(test_case.name); + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO(test_case.properties)); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + ASSERT_THAT(credentialed->SetStorageCredentials(test_case.vended), IsOk()); + + if (test_case.expired) { + EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsError(ErrorKind::kTokenExpired)); + } else { + EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsOk()); + } + } +} + +TEST_F(ArrowS3FileIOTest, FailsOnceCredentialsHaveExpiredAndCannotBeRefreshed) { + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO({})); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()).WillOnce([]() -> Result> { + return NotFound("catalog unreachable"); + }); + ASSERT_THAT( + credentialed->InitializeStorageCredentials( + ExpiringCredentials(-std::chrono::minutes(1), "expired-key", "s3://bucket/"), + provider), + IsOk()); + + EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsError(ErrorKind::kTokenExpired)); + EXPECT_THAT(io->NewOutputFile("s3://bucket/key"), IsError(ErrorKind::kTokenExpired)); + EXPECT_THAT(io->DeleteFile("s3://bucket/key"), IsError(ErrorKind::kTokenExpired)); + EXPECT_THAT(io->DeleteFiles({"s3://bucket/key"}), + HasErrorMessage("Failed to delete 1 of 1 files")); + // Locations outside the vended prefix keep the default credentials. + EXPECT_THAT(io->NewInputFile("s3://other-bucket/key"), IsOk()); +} + TEST_F(ArrowS3FileIOTest, OperationsSurviveConcurrentCredentialInstalls) { auto result = MakeS3FileIO({}); ASSERT_THAT(result, IsOk()); @@ -642,6 +1089,189 @@ TEST_F(ArrowS3FileIOTest, LongestCredentialPrefix) { IsOk()); } +TEST_F(ArrowS3FileIOTest, RefreshReachesHandlesOpenedBeforeIt) { + if (!HasIntegrationEnv()) { + GTEST_SKIP() << "Set ICEBERG_TEST_S3_URI to enable S3 IO test"; + } + + auto properties = PropertiesFromEnv(); + const auto access_key = properties.find(std::string(S3Properties::kAccessKeyId)); + const auto secret_key = properties.find(std::string(S3Properties::kSecretAccessKey)); + if (access_key == properties.end() || secret_key == properties.end()) { + GTEST_SKIP() << "Set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY to enable " + "credential refresh test"; + } + + // New keys reach a handle opened before them, from a refresh its own request + // runs, but not keys installed with other client settings. + struct Case { + std::string_view name; + bool refresh; + std::unordered_map settings; + bool handle_reads; + }; + const std::vector cases = { + {.name = "refresh", .refresh = true, .settings = {}, .handle_reads = true}, + {.name = "install with another setting", + .refresh = false, + .settings = {{std::string(S3Properties::kConnectTimeoutMs), "5000"}}, + .handle_reads = false}}; + for (const auto& test_case : cases) { + SCOPED_TRACE(test_case.name); + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO(properties)); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + const auto prefix = ObjectUri("refresh_in_place/"); + const auto uri = prefix + "object"; + constexpr std::string_view kContent = "read with keys refreshed after it was opened"; + ASSERT_THAT(io->WriteFile(uri, kContent), IsOk()); + + // The vended keys are bad and come due right after the handle opens; only + // the new ones can read. + constexpr auto kDueIn = std::chrono::seconds(2); + auto new_config = test_case.settings; + new_config.insert(*access_key); + new_config.insert(*secret_key); + const std::vector new_credentials = { + {.prefix = prefix, .config = std::move(new_config)}}; + std::shared_ptr provider; + if (test_case.refresh) { + provider = std::make_shared(); + EXPECT_CALL(*provider, Load()).WillOnce(::testing::Return(new_credentials)); + } + ASSERT_THAT( + credentialed->InitializeStorageCredentials( + {{.prefix = prefix, + .config = {{std::string(S3Properties::kAccessKeyId), "bad-access-key"}, + {std::string(S3Properties::kSecretAccessKey), "bad-secret"}, + {std::string(S3Properties::kSessionToken), "bad-token"}, + {std::string(S3Properties::kSessionTokenExpiresAtMs), + ExpiresAtMs(std::chrono::minutes(5) + kDueIn)}}}}, + provider), + IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto file, io->NewInputFile(uri, kContent.size())); + + if (test_case.refresh) { + std::this_thread::sleep_for(kDueIn + std::chrono::seconds(1)); + } else { + ASSERT_THAT(credentialed->SetStorageCredentials(new_credentials), IsOk()); + } + auto read = ReadAll(*file, kContent.size()); + if (test_case.handle_reads) { + EXPECT_THAT(read, HasValue(::testing::Eq(std::string(kContent)))); + } else { + // Still signs with the bad keys it was opened with. + EXPECT_THAT(read, ::testing::Not(IsOk())); + } + EXPECT_EQ(credentialed->credentials(), new_credentials); + EXPECT_THAT(io->DeleteFile(uri), IsOk()); + } +} + +TEST_F(ArrowS3FileIOTest, HandleKeepsItsKeysOnceTheFileIOIsGone) { + if (!HasIntegrationEnv()) { + GTEST_SKIP() << "Set ICEBERG_TEST_S3_URI to enable S3 IO test"; + } + + auto properties = PropertiesFromEnv(); + const auto access_key = properties.find(std::string(S3Properties::kAccessKeyId)); + const auto secret_key = properties.find(std::string(S3Properties::kSecretAccessKey)); + if (access_key == properties.end() || secret_key == properties.end()) { + GTEST_SKIP() << "Set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY to enable " + "credential install test"; + } + + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO(properties)); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + const auto prefix = ObjectUri("install_then_release/"); + const auto uri = prefix + "object"; + constexpr std::string_view kContent = "read after its FileIO is gone"; + ASSERT_THAT(io->WriteFile(uri, kContent), IsOk()); + + // Working keys replace bad ones with the same settings. Building their client + // must not leave it with the bad keys once the FileIO cannot supply new ones. + ASSERT_THAT(credentialed->SetStorageCredentials( + {{.prefix = prefix, .config = BadS3Credentials()}}), + IsOk()); + ASSERT_THAT(credentialed->SetStorageCredentials( + {{.prefix = prefix, .config = {*access_key, *secret_key}}}), + IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto file, io->NewInputFile(uri, kContent.size())); + io.reset(); + + EXPECT_THAT(ReadAll(*file, kContent.size()), + HasValue(::testing::Eq(std::string(kContent)))); + ICEBERG_UNWRAP_OR_FAIL(auto cleanup, MakeS3FileIO(properties)); + EXPECT_THAT(cleanup->DeleteFile(uri), IsOk()); +} + +TEST_F(ArrowS3FileIOTest, RequestWaitsForARefreshUnderWayOnceExpired) { + if (!HasIntegrationEnv()) { + GTEST_SKIP() << "Set ICEBERG_TEST_S3_URI to enable S3 IO test"; + } + + auto properties = PropertiesFromEnv(); + const auto access_key = properties.find(std::string(S3Properties::kAccessKeyId)); + const auto secret_key = properties.find(std::string(S3Properties::kSecretAccessKey)); + if (access_key == properties.end() || secret_key == properties.end()) { + GTEST_SKIP() << "Set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY to enable " + "credential refresh test"; + } + + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO(properties)); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + const auto prefix = ObjectUri("refresh_while_expired/"); + const auto uri = prefix + "object"; + constexpr std::string_view kContent = "read with keys a refresh under way fetched"; + ASSERT_THAT(io->WriteFile(uri, kContent), IsOk()); + + // Bad keys that expire right after the handle opens. The first refresh + // fails; the second, after expiry, is slow and returns working keys. + constexpr auto kExpiresIn = std::chrono::seconds(2); + std::promise refresh_started; + auto started = refresh_started.get_future(); + const std::vector refreshed = { + {.prefix = prefix, .config = {*access_key, *secret_key}}}; + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()) + .WillOnce([]() -> Result> { + return ServiceUnavailable("catalog unavailable"); + }) + .WillOnce([&]() -> Result> { + refresh_started.set_value(); + std::this_thread::sleep_for(std::chrono::milliseconds(500)); + return refreshed; + }); + ASSERT_THAT( + credentialed->InitializeStorageCredentials( + {{.prefix = prefix, + .config = {{std::string(S3Properties::kAccessKeyId), "bad-access-key"}, + {std::string(S3Properties::kSecretAccessKey), "bad-secret"}, + {std::string(S3Properties::kSessionToken), "bad-token"}, + {std::string(S3Properties::kSessionTokenExpiresAtMs), + ExpiresAtMs(kExpiresIn)}}}}, + provider), + IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto file, io->NewInputFile(uri, kContent.size())); + + // Past the expiry, and the backoff after the failed refresh. + std::this_thread::sleep_for(kExpiresIn + std::chrono::milliseconds(500)); + std::thread operation([&] { EXPECT_THAT(io->NewInputFile(uri), IsOk()); }); + // Read while that refresh runs. No fatal assertions before the join. + EXPECT_EQ(started.wait_for(std::chrono::seconds(10)), std::future_status::ready); + auto read = ReadAll(*file, kContent.size()); + operation.join(); + + EXPECT_THAT(read, HasValue(::testing::Eq(std::string(kContent)))); + EXPECT_EQ(credentialed->credentials(), refreshed); + EXPECT_THAT(io->DeleteFile(uri), IsOk()); +} + TEST_F(ArrowS3FileIOTest, DeleteFilesAttemptsEveryFile) { if (!HasIntegrationEnv()) { GTEST_SKIP() << "Set ICEBERG_TEST_S3_URI to enable S3 IO test"; @@ -692,10 +1322,8 @@ TEST_F(ArrowS3FileIOTest, DeleteFilesAttemptsEveryFile) { EXPECT_THAT(io->DeleteFiles({paths[0], paths[2]}), IsOk()); } -// The credential is vended under the oss spelling and the object addressed as -// `s3://`, so they only meet through canonicalization — and every other path -// to authentication is broken. (rest_arrow_file_io_test covers the mirrored -// direction.) +// The credential and object use the same S3-compatible scheme. The scheme is +// preserved for routing, while the underlying endpoint remains S3-compatible. TEST_F(ArrowS3FileIOTest, AppliesOssCredentialInRealRoundTrip) { if (!HasIntegrationEnv()) { GTEST_SKIP() << "Set ICEBERG_TEST_S3_URI to enable S3 IO test"; @@ -723,15 +1351,13 @@ TEST_F(ArrowS3FileIOTest, AppliesOssCredentialInRealRoundTrip) { const auto scheme_end = object_uri.find("://"); ASSERT_NE(scheme_end, std::string::npos) << "ICEBERG_TEST_S3_URI must carry a scheme"; // Both spellings are forced, so they cross whatever scheme the env URI uses. - const auto oss_spelling = std::string("oss").append(object_uri.substr(scheme_end)); - const auto oss_prefix = - oss_spelling.substr(0, oss_spelling.size() - object_name.size()); - const auto s3_uri = std::string("s3").append(object_uri.substr(scheme_end)); + const auto oss_uri = std::string("oss").append(object_uri.substr(scheme_end)); + const auto oss_prefix = oss_uri.substr(0, oss_uri.size() - object_name.size()); EXPECT_THAT(credentialed->SetStorageCredentials( {{.prefix = oss_prefix, .config = std::move(properties)}}), IsOk()); - EXPECT_THAT(CheckReadWrite(*io, s3_uri, "hello oss with vended credentials"), IsOk()); + EXPECT_THAT(CheckReadWrite(*io, oss_uri, "hello oss with vended credentials"), IsOk()); } TEST_F(ArrowS3FileIOTest, DeleteFilesReachesEveryCredentialPrefix) { diff --git a/src/iceberg/test/http_client_test.cc b/src/iceberg/test/http_client_test.cc new file mode 100644 index 000000000..4203e7426 --- /dev/null +++ b/src/iceberg/test/http_client_test.cc @@ -0,0 +1,175 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 "iceberg/catalog/rest/http_client.h" + +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#ifndef _WIN32 +# include + +# include +# include +# include +# include +#endif + +#include "iceberg/catalog/rest/auth/auth_session.h" +#include "iceberg/catalog/rest/error_handlers.h" +#include "iceberg/test/matchers.h" + +namespace iceberg::rest { + +#ifndef _WIN32 + +namespace { + +// A loopback socket listening on an ephemeral port. +struct Listener { + int fd = -1; + uint16_t port = 0; + + ~Listener() { + if (fd >= 0) { + ::close(fd); + } + } +}; + +void Listen(Listener& listener) { + listener.fd = ::socket(AF_INET, SOCK_STREAM, 0); + ASSERT_GE(listener.fd, 0); + sockaddr_in address{}; + address.sin_family = AF_INET; + address.sin_addr.s_addr = htonl(INADDR_LOOPBACK); + ASSERT_EQ(::bind(listener.fd, reinterpret_cast(&address), sizeof(address)), + 0); + ASSERT_EQ(::listen(listener.fd, 1), 0); + socklen_t length = sizeof(address); + ASSERT_EQ(::getsockname(listener.fd, reinterpret_cast(&address), &length), + 0); + listener.port = ntohs(address.sin_port); +} + +// Accepts one connection, reads the request head and sends `response`. Then +// ends the connection if `end`, or else waits for the client to close it. +void ServeOnce(int listener, std::string_view response, bool end) { + // Bound every wait, so the test cannot hang. + timeval wait{}; + wait.tv_sec = 10; + ::setsockopt(listener, SOL_SOCKET, SO_RCVTIMEO, &wait, sizeof(wait)); + const int connection = ::accept(listener, nullptr, nullptr); + if (connection < 0) { + return; + } + ::setsockopt(connection, SOL_SOCKET, SO_RCVTIMEO, &wait, sizeof(wait)); + std::string request; + std::array buffer; + while (request.find("\r\n\r\n") == std::string::npos) { + const auto received = ::recv(connection, buffer.data(), buffer.size(), 0); + if (received <= 0) { + break; + } + request.append(buffer.data(), static_cast(received)); + } + ::send(connection, response.data(), response.size(), 0); + if (end) { + ::shutdown(connection, SHUT_WR); + } + // Drain until the client closes, so ours closes cleanly. + while (::recv(connection, buffer.data(), buffer.size(), 0) > 0) { + } + ::close(connection); +} + +} // namespace + +TEST(HttpClientTest, GetFailsWhenItsTransferFails) { + // A 200 arriving first does not make a truncated response a success. + constexpr std::string_view kTruncated = + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n" + "Content-Length: 100\r\n\r\n{"; + struct Case { + std::string_view name; + std::string_view response; + bool end; + std::string_view error; + }; + for (const auto& test_case : + {Case{.name = "no response", .response = "", .end = false, .error = "timed out"}, + Case{.name = "body stalls", + .response = kTruncated, + .end = false, + .error = "timed out"}, + Case{.name = "body ends early", + .response = kTruncated, + .end = true, + .error = "failed after status 200"}}) { + SCOPED_TRACE(test_case.name); + Listener listener; + ASSERT_NO_FATAL_FAILURE(Listen(listener)); + const auto url = std::format("http://127.0.0.1:{}/v1/config", listener.port); + + HttpClient client({}, {}, std::chrono::milliseconds(500)); + auto session = auth::AuthSession::MakeDefault({}); + std::thread server( + [&] { ServeOnce(listener.fd, test_case.response, test_case.end); }); + auto response = client.Get(url, /*params=*/{}, /*headers=*/{}, + *DefaultErrorHandler::Instance(), *session); + server.join(); + EXPECT_THAT(response, IsError(ErrorKind::kIOError)); + EXPECT_THAT(response, HasErrorMessage(test_case.error)); + } +} + +TEST(HttpClientTest, PostReportsTransferFailureAfterSuccessfulStatus) { + // A write that got its status may have taken effect, so the caller must treat + // the result as unknown rather than retrying or cleaning up automatically. + constexpr std::string_view kBody = R"({"metadata-location":"s3://bucket/m.json"})"; + const auto response_bytes = std::format( + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n" + "Transfer-Encoding: chunked\r\n\r\n{:x}\r\n{}\r\n", + kBody.size(), kBody); + Listener listener; + ASSERT_NO_FATAL_FAILURE(Listen(listener)); + const auto url = std::format("http://127.0.0.1:{}/v1/commit", listener.port); + + // Outlives the response, which keeps a curl handle in its connection pool. + HttpClient client; + auto session = auth::AuthSession::MakeDefault({}); + std::thread server([&] { ServeOnce(listener.fd, response_bytes, /*end=*/true); }); + auto response = + client.Post(url, "{}", /*headers=*/{}, *DefaultErrorHandler::Instance(), *session); + server.join(); + EXPECT_THAT(response, IsError(ErrorKind::kIOError)); + EXPECT_THAT(response, HasErrorMessage("failed after status 200")); +} + +#endif + +} // namespace iceberg::rest diff --git a/src/iceberg/test/mock_io.h b/src/iceberg/test/mock_io.h index f643a852b..63abb6ea6 100644 --- a/src/iceberg/test/mock_io.h +++ b/src/iceberg/test/mock_io.h @@ -37,6 +37,7 @@ #include #include "iceberg/file_io.h" +#include "iceberg/storage_credential.h" #include "iceberg/util/macros.h" namespace iceberg { @@ -256,4 +257,9 @@ class MockFileIO : public FileIO { std::shared_ptr files_ = std::make_shared(); }; +class MockStorageCredentialProvider : public StorageCredentialProvider { + public: + MOCK_METHOD((Result>), Load, (), (override)); +}; + } // namespace iceberg diff --git a/src/iceberg/test/rest_arrow_file_io_test.cc b/src/iceberg/test/rest_arrow_file_io_test.cc index 83966dc96..7553e30d6 100644 --- a/src/iceberg/test/rest_arrow_file_io_test.cc +++ b/src/iceberg/test/rest_arrow_file_io_test.cc @@ -21,6 +21,7 @@ /// \brief Covers REST -> ResolvingFileIO -> registry -> Arrow FileIO against the /// real registered implementations, which mock delegates cannot exercise. +#include #include #include #include @@ -39,7 +40,9 @@ #include "iceberg/catalog/rest/rest_file_io.h" #include "iceberg/storage_credential.h" #include "iceberg/test/matchers.h" +#include "iceberg/test/mock_io.h" #include "iceberg/test/temp_file_test_base.h" +#include "iceberg/util/timepoint.h" namespace iceberg::rest { @@ -148,8 +151,8 @@ std::string AsOssUri(std::string_view uri) { } // Resolution, credential matching and real I/O for `oss://`. The credential -// env vars are scrubbed, so only the vended `s3`-scoped credential matching -// the canonicalized location can authenticate. +// env vars are scrubbed, so only the vended `oss`-scoped credential matching +// the location can authenticate. TEST_F(RestArrowFileIOTest, ReadsBackWhatItWroteThroughAnOssLocation) { const auto base_uri = GetEnvIfSet("ICEBERG_TEST_S3_URI"); if (!base_uri.has_value()) { @@ -176,10 +179,10 @@ TEST_F(RestArrowFileIOTest, ReadsBackWhatItWroteThroughAnOssLocation) { ScopedScrubbedAwsCredentialEnv scrubbed; - // Scoped to `s3`, while the data it grants access to is addressed as `oss://`. - auto io = MakeTableFileIO({{"warehouse", "logical_warehouse_name"}}, - /*table_config=*/{}, - {{.prefix = "s3", .config = std::move(credential_config)}}); + // The scheme is preserved for routing, while the endpoint is S3-compatible. + auto io = + MakeTableFileIO({{"warehouse", "logical_warehouse_name"}}, /*table_config=*/{}, + {{.prefix = "oss://", .config = std::move(credential_config)}}); ASSERT_THAT(io, IsOk()); auto object_uri = AsOssUri(*base_uri); @@ -195,6 +198,56 @@ TEST_F(RestArrowFileIOTest, ReadsBackWhatItWroteThroughAnOssLocation) { EXPECT_THAT(io.value()->DeleteFile(object_uri), IsOk()); } +TEST_F(RestArrowFileIOTest, RefreshesCredentialsThroughTheRealS3FileIO) { + const auto base_uri = GetEnvIfSet("ICEBERG_TEST_S3_URI"); + const auto access_key = GetEnvIfSet("AWS_ACCESS_KEY_ID"); + const auto secret_key = GetEnvIfSet("AWS_SECRET_ACCESS_KEY"); + if (!base_uri || !access_key || !secret_key) { + GTEST_SKIP() + << "Set ICEBERG_TEST_S3_URI, AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY"; + } + + std::unordered_map config = { + {"s3.access-key-id", *access_key}, {"s3.secret-access-key", *secret_key}}; + if (const auto token = GetEnvIfSet("AWS_SESSION_TOKEN")) { + config["s3.session-token"] = *token; + // A refreshed session token must say when it expires. + config["s3.session-token-expires-at-ms"] = std::to_string( + UnixMsFromTimePointMs(CurrentTimePointMs() + std::chrono::hours(1))); + } + if (const auto endpoint = GetEnvIfSet("ICEBERG_TEST_S3_ENDPOINT")) { + config["s3.endpoint"] = *endpoint; + } + if (const auto region = GetEnvIfSet("AWS_REGION")) { + config["client.region"] = *region; + } + const std::vector refreshed = {{.prefix = "s3", .config = config}}; + config["s3.access-key-id"] = "bad-access-key"; + config["s3.secret-access-key"] = "bad-secret-key"; + config["s3.session-token"] = "expired-token"; + config["s3.session-token-expires-at-ms"] = "0"; + + // Only the provider can supply credentials that authenticate this round trip. + ScopedScrubbedAwsCredentialEnv scrubbed; + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()).WillOnce(::testing::Return(refreshed)); + auto io = MakeTableFileIO({{"warehouse", "logical_warehouse_name"}}, + /*table_config=*/{}, + {{.prefix = "s3", .config = std::move(config)}}, provider); + ASSERT_THAT(io, IsOk()); + + auto object_uri = *base_uri; + if (!object_uri.ends_with('/')) { + object_uri += '/'; + } + object_uri += "iceberg_refresh_provider.txt"; + constexpr std::string_view kContent = "written after refreshing vended credentials"; + ASSERT_THAT(io.value()->WriteFile(object_uri, kContent), IsOk()); + EXPECT_THAT(io.value()->ReadFile(object_uri, std::nullopt), + HasValue(::testing::Eq(std::string(kContent)))); + EXPECT_THAT(io.value()->DeleteFile(object_uri), IsOk()); +} + #endif // ICEBERG_S3_ENABLED } // namespace diff --git a/src/iceberg/test/rest_catalog_unit_test.cc b/src/iceberg/test/rest_catalog_unit_test.cc index 97c9c3650..86865a663 100644 --- a/src/iceberg/test/rest_catalog_unit_test.cc +++ b/src/iceberg/test/rest_catalog_unit_test.cc @@ -18,6 +18,7 @@ */ #include +#include #include #include #include