feat(rest): support OAuth2 token exchange sessions

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);