feat(auth): implement AuthSessionCache for OAuth2 child sessions

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

src/iceberg/catalog/rest/auth/auth_manager.cc

@@ -19,6 +19,7 @@
 
 #include "iceberg/catalog/rest/auth/auth_manager.h"
 
+#include <algorithm>
 #include <array>
 #include <chrono>
 #include <optional>
@@ -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"
@@ -57,6 +60,43 @@ std::optional<std::pair<std::string, std::string>> 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<ChildCredential> SelectChildCredential(
+    const std::unordered_map<std::string, std::string>& 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<std::string, std::string> FilterTableSessionProperties(
     const std::unordered_map<std::string, std::string>& properties) {
   std::unordered_map<std::string, std::string> filtered;
@@ -156,13 +196,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 +214,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<AuthSession> session) {
+              if (auto oauth2 =
+                      std::dynamic_pointer_cast<internal::OAuth2Session>(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 +247,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 +280,34 @@ class OAuth2Manager : public AuthManager {
 
   Result<std::shared_ptr<AuthSession>> ContextualSession(
       const SessionContext& context, std::shared_ptr<AuthSession> parent) override {
-    // TODO(lishuxu): Add child-session caching and refresh, matching Java
-    // AuthSessionCache.
-    return MaybeCreateChildSession(context.credentials, /*allow_credential=*/true,
-                                   std::move(parent));
+    auto credential =
+        SelectChildCredential(context.credentials, /*allow_credential=*/true);
+    if (!credential.has_value()) {
+      return parent;
+    }
+    // Contextual sessions are cached by session ID, whatever credential they use.
+    std::string cache_key = context.session_id.empty() ? "" : "ctx:" + context.session_id;
+    return GetOrCreateChildSession(*credential, std::move(parent), cache_key);
   }
 
   Result<std::shared_ptr<AuthSession>> TableSession(
       [[maybe_unused]] const TableIdentifier& table,
       const std::unordered_map<std::string, std::string>& properties,
       std::shared_ptr<AuthSession> parent) override {
-    return MaybeCreateChildSession(FilterTableSessionProperties(properties),
-                                   /*allow_credential=*/false, std::move(parent));
+    auto credential = SelectChildCredential(FilterTableSessionProperties(properties),
+                                            /*allow_credential=*/false);
+    if (!credential.has_value()) {
+      return parent;
+    }
+    // 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 {
+    if (session_cache_) {
+      session_cache_->Close();
+    }
     refresh_client_.reset();
     return {};
   }
@@ -265,60 +341,96 @@ class OAuth2Manager : public AuthManager {
                                    config.optional_oauth_params(), refresh_client_);
   }
 
-  Result<std::shared_ptr<AuthSession>> MaybeCreateChildSession(
-      const std::unordered_map<std::string, std::string>& credentials,
-      bool allow_credential, std::shared_ptr<AuthSession> parent) {
-    auto token_it = credentials.find(AuthProperties::kToken.key());
-    auto credential_it = credentials.find(AuthProperties::kCredential.key());
-    auto typed_token = FindPreferredTypedToken(credentials);
-    if (token_it == credentials.end() &&
-        (!allow_credential || credential_it == credentials.end()) &&
-        !typed_token.has_value()) {
-      return parent;
-    }
-
+  Result<std::shared_ptr<AuthSession>> GetOrCreateChildSession(
+      const ChildCredential& credential, std::shared_ptr<AuthSession> parent,
+      const std::string& cache_key) {
     ICEBERG_PRECHECK(refresh_client_ != nullptr,
                      "OAuth2 catalog session must be initialized before child sessions");
     auto parent_info = parent->OAuth2Info();
     ICEBERG_PRECHECK(parent_info.has_value(),
                      "OAuth2 child session requires OAuth2 parent metadata");
 
-    if (token_it != credentials.end()) {
-      ICEBERG_ASSIGN_OR_RAISE(auto config,
-                              ChildConfig(*parent_info, parent_info->credential));
-      return MakeSession(AccessTokenResponse(token_it->second), config,
-                         /*keep_refreshed=*/false);
+    // Use cache if we have a valid cache_key
+    if (!cache_key.empty() && session_cache_) {
+      return session_cache_->Get(
+          cache_key, [&]() -> Result<std::shared_ptr<AuthSession>> {
+            return CreateChildSessionUncached(credential, *parent_info);
+          });
     }
 
-    if (allow_credential && credential_it != credentials.end()) {
-      ICEBERG_ASSIGN_OR_RAISE(auto config,
-                              ChildConfig(*parent_info, credential_it->second));
-      ICEBERG_ASSIGN_OR_RAISE(auto response,
-                              OAuth2Util::FetchToken(*refresh_client_, *parent, config));
-      return MakeSession(response, config, /*keep_refreshed=*/false);
+    // No cache, create session directly
+    return CreateChildSessionUncached(credential, *parent_info);
+  }
+
+  Result<std::shared_ptr<AuthSession>> CreateChildSessionUncached(
+      const ChildCredential& credential, const OAuth2SessionInfo& parent_info) {
+    switch (credential.kind) {
+      case ChildCredential::Kind::kToken:
+        return CreateTokenChildSession(credential.value, parent_info);
+      case ChildCredential::Kind::kCredential:
+        return CreateCredentialChildSession(credential.value, parent_info);
+      case ChildCredential::Kind::kTokenExchange:
+        return CreateTokenExchangeChildSession(credential.value, credential.token_type,
+                                               parent_info);
     }
+    std::unreachable();
+  }
+
+  Result<std::shared_ptr<AuthSession>> 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<std::shared_ptr<AuthSession>> 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<std::shared_ptr<AuthSession>> CreateTokenExchangeChildSession(
+      const std::string& subject_token, const std::string& subject_token_type,
+      const OAuth2SessionInfo& parent_info) {
     std::optional<std::string> actor_token;
     std::optional<std::string> 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,
-                                  typed_token->first, actor_token, actor_token_type,
-                                  parent_info->scope, parent_info->oauth2_server_uri,
-                                  parent_info->optional_oauth_params));
+        OAuth2Util::ExchangeToken(*refresh_client_, *temp_parent, {}, subject_token,
+                                  subject_token_type, actor_token, actor_token_type,
+                                  parent_info.scope, parent_info.oauth2_server_uri,
+                                  parent_info.optional_oauth_params));
     ICEBERG_ASSIGN_OR_RAISE(auto config,
-                            ChildConfig(*parent_info, parent_info->credential));
+                            ChildConfig(parent_info, parent_info.credential));
     return MakeSession(response, config, /*keep_refreshed=*/false);
   }
 
   /// Token response and start time captured by InitSession.
   std::optional<OAuthTokenResponse> auth_response_;
   std::optional<std::chrono::steady_clock::time_point> start_time_;
   std::shared_ptr<HttpClient> 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<internal::AuthSessionCache> session_cache_;
 };
 
 Result<std::unique_ptr<AuthManager>> MakeOAuth2Manager(

src/iceberg/catalog/rest/auth/auth_properties.h

@@ -78,10 +78,18 @@ class ICEBERG_REST_EXPORT AuthProperties : public ConfigBase<AuthProperties> {
   inline static Entry<std::string> kOAuth2ServerUri{"oauth2-server-uri",
                                                     "v1/oauth/tokens"};
   inline static Entry<bool> 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<bool> kExchangeEnabled{"token-exchange-enabled", true};
+
   inline static Entry<std::string> kAudience{"audience", ""};
   inline static Entry<std::string> 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<int64_t> kSessionTimeoutMs{"auth.session-timeout-ms", 3'600'000};
+
   // ---- OAuth2 token type constants ----
 
   inline static const std::string kAccessTokenType =

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<std::string, std::string> optional_oauth_params;
+  std::unordered_map<std::string, std::string> headers;
 };
 
 /// \brief An authentication session that can authenticate outgoing HTTP requests.