diff --git a/src/iceberg/catalog/rest/CMakeLists.txt b/src/iceberg/catalog/rest/CMakeLists.txt index a8146d468..65c037247 100644 --- a/src/iceberg/catalog/rest/CMakeLists.txt +++ b/src/iceberg/catalog/rest/CMakeLists.txt @@ -22,6 +22,7 @@ set(ICEBERG_REST_SOURCES auth/auth_managers.cc auth/auth_properties.cc auth/auth_session.cc + auth/auth_session_cache.cc auth/oauth2_util.cc auth/sigv4_manager.cc catalog_properties.cc diff --git a/src/iceberg/catalog/rest/auth/auth_manager.cc b/src/iceberg/catalog/rest/auth/auth_manager.cc index df5d7304a..1a4989a1c 100644 --- a/src/iceberg/catalog/rest/auth/auth_manager.cc +++ b/src/iceberg/catalog/rest/auth/auth_manager.cc @@ -21,18 +21,23 @@ #include #include +#include #include +#include #include #include #include "iceberg/catalog/rest/auth/auth_manager_internal.h" #include "iceberg/catalog/rest/auth/auth_properties.h" #include "iceberg/catalog/rest/auth/auth_session.h" +#include "iceberg/catalog/rest/auth/auth_session_cache_internal.h" #include "iceberg/catalog/rest/auth/auth_session_internal.h" #include "iceberg/catalog/rest/auth/oauth2_util.h" #include "iceberg/catalog/session_context.h" +#include "iceberg/logging/log_macros.h" #include "iceberg/util/base64.h" #include "iceberg/util/macros.h" +#include "iceberg/util/string_util.h" namespace iceberg::rest::auth { @@ -57,6 +62,127 @@ std::optional> FindPreferredTypedToken( return std::nullopt; } +/// \brief Credential used to create an OAuth2 child session. +struct ChildCredential { + enum class Kind { kToken, kCredential, kTokenExchange }; + + Kind kind; + std::string value; + /// Subject token type; only set for kTokenExchange. + std::string token_type; +}; + +/// \brief Select the credential for a child session. +/// +/// The order matches Java: a bearer token, then a client credential (only when +/// `allow_credential` is true), then a typed token in kTokenPreferenceOrder. +std::optional SelectChildCredential( + const std::unordered_map& credentials, + bool allow_credential) { + if (auto token_it = credentials.find(AuthProperties::kToken.key()); + token_it != credentials.end()) { + return ChildCredential{.kind = ChildCredential::Kind::kToken, + .value = token_it->second}; + } + if (allow_credential) { + if (auto credential_it = credentials.find(AuthProperties::kCredential.key()); + credential_it != credentials.end()) { + return ChildCredential{.kind = ChildCredential::Kind::kCredential, + .value = credential_it->second}; + } + } + if (auto typed_token = FindPreferredTypedToken(credentials); typed_token.has_value()) { + return ChildCredential{.kind = ChildCredential::Kind::kTokenExchange, + .value = std::move(typed_token->second), + .token_type = std::move(typed_token->first)}; + } + return std::nullopt; +} + +/// \brief Actor token sent with a token exchange. +struct ActorToken { + std::string token; + std::string token_type; +}; + +/// \brief The parent session's token, sent as the actor of a token exchange, or +/// std::nullopt if the parent has no token. +std::optional ExchangeActor(const OAuth2SessionInfo& parent_info) { + if (parent_info.token.empty()) { + return std::nullopt; + } + return ActorToken{.token = parent_info.token, + .token_type = parent_info.issued_token_type}; +} + +void AppendCacheKeyPart(std::string& key, std::string_view part) { + key += std::to_string(part.size()); + key += ':'; + key += part; +} + +/// \brief Build the cache key of a table session. +/// +/// The key covers every input that determines how the session authenticates, so +/// two table sessions are shared only when they would send the same credentials. +/// With P(s) = ":", the formats are: +/// +/// bearer token: "tbl:" + P("token") + P(token) +/// token exchange: "tbl:" + P("exchange") + P(subject type) + P(subject) +/// + P(actor type) + P(actor) +/// +/// A bearer token is sent as-is, so its key does not depend on the parent. An +/// exchanged token also depends on the actor, the parent's token; the actor parts +/// are empty when the parent has no token, matching the exchange request. The +/// length prefix keeps the key unambiguous although token types contain ':'. +/// +/// For example, bearer token "abc" gives "tbl:5:token3:abc", and JWT "eyJ.a.b" +/// exchanged under the parent access token "tok" gives (on one line): +/// +/// tbl:8:exchange36:urn:ietf:params:oauth:token-type:jwt7:eyJ.a.b +/// 45:urn:ietf:params:oauth:token-type:access_token3:tok +/// +/// An empty token gives an empty key, which disables caching. +std::string TableCacheKey(const ChildCredential& credential, + const OAuth2SessionInfo& parent_info) { + if (credential.value.empty()) { + return {}; + } + std::string key = "tbl:"; + switch (credential.kind) { + case ChildCredential::Kind::kToken: + AppendCacheKeyPart(key, "token"); + AppendCacheKeyPart(key, credential.value); + return key; + case ChildCredential::Kind::kTokenExchange: { + auto actor = ExchangeActor(parent_info); + AppendCacheKeyPart(key, "exchange"); + AppendCacheKeyPart(key, credential.token_type); + AppendCacheKeyPart(key, credential.value); + AppendCacheKeyPart(key, actor.has_value() ? actor->token_type : ""); + AppendCacheKeyPart(key, actor.has_value() ? actor->token : ""); + return key; + } + case ChildCredential::Kind::kCredential: + // Table sessions never use a client credential. + return {}; + } + std::unreachable(); +} + +/// \brief Parse auth.session-timeout-ms, returning an error for a malformed or +/// negative value rather than throwing as AuthProperties::Get() would. +Result ParseSessionTimeout(const AuthProperties& config) { + const auto& entry = AuthProperties::kSessionTimeoutMs; + int64_t timeout_ms = entry.value(); + if (auto it = config.configs().find(entry.key()); it != config.configs().end()) { + ICEBERG_ASSIGN_OR_RAISE(timeout_ms, StringUtils::ParseNumber(it->second)); + } + ICEBERG_PRECHECK(timeout_ms >= 0, "Invalid value for {}: {} (must not be negative)", + entry.key(), timeout_ms); + return std::chrono::milliseconds(timeout_ms); +} + std::unordered_map FilterTableSessionProperties( const std::unordered_map& properties) { std::unordered_map filtered; @@ -180,6 +306,30 @@ class OAuth2Manager : public AuthManager { ICEBERG_PRECHECK(shared_client != nullptr, "OAuth2 catalog session HTTP client must not be null"); refresh_client_ = std::move(shared_client); + + // Initialize catalog-level configuration + keep_refreshed_ = config.keep_refreshed(); + ICEBERG_ASSIGN_OR_RAISE(session_timeout_, ParseSessionTimeout(config)); + + // Create session cache + ICEBERG_ASSIGN_OR_RAISE( + session_cache_, + internal::AuthSessionCache::Make( + session_timeout_, [](std::shared_ptr session) { + if (auto oauth2 = + std::dynamic_pointer_cast(session)) { + oauth2->StopRefreshing(); + } + })); + + // Reclaim idle sessions in the background on the refresh executor. Without an + // executor, the cache sweeps from Get() instead. + if (refresh_executor_ != nullptr) { + ICEBERG_RETURN_UNEXPECTED(session_cache_->StartPeriodicSweep( + *refresh_executor_, + internal::AuthSessionCache::DefaultSweepInterval(session_timeout_))); + } + // Reuse the token response and start time from the init phase. if (auth_response_.has_value()) { return MakeSession(*auth_response_, config, config.keep_refreshed(), start_time_); @@ -212,21 +362,35 @@ class OAuth2Manager : public AuthManager { Result> ContextualSession( const SessionContext& context, std::shared_ptr parent) override { - // TODO(lishuxu): Add child-session caching and refresh, matching Java - // AuthSessionCache. - return MaybeCreateChildSession(context.credentials, /*allow_credential=*/true, - std::move(parent)); + auto credential = + SelectChildCredential(context.credentials, /*allow_credential=*/true); + if (!credential.has_value()) { + return parent; + } + ICEBERG_ASSIGN_OR_RAISE(auto parent_info, ParentInfo(*parent)); + // Contextual sessions are cached by session ID, whatever credential they use. + std::string cache_key = context.session_id.empty() ? "" : "ctx:" + context.session_id; + return GetOrCreateChildSession(*credential, parent_info, cache_key); } Result> TableSession( [[maybe_unused]] const TableIdentifier& table, const std::unordered_map& properties, std::shared_ptr parent) override { - return MaybeCreateChildSession(FilterTableSessionProperties(properties), - /*allow_credential=*/false, std::move(parent)); + auto credential = SelectChildCredential(FilterTableSessionProperties(properties), + /*allow_credential=*/false); + if (!credential.has_value()) { + return parent; + } + ICEBERG_ASSIGN_OR_RAISE(auto parent_info, ParentInfo(*parent)); + return GetOrCreateChildSession(*credential, parent_info, + TableCacheKey(*credential, parent_info)); } Status Close() override { + if (session_cache_) { + session_cache_->Close(); + } refresh_client_.reset(); return {}; } @@ -263,53 +427,94 @@ class OAuth2Manager : public AuthManager { token_request_started_at, refresh_executor_); } - Result> MaybeCreateChildSession( - const std::unordered_map& credentials, - bool allow_credential, std::shared_ptr parent) { - auto token_it = credentials.find(AuthProperties::kToken.key()); - auto credential_it = credentials.find(AuthProperties::kCredential.key()); - auto typed_token = FindPreferredTypedToken(credentials); - if (token_it == credentials.end() && - (!allow_credential || credential_it == credentials.end()) && - !typed_token.has_value()) { - return parent; - } - + Result ParentInfo(const AuthSession& parent) const { ICEBERG_PRECHECK(refresh_client_ != nullptr, "OAuth2 catalog session must be initialized before child sessions"); - auto parent_info = parent->OAuth2Info(); + auto parent_info = parent.OAuth2Info(); ICEBERG_PRECHECK(parent_info.has_value(), "OAuth2 child session requires OAuth2 parent metadata"); + return std::move(*parent_info); + } - if (token_it != credentials.end()) { - ICEBERG_ASSIGN_OR_RAISE(auto config, - ChildConfig(*parent_info, parent_info->credential)); - return MakeSession(AccessTokenResponse(token_it->second), config, - /*keep_refreshed=*/false); + Result> GetOrCreateChildSession( + const ChildCredential& credential, const OAuth2SessionInfo& parent_info, + const std::string& cache_key) { + // Use cache if we have a valid cache_key + if (!cache_key.empty() && session_cache_) { + return session_cache_->Get( + cache_key, [&]() -> Result> { + return CreateChildSessionUncached(credential, parent_info); + }); } - if (allow_credential && credential_it != credentials.end()) { - ICEBERG_ASSIGN_OR_RAISE(auto config, - ChildConfig(*parent_info, credential_it->second)); - ICEBERG_ASSIGN_OR_RAISE(auto response, - OAuth2Util::FetchToken(*refresh_client_, *parent, config)); - return MakeSession(response, config, /*keep_refreshed=*/false); + // No cache, create session directly + return CreateChildSessionUncached(credential, parent_info); + } + + Result> CreateChildSessionUncached( + const ChildCredential& credential, const OAuth2SessionInfo& parent_info) { + switch (credential.kind) { + case ChildCredential::Kind::kToken: + return CreateTokenChildSession(credential.value, parent_info); + case ChildCredential::Kind::kCredential: + return CreateCredentialChildSession(credential.value, parent_info); + case ChildCredential::Kind::kTokenExchange: + return CreateTokenExchangeChildSession(credential.value, credential.token_type, + parent_info); } + std::unreachable(); + } + Result> CreateTokenChildSession( + const std::string& token, const OAuth2SessionInfo& parent_info) { + // Token child session: don't include parent credential in config. + // Token-based sessions are not refreshed; when token exchange is enabled, + // refresh will use the token itself rather than the parent credential. + auto properties = parent_info.optional_oauth_params; + properties[AuthProperties::kScope.key()] = parent_info.scope; + properties[AuthProperties::kOAuth2ServerUri.key()] = parent_info.oauth2_server_uri; + ICEBERG_ASSIGN_OR_RAISE(auto config, AuthProperties::FromProperties(properties)); + return MakeSession(AccessTokenResponse(token), config, /*keep_refreshed=*/false); + } + + Result> CreateCredentialChildSession( + const std::string& credential, const OAuth2SessionInfo& parent_info) { + ICEBERG_ASSIGN_OR_RAISE(auto config, ChildConfig(parent_info, credential)); + // Use parent session headers for FetchToken authentication + auto temp_parent = AuthSession::MakeDefault(parent_info.headers); + ICEBERG_ASSIGN_OR_RAISE( + auto response, OAuth2Util::FetchToken(*refresh_client_, *temp_parent, config)); + // Credential child sessions follow the catalog's keep_refreshed, but refreshing + // needs an executor. Without one, the child keeps its initial token. + if (keep_refreshed_ && refresh_executor_ == nullptr) { + std::call_once(no_refresh_executor_warning_, [] { + ICEBERG_LOG_WARN( + "OAuth2 credential child sessions will not refresh their tokens because " + "no auth executor was provided"); + }); + } + return MakeSession(response, config, keep_refreshed_ && refresh_executor_ != nullptr); + } + + Result> CreateTokenExchangeChildSession( + const std::string& subject_token, const std::string& subject_token_type, + const OAuth2SessionInfo& parent_info) { std::optional actor_token; std::optional actor_token_type; - if (!parent_info->token.empty()) { - actor_token = parent_info->token; - actor_token_type = parent_info->issued_token_type; + if (auto actor = ExchangeActor(parent_info); actor.has_value()) { + actor_token = std::move(actor->token); + actor_token_type = std::move(actor->token_type); } + // Use parent session headers for ExchangeToken authentication + auto temp_parent = AuthSession::MakeDefault(parent_info.headers); ICEBERG_ASSIGN_OR_RAISE( auto response, - OAuth2Util::ExchangeToken(*refresh_client_, *parent, {}, typed_token->second, - typed_token->first, actor_token, actor_token_type, - parent_info->scope, parent_info->oauth2_server_uri, - parent_info->optional_oauth_params)); + OAuth2Util::ExchangeToken(*refresh_client_, *temp_parent, {}, subject_token, + subject_token_type, actor_token, actor_token_type, + parent_info.scope, parent_info.oauth2_server_uri, + parent_info.optional_oauth_params)); ICEBERG_ASSIGN_OR_RAISE(auto config, - ChildConfig(*parent_info, parent_info->credential)); + ChildConfig(parent_info, parent_info.credential)); return MakeSession(response, config, /*keep_refreshed=*/false); } @@ -318,6 +523,12 @@ class OAuth2Manager : public AuthManager { std::optional start_time_; std::shared_ptr refresh_client_; Executor* refresh_executor_ = nullptr; + + // Catalog-level configuration initialized in CatalogSession + bool keep_refreshed_ = true; + std::chrono::milliseconds session_timeout_{3'600'000}; + std::shared_ptr session_cache_; + std::once_flag no_refresh_executor_warning_; }; Result> MakeOAuth2Manager( diff --git a/src/iceberg/catalog/rest/auth/auth_properties.h b/src/iceberg/catalog/rest/auth/auth_properties.h index 8784194cc..0cf9be5e4 100644 --- a/src/iceberg/catalog/rest/auth/auth_properties.h +++ b/src/iceberg/catalog/rest/auth/auth_properties.h @@ -78,10 +78,18 @@ class ICEBERG_REST_EXPORT AuthProperties : public ConfigBase { inline static Entry kOAuth2ServerUri{"oauth2-server-uri", "v1/oauth/tokens"}; inline static Entry kKeepRefreshed{"token-refresh-enabled", true}; + + /// Whether to use token exchange for refreshing OAuth2 sessions. + /// NOTE: Not yet implemented; the value is ignored. inline static Entry kExchangeEnabled{"token-exchange-enabled", true}; + inline static Entry kAudience{"audience", ""}; inline static Entry kResource{"resource", ""}; + /// Session cache timeout in milliseconds. Sessions will be eligible for eviction + /// after this duration of inactivity. Default is 1 hour (3,600,000 ms). + inline static Entry kSessionTimeoutMs{"auth.session-timeout-ms", 3'600'000}; + // ---- OAuth2 token type constants ---- inline static const std::string kAccessTokenType = diff --git a/src/iceberg/catalog/rest/auth/auth_session.h b/src/iceberg/catalog/rest/auth/auth_session.h index 7e330706c..393418e2a 100644 --- a/src/iceberg/catalog/rest/auth/auth_session.h +++ b/src/iceberg/catalog/rest/auth/auth_session.h @@ -43,6 +43,7 @@ struct ICEBERG_REST_EXPORT OAuth2SessionInfo { std::string scope; std::string oauth2_server_uri; std::unordered_map optional_oauth_params; + std::unordered_map headers; }; /// \brief An authentication session that can authenticate outgoing HTTP requests. diff --git a/src/iceberg/catalog/rest/auth/auth_session_cache.cc b/src/iceberg/catalog/rest/auth/auth_session_cache.cc new file mode 100644 index 000000000..aa62a8ac9 --- /dev/null +++ b/src/iceberg/catalog/rest/auth/auth_session_cache.cc @@ -0,0 +1,320 @@ +/* + * 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 +#include +#include + +#include "iceberg/catalog/rest/auth/auth_session_cache_internal.h" +#include "iceberg/catalog/rest/auth/auth_session_internal.h" +#include "iceberg/logging/log_macros.h" +#include "iceberg/util/macros.h" + +namespace iceberg::rest::auth::internal { + +namespace { + +// Bounds of DefaultSweepInterval(): sweeping at least once a minute reclaims idle +// sessions promptly, and at most once a second keeps a zero timeout from sweeping +// continuously. +constexpr auto kMinSweepInterval = std::chrono::seconds(1); +constexpr auto kMaxSweepInterval = std::chrono::minutes(1); + +// The loader may call user-provided code such as an HttpClient, which can throw. +// Turning exceptions into errors lets Get() always clean up the loading slot. +Result> RunLoader(const AuthSessionCache::Loader& loader) { + try { + return loader(); + } catch (const std::exception& e) { + return UnknownError("Auth session loader threw: {}", e.what()); + } catch (...) { + return UnknownError("Auth session loader threw a non-standard exception"); + } +} + +} // namespace + +Result> AuthSessionCache::Make( + std::chrono::milliseconds session_timeout, RemovalListener on_removal, Clock clock) { + ICEBERG_PRECHECK(session_timeout >= std::chrono::milliseconds::zero(), + "Auth session timeout must not be negative: {} ms", + session_timeout.count()); + ICEBERG_PRECHECK(clock != nullptr, "Auth session cache clock must not be null"); + return std::shared_ptr( + new AuthSessionCache(session_timeout, std::move(on_removal), std::move(clock))); +} + +std::chrono::milliseconds AuthSessionCache::DefaultSweepInterval( + std::chrono::milliseconds session_timeout) { + return std::clamp(session_timeout, kMinSweepInterval, + kMaxSweepInterval); +} + +AuthSessionCache::AuthSessionCache(std::chrono::milliseconds session_timeout, + RemovalListener on_removal, Clock clock) + : session_timeout_(session_timeout), + on_removal_(std::move(on_removal)), + clock_(std::move(clock)) {} + +AuthSessionCache::~AuthSessionCache() { Close(); } + +Result> AuthSessionCache::Get(const std::string& key, + const Loader& loader) { + // Phase 1: Look up cache (under lock) + std::shared_ptr slot; + std::vector> removed; + { + std::unique_lock lock(mutex_); + while (true) { + if (closed_) { + return Invalid("Auth session cache is closed"); + } + auto it = slots_.find(key); + if (it == slots_.end()) { + break; + } + auto existing = it->second; + if (existing->loading) { + // Wait for another thread to finish loading. Loop again after waking + // because a failed load removes the slot. + load_finished_.wait(lock, [&] { return closed_ || !existing->loading; }); + continue; + } + auto now = clock_(); + if (!IsExpired(*existing, now)) { + TouchLocked(it, now); + return existing->session; + } + EraseLocked(it, removed); + break; + } + if (!sweep_scheduler_.has_value()) { + // Without a periodic sweep, keys that are never accessed again would stay + // until Close(). Reclaim idle entries before adding one, so the cache grows + // with recently used keys rather than with every key ever used. + EvictIdleLocked(clock_(), removed); + } + // Create loading slot. Other threads calling Get(key) will wait. + slot = std::make_shared(); + slots_.emplace(key, slot); + } + + // Phase 2: Call loader (lock released) + NotifyRemoved(removed); + auto result = RunLoader(loader); + if (result.has_value() && result.value() == nullptr) { + result = Invalid("Auth session loader returned a null session"); + } + + // Phase 3: Update cache (re-acquire lock) + std::shared_ptr orphan; + { + std::lock_guard lock(mutex_); + slot->loading = false; + auto it = slots_.find(key); + bool current = it != slots_.end() && it->second == slot; + + if (!result.has_value()) { + // Failed loads are not cached. Remove the slot so next Get() can retry; a + // loading slot is not in access_order_ yet. + if (current) { + slots_.erase(it); + } + } else if (closed_ || !current) { + // Cache closed during load, or slot replaced. Mark as orphan. + orphan = result.value(); + } else { + slot->session = result.value(); + TouchLocked(it, clock_()); + } + } + load_finished_.notify_all(); + + if (orphan != nullptr) { + if (on_removal_) { + on_removal_(orphan); + } + return Invalid("Auth session cache is closed"); + } + return result; +} + +void AuthSessionCache::Sweep() { + std::vector> removed; + { + std::lock_guard lock(mutex_); + EvictIdleLocked(clock_(), removed); + } + NotifyRemoved(removed); +} + +void AuthSessionCache::TouchLocked(SlotMap::iterator it, + std::chrono::steady_clock::time_point now) { + // Every last_access is read from the clock under the lock and its entry moved to + // the back, so access_order_ stays sorted by last access. + Slot& slot = *it->second; + slot.last_access = now; + if (slot.access_pos.has_value()) { + access_order_.splice(access_order_.end(), access_order_, *slot.access_pos); + } else { + slot.access_pos = access_order_.insert(access_order_.end(), + AccessEntry{.key = &it->first, .slot = &slot}); + } +} + +void AuthSessionCache::EraseLocked(SlotMap::iterator it, + std::vector>& removed) { + Slot& slot = *it->second; + if (slot.access_pos.has_value()) { + access_order_.erase(*slot.access_pos); + } + removed.push_back(std::move(slot.session)); + slots_.erase(it); +} + +void AuthSessionCache::EvictIdleLocked( + std::chrono::steady_clock::time_point now, + std::vector>& removed) { + // Every entry after the first one that is not idle was accessed more recently. + while (!access_order_.empty()) { + const AccessEntry& oldest = access_order_.front(); + if (now - oldest.slot->last_access < session_timeout_) { + break; + } + EraseLocked(slots_.find(*oldest.key), removed); + } +} + +Status AuthSessionCache::StartPeriodicSweep(Executor& executor, + std::chrono::milliseconds interval) { + ICEBERG_PRECHECK(interval > std::chrono::milliseconds::zero(), + "Auth session sweep interval must be positive: {} ms", + interval.count()); + std::lock_guard lock(mutex_); + if (closed_ || sweep_scheduler_.has_value()) { + return {}; + } + sweep_scheduler_.emplace(executor); + sweep_interval_ = interval; + ScheduleSweepLocked(); + return {}; +} + +void AuthSessionCache::Close() { + std::vector> removed; + { + std::lock_guard lock(mutex_); + if (closed_) { + return; + } + closed_ = true; + for (auto& [key, slot] : slots_) { + // Sessions still loading are handed to the listener when their load ends. + if (!slot->loading) { + removed.push_back(std::move(slot->session)); + } + } + slots_.clear(); + access_order_.clear(); + } + load_finished_.notify_all(); + // Shut down outside the lock: Shutdown() waits for an in-progress submission, + // whose callback may take the lock. No sweep is scheduled once closed_ is set. + if (sweep_scheduler_.has_value()) { + sweep_scheduler_->Shutdown(); + } + NotifyRemoved(removed); +} + +size_t AuthSessionCache::size() const { + std::lock_guard lock(mutex_); + return slots_.size(); +} + +bool AuthSessionCache::IsExpired(const Slot& slot, + std::chrono::steady_clock::time_point now) const { + // Check 1: Idle timeout (expireAfterAccess semantics) + if (now - slot.last_access >= session_timeout_) { + return true; + } + + // Check 2: Token expiration, for sessions that will not refresh + // This covers sessions without refresh and sessions whose refresh gave up after + // repeated failures. Such a session is useless once its token expires, and + // evicting it is harmless since stopping its refresh is a no-op; the next Get() + // loads a new one. A session that is still refreshing is left alone: its refresh + // may be in flight, and evicting it would stop the refresh for callers that + // still hold it. + auto oauth2 = std::dynamic_pointer_cast(slot.session); + if (oauth2 && !oauth2->IsRefreshable()) { + auto expires_at = oauth2->ExpiresAt(); + return expires_at.has_value() && now >= *expires_at; + } + + return false; +} + +void AuthSessionCache::ScheduleSweepLocked() { + // Scheduling under the cache lock ensures Close() either cancels this task or + // sees closed_ and keeps it from being scheduled. + std::weak_ptr weak_self = weak_from_this(); + sweep_scheduler_->Schedule( + sweep_interval_, ExecutorTask([weak_self] { + if (auto self = weak_self.lock()) { + self->RunScheduledSweep(); + } + }), + ::iceberg::internal::TaskSubmissionCallback([weak_self](Status status) { + if (status.has_value()) { + return; + } + // The sweep did not run; keep sweeping so idle sessions are still reclaimed. + ICEBERG_LOG_WARN("Failed to submit auth session cache sweep: {}", + status.error().message); + if (auto self = weak_self.lock()) { + self->ScheduleNextSweep(); + } + })); +} + +void AuthSessionCache::ScheduleNextSweep() { + std::lock_guard lock(mutex_); + if (!closed_) { + ScheduleSweepLocked(); + } +} + +void AuthSessionCache::RunScheduledSweep() { + Sweep(); + ScheduleNextSweep(); +} + +void AuthSessionCache::NotifyRemoved( + std::vector>& sessions) const { + if (!on_removal_) { + return; + } + for (auto& session : sessions) { + if (session != nullptr) { + on_removal_(session); + } + } +} + +} // namespace iceberg::rest::auth::internal diff --git a/src/iceberg/catalog/rest/auth/auth_session_cache_internal.h b/src/iceberg/catalog/rest/auth/auth_session_cache_internal.h new file mode 100644 index 000000000..aee88147e --- /dev/null +++ b/src/iceberg/catalog/rest/auth/auth_session_cache_internal.h @@ -0,0 +1,161 @@ +/* + * 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. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "iceberg/catalog/rest/iceberg_rest_export.h" +#include "iceberg/catalog/rest/type_fwd.h" +#include "iceberg/result.h" +#include "iceberg/util/executor.h" +#include "iceberg/util/scheduled_executor_internal.h" + +/// \file iceberg/catalog/rest/auth/auth_session_cache_internal.h +/// \brief Cache of child authentication sessions, keyed by caller identity. + +namespace iceberg::rest::auth::internal { + +/// \brief Caches child auth sessions and expires them after a period of inactivity. +/// +/// Entries expire after `session_timeout` without access. OAuth2 sessions that +/// will not refresh also expire once their token does; such an entry is replaced +/// when its key is next accessed. Each session that leaves the cache is handed to +/// the removal listener exactly once. +/// +/// Idle entries are reclaimed in access order, least recently used first, so the +/// cost is proportional to the number of entries removed. Sweep() does this, and +/// StartPeriodicSweep() runs it in the background; without a periodic sweep, Get() +/// does it before adding an entry. +/// +/// Thread safety: all public methods are thread-safe. Loaders and the removal +/// listener are always invoked without holding the cache lock. +class ICEBERG_REST_EXPORT AuthSessionCache + : public std::enable_shared_from_this { + public: + using Clock = std::function; + using Loader = std::function>()>; + /// Called once for every session that leaves the cache. Must not block or throw. + using RemovalListener = std::function&)>; + + /// \brief Create a cache. + /// + /// \param session_timeout Idle time after which an entry expires; must not be + /// negative. Zero disables caching in effect. + /// \param on_removal Listener invoked for each session that leaves the cache. + /// \param clock Time source; tests can inject a fake clock. + static Result> Make( + std::chrono::milliseconds session_timeout, RemovalListener on_removal, + Clock clock = std::chrono::steady_clock::now); + + ~AuthSessionCache(); + + AuthSessionCache(const AuthSessionCache&) = delete; + AuthSessionCache& operator=(const AuthSessionCache&) = delete; + + /// \brief Suggested StartPeriodicSweep() interval for `session_timeout`: the + /// timeout clamped to [1s, 60s]. + static std::chrono::milliseconds DefaultSweepInterval( + std::chrono::milliseconds session_timeout); + + /// \brief Return the cached session for `key`, loading it on a miss. + /// + /// Concurrent calls for the same key run `loader` at most once; the others wait + /// for its result. Failed loads are not cached; an exception thrown by `loader` + /// is returned as an UnknownError. Returns an error once the cache + /// is closed; a session loaded after Close() is handed to the removal listener. + Result> Get(const std::string& key, const Loader& loader); + + /// \brief Remove entries idle for at least the session timeout and notify the + /// removal listener. Entries whose token expired are replaced on their next + /// access instead. + void Sweep(); + + /// \brief Run Sweep() on `executor` every `interval` until Close(). + /// + /// `executor` is borrowed and must outlive this cache and its submitted tasks. + /// `interval` must be positive. Calling this more than once, or after Close(), + /// has no effect. + Status StartPeriodicSweep(Executor& executor, std::chrono::milliseconds interval); + + /// \brief Stop periodic sweeping and remove every entry. Idempotent. + void Close(); + + /// \brief Number of entries, including loads in progress. + size_t size() const; + + private: + struct Slot; + + /// One element per loaded slot. `key` points at the key stored in slots_, which + /// stays valid until that element is erased, even when the map rehashes. + struct AccessEntry { + const std::string* key; + Slot* slot; + }; + using AccessOrder = std::list; + + struct Slot { + bool loading = true; + std::shared_ptr session; + std::chrono::steady_clock::time_point last_access; + /// Position in access_order_; set once the load succeeds. + std::optional access_pos; + }; + using SlotMap = std::unordered_map>; + + AuthSessionCache(std::chrono::milliseconds session_timeout, RemovalListener on_removal, + Clock clock); + + bool IsExpired(const Slot& slot, std::chrono::steady_clock::time_point now) const; + void TouchLocked(SlotMap::iterator it, std::chrono::steady_clock::time_point now); + void EraseLocked(SlotMap::iterator it, + std::vector>& removed); + void EvictIdleLocked(std::chrono::steady_clock::time_point now, + std::vector>& removed); + void ScheduleSweepLocked(); + void ScheduleNextSweep(); + void RunScheduledSweep(); + void NotifyRemoved(std::vector>& sessions) const; + + const std::chrono::milliseconds session_timeout_; + const RemovalListener on_removal_; + const Clock clock_; + + mutable std::mutex mutex_; + std::condition_variable load_finished_; + SlotMap slots_; + /// Loaded slots ordered by last access, least recent first. + AccessOrder access_order_; + bool closed_ = false; + std::optional<::iceberg::internal::ScheduledExecutor> sweep_scheduler_; + std::chrono::milliseconds sweep_interval_{0}; +}; + +} // namespace iceberg::rest::auth::internal diff --git a/src/iceberg/catalog/rest/auth/auth_session_internal.h b/src/iceberg/catalog/rest/auth/auth_session_internal.h index fa2f7ede9..a4454aafc 100644 --- a/src/iceberg/catalog/rest/auth/auth_session_internal.h +++ b/src/iceberg/catalog/rest/auth/auth_session_internal.h @@ -102,11 +102,46 @@ class OAuth2Session final : public AuthSession, .scope = config_.scope, .oauth2_server_uri = config_.token_endpoint, .optional_oauth_params = config_.optional_oauth_params, + .headers = headers_, }; } Status Close() override { return CloseImpl(); } + /// \brief Stop scheduling token refreshes without waiting. + /// + /// Unlike Close(), this does not wait for an in-flight refresh to finish, so it + /// is safe to call from any thread, including refresh executor workers. The + /// session keeps its current token and remains usable; an in-flight refresh may + /// still update the token but will not schedule further refreshes. + void StopRefreshing() { + { + std::lock_guard lock(refresh_mutex_); + closed_ = true; + } + if (scheduler_) { + scheduler_->Shutdown(); + } + } + + /// \brief Expiration time of the current token, or std::nullopt if unknown. + std::optional ExpiresAt() const { + std::shared_lock lock(mutex_); + if (expires_at_ == std::chrono::steady_clock::time_point{}) { + return std::nullopt; + } + return expires_at_; + } + + /// \brief Whether this session will keep refreshing its token. + /// + /// False if refresh is disabled, was stopped by StopRefreshing() or Close(), or + /// gave up after repeated failures. + bool IsRefreshable() const { + std::lock_guard lock(refresh_mutex_); + return scheduler_.has_value() && !closed_ && !refresh_stopped_; + } + ~OAuth2Session() override { std::ignore = CloseImpl(); } private: @@ -221,11 +256,18 @@ class OAuth2Session final : public AuthSession, }); if (!result.has_value()) { ICEBERG_LOG_WARN("Failed to refresh OAuth2 token: {}", result.error().message); + MarkRefreshStopped(); return std::nullopt; } return *result; } + // Record that no further refresh will be scheduled after a failure. + void MarkRefreshStopped() { + std::lock_guard lock(refresh_mutex_); + refresh_stopped_ = true; + } + Result> RefreshToken() { { std::lock_guard lock(refresh_mutex_); @@ -276,6 +318,9 @@ class OAuth2Session final : public AuthSession, ICEBERG_LOG_WARN( "Stopping OAuth2 token refresh after repeated submission failures: {}", status.error().message); + if (auto self = weak_self.lock()) { + self->MarkRefreshStopped(); + } return; } if (auto self = weak_self.lock()) { @@ -313,7 +358,9 @@ class OAuth2Session final : public AuthSession, AuthProperties refresh_properties_; std::shared_ptr client_; bool closed_ = false; - std::mutex refresh_mutex_; + // Set when refresh gives up after repeated failures; guarded by refresh_mutex_. + bool refresh_stopped_ = false; + mutable std::mutex refresh_mutex_; std::condition_variable refresh_cv_; int active_refresh_count_ = 0; std::optional<::iceberg::internal::ScheduledExecutor> scheduler_; diff --git a/src/iceberg/catalog/session_context.h b/src/iceberg/catalog/session_context.h index 19eb8fc5f..1a6bc819c 100644 --- a/src/iceberg/catalog/session_context.h +++ b/src/iceberg/catalog/session_context.h @@ -41,6 +41,9 @@ namespace iceberg { /// silently overriding credentials. struct ICEBERG_EXPORT SessionContext { /// Unique session identifier. Explicit contexts must provide a non-empty ID. + /// + /// A session ID identifies a stable caller context. Its identity and credentials + /// must not change while the ID is in use. std::string session_id; /// Caller identity. Empty means no identity was supplied. This value is diff --git a/src/iceberg/test/CMakeLists.txt b/src/iceberg/test/CMakeLists.txt index 37e20e339..7fd77ff99 100644 --- a/src/iceberg/test/CMakeLists.txt +++ b/src/iceberg/test/CMakeLists.txt @@ -335,6 +335,7 @@ if(ICEBERG_BUILD_REST) add_rest_iceberg_test(rest_catalog_test SOURCES auth_manager_test.cc + auth_session_cache_test.cc catalog_properties_test.cc error_handlers_test.cc endpoint_test.cc diff --git a/src/iceberg/test/auth_manager_test.cc b/src/iceberg/test/auth_manager_test.cc index f6f1400b2..7b6f0f3a6 100644 --- a/src/iceberg/test/auth_manager_test.cc +++ b/src/iceberg/test/auth_manager_test.cc @@ -433,6 +433,35 @@ TEST_F(AuthManagerTest, OAuth2StaticToken) { "catalog-resource"); } +TEST_F(AuthManagerTest, OAuth2ZeroSessionTimeoutCreatesCatalogSession) { + std::unordered_map properties = { + {AuthProperties::kAuthType, "oauth2"}, + {AuthProperties::kToken.key(), "my-static-token"}, + {AuthProperties::kSessionTimeoutMs.key(), "0"}, + }; + + ICEBERG_UNWRAP_OR_FAIL(auto manager, AuthManagers::Load("test-catalog", properties)); + ICEBERG_UNWRAP_OR_FAIL(auto session, manager->CatalogSession(client_, properties)); + + EXPECT_NE(session, nullptr); + EXPECT_THAT(manager->Close(), IsOk()); +} + +TEST_F(AuthManagerTest, OAuth2InvalidSessionTimeoutReturnsError) { + for (const std::string value : {"1h", "abc", "-1"}) { + std::unordered_map properties = { + {AuthProperties::kAuthType, "oauth2"}, + {AuthProperties::kToken.key(), "my-static-token"}, + {AuthProperties::kSessionTimeoutMs.key(), value}, + }; + + ICEBERG_UNWRAP_OR_FAIL(auto manager, AuthManagers::Load("test-catalog", properties)); + EXPECT_THAT(manager->CatalogSession(client_, properties), + IsError(ErrorKind::kInvalidArgument)) + << value; + } +} + // Verifies OAuth2 type is inferred from token property TEST_F(AuthManagerTest, OAuth2InferredFromToken) { std::unordered_map properties = { @@ -554,6 +583,29 @@ TEST_F(AuthManagerTest, OAuth2TableTokenCreatesChild) { EXPECT_TRUE(info->credential.empty()); } +TEST_F(AuthManagerTest, OAuth2TableTokenSessionIsCachedByToken) { + std::unordered_map properties = { + {AuthProperties::kAuthType, "oauth2"}, + }; + ICEBERG_UNWRAP_OR_FAIL(auto manager, AuthManagers::Load("test-catalog", properties)); + ICEBERG_UNWRAP_OR_FAIL(auto parent, manager->CatalogSession(client_, properties)); + TableIdentifier table{.ns = Namespace{{"db"}}, .name = "table"}; + + ICEBERG_UNWRAP_OR_FAIL( + auto first, + manager->TableSession(table, {{AuthProperties::kToken.key(), "token-a"}}, parent)); + ICEBERG_UNWRAP_OR_FAIL( + auto second, + manager->TableSession(table, {{AuthProperties::kToken.key(), "token-a"}}, parent)); + ICEBERG_UNWRAP_OR_FAIL( + auto other, + manager->TableSession(table, {{AuthProperties::kToken.key(), "token-b"}}, parent)); + + EXPECT_EQ(first, second); + EXPECT_NE(first, other); + EXPECT_THAT(manager->Close(), IsOk()); +} + TEST_F(AuthManagerTest, OAuth2TableIgnoresCredential) { std::unordered_map properties = { {AuthProperties::kAuthType, "oauth2"}, @@ -874,18 +926,75 @@ TEST(OAuth2SessionTest, StopsAfterHttpRetryBudget) { }); auto client = std::make_shared(); ICEBERG_UNWRAP_OR_FAIL(auto session, RefreshingSession(client, executor)); + auto oauth2 = std::dynamic_pointer_cast(session); + ASSERT_NE(oauth2, nullptr); ASSERT_EQ(task.wait_for(std::chrono::seconds(5)), std::future_status::ready); + EXPECT_TRUE(oauth2->IsRefreshable()); test::FakeRetryEnvironment fake_retry; ScopedRetryTestHooks hooks(fake_retry.hooks()); std::move(task.get())(); EXPECT_EQ(client->requests, 6); EXPECT_THAT(fake_retry.observed_base_delays_ms(), ::testing::ElementsAre(100, 200, 400, 800, 1600)); + // No further refresh is scheduled, so the session no longer counts as refreshable. + EXPECT_FALSE(oauth2->IsRefreshable()); ICEBERG_UNWRAP_OR_FAIL(auto request, session->Authenticate({})); EXPECT_EQ(request.headers.at("Authorization"), "Bearer initial"); ASSERT_THAT(session->Close(), IsOk()); } +TEST(OAuth2SessionTest, StopRefreshingKeepsSessionUsable) { + MockExecutor executor; + EXPECT_CALL(executor, Submit).Times(0); + OAuthTokenResponse token_response{ + .access_token = "token", + .token_type = "bearer", + .expires_in_secs = 3600, + }; + ICEBERG_UNWRAP_OR_FAIL( + auto session, + internal::OAuth2Session::Make(token_response, + {.token_endpoint = "http://localhost/oauth/tokens", + .client_id = "client_id", + .client_secret = "client_secret", + .scope = "catalog", + .refresh_executor = &executor}, + std::make_shared(), std::nullopt)); + EXPECT_TRUE(session->IsRefreshable()); + + session->StopRefreshing(); + session->StopRefreshing(); + EXPECT_FALSE(session->IsRefreshable()); + + ICEBERG_UNWRAP_OR_FAIL(auto request, session->Authenticate({})); + EXPECT_EQ(request.headers.at("Authorization"), "Bearer token"); + EXPECT_THAT(session->Close(), IsOk()); +} + +TEST(OAuth2SessionTest, ExpiresAtReflectsTokenLifetime) { + auto make_session = [](std::optional expires_in_secs) { + OAuthTokenResponse token_response{ + .access_token = "opaque-token", + .token_type = "bearer", + .expires_in_secs = expires_in_secs, + }; + return internal::OAuth2Session::Make( + token_response, {.token_endpoint = "http://localhost/oauth/tokens"}, + std::make_shared(), std::nullopt); + }; + + auto before = std::chrono::steady_clock::now(); + ICEBERG_UNWRAP_OR_FAIL(auto expiring, make_session(60)); + auto after = std::chrono::steady_clock::now(); + auto expires_at = expiring->ExpiresAt(); + ASSERT_TRUE(expires_at.has_value()); + EXPECT_GE(*expires_at, before + std::chrono::seconds(60)); + EXPECT_LE(*expires_at, after + std::chrono::seconds(60)); + + ICEBERG_UNWRAP_OR_FAIL(auto unknown, make_session(std::nullopt)); + EXPECT_FALSE(unknown->ExpiresAt().has_value()); +} + TEST(OAuth2SessionTest, InitTokenExpirationUsesRequestStartTime) { OAuthTokenResponse token_response{ .access_token = "opaque-token", @@ -917,4 +1026,214 @@ TEST(OAuth2SessionTest, JwtExpirationTakesPriorityOverExpiresIn) { EXPECT_EQ(*expiration, std::chrono::steady_clock::time_point(std::chrono::seconds(70))); } +namespace { + +/// Creates a credential child session from a static-token catalog whose token +/// endpoint always succeeds. +Result> CredentialChildSession( + AuthManager& manager) { + auto client = std::make_shared(); + client->response = []() -> Result { + return HttpResponse( + 200, R"({"access_token":"child-token","token_type":"bearer","expires_in":3600})"); + }; + std::unordered_map properties = { + {AuthProperties::kAuthType, "oauth2"}, + {AuthProperties::kToken.key(), "catalog-token"}, + }; + ICEBERG_ASSIGN_OR_RAISE(auto parent, manager.CatalogSession(client, properties)); + SessionContext context{ + .session_id = "tenant-a", + .credentials = {{AuthProperties::kCredential.key(), "client:secret"}}, + }; + ICEBERG_ASSIGN_OR_RAISE(auto child, manager.ContextualSession(context, parent)); + auto oauth2 = std::dynamic_pointer_cast(child); + ICEBERG_PRECHECK(oauth2 != nullptr, "child session is not an OAuth2 session"); + return oauth2; +} + +} // namespace + +TEST(OAuth2ChildSessionTest, CredentialChildRefreshesWithExecutor) { + // The refresh is due long after this test ends, so nothing is submitted. + MockExecutor executor; + EXPECT_CALL(executor, Submit).Times(0); + ICEBERG_UNWRAP_OR_FAIL( + auto manager, + AuthManagers::Load("test-catalog", {{AuthProperties::kAuthType, "oauth2"}})); + manager->SetAuthRefreshExecutor(&executor); + + ICEBERG_UNWRAP_OR_FAIL(auto child, CredentialChildSession(*manager)); + + EXPECT_TRUE(child->IsRefreshable()); + ICEBERG_UNWRAP_OR_FAIL(auto request, child->Authenticate({})); + EXPECT_EQ(request.headers.at("Authorization"), "Bearer child-token"); + EXPECT_THAT(manager->Close(), IsOk()); +} + +TEST(OAuth2ChildSessionTest, CredentialChildWithoutExecutorDoesNotRefresh) { + ICEBERG_UNWRAP_OR_FAIL( + auto manager, + AuthManagers::Load("test-catalog", {{AuthProperties::kAuthType, "oauth2"}})); + + ICEBERG_UNWRAP_OR_FAIL(auto child, CredentialChildSession(*manager)); + + EXPECT_FALSE(child->IsRefreshable()); + ICEBERG_UNWRAP_OR_FAIL(auto request, child->Authenticate({})); + EXPECT_EQ(request.headers.at("Authorization"), "Bearer child-token"); + EXPECT_THAT(manager->Close(), IsOk()); +} + +TEST(OAuth2ChildSessionTest, ContextSessionIsReloadedAfterRefreshGivesUp) { + MockExecutor executor; + std::promise queued; + auto task = queued.get_future(); + EXPECT_CALL(executor, Submit).WillOnce([&](ExecutorTask work) -> Status { + queued.set_value(std::move(work)); + return {}; + }); + // The token endpoint first issues an already expired token, then fails until it + // recovers. + auto client = std::make_shared(); + std::atomic recovered = false; + client->response = [&]() -> Result { + if (client->requests == 1) { + return HttpResponse( + 200, R"({"access_token":"expired","token_type":"bearer","expires_in":0})"); + } + if (!recovered) { + return IOError("token endpoint unavailable"); + } + return HttpResponse( + 200, R"({"access_token":"fresh","token_type":"bearer","expires_in":3600})"); + }; + std::unordered_map properties = { + {AuthProperties::kAuthType, "oauth2"}, + {AuthProperties::kToken.key(), "catalog-token"}, + }; + ICEBERG_UNWRAP_OR_FAIL(auto manager, AuthManagers::Load("test-catalog", properties)); + manager->SetAuthRefreshExecutor(&executor); + ICEBERG_UNWRAP_OR_FAIL(auto parent, manager->CatalogSession(client, properties)); + SessionContext context{ + .session_id = "tenant-a", + .credentials = {{AuthProperties::kCredential.key(), "client:secret"}}, + }; + + ICEBERG_UNWRAP_OR_FAIL(auto first, manager->ContextualSession(context, parent)); + ASSERT_EQ(task.wait_for(std::chrono::seconds(5)), std::future_status::ready); + // While its refresh is pending, the expired session stays cached. + ICEBERG_UNWRAP_OR_FAIL(auto pending, manager->ContextualSession(context, parent)); + EXPECT_EQ(pending, first); + + { + test::FakeRetryEnvironment fake_retry; + ScopedRetryTestHooks hooks(fake_retry.hooks()); + std::move(task.get())(); // every refresh attempt fails + } + recovered = true; + + // The refresh gave up and the token has expired, so a new session is loaded. + ICEBERG_UNWRAP_OR_FAIL(auto reloaded, manager->ContextualSession(context, parent)); + EXPECT_NE(reloaded, first); + ICEBERG_UNWRAP_OR_FAIL(auto request, reloaded->Authenticate({})); + EXPECT_EQ(request.headers.at("Authorization"), "Bearer fresh"); + EXPECT_THAT(manager->Close(), IsOk()); +} + +namespace { + +/// Answers token exchanges with a token that records the subject and actor. +class ExchangeHttpClient : public HttpClient { + public: + Result PostForm(const std::string&, + const std::unordered_map& form, + const std::unordered_map&, + const ErrorHandler&, AuthSession&) override { + ++requests; + auto actor = form.find("actor_token"); + std::string token = "sub=" + form.at("subject_token") + ",act=" + + (actor == form.end() ? std::string("none") : actor->second); + return HttpResponse(200, R"({"access_token":")" + token + + R"(","token_type":"bearer","issued_token_type":")" + + AuthProperties::kAccessTokenType + R"("})"); + } + std::atomic requests{0}; +}; + +class TableSessionCacheKeyTest : public ::testing::Test { + protected: + void SetUp() override { + ICEBERG_UNWRAP_OR_FAIL( + manager_, + AuthManagers::Load("test-catalog", {{AuthProperties::kAuthType, "oauth2"}})); + ICEBERG_UNWRAP_OR_FAIL( + catalog_, manager_->CatalogSession( + client_, {{AuthProperties::kAuthType, "oauth2"}, + {AuthProperties::kToken.key(), "catalog-token"}})); + } + + void TearDown() override { EXPECT_THAT(manager_->Close(), IsOk()); } + + /// A contextual session holding `token`, used as a table session parent. + Result> Parent(const std::string& session_id, + const std::string& token) { + return manager_->ContextualSession( + {.session_id = session_id, + .credentials = {{AuthProperties::kToken.key(), token}}}, + catalog_); + } + + Result> Table( + const std::unordered_map& config, + std::shared_ptr parent) { + return manager_->TableSession(table_, config, std::move(parent)); + } + + static std::string Authorization(AuthSession& session) { + auto request = session.Authenticate({}); + return request.has_value() ? request->headers.at("Authorization") : ""; + } + + std::shared_ptr client_ = std::make_shared(); + std::unique_ptr manager_; + std::shared_ptr catalog_; + TableIdentifier table_{.ns = Namespace{{"db"}}, .name = "orders"}; +}; + +} // namespace + +TEST_F(TableSessionCacheKeyTest, ExchangedSessionsAreKeyedByActor) { + ICEBERG_UNWRAP_OR_FAIL(auto alice, Parent("alice", "alice-token")); + ICEBERG_UNWRAP_OR_FAIL(auto bob, Parent("bob", "bob-token")); + std::unordered_map config = { + {AuthProperties::kJwtTokenType, "table-jwt"}}; + + ICEBERG_UNWRAP_OR_FAIL(auto alice_first, Table(config, alice)); + ICEBERG_UNWRAP_OR_FAIL(auto alice_again, Table(config, alice)); + ICEBERG_UNWRAP_OR_FAIL(auto bob_session, Table(config, bob)); + + EXPECT_EQ(alice_again, alice_first); + EXPECT_NE(bob_session, alice_first); + EXPECT_EQ(client_->requests, 2); + EXPECT_EQ(Authorization(*alice_first), "Bearer sub=table-jwt,act=alice-token"); + EXPECT_EQ(Authorization(*bob_session), "Bearer sub=table-jwt,act=bob-token"); +} + +TEST_F(TableSessionCacheKeyTest, SessionsAreKeyedByFlowAndTokenType) { + ICEBERG_UNWRAP_OR_FAIL(auto alice, Parent("alice", "alice-token")); + + ICEBERG_UNWRAP_OR_FAIL(auto bearer, + Table({{AuthProperties::kToken.key(), "same-value"}}, alice)); + ICEBERG_UNWRAP_OR_FAIL(auto jwt, + Table({{AuthProperties::kJwtTokenType, "same-value"}}, alice)); + ICEBERG_UNWRAP_OR_FAIL(auto saml2, + Table({{AuthProperties::kSaml2TokenType, "same-value"}}, alice)); + + EXPECT_NE(jwt, bearer); + EXPECT_NE(saml2, jwt); + EXPECT_EQ(client_->requests, 2); + EXPECT_EQ(Authorization(*bearer), "Bearer same-value"); + EXPECT_EQ(Authorization(*jwt), "Bearer sub=same-value,act=alice-token"); +} + } // namespace iceberg::rest::auth diff --git a/src/iceberg/test/auth_session_cache_test.cc b/src/iceberg/test/auth_session_cache_test.cc new file mode 100644 index 000000000..074459e73 --- /dev/null +++ b/src/iceberg/test/auth_session_cache_test.cc @@ -0,0 +1,582 @@ +/* + * 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "iceberg/catalog/rest/auth/auth_session.h" +#include "iceberg/catalog/rest/auth/auth_session_cache_internal.h" +#include "iceberg/catalog/rest/auth/auth_session_internal.h" +#include "iceberg/catalog/rest/http_client.h" +#include "iceberg/test/executor.h" +#include "iceberg/test/matchers.h" + +namespace iceberg::rest::auth { + +namespace { + +using internal::AuthSessionCache; +using namespace std::chrono_literals; + +constexpr auto kTimeout = std::chrono::milliseconds(1000); +constexpr auto kTokenLifetime = std::chrono::seconds(60); + +/// Starts at the real steady clock so it is comparable with OAuth2Session::ExpiresAt(). +class FakeClock { + public: + std::chrono::steady_clock::time_point Now() const { + return std::chrono::steady_clock::time_point(std::chrono::milliseconds(now_ms_)); + } + + AuthSessionCache::Clock AsClock() { + return [this] { return Now(); }; + } + + void Advance(std::chrono::milliseconds delta) { now_ms_ += delta.count(); } + + private: + std::atomic now_ms_{std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count()}; +}; + +class RemovalRecorder { + public: + AuthSessionCache::RemovalListener Listener() { + return [this](const std::shared_ptr& session) { + std::lock_guard lock(mutex_); + removed_.push_back(session); + }; + } + + std::vector> removed() const { + std::lock_guard lock(mutex_); + return removed_; + } + + private: + mutable std::mutex mutex_; + std::vector> removed_; +}; + +/// One-shot latch used to hold a loader inside Get(). +class Gate { + public: + void Open() { + { + std::lock_guard lock(mutex_); + open_ = true; + } + cv_.notify_all(); + } + + void Wait() { + std::unique_lock lock(mutex_); + cv_.wait(lock, [this] { return open_; }); + } + + private: + std::mutex mutex_; + std::condition_variable cv_; + bool open_ = false; +}; + +AuthSessionCache::Loader CountingLoader(std::atomic& calls) { + return [&calls]() -> Result> { + ++calls; + return AuthSession::MakeDefault({}); + }; +} + +/// Loads an OAuth2 session whose token expires kTokenLifetime after `clock` reads. +/// The session refreshes on `refresh_executor` when it is not null. +AuthSessionCache::Loader OAuth2Loader(std::atomic& calls, const FakeClock& clock, + Executor* refresh_executor) { + return [&calls, &clock, refresh_executor]() -> Result> { + ++calls; + OAuthTokenResponse token_response{ + .access_token = "opaque-token", + .token_type = "bearer", + .expires_in_secs = kTokenLifetime.count(), + }; + ICEBERG_ASSIGN_OR_RAISE( + auto session, + internal::OAuth2Session::Make(token_response, + {.token_endpoint = "http://localhost/oauth/tokens", + .refresh_executor = refresh_executor}, + std::make_shared(), clock.Now())); + return session; + }; +} + +} // namespace + +class AuthSessionCacheTest : public ::testing::Test { + protected: + std::shared_ptr MakeCache( + std::chrono::milliseconds timeout = kTimeout) { + auto cache = AuthSessionCache::Make(timeout, recorder_.Listener(), clock_.AsClock()); + EXPECT_THAT(cache, IsOk()); + return cache.value(); + } + + // Declared before any cache so they outlive it. + FakeClock clock_; + RemovalRecorder recorder_; +}; + +TEST_F(AuthSessionCacheTest, RejectsNegativeTimeout) { + EXPECT_THAT(AuthSessionCache::Make(-1ms, recorder_.Listener()), + IsError(ErrorKind::kInvalidArgument)); +} + +TEST_F(AuthSessionCacheTest, ReturnsCachedSessionOnHit) { + auto cache = MakeCache(); + std::atomic calls = 0; + + ICEBERG_UNWRAP_OR_FAIL(auto first, cache->Get("tenant-a", CountingLoader(calls))); + ICEBERG_UNWRAP_OR_FAIL(auto second, cache->Get("tenant-a", CountingLoader(calls))); + + EXPECT_EQ(first, second); + EXPECT_EQ(calls, 1); + EXPECT_EQ(cache->size(), 1); +} + +TEST_F(AuthSessionCacheTest, KeysAreIsolated) { + auto cache = MakeCache(); + std::atomic calls = 0; + + ICEBERG_UNWRAP_OR_FAIL(auto a, cache->Get("tenant-a", CountingLoader(calls))); + ICEBERG_UNWRAP_OR_FAIL(auto b, cache->Get("tenant-b", CountingLoader(calls))); + + EXPECT_NE(a, b); + EXPECT_EQ(calls, 2); +} + +TEST_F(AuthSessionCacheTest, FailedLoadIsNotCached) { + auto cache = MakeCache(); + auto failing = []() -> Result> { + return AuthenticationFailed("token endpoint unavailable"); + }; + std::atomic calls = 0; + + EXPECT_THAT(cache->Get("tenant-a", failing), IsError(ErrorKind::kAuthenticationFailed)); + EXPECT_EQ(cache->size(), 0); + EXPECT_THAT(cache->Get("tenant-a", CountingLoader(calls)), IsOk()); + EXPECT_EQ(calls, 1); +} + +TEST_F(AuthSessionCacheTest, NullSessionIsRejected) { + auto cache = MakeCache(); + auto null_loader = []() -> Result> { return nullptr; }; + + EXPECT_THAT(cache->Get("tenant-a", null_loader), IsError(ErrorKind::kInvalid)); + EXPECT_EQ(cache->size(), 0); +} + +TEST_F(AuthSessionCacheTest, ExpiresAfterIdleTimeout) { + auto cache = MakeCache(); + std::atomic calls = 0; + ICEBERG_UNWRAP_OR_FAIL(auto first, cache->Get("tenant-a", CountingLoader(calls))); + + // Each access resets the idle timer, so the entry outlives the timeout in total. + clock_.Advance(kTimeout - 1ms); + ICEBERG_UNWRAP_OR_FAIL(auto hit, cache->Get("tenant-a", CountingLoader(calls))); + clock_.Advance(kTimeout - 1ms); + ICEBERG_UNWRAP_OR_FAIL(auto still_hit, cache->Get("tenant-a", CountingLoader(calls))); + EXPECT_EQ(hit, first); + EXPECT_EQ(still_hit, first); + EXPECT_EQ(calls, 1); + + clock_.Advance(kTimeout); + ICEBERG_UNWRAP_OR_FAIL(auto reloaded, cache->Get("tenant-a", CountingLoader(calls))); + EXPECT_NE(reloaded, first); + EXPECT_EQ(calls, 2); + EXPECT_THAT(recorder_.removed(), ::testing::ElementsAre(first)); +} + +TEST_F(AuthSessionCacheTest, SweepRemovesOnlyExpiredEntries) { + auto cache = MakeCache(); + std::atomic calls = 0; + ICEBERG_UNWRAP_OR_FAIL(auto idle, cache->Get("idle", CountingLoader(calls))); + ICEBERG_UNWRAP_OR_FAIL(auto active, cache->Get("active", CountingLoader(calls))); + + clock_.Advance(kTimeout - 1ms); + ASSERT_THAT(cache->Get("active", CountingLoader(calls)), IsOk()); + clock_.Advance(1ms); + cache->Sweep(); + + EXPECT_EQ(cache->size(), 1); + EXPECT_THAT(recorder_.removed(), ::testing::ElementsAre(idle)); + ICEBERG_UNWRAP_OR_FAIL(auto still_active, cache->Get("active", CountingLoader(calls))); + EXPECT_EQ(still_active, active); +} + +TEST_F(AuthSessionCacheTest, NonRefreshableSessionExpiresWithToken) { + auto cache = MakeCache(std::chrono::hours(1)); + std::atomic calls = 0; + auto loader = OAuth2Loader(calls, clock_, /*refresh_executor=*/nullptr); + ICEBERG_UNWRAP_OR_FAIL(auto first, cache->Get("tenant-a", loader)); + + clock_.Advance(kTokenLifetime - 1ms); + ICEBERG_UNWRAP_OR_FAIL(auto hit, cache->Get("tenant-a", loader)); + EXPECT_EQ(hit, first); + + // Sweep() only reclaims idle entries; the expired token is replaced on access. + clock_.Advance(1ms); + cache->Sweep(); + EXPECT_EQ(cache->size(), 1); + EXPECT_TRUE(recorder_.removed().empty()); + + ICEBERG_UNWRAP_OR_FAIL(auto reloaded, cache->Get("tenant-a", loader)); + EXPECT_NE(reloaded, first); + EXPECT_EQ(calls, 2); + EXPECT_THAT(recorder_.removed(), ::testing::ElementsAre(first)); +} + +TEST_F(AuthSessionCacheTest, RefreshableSessionOutlivesToken) { + // Must outlive the cache. The refresh is due long after this test ends, so + // nothing is submitted. + test::MockExecutor executor; + EXPECT_CALL(executor, Submit).Times(0); + auto cache = MakeCache(std::chrono::hours(1)); + std::atomic calls = 0; + auto loader = OAuth2Loader(calls, clock_, &executor); + ICEBERG_UNWRAP_OR_FAIL(auto first, cache->Get("tenant-a", loader)); + + // Its refresh may still be in flight; eviction would stop it for other holders. + clock_.Advance(kTokenLifetime + 1s); + cache->Sweep(); + ICEBERG_UNWRAP_OR_FAIL(auto hit, cache->Get("tenant-a", loader)); + + EXPECT_EQ(hit, first); + EXPECT_EQ(calls, 1); + EXPECT_TRUE(recorder_.removed().empty()); +} + +TEST_F(AuthSessionCacheTest, GetEvictsIdleEntriesWithoutPeriodicSweep) { + auto cache = MakeCache(); + std::atomic calls = 0; + ICEBERG_UNWRAP_OR_FAIL(auto a, cache->Get("a", CountingLoader(calls))); + ICEBERG_UNWRAP_OR_FAIL(auto b, cache->Get("b", CountingLoader(calls))); + ICEBERG_UNWRAP_OR_FAIL(auto c, cache->Get("c", CountingLoader(calls))); + + // Keys that are never accessed again are reclaimed when another key is added. + clock_.Advance(kTimeout); + ASSERT_THAT(cache->Get("new", CountingLoader(calls)), IsOk()); + + EXPECT_EQ(cache->size(), 1); + EXPECT_THAT(recorder_.removed(), ::testing::UnorderedElementsAre(a, b, c)); +} + +TEST_F(AuthSessionCacheTest, EvictionFollowsAccessOrder) { + auto cache = MakeCache(); + std::atomic calls = 0; + ICEBERG_UNWRAP_OR_FAIL(auto a, cache->Get("a", CountingLoader(calls))); + ICEBERG_UNWRAP_OR_FAIL(auto b, cache->Get("b", CountingLoader(calls))); + ICEBERG_UNWRAP_OR_FAIL(auto c, cache->Get("c", CountingLoader(calls))); + + // Accessing a moves it behind b and c, so only b and c become idle. + clock_.Advance(kTimeout / 2); + ASSERT_THAT(cache->Get("a", CountingLoader(calls)), IsOk()); + clock_.Advance(kTimeout * 6 / 10); + cache->Sweep(); + + EXPECT_EQ(cache->size(), 1); + EXPECT_THAT(recorder_.removed(), ::testing::ElementsAre(b, c)); + ICEBERG_UNWRAP_OR_FAIL(auto still_cached, cache->Get("a", CountingLoader(calls))); + EXPECT_EQ(still_cached, a); + EXPECT_EQ(calls, 3); +} + +TEST_F(AuthSessionCacheTest, GetDoesNotEvictWithPeriodicSweep) { + test::ThreadExecutor executor; // must outlive the cache + auto cache = MakeCache(); + ASSERT_THAT(cache->StartPeriodicSweep(executor, std::chrono::hours(1)), IsOk()); + std::atomic calls = 0; + ASSERT_THAT(cache->Get("a", CountingLoader(calls)), IsOk()); + + // Idle entries are left to the periodic sweep. + clock_.Advance(kTimeout * 2); + ASSERT_THAT(cache->Get("b", CountingLoader(calls)), IsOk()); + + EXPECT_EQ(cache->size(), 2); + EXPECT_TRUE(recorder_.removed().empty()); + cache->Close(); +} + +TEST_F(AuthSessionCacheTest, ConcurrentGetsForSameKeyLoadOnce) { + auto cache = MakeCache(); + std::atomic calls = 0; + Gate entered; + Gate release; + auto blocking_loader = [&]() -> Result> { + ++calls; + entered.Open(); + release.Wait(); + return AuthSession::MakeDefault({}); + }; + + constexpr int kThreads = 8; + std::vector> results(kThreads); + std::vector threads; + for (int i = 0; i < kThreads; ++i) { + threads.emplace_back([&, i] { + auto session = cache->Get("tenant-a", blocking_loader); + if (session.has_value()) { + results[i] = session.value(); + } + }); + } + entered.Wait(); + std::this_thread::sleep_for(20ms); // let the other callers queue up + release.Open(); + for (auto& thread : threads) { + thread.join(); + } + + EXPECT_EQ(calls, 1); + ASSERT_NE(results[0], nullptr); + for (const auto& session : results) { + EXPECT_EQ(session, results[0]); + } +} + +TEST_F(AuthSessionCacheTest, WaitersRetryAfterFailedLoad) { + auto cache = MakeCache(); + std::atomic calls = 0; + Gate entered; + Gate release; + auto loader = [&]() -> Result> { + if (++calls == 1) { + entered.Open(); + release.Wait(); + return AuthenticationFailed("first attempt fails"); + } + return AuthSession::MakeDefault({}); + }; + + Result> first_result = nullptr; + std::thread first([&] { first_result = cache->Get("tenant-a", loader); }); + entered.Wait(); + Result> waiter_result = nullptr; + std::thread waiter([&] { waiter_result = cache->Get("tenant-a", loader); }); + std::this_thread::sleep_for(20ms); + release.Open(); + first.join(); + waiter.join(); + + EXPECT_THAT(first_result, IsError(ErrorKind::kAuthenticationFailed)); + EXPECT_THAT(waiter_result, IsOk()); + EXPECT_EQ(calls, 2); + EXPECT_EQ(cache->size(), 1); +} + +TEST_F(AuthSessionCacheTest, LoaderExceptionDoesNotBlockLaterGets) { + auto cache = MakeCache(); + std::atomic calls = 0; + Gate entered; + Gate release; + auto loader = [&]() -> Result> { + if (++calls == 1) { + entered.Open(); + release.Wait(); + throw std::runtime_error("HTTP library failure"); + } + return AuthSession::MakeDefault({}); + }; + + auto first = + std::async(std::launch::async, [&] { return cache->Get("tenant-a", loader); }); + entered.Wait(); + auto waiter = + std::async(std::launch::async, [&] { return cache->Get("tenant-a", loader); }); + std::this_thread::sleep_for(20ms); + release.Open(); + + if (waiter.wait_for(5s) != std::future_status::ready) { + cache->Close(); // Wake the waiter so the test fails instead of hanging. + FAIL() << "a waiter stayed blocked after the loader threw"; + } + auto first_result = first.get(); + ASSERT_THAT(first_result, IsError(ErrorKind::kUnknownError)); + EXPECT_THAT(first_result.error().message, ::testing::HasSubstr("HTTP library failure")); + ICEBERG_UNWRAP_OR_FAIL(auto loaded, waiter.get()); + + ICEBERG_UNWRAP_OR_FAIL(auto hit, cache->Get("tenant-a", loader)); + EXPECT_EQ(hit, loaded); + EXPECT_EQ(calls, 2); + EXPECT_EQ(cache->size(), 1); +} + +TEST_F(AuthSessionCacheTest, CloseRemovesAllEntriesOnce) { + auto cache = MakeCache(); + std::atomic calls = 0; + ICEBERG_UNWRAP_OR_FAIL(auto a, cache->Get("tenant-a", CountingLoader(calls))); + ICEBERG_UNWRAP_OR_FAIL(auto b, cache->Get("tenant-b", CountingLoader(calls))); + + cache->Close(); + cache->Close(); + + EXPECT_EQ(cache->size(), 0); + EXPECT_THAT(recorder_.removed(), ::testing::UnorderedElementsAre(a, b)); + EXPECT_THAT(cache->Get("tenant-a", CountingLoader(calls)), + IsError(ErrorKind::kInvalid)); + EXPECT_EQ(calls, 2); +} + +TEST_F(AuthSessionCacheTest, LoadFinishingAfterCloseIsRemovedAndWaitersFail) { + auto cache = MakeCache(); + Gate entered; + Gate release; + std::shared_ptr loaded; + auto loader = [&]() -> Result> { + entered.Open(); + release.Wait(); + loaded = AuthSession::MakeDefault({}); + return loaded; + }; + + Result> loader_result = nullptr; + std::thread loading([&] { loader_result = cache->Get("tenant-a", loader); }); + entered.Wait(); + Result> waiter_result = nullptr; + std::thread waiter([&] { waiter_result = cache->Get("tenant-a", loader); }); + std::this_thread::sleep_for(20ms); + + cache->Close(); + waiter.join(); // Close() wakes waiters before the load finishes. + EXPECT_THAT(waiter_result, IsError(ErrorKind::kInvalid)); + + release.Open(); + loading.join(); + EXPECT_THAT(loader_result, IsError(ErrorKind::kInvalid)); + EXPECT_EQ(cache->size(), 0); + EXPECT_THAT(recorder_.removed(), ::testing::ElementsAre(loaded)); +} + +TEST_F(AuthSessionCacheTest, EverySessionIsRemovedExactlyOnceUnderContention) { + auto cache = MakeCache(); + std::mutex loaded_mutex; + std::set> loaded; + auto loader = [&]() -> Result> { + auto session = AuthSession::MakeDefault({}); + std::lock_guard lock(loaded_mutex); + loaded.insert(session); + return session; + }; + + std::atomic stop = false; + std::thread sweeper([&] { + while (!stop) { + cache->Sweep(); + std::this_thread::yield(); + } + }); + std::vector workers; + for (int t = 0; t < 4; ++t) { + workers.emplace_back([&, t] { + for (int i = 0; i < 500; ++i) { + EXPECT_THAT(cache->Get("key-" + std::to_string((i + t) % 5), loader), IsOk()); + if (i % 10 == 0) { + clock_.Advance(kTimeout / 3); + } + } + }); + } + for (auto& worker : workers) { + worker.join(); + } + stop = true; + sweeper.join(); + cache->Close(); + + auto removed = recorder_.removed(); + std::set> unique_removed(removed.begin(), removed.end()); + EXPECT_EQ(unique_removed.size(), removed.size()) << "a session was removed twice"; + EXPECT_EQ(unique_removed, loaded); +} + +TEST_F(AuthSessionCacheTest, PeriodicSweepRemovesExpiredEntries) { + test::ThreadExecutor executor; // must outlive the cache + auto cache = MakeCache(); + std::atomic calls = 0; + ICEBERG_UNWRAP_OR_FAIL(auto idle, cache->Get("idle", CountingLoader(calls))); + clock_.Advance(kTimeout); + + ASSERT_THAT(cache->StartPeriodicSweep(executor, 10ms), IsOk()); + + // Sweep invokes the listener after erasing the entry, so wait for the listener. + auto deadline = std::chrono::steady_clock::now() + 5s; + while (recorder_.removed().empty() && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(5ms); + } + EXPECT_EQ(cache->size(), 0); + EXPECT_THAT(recorder_.removed(), ::testing::ElementsAre(idle)); +} + +TEST_F(AuthSessionCacheTest, PeriodicSweepContinuesAfterSubmissionFailure) { + test::ThreadExecutor executor(IOError("executor rejected the task")); + auto cache = MakeCache(); + ASSERT_THAT(cache->StartPeriodicSweep(executor, 10ms), IsOk()); + + auto deadline = std::chrono::steady_clock::now() + 5s; + while (executor.submit_count() < 3 && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(5ms); + } + EXPECT_GE(executor.submit_count(), 3); + cache->Close(); +} + +TEST_F(AuthSessionCacheTest, PeriodicSweepArgumentsAndLifecycle) { + test::ThreadExecutor executor; + auto cache = MakeCache(); + + EXPECT_THAT(cache->StartPeriodicSweep(executor, 0ms), + IsError(ErrorKind::kInvalidArgument)); + EXPECT_THAT(cache->StartPeriodicSweep(executor, 10ms), IsOk()); + EXPECT_THAT(cache->StartPeriodicSweep(executor, 10ms), IsOk()); // no-op + + // Destroying the cache cancels the pending sweep. + cache.reset(); + int submitted = executor.submit_count(); + ASSERT_TRUE(test::WaitForSharedTimer(30ms)); + EXPECT_EQ(executor.submit_count(), submitted); + + auto closed = MakeCache(); + closed->Close(); + EXPECT_THAT(closed->StartPeriodicSweep(executor, 10ms), IsOk()); // no-op + ASSERT_TRUE(test::WaitForSharedTimer(30ms)); + EXPECT_EQ(executor.submit_count(), submitted); +} + +} // namespace iceberg::rest::auth diff --git a/src/iceberg/test/rest_catalog_integration_test.cc b/src/iceberg/test/rest_catalog_integration_test.cc index c03c9c271..0e3ba6a3d 100644 --- a/src/iceberg/test/rest_catalog_integration_test.cc +++ b/src/iceberg/test/rest_catalog_integration_test.cc @@ -313,6 +313,13 @@ TEST_F(RestCatalogIntegrationTest, OAuthTableTypedTokenEndToEnd) { auto info = child->OAuth2Info(); ASSERT_TRUE(info.has_value()); EXPECT_EQ(info->issued_token_type, auth::AuthProperties::kAccessTokenType); + + // The same typed token reuses the cached session instead of exchanging again. + ICEBERG_UNWRAP_OR_FAIL( + auto cached, + manager->TableSession( + table, {{auth::AuthProperties::kJwtTokenType, "table-jwt-token"}}, parent)); + EXPECT_EQ(cached, child); } TEST_F(RestCatalogIntegrationTest, OAuthTokenExchangeWithoutActorEndToEnd) {