From af1d15b582d71c874a9362c5de0c6bea8b4925d8 Mon Sep 17 00:00:00 2001 From: "shuxu.li" Date: Mon, 5 Oct 2026 21:50:42 +0800 Subject: [PATCH 1/2] feat(auth): implement AuthSessionCache for OAuth2 child sessions This commit implements Phase 2 of the OAuth2 enhancement (#962), adding session caching and lifecycle management for contextual and table-scoped child sessions. **Core Changes:** 1. **AuthSessionCache implementation** - Guava-like cache with expireAfterAccess semantics - Loading deduplication: concurrent Get() calls for the same key load once - Dual expiration: idle timeout + OAuth2 token expiration - Thread-safe with fine-grained locking (loader runs outside lock) - Periodic sweep for expired session cleanup - Removal listener integration for graceful session shutdown 2. **OAuth2Session enhancements** - Add `StopRefreshing()`: non-blocking refresh cancellation for cache eviction - Add `ExpiresAt()`: expose token expiration time for cache expiry checks - Add `IsRefreshable()`: indicates whether auto-refresh is enabled - Extend `OAuth2Info()` to include `headers` (fixes child auth Issue 2) - Improve `CloseImpl()`: always wait for active refreshes, defensive cancel - Add `StoreScheduledTask()`: handle Close/StopRefreshing race conditions 3. **AuthManager child session handling** - Cache contextual sessions by `session_id` (key: "ctx:") - Cache table sessions by `token` value (key: "tbl:") - Fix Issue 2: child sessions now use parent's headers for authentication - Credential-based child sessions inherit catalog-level `keep_refreshed` - Token/exchange child sessions remain non-refreshed (Phase 2 scope) 4. **Configuration** - Add `auth.session-timeout-ms` (default: 1 hour) - Document `token-exchange-enabled` as not yet implemented - Sweep interval: clamped between 1s and 60s based on timeout 5. **SessionContext extension** - Add `session_id` field for stable caller context identification - Add `identity` field for descriptive caller identity - Add `credentials` map for context-specific auth material **Testing:** - 14 new cache tests covering concurrency, expiration, and lifecycle - Extended auth_manager tests for child session creation **Notes:** - Token exchange refresh is not yet implemented (see TODO in DoRefreshAttempt) - Failed loads are not cached (retry on next Get) - Orphan sessions (loaded after Close) are stopped immediately --- src/iceberg/catalog/rest/CMakeLists.txt | 1 + src/iceberg/catalog/rest/auth/auth_manager.cc | 117 ++++- .../catalog/rest/auth/auth_properties.h | 8 + src/iceberg/catalog/rest/auth/auth_session.h | 1 + .../catalog/rest/auth/auth_session_cache.cc | 240 +++++++++++ .../rest/auth/auth_session_cache_internal.h | 129 ++++++ .../catalog/rest/auth/auth_session_internal.h | 58 ++- .../rest/auth/token_refresh_scheduler.cc | 4 +- .../rest/auth/token_refresh_scheduler.h | 3 + src/iceberg/catalog/session_context.h | 3 + src/iceberg/test/CMakeLists.txt | 1 + src/iceberg/test/auth_manager_test.cc | 63 +++ src/iceberg/test/auth_session_cache_test.cc | 403 ++++++++++++++++++ 13 files changed, 996 insertions(+), 35 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 0736290bc5..8e7bfe3347 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 auth/token_refresh_scheduler.cc diff --git a/src/iceberg/catalog/rest/auth/auth_manager.cc b/src/iceberg/catalog/rest/auth/auth_manager.cc index 5facfa3eb5..3aa78345e8 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,8 +29,10 @@ #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/rest/auth/token_refresh_scheduler.h" #include "iceberg/catalog/session_context.h" #include "iceberg/util/base64.h" #include "iceberg/util/macros.h" @@ -156,13 +159,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())); } @@ -176,6 +177,31 @@ 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. Sweep at most once a minute and at least + // once a second, so a zero timeout still gets periodic cleanup. + auto sweep_interval = std::clamp(session_timeout_, std::chrono::milliseconds(1000), + std::chrono::milliseconds(60000)); + ICEBERG_RETURN_UNEXPECTED(session_cache_->StartPeriodicSweep( + TokenRefreshScheduler::Instance(), sweep_interval)); + // Reuse the token response and start time from the init phase. if (auth_response_.has_value()) { return internal::MakeOAuth2Session( @@ -184,8 +210,8 @@ class OAuth2Manager : public AuthManager { config.optional_oauth_params(), refresh_client_, 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(), @@ -217,21 +243,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 {}; } @@ -267,7 +303,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); @@ -283,35 +320,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); } @@ -319,6 +386,12 @@ class OAuth2Manager : public AuthManager { std::optional auth_response_; std::optional start_time_; std::shared_ptr refresh_client_; + + // 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 bdc77ebc4c..71b25959b1 100644 --- a/src/iceberg/catalog/rest/auth/auth_session.h +++ b/src/iceberg/catalog/rest/auth/auth_session.h @@ -42,6 +42,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..30206f8a5d --- /dev/null +++ b/src/iceberg/catalog/rest/auth/auth_session_cache.cc @@ -0,0 +1,240 @@ +/* + * 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/catalog/rest/auth/token_refresh_scheduler.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(TokenRefreshScheduler& scheduler, + 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_ != nullptr) { + return {}; + } + sweep_scheduler_ = &scheduler; + sweep_interval_ = interval; + ScheduleSweepLocked(); + return {}; +} + +void AuthSessionCache::Close() { + std::vector> removed; + TokenRefreshScheduler* scheduler = nullptr; + uint64_t sweep_task_id = 0; + { + 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(); + scheduler = sweep_scheduler_; + sweep_task_id = std::exchange(sweep_task_id_, 0); + } + load_finished_.notify_all(); + if (scheduler != nullptr) { + scheduler->Cancel(sweep_task_id); + } + 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 sees this task's id + // or prevents it from being scheduled. + sweep_task_id_ = + sweep_scheduler_->Schedule(sweep_interval_, [weak_self = weak_from_this()] { + if (auto self = weak_self.lock()) { + self->RunScheduledSweep(); + } + }); +} + +void AuthSessionCache::RunScheduledSweep() { + Sweep(); + std::lock_guard lock(mutex_); + if (!closed_) { + ScheduleSweepLocked(); + } +} + +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..3bdcffd793 --- /dev/null +++ b/src/iceberg/catalog/rest/auth/auth_session_cache_internal.h @@ -0,0 +1,129 @@ +/* + * 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" + +/// \file iceberg/catalog/rest/auth/auth_session_cache_internal.h +/// \brief Cache of child authentication sessions, keyed by caller identity. + +namespace iceberg::rest::auth { +class TokenRefreshScheduler; +} // namespace iceberg::rest::auth + +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() every `interval` on `scheduler` until Close(). + /// + /// `scheduler` must outlive this cache or the call to Close(). `interval` must + /// be positive. Calling this more than once, or after Close(), has no effect. + Status StartPeriodicSweep(TokenRefreshScheduler& scheduler, + 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 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; + TokenRefreshScheduler* sweep_scheduler_ = nullptr; + std::chrono::milliseconds sweep_interval_{0}; + uint64_t sweep_task_id_ = 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 b8db3836e0..22b4e3f33d 100644 --- a/src/iceberg/catalog/rest/auth/auth_session_internal.h +++ b/src/iceberg/catalog/rest/auth/auth_session_internal.h @@ -99,11 +99,36 @@ 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 the refresh scheduler thread. 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() { + closed_.store(true); + TokenRefreshScheduler::Instance().Cancel( + scheduled_task_id_.exchange(kInvalidTaskHandle)); + } + + /// \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 config_.keep_refreshed; } + ~OAuth2Session() override { std::ignore = CloseImpl(); } private: @@ -114,14 +139,16 @@ class OAuth2Session final : public AuthSession, client_(std::move(client)) {} Status CloseImpl() { - bool expected = false; - if (!closed_.compare_exchange_strong(expected, true)) { - return {}; - } - TokenRefreshScheduler::Instance().Cancel(scheduled_task_id_.exchange(0)); + closed_.store(true); + TokenRefreshScheduler::Instance().Cancel( + scheduled_task_id_.exchange(kInvalidTaskHandle)); std::unique_lock lock(refresh_mutex_); refresh_cv_.wait(lock, [this] { return active_refresh_count_ == 0; }); - TokenRefreshScheduler::Instance().Cancel(scheduled_task_id_.exchange(0)); + // Defensive: cancel again in case a refresh callback scheduled a retry between + // the first cancel and acquiring refresh_mutex_. With StoreScheduledTask()'s + // closed_ recheck, this should be a no-op, but it's cheap insurance. + TokenRefreshScheduler::Instance().Cancel( + scheduled_task_id_.exchange(kInvalidTaskHandle)); return {}; } @@ -200,8 +227,7 @@ class OAuth2Session final : public AuthSession, if (closed_.load()) return; auto empty_session = AuthSession::MakeDefault({}); - // TODO(lishuxu): Honor token-exchange-enabled and refresh via token exchange, - // matching Java. + // TODO(lishuxu): Use token exchange to refresh exchanged tokens when enabled. auto result = OAuth2Util::FetchToken(*client_, *empty_session, refresh_properties_); if (result.has_value()) { auto& response = result.value(); @@ -225,7 +251,7 @@ class OAuth2Session final : public AuthSession, self->DoRefreshAttempt(next_attempt, next_backoff, refresh_started_at); } }); - scheduled_task_id_.store(retry_id); + StoreScheduledTask(retry_id); } } @@ -239,7 +265,17 @@ class OAuth2Session final : public AuthSession, delay, [weak_self = std::move(weak_self)] { if (auto self = weak_self.lock()) self->DoRefresh(); }); - scheduled_task_id_.store(new_id); + StoreScheduledTask(new_id); + } + + void StoreScheduledTask(uint64_t task_id) { + scheduled_task_id_.store(task_id); + // StopRefreshing() may run between the caller's closed_ check and the store + // above; cancel the task it could not see. + if (closed_.load()) { + TokenRefreshScheduler::Instance().Cancel( + scheduled_task_id_.exchange(kInvalidTaskHandle)); + } } std::chrono::milliseconds CalculateRefreshDelay() const { @@ -264,7 +300,7 @@ class OAuth2Session final : public AuthSession, Config config_; AuthProperties refresh_properties_; std::shared_ptr client_; - std::atomic scheduled_task_id_{0}; + std::atomic scheduled_task_id_{kInvalidTaskHandle}; std::atomic closed_{false}; std::mutex refresh_mutex_; std::condition_variable refresh_cv_; diff --git a/src/iceberg/catalog/rest/auth/token_refresh_scheduler.cc b/src/iceberg/catalog/rest/auth/token_refresh_scheduler.cc index d66ed7d906..3c9da34a02 100644 --- a/src/iceberg/catalog/rest/auth/token_refresh_scheduler.cc +++ b/src/iceberg/catalog/rest/auth/token_refresh_scheduler.cc @@ -38,7 +38,7 @@ uint64_t TokenRefreshScheduler::Schedule(std::chrono::milliseconds delay, std::function callback) { std::lock_guard lock(mutex_); if (shutdown_) { - return 0; + return kInvalidTaskHandle; } uint64_t id = next_id_++; tasks_.push_back(Task{.id = id, @@ -49,7 +49,7 @@ uint64_t TokenRefreshScheduler::Schedule(std::chrono::milliseconds delay, } void TokenRefreshScheduler::Cancel(uint64_t handle) { - if (handle == 0) return; + if (handle == kInvalidTaskHandle) return; std::lock_guard lock(mutex_); std::erase_if(tasks_, [handle](const Task& t) { return t.id == handle; }); } diff --git a/src/iceberg/catalog/rest/auth/token_refresh_scheduler.h b/src/iceberg/catalog/rest/auth/token_refresh_scheduler.h index 5ef20ed94a..3d3841c0a1 100644 --- a/src/iceberg/catalog/rest/auth/token_refresh_scheduler.h +++ b/src/iceberg/catalog/rest/auth/token_refresh_scheduler.h @@ -44,6 +44,9 @@ namespace iceberg::rest::auth { /// /// TODO(lishuxu): Migrate to the shared thread pool abstraction once available /// (see https://github.com/apache/iceberg-cpp/pull/646#discussion_r3304315308). +/// Invalid task handle returned when scheduling fails or used to represent "no task". +constexpr uint64_t kInvalidTaskHandle = 0; + class ICEBERG_REST_EXPORT TokenRefreshScheduler { public: /// \brief Get the global singleton instance. 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 891ce374a5..391574197e 100644 --- a/src/iceberg/test/CMakeLists.txt +++ b/src/iceberg/test/CMakeLists.txt @@ -327,6 +327,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 fd9605fc31..541c9b9316 100644 --- a/src/iceberg/test/auth_manager_test.cc +++ b/src/iceberg/test/auth_manager_test.cc @@ -432,6 +432,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"}, @@ -791,6 +805,55 @@ TEST(OAuth2SessionTest, InitialTokenIsUsed) { ASSERT_THAT(session->Close(), IsOk()); } +TEST(OAuth2SessionTest, StopRefreshingKeepsSessionUsable) { + 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", + .keep_refreshed = true}, + std::make_shared(), std::nullopt)); + + 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", .keep_refreshed = false}, + 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..0ec2e55a55 --- /dev/null +++ b/src/iceberg/test/auth_session_cache_test.cc @@ -0,0 +1,403 @@ +/* + * 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/catalog/rest/auth/token_refresh_scheduler.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) { + TokenRefreshScheduler scheduler; // 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(scheduler, 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, PeriodicSweepArgumentsAndLifecycle) { + TokenRefreshScheduler scheduler; + auto cache = MakeCache(); + + EXPECT_THAT(cache->StartPeriodicSweep(scheduler, 0ms), + IsError(ErrorKind::kInvalidArgument)); + EXPECT_THAT(cache->StartPeriodicSweep(scheduler, 10ms), IsOk()); + EXPECT_THAT(cache->StartPeriodicSweep(scheduler, 10ms), IsOk()); // no-op + + // Destroying the cache cancels the sweep; the scheduler keeps running safely. + cache.reset(); + std::this_thread::sleep_for(30ms); + + auto closed = MakeCache(); + closed->Close(); + EXPECT_THAT(closed->StartPeriodicSweep(scheduler, 10ms), IsOk()); // no-op +} + +} // namespace iceberg::rest::auth From 870be7db306b2fcae83ffa7c1a32c5dbef0abd67 Mon Sep 17 00:00:00 2001 From: "shuxu.li" Date: Tue, 20 Oct 2026 17:31:06 +0800 Subject: [PATCH 2/2] 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 | 72 ++++++++- .../test/rest_catalog_integration_test.cc | 7 + 6 files changed, 206 insertions(+), 74 deletions(-) diff --git a/src/iceberg/catalog/rest/auth/auth_manager.cc b/src/iceberg/catalog/rest/auth/auth_manager.cc index 3aa78345e8..d9bb95c9e1 100644 --- a/src/iceberg/catalog/rest/auth/auth_manager.cc +++ b/src/iceberg/catalog/rest/auth/auth_manager.cc @@ -60,6 +60,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; @@ -243,25 +280,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 { @@ -301,19 +341,9 @@ class OAuth2Manager : public AuthManager { config.optional_oauth_params(), refresh_client_); } - 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(); @@ -322,47 +352,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()) { @@ -373,8 +412,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 30206f8a5d..661638334a 100644 --- a/src/iceberg/catalog/rest/auth/auth_session_cache.cc +++ b/src/iceberg/catalog/rest/auth/auth_session_cache.cc @@ -190,17 +190,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 3bdcffd793..86ddedbbde 100644 --- a/src/iceberg/catalog/rest/auth/auth_session_cache_internal.h +++ b/src/iceberg/catalog/rest/auth/auth_session_cache_internal.h @@ -45,10 +45,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 541c9b9316..a016b393fd 100644 --- a/src/iceberg/test/auth_manager_test.cc +++ b/src/iceberg/test/auth_manager_test.cc @@ -566,6 +566,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 0ec2e55a55..d3cfb841af 100644 --- a/src/iceberg/test/auth_session_cache_test.cc +++ b/src/iceberg/test/auth_session_cache_test.cc @@ -33,7 +33,9 @@ #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/token_refresh_scheduler.h" +#include "iceberg/catalog/rest/http_client.h" #include "iceberg/test/matchers.h" namespace iceberg::rest::auth { @@ -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,26 @@ AuthSessionCache::Loader CountingLoader(std::atomic& calls) { }; } +/// Loads an OAuth2 session whose token expires kTokenLifetime after `clock` reads. +AuthSessionCache::Loader OAuth2Loader(std::atomic& calls, const FakeClock& clock, + bool keep_refreshed) { + return [&calls, &clock, keep_refreshed]() -> 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", + .keep_refreshed = keep_refreshed}, + std::make_shared(), clock.Now())); + return session; + }; +} + } // namespace class AuthSessionCacheTest : public ::testing::Test { @@ -210,6 +238,42 @@ 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_, /*keep_refreshed=*/false); + 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) { + auto cache = MakeCache(std::chrono::hours(1)); + std::atomic calls = 0; + auto loader = OAuth2Loader(calls, clock_, /*keep_refreshed=*/true); + 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) {