feat(rest): support OAuth2 token exchange sessions (#867)

Add RFC 8693 token exchange support, including OAuth2 endpoint handling and response parsing. Create contextual and table child sessions from direct tokens, credentials, or typed tokens.

Keep OAuth2 metadata in one synchronized session snapshot. Use shared `HttpClient` ownership so refresh tasks can safely keep the client alive. Align token expiry with Java by preferring JWT `exp` and otherwise using the request start time plus `expires_in`.

Move OAuth2 helpers into `OAuth2Util` and update unit and integration tests. Leave catalog token refresh, token exchange during refresh, and child session caching as follow-up work.

Author: lishuxu

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

@@ -19,22 +19,67 @@
 
 #include "iceberg/catalog/rest/auth/auth_manager.h"
 
+#include <array>
+#include <chrono>
 #include <optional>
+#include <string_view>
+#include <utility>
 
 #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_internal.h"
 #include "iceberg/catalog/rest/auth/oauth2_util.h"
+#include "iceberg/catalog/session_context.h"
 #include "iceberg/util/base64.h"
 #include "iceberg/util/macros.h"
 
 namespace iceberg::rest::auth {
 
+namespace {
+
+constexpr std::string_view kAuthorizationHeader = "Authorization";
+
+const std::array<std::string_view, 5> kTokenPreferenceOrder = {
+    AuthProperties::kIdTokenType,    AuthProperties::kAccessTokenType,
+    AuthProperties::kJwtTokenType,   AuthProperties::kSaml2TokenType,
+    AuthProperties::kSaml1TokenType,
+};
+
+std::optional<std::pair<std::string, std::string>> FindPreferredTypedToken(
+    const std::unordered_map<std::string, std::string>& credentials) {
+  for (std::string_view token_type : kTokenPreferenceOrder) {
+    auto token_it = credentials.find(std::string(token_type));
+    if (token_it != credentials.end()) {
+      return std::pair{token_it->first, token_it->second};
+    }
+  }
+  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;
+  if (auto token_it = properties.find(AuthProperties::kToken.key());
+      token_it != properties.end()) {
+    filtered.emplace(token_it->first, token_it->second);
+  }
+  for (std::string_view token_type : kTokenPreferenceOrder) {
+    auto token_it = properties.find(std::string(token_type));
+    if (token_it != properties.end()) {
+      filtered.emplace(token_it->first, token_it->second);
+    }
+  }
+  return filtered;
+}
+
+}  // namespace
+
 Result<std::shared_ptr<AuthSession>> AuthManager::InitSession(
-    HttpClient& init_client,
+    std::shared_ptr<HttpClient> init_client,
     const std::unordered_map<std::string, std::string>& properties) {
   // By default, use the catalog session for initialization
-  return CatalogSession(init_client, properties);
+  return CatalogSession(std::move(init_client), properties);
 }
 
 Result<std::shared_ptr<AuthSession>> AuthManager::ContextualSession(
@@ -55,7 +100,7 @@ Result<std::shared_ptr<AuthSession>> AuthManager::TableSession(
 class NoopAuthManager : public AuthManager {
  public:
   Result<std::shared_ptr<AuthSession>> CatalogSession(
-      [[maybe_unused]] HttpClient& client,
+      [[maybe_unused]] std::shared_ptr<HttpClient> client,
       [[maybe_unused]] const std::unordered_map<std::string, std::string>& properties)
       override {
     return AuthSession::MakeDefault({});
@@ -72,7 +117,7 @@ Result<std::unique_ptr<AuthManager>> MakeNoopAuthManager(
 class BasicAuthManager : public AuthManager {
  public:
   Result<std::shared_ptr<AuthSession>> CatalogSession(
-      [[maybe_unused]] HttpClient& client,
+      [[maybe_unused]] std::shared_ptr<HttpClient> client,
       const std::unordered_map<std::string, std::string>& properties) override {
     auto username_it = properties.find(AuthProperties::kBasicUsername);
     ICEBERG_PRECHECK(username_it != properties.end() && !username_it->second.empty(),
@@ -96,67 +141,184 @@ Result<std::unique_ptr<AuthManager>> MakeBasicAuthManager(
 class OAuth2Manager : public AuthManager {
  public:
   Result<std::shared_ptr<AuthSession>> InitSession(
-      HttpClient& init_client,
+      std::shared_ptr<HttpClient> init_client,
       const std::unordered_map<std::string, std::string>& properties) override {
+    ICEBERG_PRECHECK(init_client != nullptr,
+                     "OAuth2 initialization HTTP client must not be null");
     ICEBERG_ASSIGN_OR_RAISE(auto config, AuthProperties::FromProperties(properties));
     // No token refresh during init (short-lived session).
     config.Set(AuthProperties::kKeepRefreshed, false);
 
     // Credential takes priority: fetch a fresh token for the config request.
     if (!config.credential().empty()) {
-      auto init_session = AuthSession::MakeDefault(AuthHeaders(config.token()));
-      ICEBERG_ASSIGN_OR_RAISE(init_token_response_,
-                              FetchToken(init_client, *init_session, config));
-      return AuthSession::MakeDefault(AuthHeaders(init_token_response_->access_token));
+      auto init_session =
+          AuthSession::MakeDefault(OAuth2Util::AuthHeaders(config.token()));
+      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()) {
-      return AuthSession::MakeDefault(AuthHeaders(config.token()));
+      // TODO(lishuxu): Match Java OAuth2Util.AuthSession.fromAccessToken here.
+      return AuthSession::MakeDefault(OAuth2Util::AuthHeaders(config.token()));
     }
 
     return AuthSession::MakeDefault({});
   }
 
   Result<std::shared_ptr<AuthSession>> CatalogSession(
-      HttpClient& client,
+      std::shared_ptr<HttpClient> shared_client,
       const std::unordered_map<std::string, std::string>& properties) override {
     ICEBERG_ASSIGN_OR_RAISE(auto config, AuthProperties::FromProperties(properties));
-
-    // Reuse token from init phase.
-    if (init_token_response_.has_value()) {
-      auto token_response = std::move(*init_token_response_);
-      init_token_response_.reset();
-      return AuthSession::MakeOAuth2(token_response, config.oauth2_server_uri(),
-                                     config.client_id(), config.client_secret(),
-                                     config.scope(), config.keep_refreshed(),
-                                     config.optional_oauth_params(), client);
+    ICEBERG_PRECHECK(shared_client != nullptr,
+                     "OAuth2 catalog session HTTP client must not be null");
+    refresh_client_ = std::move(shared_client);
+    // Reuse the token response and start time from the init phase.
+    if (auth_response_.has_value()) {
+      return internal::MakeOAuth2Session(
+          *auth_response_, config.oauth2_server_uri(), config.client_id(),
+          config.client_secret(), config.scope(), config.keep_refreshed(),
+          config.optional_oauth_params(), refresh_client_, start_time_);
     }
 
-    // If token is provided, use it directly.
+    // TODO(lishuxu): Honor token-refresh-enabled for catalog bearer tokens, matching
+    // Java. If token is provided, use it directly.
     if (!config.token().empty()) {
-      return AuthSession::MakeDefault(AuthHeaders(config.token()));
+      OAuthTokenResponse token_response{
+          .access_token = config.token(),
+          .token_type = "bearer",
+          .issued_token_type = AuthProperties::kAccessTokenType,
+      };
+      return AuthSession::MakeOAuth2(token_response, config.oauth2_server_uri(),
+                                     config.client_id(), config.client_secret(),
+                                     config.scope(), /*keep_refreshed=*/false,
+                                     config.optional_oauth_params(), refresh_client_);
     }
 
     // Fetch a new token using client_credentials grant.
     if (!config.credential().empty()) {
-      auto base_session = AuthSession::MakeDefault(AuthHeaders(config.token()));
+      auto base_session =
+          AuthSession::MakeDefault(OAuth2Util::AuthHeaders(config.token()));
       OAuthTokenResponse token_response;
-      ICEBERG_ASSIGN_OR_RAISE(token_response, FetchToken(client, *base_session, config));
+      ICEBERG_ASSIGN_OR_RAISE(
+          token_response,
+          OAuth2Util::FetchToken(*refresh_client_, *base_session, config));
       return AuthSession::MakeOAuth2(token_response, config.oauth2_server_uri(),
                                      config.client_id(), config.client_secret(),
                                      config.scope(), config.keep_refreshed(),
-                                     config.optional_oauth_params(), client);
+                                     config.optional_oauth_params(), refresh_client_);
     }
 
-    return AuthSession::MakeDefault({});
+    return MakeSession(AccessTokenResponse(""), config, /*keep_refreshed=*/false);
   }
 
-  // TODO(lishuxu): Override TableSession() for token exchange (RFC 8693).
-  // TODO(lishuxu): Override ContextualSession() for per-context exchange.
+  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));
+  }
+
+  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));
+  }
+
+  Status Close() override {
+    refresh_client_.reset();
+    return {};
+  }
 
  private:
-  /// Cached token from InitSession
-  std::optional<OAuthTokenResponse> init_token_response_;
+  static OAuthTokenResponse AccessTokenResponse(std::string token) {
+    return {
+        .access_token = std::move(token),
+        .token_type = "bearer",
+        .issued_token_type = AuthProperties::kAccessTokenType,
+    };
+  }
+
+  static Result<AuthProperties> ChildConfig(const OAuth2SessionInfo& parent_info,
+                                            const std::string& credential) {
+    auto properties = parent_info.optional_oauth_params;
+    properties[AuthProperties::kCredential.key()] = credential;
+    properties[AuthProperties::kScope.key()] = parent_info.scope;
+    properties[AuthProperties::kOAuth2ServerUri.key()] = parent_info.oauth2_server_uri;
+    return AuthProperties::FromProperties(properties);
+  }
+
+  Result<std::shared_ptr<AuthSession>> MakeSession(
+      const OAuthTokenResponse& token_response, const AuthProperties& config,
+      bool keep_refreshed) const {
+    ICEBERG_PRECHECK(refresh_client_ != nullptr,
+                     "OAuth2 catalog session must be initialized before child sessions");
+    return AuthSession::MakeOAuth2(token_response, config.oauth2_server_uri(),
+                                   config.client_id(), config.client_secret(),
+                                   config.scope(), keep_refreshed,
+                                   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;
+    }
+
+    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);
+    }
+
+    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);
+    }
+
+    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;
+    }
+    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));
+    ICEBERG_ASSIGN_OR_RAISE(auto config,
+                            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_;
 };
 
 Result<std::unique_ptr<AuthManager>> MakeOAuth2Manager(

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

@@ -48,7 +48,7 @@ class ICEBERG_REST_EXPORT AuthManager {
   /// \param properties Client configuration supplied by the catalog.
   /// \return Session for initialization or an error if credentials cannot be acquired.
   virtual Result<std::shared_ptr<AuthSession>> InitSession(
-      HttpClient& init_client,
+      std::shared_ptr<HttpClient> init_client,
       const std::unordered_map<std::string, std::string>& properties);
 
   /// \brief Create the long-lived catalog session that acts as the parent session.
@@ -62,7 +62,7 @@ class ICEBERG_REST_EXPORT AuthManager {
   /// \return Session for catalog operations or an error if authentication cannot be set
   /// up.
   virtual Result<std::shared_ptr<AuthSession>> CatalogSession(
-      HttpClient& shared_client,
+      std::shared_ptr<HttpClient> shared_client,
       const std::unordered_map<std::string, std::string>& properties) = 0;
 
   /// \brief Create or reuse a session for a specific context.

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

@@ -22,6 +22,7 @@
 #include <utility>
 
 #include "iceberg/catalog/rest/catalog_properties.h"
+#include "iceberg/catalog/rest/rest_util.h"
 
 namespace iceberg::rest::auth {
 
@@ -35,6 +36,31 @@ std::pair<std::string, std::string> ParseCredential(const std::string& credentia
   return {credential.substr(0, colon_pos), credential.substr(colon_pos + 1)};
 }
 
+Result<std::string> ResolveOAuth2ServerUri(
+    const std::unordered_map<std::string, std::string>& properties) {
+  auto endpoint_it = properties.find(AuthProperties::kOAuth2ServerUri.key());
+  std::string endpoint = endpoint_it == properties.end()
+                             ? AuthProperties::kOAuth2ServerUri.value()
+                             : endpoint_it->second;
+
+  if (endpoint.starts_with("http://") || endpoint.starts_with("https://")) {
+    return endpoint;
+  }
+  if (endpoint.empty()) {
+    return endpoint;
+  }
+  auto uri_it = properties.find(RestCatalogProperties::kUri.key());
+  if (uri_it == properties.end() || uri_it->second.empty()) {
+    return endpoint;
+  }
+
+  auto base_uri = std::string(TrimTrailingSlash(uri_it->second));
+  if (endpoint.starts_with('/')) {
+    return base_uri + endpoint;
+  }
+  return base_uri + "/" + endpoint;
+}
+
 }  // namespace
 
 std::unordered_map<std::string, std::string> AuthProperties::optional_oauth_params()
@@ -61,19 +87,8 @@ Result<AuthProperties> AuthProperties::FromProperties(
     config.client_secret_ = std::move(secret);
   }
 
-  // Resolve token endpoint: if not explicitly set, derive from catalog URI
-  if (properties.find(kOAuth2ServerUri.key()) == properties.end() ||
-      properties.at(kOAuth2ServerUri.key()).empty()) {
-    auto uri_it = properties.find(RestCatalogProperties::kUri.key());
-    if (uri_it != properties.end() && !uri_it->second.empty()) {
-      std::string_view base = uri_it->second;
-      while (!base.empty() && base.back() == '/') {
-        base.remove_suffix(1);
-      }
-      config.Set(kOAuth2ServerUri,
-                 std::string(base) + "/" + std::string(kOAuth2ServerUri.value()));
-    }
-  }
+  ICEBERG_ASSIGN_OR_RAISE(auto oauth2_server_uri, ResolveOAuth2ServerUri(properties));
+  config.Set(kOAuth2ServerUri, std::move(oauth2_server_uri));
 
   // TODO(lishuxu): Parse JWT exp claim from token to set expires_at_millis_.
 

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

@@ -82,6 +82,20 @@ class ICEBERG_REST_EXPORT AuthProperties : public ConfigBase<AuthProperties> {
   inline static Entry<std::string> kAudience{"audience", ""};
   inline static Entry<std::string> kResource{"resource", ""};
 
+  // ---- OAuth2 token type constants ----
+
+  inline static const std::string kAccessTokenType =
+      "urn:ietf:params:oauth:token-type:access_token";
+  inline static const std::string kRefreshTokenType =
+      "urn:ietf:params:oauth:token-type:refresh_token";
+  inline static const std::string kIdTokenType =
+      "urn:ietf:params:oauth:token-type:id_token";
+  inline static const std::string kSaml1TokenType =
+      "urn:ietf:params:oauth:token-type:saml1";
+  inline static const std::string kSaml2TokenType =
+      "urn:ietf:params:oauth:token-type:saml2";
+  inline static const std::string kJwtTokenType = "urn:ietf:params:oauth:token-type:jwt";
+
   /// \brief Build an AuthProperties from a properties map.
   static Result<AuthProperties> FromProperties(
       const std::unordered_map<std::string, std::string>& properties);