From de528cada7b42765c5c595e06a60cc3374f44572 Mon Sep 17 00:00:00 2001 From: "shuxu.li" Date: Mon, 5 Oct 2026 21:50:42 +0800 Subject: [PATCH 1/5] feat(auth): implement AuthSessionCache for OAuth2 child sessions Cache contextual and table-scoped OAuth2 child sessions so that context credentials do not trigger a token request on every catalog call (#962). AuthSessionCache: - Expire entries after auth.session-timeout-ms without access (default 1 hour), and when the token of an OAuth2 session expires. - Load each key once under concurrent Get() calls, with the loader running outside the lock. Failed loads are not cached. - Hand every removed session to the removal listener exactly once. A session loaded after Close() is removed immediately. - Sweep expired entries periodically on the auth refresh executor, every session timeout clamped to [1s, 60s]. Without an executor, entries are reclaimed when accessed again or on Close(). OAuth2Manager: - Cache contextual sessions by session_id and table sessions by token. - Authenticate child token requests with the parent session's headers. - Credential child sessions follow token-refresh-enabled and, with refresh enabled, require the auth executor. Token and exchange child sessions do not refresh until refresh via token exchange is supported. - Token child sessions no longer carry the parent's credential. OAuth2Session: - Add StopRefreshing(), a non-blocking stop used on cache removal. - Add ExpiresAt() and IsRefreshable(), and include the session headers in OAuth2Info(). --- src/iceberg/catalog/rest/CMakeLists.txt | 1 + src/iceberg/catalog/rest/auth/auth_manager.cc | 120 ++++- .../catalog/rest/auth/auth_properties.h | 8 + src/iceberg/catalog/rest/auth/auth_session.h | 1 + .../catalog/rest/auth/auth_session_cache.cc | 254 +++++++++++ .../rest/auth/auth_session_cache_internal.h | 127 ++++++ .../catalog/rest/auth/auth_session_internal.h | 29 ++ src/iceberg/catalog/session_context.h | 3 + src/iceberg/test/CMakeLists.txt | 1 + src/iceberg/test/auth_manager_test.cc | 65 +++ src/iceberg/test/auth_session_cache_test.cc | 420 ++++++++++++++++++ 11 files changed, 1007 insertions(+), 22 deletions(-) create mode 100644 src/iceberg/catalog/rest/auth/auth_session_cache.cc create mode 100644 src/iceberg/catalog/rest/auth/auth_session_cache_internal.h create mode 100644 src/iceberg/test/auth_session_cache_test.cc diff --git a/src/iceberg/catalog/rest/CMakeLists.txt b/src/iceberg/catalog/rest/CMakeLists.txt index a8146d468a..65c037247f 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 df5d7304a8..e1ff2014b9 100644 --- a/src/iceberg/catalog/rest/auth/auth_manager.cc +++ b/src/iceberg/catalog/rest/auth/auth_manager.cc @@ -19,6 +19,7 @@ #include "iceberg/catalog/rest/auth/auth_manager.h" +#include #include #include #include @@ -28,6 +29,7 @@ #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" @@ -160,13 +162,11 @@ class OAuth2Manager : public AuthManager { start_time_ = std::chrono::steady_clock::now(); ICEBERG_ASSIGN_OR_RAISE( auth_response_, OAuth2Util::FetchToken(*init_client, *init_session, config)); - // TODO(lishuxu): Match Java OAuth2Util.AuthSession.fromTokenResponse here. return AuthSession::MakeDefault( OAuth2Util::AuthHeaders(auth_response_->access_token)); } if (!config.token().empty()) { - // TODO(lishuxu): Match Java OAuth2Util.AuthSession.fromAccessToken here. return AuthSession::MakeDefault(OAuth2Util::AuthHeaders(config.token())); } @@ -180,13 +180,42 @@ 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(); + exchange_enabled_ = config.exchange_enabled(); + session_timeout_ = + std::chrono::milliseconds(config.Get(AuthProperties::kSessionTimeoutMs)); + + // 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(); + } + })); + + // Periodically reclaim idle sessions on the refresh executor. Sweep at most once + // a minute and at least once a second, so a zero timeout still gets periodic + // cleanup. Without an executor, expired sessions are reclaimed only when their key + // is accessed again or on Close(). + if (refresh_executor_ != nullptr) { + auto sweep_interval = std::clamp(session_timeout_, std::chrono::milliseconds(1000), + std::chrono::milliseconds(60000)); + ICEBERG_RETURN_UNEXPECTED( + session_cache_->StartPeriodicSweep(*refresh_executor_, sweep_interval)); + } + // 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_); } - // TODO(lishuxu): Honor token-refresh-enabled for catalog bearer tokens, matching - // Java. If token is provided, use it directly. + // TODO(lishuxu): Honor token-refresh-enabled for catalog bearer tokens. + // If token is provided, use it directly. if (!config.token().empty()) { OAuthTokenResponse token_response{ .access_token = config.token(), @@ -212,21 +241,31 @@ 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. + // Use session_id as cache key for contextual sessions + std::string cache_key = context.session_id.empty() ? "" : "ctx:" + context.session_id; return MaybeCreateChildSession(context.credentials, /*allow_credential=*/true, - std::move(parent)); + std::move(parent), cache_key); } Result> TableSession( [[maybe_unused]] const TableIdentifier& table, const std::unordered_map& properties, std::shared_ptr parent) override { + // Use token value as cache key for table sessions + auto token_it = properties.find(AuthProperties::kToken.key()); + std::string cache_key; + if (token_it != properties.end() && !token_it->second.empty()) { + cache_key = "tbl:" + token_it->second; + } return MaybeCreateChildSession(FilterTableSessionProperties(properties), - /*allow_credential=*/false, std::move(parent)); + /*allow_credential=*/false, std::move(parent), + cache_key); } Status Close() override { + if (session_cache_) { + session_cache_->Close(); + } refresh_client_.reset(); return {}; } @@ -265,7 +304,8 @@ class OAuth2Manager : public AuthManager { Result> MaybeCreateChildSession( const std::unordered_map& credentials, - bool allow_credential, std::shared_ptr parent) { + bool allow_credential, std::shared_ptr parent, + const std::string& cache_key) { auto token_it = credentials.find(AuthProperties::kToken.key()); auto credential_it = credentials.find(AuthProperties::kCredential.key()); auto typed_token = FindPreferredTypedToken(credentials); @@ -281,35 +321,65 @@ class OAuth2Manager : public AuthManager { ICEBERG_PRECHECK(parent_info.has_value(), "OAuth2 child session requires OAuth2 parent metadata"); + // Use cache if we have a valid cache_key + if (!cache_key.empty() && session_cache_) { + return session_cache_->Get(cache_key, + [&]() -> Result> { + return CreateChildSessionUncached( + credentials, allow_credential, *parent_info); + }); + } + + // No cache, create session directly + return CreateChildSessionUncached(credentials, allow_credential, *parent_info); + } + + Result> CreateChildSessionUncached( + const std::unordered_map& credentials, + bool allow_credential, const OAuth2SessionInfo& parent_info) { + 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()) { - ICEBERG_ASSIGN_OR_RAISE(auto config, - ChildConfig(*parent_info, parent_info->credential)); + // 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_it->second), config, /*keep_refreshed=*/false); } 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); + ChildConfig(parent_info, credential_it->second)); + // 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 use keep_refreshed from catalog level + return MakeSession(response, config, keep_refreshed_); } 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 (!parent_info.token.empty()) { + actor_token = parent_info.token; + actor_token_type = parent_info.issued_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, + OAuth2Util::ExchangeToken(*refresh_client_, *temp_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)); + 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 +388,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; + bool exchange_enabled_ = true; + std::chrono::milliseconds session_timeout_{3'600'000}; + std::shared_ptr session_cache_; }; Result> MakeOAuth2Manager( diff --git a/src/iceberg/catalog/rest/auth/auth_properties.h b/src/iceberg/catalog/rest/auth/auth_properties.h index 8784194cce..730b32d8ad 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. Config is read but not used in refresh logic. 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 7e330706c2..393418e2a4 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 0000000000..c365ab1ca2 --- /dev/null +++ b/src/iceberg/catalog/rest/auth/auth_session_cache.cc @@ -0,0 +1,254 @@ +/* + * 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 "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 { + +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))); +} + +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)) { + existing->last_access = now; + return existing->session; + } + removed.push_back(std::move(existing->session)); + slots_.erase(it); + break; + } + // 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 = 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. + 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(); + slot->last_access = 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_); + auto now = clock_(); + for (auto it = slots_.begin(); it != slots_.end();) { + if (!it->second->loading && IsExpired(*it->second, now)) { + removed.push_back(std::move(it->second->session)); + it = slots_.erase(it); + } else { + ++it; + } + } + } + NotifyRemoved(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(); + } + 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 + // Even for refreshable sessions, check token expiration to handle cases where: + // - Refresh has permanently failed (after max retries) + // - Refresh is delayed or not scheduled + // - Token was obtained from external source (table token) + auto oauth2 = std::dynamic_pointer_cast(slot.session); + if (oauth2) { + auto expires_at = oauth2->ExpiresAt(); + if (expires_at.has_value() && now >= *expires_at) { + return true; + } + } + + 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 0000000000..32ac29ed77 --- /dev/null +++ b/src/iceberg/catalog/rest/auth/auth_session_cache_internal.h @@ -0,0 +1,127 @@ +/* + * 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 "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, and each session that +/// leaves the cache is handed to the removal listener exactly once. Idle entries +/// are reclaimed when accessed or when Sweep() runs; StartPeriodicSweep() can +/// schedule Sweep() to run periodically. +/// +/// 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. + 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 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. 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 all expired entries and notify the removal listener. + 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 { + bool loading = true; + std::shared_ptr session; + std::chrono::steady_clock::time_point last_access; + }; + + 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 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_; + std::unordered_map> slots_; + 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 fa2f7ede99..9560467471 100644 --- a/src/iceberg/catalog/rest/auth/auth_session_internal.h +++ b/src/iceberg/catalog/rest/auth/auth_session_internal.h @@ -102,11 +102,40 @@ 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 can refresh its token automatically. + bool IsRefreshable() const { return scheduler_.has_value(); } + ~OAuth2Session() override { std::ignore = CloseImpl(); } private: diff --git a/src/iceberg/catalog/session_context.h b/src/iceberg/catalog/session_context.h index 19eb8fc5f8..1a6bc819cb 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 37e20e3391..7fd77ff991 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 f6f1400b2b..b7497ec601 100644 --- a/src/iceberg/test/auth_manager_test.cc +++ b/src/iceberg/test/auth_manager_test.cc @@ -434,6 +434,20 @@ TEST_F(AuthManagerTest, OAuth2StaticToken) { } // Verifies OAuth2 type is inferred from token property +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, OAuth2InferredFromToken) { std::unordered_map properties = { {AuthProperties::kToken.key(), "inferred-token"}, @@ -886,6 +900,57 @@ TEST(OAuth2SessionTest, StopsAfterHttpRetryBudget) { 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(); + + 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", 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 0000000000..3f806c6aea --- /dev/null +++ b/src/iceberg/test/auth_session_cache_test.cc @@ -0,0 +1,420 @@ +/* + * 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 "iceberg/catalog/rest/auth/auth_session.h" +#include "iceberg/catalog/rest/auth/auth_session_cache_internal.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); + +class FakeClock { + public: + AuthSessionCache::Clock AsClock() { + return [this] { + return std::chrono::steady_clock::time_point(std::chrono::milliseconds(now_ms_)); + }; + } + + void Advance(std::chrono::milliseconds delta) { now_ms_ += delta.count(); } + + private: + std::atomic now_ms_{0}; +}; + +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({}); + }; +} + +} // 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, 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, 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()); + + auto deadline = std::chrono::steady_clock::now() + 5s; + while (cache->size() != 0 && 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 From 21991c011a2306c127d306f1a70a7d289a064e86 Mon Sep 17 00:00:00 2001 From: "shuxu.li" Date: Tue, 20 Oct 2026 17:31:06 +0800 Subject: [PATCH 2/5] fix(auth): address AuthSessionCache review comments - Cache table sessions that use typed tokens. The child credential is now selected once by SelectChildCredential, and the table cache key is the value of the token actually used. Previously only the plain token property produced a key, so typed tokens (jwt, id_token, ...) skipped the cache and ran a token exchange on every table load. - Evict only non-refreshable sessions when their token expires. Evicting a refreshable session could land before an in-flight refresh completed and stop refreshing for callers still holding it, such as loaded tables. Refreshable sessions now expire only after the idle timeout. --- src/iceberg/catalog/rest/auth/auth_manager.cc | 153 +++++++++++------- .../catalog/rest/auth/auth_session_cache.cc | 16 +- .../rest/auth/auth_session_cache_internal.h | 9 +- src/iceberg/test/auth_manager_test.cc | 23 +++ src/iceberg/test/auth_session_cache_test.cc | 77 ++++++++- .../test/rest_catalog_integration_test.cc | 7 + 6 files changed, 211 insertions(+), 74 deletions(-) diff --git a/src/iceberg/catalog/rest/auth/auth_manager.cc b/src/iceberg/catalog/rest/auth/auth_manager.cc index e1ff2014b9..ce6b459604 100644 --- a/src/iceberg/catalog/rest/auth/auth_manager.cc +++ b/src/iceberg/catalog/rest/auth/auth_manager.cc @@ -59,6 +59,43 @@ 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; +} + std::unordered_map FilterTableSessionProperties( const std::unordered_map& properties) { std::unordered_map filtered; @@ -241,25 +278,28 @@ class OAuth2Manager : public AuthManager { Result> ContextualSession( const SessionContext& context, std::shared_ptr parent) override { - // Use session_id as cache key for contextual sessions + auto credential = + SelectChildCredential(context.credentials, /*allow_credential=*/true); + if (!credential.has_value()) { + return 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 MaybeCreateChildSession(context.credentials, /*allow_credential=*/true, - std::move(parent), cache_key); + return GetOrCreateChildSession(*credential, std::move(parent), cache_key); } Result> TableSession( [[maybe_unused]] const TableIdentifier& table, const std::unordered_map& properties, std::shared_ptr parent) override { - // Use token value as cache key for table sessions - auto token_it = properties.find(AuthProperties::kToken.key()); - std::string cache_key; - if (token_it != properties.end() && !token_it->second.empty()) { - cache_key = "tbl:" + token_it->second; + auto credential = SelectChildCredential(FilterTableSessionProperties(properties), + /*allow_credential=*/false); + if (!credential.has_value()) { + return parent; } - return MaybeCreateChildSession(FilterTableSessionProperties(properties), - /*allow_credential=*/false, std::move(parent), - cache_key); + // Table sessions are cached by the value of the token they use. + std::string cache_key = credential->value.empty() ? "" : "tbl:" + credential->value; + return GetOrCreateChildSession(*credential, std::move(parent), cache_key); } Status Close() override { @@ -302,19 +342,9 @@ class OAuth2Manager : public AuthManager { token_request_started_at, refresh_executor_); } - Result> MaybeCreateChildSession( - const std::unordered_map& credentials, - bool allow_credential, std::shared_ptr parent, + Result> GetOrCreateChildSession( + const ChildCredential& credential, std::shared_ptr parent, const std::string& cache_key) { - 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; - } - ICEBERG_PRECHECK(refresh_client_ != nullptr, "OAuth2 catalog session must be initialized before child sessions"); auto parent_info = parent->OAuth2Info(); @@ -323,47 +353,56 @@ class OAuth2Manager : public AuthManager { // Use cache if we have a valid cache_key if (!cache_key.empty() && session_cache_) { - return session_cache_->Get(cache_key, - [&]() -> Result> { - return CreateChildSessionUncached( - credentials, allow_credential, *parent_info); - }); + return session_cache_->Get( + cache_key, [&]() -> Result> { + return CreateChildSessionUncached(credential, *parent_info); + }); } // No cache, create session directly - return CreateChildSessionUncached(credentials, allow_credential, *parent_info); + return CreateChildSessionUncached(credential, *parent_info); } Result> CreateChildSessionUncached( - const std::unordered_map& credentials, - bool allow_credential, const OAuth2SessionInfo& parent_info) { - 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()) { - // 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_it->second), config, - /*keep_refreshed=*/false); + 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(); + } - if (allow_credential && credential_it != credentials.end()) { - ICEBERG_ASSIGN_OR_RAISE(auto config, - ChildConfig(parent_info, credential_it->second)); - // 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 use keep_refreshed from catalog level - return MakeSession(response, config, keep_refreshed_); - } + 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 use keep_refreshed from catalog level + return MakeSession(response, config, keep_refreshed_); + } + 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()) { @@ -374,8 +413,8 @@ class OAuth2Manager : public AuthManager { auto temp_parent = AuthSession::MakeDefault(parent_info.headers); ICEBERG_ASSIGN_OR_RAISE( auto response, - OAuth2Util::ExchangeToken(*refresh_client_, *temp_parent, {}, typed_token->second, - typed_token->first, actor_token, actor_token_type, + 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, diff --git a/src/iceberg/catalog/rest/auth/auth_session_cache.cc b/src/iceberg/catalog/rest/auth/auth_session_cache.cc index c365ab1ca2..8868e8936a 100644 --- a/src/iceberg/catalog/rest/auth/auth_session_cache.cc +++ b/src/iceberg/catalog/rest/auth/auth_session_cache.cc @@ -188,17 +188,15 @@ bool AuthSessionCache::IsExpired(const Slot& slot, return true; } - // Check 2: Token expiration - // Even for refreshable sessions, check token expiration to handle cases where: - // - Refresh has permanently failed (after max retries) - // - Refresh is delayed or not scheduled - // - Token was obtained from external source (table token) + // Check 2: Token expiration, for sessions that cannot refresh + // Such a session is useless once its token expires, and evicting it is harmless + // since stopping its refresh is a no-op. A refreshable session is left alone: + // its refresh may still 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) { + if (oauth2 && !oauth2->IsRefreshable()) { auto expires_at = oauth2->ExpiresAt(); - if (expires_at.has_value() && now >= *expires_at) { - return true; - } + return expires_at.has_value() && now >= *expires_at; } return false; diff --git a/src/iceberg/catalog/rest/auth/auth_session_cache_internal.h b/src/iceberg/catalog/rest/auth/auth_session_cache_internal.h index 32ac29ed77..7818cc7cf2 100644 --- a/src/iceberg/catalog/rest/auth/auth_session_cache_internal.h +++ b/src/iceberg/catalog/rest/auth/auth_session_cache_internal.h @@ -43,10 +43,11 @@ 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, and each session that -/// leaves the cache is handed to the removal listener exactly once. Idle entries -/// are reclaimed when accessed or when Sweep() runs; StartPeriodicSweep() can -/// schedule Sweep() to run periodically. +/// Entries expire after `session_timeout` without access. OAuth2 sessions that +/// cannot refresh also expire once their token does. Each session that leaves the +/// cache is handed to the removal listener exactly once. Expired entries are +/// reclaimed when accessed or when Sweep() runs; StartPeriodicSweep() can schedule +/// Sweep() to run periodically. /// /// Thread safety: all public methods are thread-safe. Loaders and the removal /// listener are always invoked without holding the cache lock. diff --git a/src/iceberg/test/auth_manager_test.cc b/src/iceberg/test/auth_manager_test.cc index b7497ec601..fae2bd5971 100644 --- a/src/iceberg/test/auth_manager_test.cc +++ b/src/iceberg/test/auth_manager_test.cc @@ -568,6 +568,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"}, diff --git a/src/iceberg/test/auth_session_cache_test.cc b/src/iceberg/test/auth_session_cache_test.cc index 3f806c6aea..6dc2c1cb6c 100644 --- a/src/iceberg/test/auth_session_cache_test.cc +++ b/src/iceberg/test/auth_session_cache_test.cc @@ -33,6 +33,8 @@ #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" @@ -44,19 +46,25 @@ 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 std::chrono::steady_clock::time_point(std::chrono::milliseconds(now_ms_)); - }; + return [this] { return Now(); }; } void Advance(std::chrono::milliseconds delta) { now_ms_ += delta.count(); } private: - std::atomic now_ms_{0}; + std::atomic now_ms_{std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count()}; }; class RemovalRecorder { @@ -107,6 +115,27 @@ AuthSessionCache::Loader CountingLoader(std::atomic& calls) { }; } +/// 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 { @@ -210,6 +239,46 @@ TEST_F(AuthSessionCacheTest, SweepRemovesOnlyExpiredEntries) { 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); + + clock_.Advance(1ms); + cache->Sweep(); + EXPECT_EQ(cache->size(), 0); + EXPECT_THAT(recorder_.removed(), ::testing::ElementsAre(first)); + + ICEBERG_UNWRAP_OR_FAIL(auto reloaded, cache->Get("tenant-a", loader)); + EXPECT_NE(reloaded, first); + EXPECT_EQ(calls, 2); +} + +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, ConcurrentGetsForSameKeyLoadOnce) { auto cache = MakeCache(); std::atomic calls = 0; diff --git a/src/iceberg/test/rest_catalog_integration_test.cc b/src/iceberg/test/rest_catalog_integration_test.cc index c03c9c2716..0e3ba6a3d3 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) { From e60ad84856d74d9fc1d5695d4efa8214ba254737 Mon Sep 17 00:00:00 2001 From: "shuxu.li" Date: Sat, 24 Oct 2026 19:11:23 +0800 Subject: [PATCH 3/5] fix(auth): refresh credential child sessions only with an auth executor Since #1033, refreshing a session requires an auth executor. Credential child sessions followed token-refresh-enabled unconditionally, so a static-token catalog without an executor failed to create them. They now refresh only when an executor is provided and otherwise keep their initial token, as before, with a one-time warning. --- src/iceberg/catalog/rest/auth/auth_manager.cc | 15 ++++- src/iceberg/test/auth_manager_test.cc | 58 +++++++++++++++++++ 2 files changed, 71 insertions(+), 2 deletions(-) diff --git a/src/iceberg/catalog/rest/auth/auth_manager.cc b/src/iceberg/catalog/rest/auth/auth_manager.cc index ce6b459604..8ee36c920a 100644 --- a/src/iceberg/catalog/rest/auth/auth_manager.cc +++ b/src/iceberg/catalog/rest/auth/auth_manager.cc @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -33,6 +34,7 @@ #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" @@ -396,8 +398,16 @@ class OAuth2Manager : public AuthManager { auto temp_parent = AuthSession::MakeDefault(parent_info.headers); ICEBERG_ASSIGN_OR_RAISE( auto response, OAuth2Util::FetchToken(*refresh_client_, *temp_parent, config)); - // Credential child sessions use keep_refreshed from catalog level - return MakeSession(response, config, keep_refreshed_); + // 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( @@ -433,6 +443,7 @@ class OAuth2Manager : public AuthManager { bool exchange_enabled_ = 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/test/auth_manager_test.cc b/src/iceberg/test/auth_manager_test.cc index fae2bd5971..4850400732 100644 --- a/src/iceberg/test/auth_manager_test.cc +++ b/src/iceberg/test/auth_manager_test.cc @@ -1005,4 +1005,62 @@ 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()); +} + } // namespace iceberg::rest::auth From a44de986b5817e61bd788d8204193a40503bdbcf Mon Sep 17 00:00:00 2001 From: "shuxu.li" Date: Mon, 2 Nov 2026 18:23:11 +0800 Subject: [PATCH 4/5] fix(auth): address AuthSessionCache review findings - Key table sessions by flow, token type and actor. An exchanged table session depends on the parent token sent as the actor, but the key held only the token value, so different parents shared one session. - Reload a cached session whose refresh gave up once its token expires. Such a session still counted as refreshable, so a frequently used key kept returning an expired token after the token endpoint recovered. - Evict idle sessions in access order, in time proportional to the number removed, and from Get() when no periodic sweep runs. Without an auth executor, keys that were never used again stayed cached until Close. Sweep() now evicts idle sessions only; one whose token expired is replaced on its next access. - Clean up the loading slot when a session loader throws, so later Get calls for the key no longer block until Close. - Return an error for a malformed auth.session-timeout-ms instead of throwing. - Make the periodic sweep test wait for the removal listener. --- src/iceberg/catalog/rest/auth/auth_manager.cc | 141 +++++++++++--- .../catalog/rest/auth/auth_properties.h | 2 +- .../catalog/rest/auth/auth_session_cache.cc | 108 +++++++++-- .../rest/auth/auth_session_cache_internal.h | 49 ++++- .../catalog/rest/auth/auth_session_internal.h | 24 ++- src/iceberg/test/auth_manager_test.cc | 175 +++++++++++++++++- src/iceberg/test/auth_session_cache_test.cc | 99 +++++++++- 7 files changed, 534 insertions(+), 64 deletions(-) diff --git a/src/iceberg/catalog/rest/auth/auth_manager.cc b/src/iceberg/catalog/rest/auth/auth_manager.cc index 8ee36c920a..1a4989a1c7 100644 --- a/src/iceberg/catalog/rest/auth/auth_manager.cc +++ b/src/iceberg/catalog/rest/auth/auth_manager.cc @@ -19,11 +19,11 @@ #include "iceberg/catalog/rest/auth/auth_manager.h" -#include #include #include #include #include +#include #include #include @@ -37,6 +37,7 @@ #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 { @@ -98,6 +99,90 @@ std::optional SelectChildCredential( 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; @@ -201,11 +286,13 @@ class OAuth2Manager : public AuthManager { start_time_ = std::chrono::steady_clock::now(); ICEBERG_ASSIGN_OR_RAISE( auth_response_, OAuth2Util::FetchToken(*init_client, *init_session, config)); + // TODO(lishuxu): Match Java OAuth2Util.AuthSession.fromTokenResponse here. return AuthSession::MakeDefault( OAuth2Util::AuthHeaders(auth_response_->access_token)); } if (!config.token().empty()) { + // TODO(lishuxu): Match Java OAuth2Util.AuthSession.fromAccessToken here. return AuthSession::MakeDefault(OAuth2Util::AuthHeaders(config.token())); } @@ -222,9 +309,7 @@ class OAuth2Manager : public AuthManager { // Initialize catalog-level configuration keep_refreshed_ = config.keep_refreshed(); - exchange_enabled_ = config.exchange_enabled(); - session_timeout_ = - std::chrono::milliseconds(config.Get(AuthProperties::kSessionTimeoutMs)); + ICEBERG_ASSIGN_OR_RAISE(session_timeout_, ParseSessionTimeout(config)); // Create session cache ICEBERG_ASSIGN_OR_RAISE( @@ -237,15 +322,12 @@ class OAuth2Manager : public AuthManager { } })); - // Periodically reclaim idle sessions on the refresh executor. Sweep at most once - // a minute and at least once a second, so a zero timeout still gets periodic - // cleanup. Without an executor, expired sessions are reclaimed only when their key - // is accessed again or on Close(). + // Reclaim idle sessions in the background on the refresh executor. Without an + // executor, the cache sweeps from Get() instead. if (refresh_executor_ != nullptr) { - auto sweep_interval = std::clamp(session_timeout_, std::chrono::milliseconds(1000), - std::chrono::milliseconds(60000)); - ICEBERG_RETURN_UNEXPECTED( - session_cache_->StartPeriodicSweep(*refresh_executor_, sweep_interval)); + ICEBERG_RETURN_UNEXPECTED(session_cache_->StartPeriodicSweep( + *refresh_executor_, + internal::AuthSessionCache::DefaultSweepInterval(session_timeout_))); } // Reuse the token response and start time from the init phase. @@ -253,8 +335,8 @@ class OAuth2Manager : public AuthManager { return MakeSession(*auth_response_, config, config.keep_refreshed(), start_time_); } - // TODO(lishuxu): Honor token-refresh-enabled for catalog bearer tokens. - // If token is provided, use it directly. + // TODO(lishuxu): Honor token-refresh-enabled for catalog bearer tokens, matching + // Java. If token is provided, use it directly. if (!config.token().empty()) { OAuthTokenResponse token_response{ .access_token = config.token(), @@ -285,9 +367,10 @@ class OAuth2Manager : public AuthManager { 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, std::move(parent), cache_key); + return GetOrCreateChildSession(*credential, parent_info, cache_key); } Result> TableSession( @@ -299,9 +382,9 @@ class OAuth2Manager : public AuthManager { if (!credential.has_value()) { return parent; } - // Table sessions are cached by the value of the token they use. - std::string cache_key = credential->value.empty() ? "" : "tbl:" + credential->value; - return GetOrCreateChildSession(*credential, std::move(parent), cache_key); + ICEBERG_ASSIGN_OR_RAISE(auto parent_info, ParentInfo(*parent)); + return GetOrCreateChildSession(*credential, parent_info, + TableCacheKey(*credential, parent_info)); } Status Close() override { @@ -344,25 +427,28 @@ class OAuth2Manager : public AuthManager { token_request_started_at, refresh_executor_); } - Result> GetOrCreateChildSession( - const ChildCredential& credential, std::shared_ptr parent, - const std::string& cache_key) { + 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); + } + 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); + return CreateChildSessionUncached(credential, parent_info); }); } // No cache, create session directly - return CreateChildSessionUncached(credential, *parent_info); + return CreateChildSessionUncached(credential, parent_info); } Result> CreateChildSessionUncached( @@ -415,9 +501,9 @@ class OAuth2Manager : public AuthManager { 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); @@ -440,7 +526,6 @@ class OAuth2Manager : public AuthManager { // Catalog-level configuration initialized in CatalogSession bool keep_refreshed_ = true; - bool exchange_enabled_ = true; std::chrono::milliseconds session_timeout_{3'600'000}; std::shared_ptr session_cache_; std::once_flag no_refresh_executor_warning_; diff --git a/src/iceberg/catalog/rest/auth/auth_properties.h b/src/iceberg/catalog/rest/auth/auth_properties.h index 730b32d8ad..0cf9be5e44 100644 --- a/src/iceberg/catalog/rest/auth/auth_properties.h +++ b/src/iceberg/catalog/rest/auth/auth_properties.h @@ -80,7 +80,7 @@ class ICEBERG_REST_EXPORT AuthProperties : public ConfigBase { inline static Entry kKeepRefreshed{"token-refresh-enabled", true}; /// Whether to use token exchange for refreshing OAuth2 sessions. - /// NOTE: Not yet implemented. Config is read but not used in refresh logic. + /// NOTE: Not yet implemented; the value is ignored. inline static Entry kExchangeEnabled{"token-exchange-enabled", true}; inline static Entry kAudience{"audience", ""}; diff --git a/src/iceberg/catalog/rest/auth/auth_session_cache.cc b/src/iceberg/catalog/rest/auth/auth_session_cache.cc index 8868e8936a..53521fc87f 100644 --- a/src/iceberg/catalog/rest/auth/auth_session_cache.cc +++ b/src/iceberg/catalog/rest/auth/auth_session_cache.cc @@ -17,6 +17,8 @@ * under the License. */ +#include +#include #include #include "iceberg/catalog/rest/auth/auth_session_cache_internal.h" @@ -26,6 +28,28 @@ 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(), @@ -36,6 +60,12 @@ Result> AuthSessionCache::Make( 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), @@ -68,13 +98,18 @@ Result> AuthSessionCache::Get(const std::string& ke } auto now = clock_(); if (!IsExpired(*existing, now)) { - existing->last_access = now; + TouchLocked(it, now); return existing->session; } - removed.push_back(std::move(existing->session)); - slots_.erase(it); + 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); @@ -82,7 +117,7 @@ Result> AuthSessionCache::Get(const std::string& ke // Phase 2: Call loader (lock released) NotifyRemoved(removed); - auto result = loader(); + auto result = RunLoader(loader); if (result.has_value() && result.value() == nullptr) { result = Invalid("Auth session loader returned a null session"); } @@ -96,7 +131,8 @@ Result> AuthSessionCache::Get(const std::string& ke 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. + // 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); } @@ -105,7 +141,7 @@ Result> AuthSessionCache::Get(const std::string& ke orphan = result.value(); } else { slot->session = result.value(); - slot->last_access = clock_(); + TouchLocked(it, clock_()); } } load_finished_.notify_all(); @@ -123,19 +159,48 @@ void AuthSessionCache::Sweep() { std::vector> removed; { std::lock_guard lock(mutex_); - auto now = clock_(); - for (auto it = slots_.begin(); it != slots_.end();) { - if (!it->second->loading && IsExpired(*it->second, now)) { - removed.push_back(std::move(it->second->session)); - it = slots_.erase(it); - } else { - ++it; - } - } + 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{&it->first, &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(), @@ -166,6 +231,7 @@ void AuthSessionCache::Close() { } } slots_.clear(); + access_order_.clear(); } load_finished_.notify_all(); // Shut down outside the lock: Shutdown() waits for an in-progress submission, @@ -188,11 +254,13 @@ bool AuthSessionCache::IsExpired(const Slot& slot, return true; } - // Check 2: Token expiration, for sessions that cannot refresh - // Such a session is useless once its token expires, and evicting it is harmless - // since stopping its refresh is a no-op. A refreshable session is left alone: - // its refresh may still be in flight, and evicting it would stop the refresh - // for callers that still hold it. + // 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(); diff --git a/src/iceberg/catalog/rest/auth/auth_session_cache_internal.h b/src/iceberg/catalog/rest/auth/auth_session_cache_internal.h index 7818cc7cf2..aee88147e1 100644 --- a/src/iceberg/catalog/rest/auth/auth_session_cache_internal.h +++ b/src/iceberg/catalog/rest/auth/auth_session_cache_internal.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -44,10 +45,14 @@ 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 -/// cannot refresh also expire once their token does. Each session that leaves the -/// cache is handed to the removal listener exactly once. Expired entries are -/// reclaimed when accessed or when Sweep() runs; StartPeriodicSweep() can schedule -/// Sweep() to run periodically. +/// 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. @@ -56,7 +61,7 @@ class ICEBERG_REST_EXPORT AuthSessionCache public: using Clock = std::function; using Loader = std::function>()>; - /// Called once for every session that leaves the cache. Must not block. + /// Called once for every session that leaves the cache. Must not block or throw. using RemovalListener = std::function&)>; /// \brief Create a cache. @@ -74,14 +79,22 @@ class ICEBERG_REST_EXPORT 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. Returns an error once the cache + /// 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 all expired entries and notify the removal listener. + /// \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(). @@ -98,16 +111,34 @@ class ICEBERG_REST_EXPORT AuthSessionCache 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(); @@ -119,7 +150,9 @@ class ICEBERG_REST_EXPORT AuthSessionCache mutable std::mutex mutex_; std::condition_variable load_finished_; - std::unordered_map> slots_; + 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}; diff --git a/src/iceberg/catalog/rest/auth/auth_session_internal.h b/src/iceberg/catalog/rest/auth/auth_session_internal.h index 9560467471..a4454aafc9 100644 --- a/src/iceberg/catalog/rest/auth/auth_session_internal.h +++ b/src/iceberg/catalog/rest/auth/auth_session_internal.h @@ -133,8 +133,14 @@ class OAuth2Session final : public AuthSession, return expires_at_; } - /// \brief Whether this session can refresh its token automatically. - bool IsRefreshable() const { return scheduler_.has_value(); } + /// \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(); } @@ -250,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_); @@ -305,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()) { @@ -342,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/test/auth_manager_test.cc b/src/iceberg/test/auth_manager_test.cc index 4850400732..7b6f0f3a62 100644 --- a/src/iceberg/test/auth_manager_test.cc +++ b/src/iceberg/test/auth_manager_test.cc @@ -433,7 +433,6 @@ TEST_F(AuthManagerTest, OAuth2StaticToken) { "catalog-resource"); } -// Verifies OAuth2 type is inferred from token property TEST_F(AuthManagerTest, OAuth2ZeroSessionTimeoutCreatesCatalogSession) { std::unordered_map properties = { {AuthProperties::kAuthType, "oauth2"}, @@ -448,6 +447,22 @@ TEST_F(AuthManagerTest, OAuth2ZeroSessionTimeoutCreatesCatalogSession) { 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 = { {AuthProperties::kToken.key(), "inferred-token"}, @@ -911,13 +926,18 @@ 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()); @@ -944,6 +964,7 @@ TEST(OAuth2SessionTest, StopRefreshingKeepsSessionUsable) { session->StopRefreshing(); session->StopRefreshing(); + EXPECT_FALSE(session->IsRefreshable()); ICEBERG_UNWRAP_OR_FAIL(auto request, session->Authenticate({})); EXPECT_EQ(request.headers.at("Authorization"), "Bearer token"); @@ -1063,4 +1084,156 @@ TEST(OAuth2ChildSessionTest, CredentialChildWithoutExecutorDoesNotRefresh) { 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 index 6dc2c1cb6c..074459e738 100644 --- a/src/iceberg/test/auth_session_cache_test.cc +++ b/src/iceberg/test/auth_session_cache_test.cc @@ -21,9 +21,11 @@ #include #include #include +#include #include #include #include +#include #include #include #include @@ -249,14 +251,16 @@ TEST_F(AuthSessionCacheTest, NonRefreshableSessionExpiresWithToken) { 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(), 0); - EXPECT_THAT(recorder_.removed(), ::testing::ElementsAre(first)); + 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) { @@ -279,6 +283,57 @@ TEST_F(AuthSessionCacheTest, RefreshableSessionOutlivesToken) { 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; @@ -346,6 +401,43 @@ TEST_F(AuthSessionCacheTest, WaitersRetryAfterFailedLoad) { 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; @@ -443,8 +535,9 @@ TEST_F(AuthSessionCacheTest, PeriodicSweepRemovesExpiredEntries) { 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 (cache->size() != 0 && std::chrono::steady_clock::now() < deadline) { + while (recorder_.removed().empty() && std::chrono::steady_clock::now() < deadline) { std::this_thread::sleep_for(5ms); } EXPECT_EQ(cache->size(), 0); From 41dbdb32b22a78325dbaa067cfdba32dcb7a876c Mon Sep 17 00:00:00 2001 From: "shuxu.li" Date: Mon, 2 Nov 2026 19:44:37 +0800 Subject: [PATCH 5/5] style(auth): use a designated initializer for AccessEntry --- src/iceberg/catalog/rest/auth/auth_session_cache.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/iceberg/catalog/rest/auth/auth_session_cache.cc b/src/iceberg/catalog/rest/auth/auth_session_cache.cc index 53521fc87f..aa62a8ac91 100644 --- a/src/iceberg/catalog/rest/auth/auth_session_cache.cc +++ b/src/iceberg/catalog/rest/auth/auth_session_cache.cc @@ -173,8 +173,8 @@ void AuthSessionCache::TouchLocked(SlotMap::iterator it, 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{&it->first, &slot}); + slot.access_pos = access_order_.insert(access_order_.end(), + AccessEntry{.key = &it->first, .slot = &slot}); } }