From de528cada7b42765c5c595e06a60cc3374f44572 Mon Sep 17 00:00:00 2001 From: "shuxu.li" Date: Mon, 5 Oct 2026 21:50:42 +0800 Subject: [PATCH 1/3] 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/3] 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/3] 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