feat(rest): support OAuth2 token exchange refresh and recovery

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

@@ -309,6 +309,7 @@ class OAuth2Manager : public AuthManager {
 
     // Initialize catalog-level configuration
     keep_refreshed_ = config.keep_refreshed();
+    exchange_enabled_ = config.exchange_enabled();
     ICEBERG_ASSIGN_OR_RAISE(session_timeout_, ParseSessionTimeout(config));
 
     // Create session cache
@@ -420,11 +421,10 @@ class OAuth2Manager : public AuthManager {
           std::nullopt) const {
     ICEBERG_PRECHECK(refresh_client_ != nullptr,
                      "OAuth2 catalog session must be initialized before child sessions");
-    return internal::MakeOAuth2Session(token_response, config.oauth2_server_uri(),
-                                       config.client_id(), config.client_secret(),
-                                       config.scope(), keep_refreshed,
-                                       config.optional_oauth_params(), refresh_client_,
-                                       token_request_started_at, refresh_executor_);
+    return internal::MakeOAuth2Session(
+        token_response, config.oauth2_server_uri(), config.credential(), config.scope(),
+        keep_refreshed, config.optional_oauth_params(), refresh_client_,
+        token_request_started_at, refresh_executor_, exchange_enabled_);
   }
 
   Result<OAuth2SessionInfo> ParentInfo(const AuthSession& parent) const {
@@ -530,6 +530,7 @@ class OAuth2Manager : public AuthManager {
 
   // Catalog-level configuration initialized in CatalogSession
   bool keep_refreshed_ = true;
+  bool exchange_enabled_ = true;
   std::chrono::milliseconds session_timeout_{3'600'000};
   std::shared_ptr<internal::AuthSessionCache> session_cache_;
   std::once_flag no_refresh_executor_warning_;

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

@@ -79,8 +79,8 @@ class ICEBERG_REST_EXPORT AuthProperties : public ConfigBase<AuthProperties> {
                                                     "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; the value is ignored.
+  /// Whether to use token exchange when refreshing credential-based OAuth2 sessions.
+  /// When false, refresh uses client credentials.
   inline static Entry<bool> kExchangeEnabled{"token-exchange-enabled", true};
 
   inline static Entry<std::string> kAudience{"audience", ""};

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

@@ -64,28 +64,30 @@ Result<std::shared_ptr<AuthSession>> AuthSession::MakeOAuth2(
     const std::unordered_map<std::string, std::string>& optional_oauth_params,
     std::shared_ptr<HttpClient> client, Executor* auth_executor) {
   return internal::MakeOAuth2Session(
-      initial_token, token_endpoint, client_id, client_secret, scope, keep_refreshed,
-      optional_oauth_params, std::move(client), std::nullopt, auth_executor);
+      initial_token, token_endpoint,
+      client_id.empty() ? client_secret : client_id + ":" + client_secret, scope,
+      keep_refreshed, optional_oauth_params, std::move(client), std::nullopt,
+      auth_executor,
+      /*exchange_enabled=*/false);
 }
 
 Result<std::shared_ptr<AuthSession>> internal::MakeOAuth2Session(
     const OAuthTokenResponse& initial_token, const std::string& token_endpoint,
-    const std::string& client_id, const std::string& client_secret,
-    const std::string& scope, bool keep_refreshed,
+    const std::string& credential, const std::string& scope, bool keep_refreshed,
     const std::unordered_map<std::string, std::string>& optional_oauth_params,
     std::shared_ptr<HttpClient> client,
     std::optional<std::chrono::steady_clock::time_point> token_request_started_at,
-    Executor* refresh_executor) {
+    Executor* refresh_executor, bool exchange_enabled) {
   ICEBERG_PRECHECK(!keep_refreshed || refresh_executor != nullptr,
                    "OAuth2 token refresh is enabled but no auth executor was provided; "
                    "pass an executor or set {}=false",
                    AuthProperties::kKeepRefreshed.key());
   internal::OAuth2Session::Config config{
       .token_endpoint = token_endpoint,
-      .client_id = client_id,
-      .client_secret = client_secret,
+      .credential = credential,
       .scope = scope,
       .optional_oauth_params = optional_oauth_params,
+      .exchange_enabled = exchange_enabled,
       .refresh_executor = keep_refreshed ? refresh_executor : nullptr,
   };
   ICEBERG_ASSIGN_OR_RAISE(auto session, internal::OAuth2Session::Make(

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

@@ -88,6 +88,7 @@ class ICEBERG_REST_EXPORT AuthSession {
   /// \brief Create an OAuth2 session with automatic token refresh.
   ///
   /// Refreshes expiring tokens in the background when keep_refreshed is true.
+  /// This factory uses the client-credentials grant for refresh.
   /// Authenticate() uses the cached token without refreshing it.
   ///
   /// \param initial_token The initial token response from FetchToken().

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

@@ -23,6 +23,7 @@
 #include <chrono>
 #include <condition_variable>
 #include <exception>
+#include <functional>
 #include <memory>
 #include <mutex>
 #include <optional>
@@ -36,6 +37,7 @@
 #include "iceberg/catalog/rest/auth/oauth2_util.h"
 #include "iceberg/catalog/rest/http_client.h"
 #include "iceberg/logging/log_macros.h"
+#include "iceberg/util/base64.h"
 #include "iceberg/util/macros.h"
 #include "iceberg/util/retry_util.h"
 #include "iceberg/util/scheduled_executor_internal.h"
@@ -65,12 +67,16 @@ class OAuth2Session final : public AuthSession,
  public:
   struct Config {
     std::string token_endpoint;
-    std::string client_id;
-    std::string client_secret;
+    std::string credential;
     std::string scope;
     std::unordered_map<std::string, std::string> optional_oauth_params;
+    /// The legacy public factory keeps client-credentials refresh.
+    bool exchange_enabled = false;
     /// Borrowed executor for background refresh. Null disables refresh.
     Executor* refresh_executor = nullptr;
+    /// Internal clock override for deterministic expiration tests.
+    std::function<std::chrono::steady_clock::time_point()> now =
+        std::chrono::steady_clock::now;
   };
 
   static Result<std::shared_ptr<OAuth2Session>> Make(
@@ -98,7 +104,7 @@ class OAuth2Session final : public AuthSession,
     return OAuth2SessionInfo{
         .token = token_,
         .issued_token_type = issued_token_type_,
-        .credential = Credential(config_),
+        .credential = config_.credential,
         .scope = config_.scope,
         .oauth2_server_uri = config_.token_endpoint,
         .optional_oauth_params = config_.optional_oauth_params,
@@ -135,8 +141,8 @@ class OAuth2Session final : public AuthSession,
 
   /// \brief Whether this session will keep refreshing its token.
   ///
-  /// False if refresh is disabled, was stopped by StopRefreshing() or Close(), or
-  /// gave up after repeated failures.
+  /// False if expiration is unknown, refresh is disabled, was stopped by
+  /// StopRefreshing() or Close(), or gave up after repeated failures.
   bool IsRefreshable() const {
     std::lock_guard lock(refresh_mutex_);
     return scheduler_.has_value() && !closed_ && !refresh_stopped_;
@@ -169,15 +175,10 @@ class OAuth2Session final : public AuthSession,
     return {};
   }
 
-  static std::string Credential(const Config& config) {
-    return config.client_id.empty() ? config.client_secret
-                                    : config.client_id + ":" + config.client_secret;
-  }
-
   static Result<AuthProperties> MakeRefreshProperties(const Config& config) {
     std::unordered_map<std::string, std::string> properties =
         config.optional_oauth_params;
-    properties[AuthProperties::kCredential.key()] = Credential(config);
+    properties[AuthProperties::kCredential.key()] = config.credential;
     properties[AuthProperties::kScope.key()] = config.scope;
     properties[AuthProperties::kOAuth2ServerUri.key()] = config.token_endpoint;
     return AuthProperties::FromProperties(properties);
@@ -215,32 +216,31 @@ class OAuth2Session final : public AuthSession,
   };
 
   void UpdateTokenState(const OAuthTokenResponse& token_response,
-                        std::optional<std::chrono::steady_clock::time_point>
-                            token_request_started_at = std::nullopt) {
+                        std::optional<std::chrono::steady_clock::time_point> expiration) {
     token_ = token_response.access_token;
     issued_token_type_ = token_response.issued_token_type.empty()
                              ? AuthProperties::kAccessTokenType
                              : token_response.issued_token_type;
     headers_ = OAuth2Util::AuthHeaders(token_);
-    expires_at_ = std::chrono::steady_clock::time_point{};
-    auto request_started_at =
-        token_request_started_at.value_or(std::chrono::steady_clock::now());
-    if (auto expiration = TokenExpirationTime(token_response, request_started_at);
-        expiration.has_value()) {
-      expires_at_ = *expiration;
-    }
+    expires_at_ = expiration.value_or(std::chrono::steady_clock::time_point{});
   }
 
   void SetInitialToken(
       const OAuthTokenResponse& token_response,
       std::optional<std::chrono::steady_clock::time_point> token_request_started_at) {
-    UpdateTokenState(token_response, token_request_started_at);
+    auto request_started_at = token_request_started_at.value_or(config_.now());
+    UpdateTokenState(
+        token_response,
+        TokenExpirationTime(token_response, request_started_at,
+                            std::chrono::system_clock::now(), config_.now()));
     if (!scheduler_) {
       return;
     }
     // Refresh an already-expired initial token immediately.
     if (auto delay = CalculateRefreshDelay()) {
       ScheduleRefresh(*delay);
+    } else {
+      MarkRefreshStopped();
     }
   }
 
@@ -254,43 +254,95 @@ class OAuth2Session final : public AuthSession,
     auto result = RetryRunner<retry::StopRetryOn<>>({.num_retries = 5}).Run([this] {
       return RefreshToken();
     });
+    if (!result.has_value()) {
+      // One final credential-backed attempt, within the same close guard.
+      result = RefreshToken(/*recover=*/true);
+    }
     if (!result.has_value()) {
       ICEBERG_LOG_WARN("Failed to refresh OAuth2 token: {}", result.error().message);
       MarkRefreshStopped();
       return std::nullopt;
     }
+    if (!result->has_value()) {
+      MarkRefreshStopped();
+    }
     return *result;
   }
 
-  // Record that no further refresh will be scheduled after a failure.
+  // Record that no further refresh will be scheduled.
   void MarkRefreshStopped() {
     std::lock_guard lock(refresh_mutex_);
     refresh_stopped_ = true;
   }
 
-  Result<std::optional<std::chrono::milliseconds>> RefreshToken() {
+  class ClientCredentialsSession final : public AuthSession {
+   public:
+    Result<HttpRequest> Authenticate(HttpRequest request) override {
+      // HttpClient merges its defaults before invoking the session.
+      request.headers.erase("Authorization");
+      return request;
+    }
+  };
+
+  Result<OAuthTokenResponse> FetchRefreshToken(const std::string& token,
+                                               const std::string& issued_token_type,
+                                               bool recover) {
+    if (!config_.exchange_enabled) {
+      ClientCredentialsSession session;
+      return OAuth2Util::FetchToken(*client_, session, refresh_properties_);
+    }
+    auto empty_session = AuthSession::MakeDefault({});
+    auto headers =
+        recover
+            ? std::unordered_map<std::string, std::string>{{"Authorization",
+                                                            "Basic " +
+                                                                Base64::Encode(
+                                                                    config_.credential)}}
+            : OAuth2Util::AuthHeaders(token);
+    // Explicit request headers override any default Authorization on the client.
+    // Keep the current token as subject, without a parent actor.
+    return OAuth2Util::ExchangeToken(*client_, *empty_session, headers, token,
+                                     issued_token_type, std::nullopt, std::nullopt,
+                                     config_.scope, config_.token_endpoint,
+                                     config_.optional_oauth_params);
+  }
+
+  Result<std::optional<std::chrono::milliseconds>> RefreshToken(bool recover = false) {
     {
       std::lock_guard lock(refresh_mutex_);
       if (closed_) {
         return std::nullopt;
       }
     }
     try {
-      auto request_started_at = std::chrono::steady_clock::now();
-      auto empty_session = AuthSession::MakeDefault({});
-      // TODO(lishuxu): Honor token-exchange-enabled when refreshing tokens.
-      ICEBERG_ASSIGN_OR_RAISE(
-          auto response,
-          OAuth2Util::FetchToken(*client_, *empty_session, refresh_properties_));
+      std::string token;
+      std::string issued_token_type;
       {
-        std::unique_lock lock(mutex_);
-        UpdateTokenState(response, request_started_at);
+        std::shared_lock lock(mutex_);
+        token = token_;
+        issued_token_type = issued_token_type_;
+        recover = recover || (expires_at_ != std::chrono::steady_clock::time_point{} &&
+                              expires_at_ <= config_.now());
       }
-      auto delay = CalculateRefreshDelay();
-      if (delay.has_value() && *delay <= std::chrono::milliseconds::zero()) {
-        return TokenExpired("Refreshed OAuth2 token is already expired");
+      if ((!config_.exchange_enabled || recover) && config_.credential.empty()) {
+        ICEBERG_LOG_WARN("Stopping OAuth2 token refresh: no credential available");
+        return std::nullopt;
+      }
+      auto request_started_at = config_.now();
+      ICEBERG_ASSIGN_OR_RAISE(auto response,
+                              FetchRefreshToken(token, issued_token_type, recover));
+      auto expiration = TokenExpirationTime(
+          response, request_started_at, std::chrono::system_clock::now(), config_.now());
+      {
+        std::unique_lock lock(mutex_);
+        // Keep the current token available for requests and exchange retries until
+        // the replacement has passed expiration validation.
+        if (expiration.has_value() && *expiration <= config_.now()) {
+          return TokenExpired("Refreshed OAuth2 token is already expired");
+        }
+        UpdateTokenState(response, expiration);
       }
-      return delay;
+      return CalculateRefreshDelay();
     } catch (const std::exception& e) {
       return UnknownError("OAuth2 token refresh threw: {}", e.what());
     } catch (...) {
@@ -338,7 +390,7 @@ class OAuth2Session final : public AuthSession,
     if (expires_at_ == std::chrono::steady_clock::time_point{}) {
       return std::nullopt;
     }
-    auto now = std::chrono::steady_clock::now();
+    auto now = config_.now();
     if (expires_at_ <= now) {
       return std::chrono::milliseconds::zero();
     }
@@ -369,11 +421,10 @@ class OAuth2Session final : public AuthSession,
 /// Returns InvalidArgument if keep_refreshed is true without an executor.
 Result<std::shared_ptr<AuthSession>> MakeOAuth2Session(
     const OAuthTokenResponse& initial_token, const std::string& token_endpoint,
-    const std::string& client_id, const std::string& client_secret,
-    const std::string& scope, bool keep_refreshed,
+    const std::string& credential, const std::string& scope, bool keep_refreshed,
     const std::unordered_map<std::string, std::string>& optional_oauth_params,
     std::shared_ptr<HttpClient> client,
     std::optional<std::chrono::steady_clock::time_point> token_request_started_at,
-    Executor* refresh_executor);
+    Executor* refresh_executor, bool exchange_enabled);
 
 }  // namespace iceberg::rest::auth::internal