Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
112 changes: 108 additions & 4 deletions src/iceberg/catalog/rest/auth/auth_manager.cc
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
#include "iceberg/catalog/rest/auth/auth_properties.h"
#include "iceberg/catalog/rest/auth/auth_session.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"

Expand Down Expand Up @@ -121,6 +122,7 @@ class OAuth2Manager : public AuthManager {
HttpClient& client,
const std::unordered_map<std::string, std::string>& properties) override {
ICEBERG_ASSIGN_OR_RAISE(auto config, AuthProperties::FromProperties(properties));
shared_client_ = &client;

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

shared_client_ is a borrowed HttpClient*, but child sessions use it after this call. Please pass a std::shared_ptr<HttpClient> through the manager/session API so the client lifetime is explicit.


// Reuse token from init phase.
if (init_token_response_.has_value()) {
Expand All @@ -134,7 +136,15 @@ class OAuth2Manager : public AuthManager {

// 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(), client);
}

// Fetch a new token using client_credentials grant.
Expand All @@ -148,15 +158,109 @@ class OAuth2Manager : public AuthManager {
config.optional_oauth_params(), client);
}

return AuthSession::MakeDefault({});
return MakeSession(AccessTokenResponse(""), config, /*keep_refreshed=*/false);
}

Result<std::shared_ptr<AuthSession>> ContextualSession(
const SessionContext& context, std::shared_ptr<AuthSession> parent) override {
return MaybeCreateChildSession(context.credentials, /*allow_credential=*/true,
Comment thread
wgtmac marked this conversation as resolved.
std::move(parent));
}

// TODO(lishuxu): Override TableSession() for token exchange (RFC 8693).
// TODO(lishuxu): Override ContextualSession() for per-context exchange.
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));
}

private:
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(shared_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(), *shared_client_);
}

Result<std::shared_ptr<AuthSession>> MaybeCreateChildSession(
const std::unordered_map<std::string, std::string>& credentials,
bool allow_credential, std::shared_ptr<AuthSession> parent) const {
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(shared_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,
Comment thread
wgtmac marked this conversation as resolved.
/*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,
FetchToken(*shared_client_, *parent, config));
return MakeSession(response, config, /*keep_refreshed=*/false);
}

std::optional<OAuth2Token> actor;
if (!parent_info->token.empty()) {
actor = OAuth2Token{
.token_type = parent_info->issued_token_type,
.token = parent_info->token,
};
}
TokenExchangeRequest request{
.oauth2_server_uri = parent_info->oauth2_server_uri,
.subject = std::move(*typed_token),
.actor = std::move(actor),
.scope = parent_info->scope,
.optional_oauth_params = parent_info->optional_oauth_params,
};
ICEBERG_ASSIGN_OR_RAISE(auto response,
ExchangeToken(*shared_client_, *parent, {}, request));
ICEBERG_ASSIGN_OR_RAISE(auto config,
ChildConfig(*parent_info, parent_info->credential));
return MakeSession(response, config, /*keep_refreshed=*/false);
}

/// Cached token from InitSession
std::optional<OAuthTokenResponse> init_token_response_;
HttpClient* shared_client_ = nullptr;
};

Result<std::unique_ptr<AuthManager>> MakeOAuth2Manager(
Expand Down
38 changes: 25 additions & 13 deletions src/iceberg/catalog/rest/auth/auth_properties.cc
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
#include <utility>

#include "iceberg/catalog/rest/catalog_properties.h"
#include "iceberg/catalog/rest/rest_util.h"

namespace iceberg::rest::auth {

Expand All @@ -35,6 +36,28 @@ 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() || endpoint_it->second.empty()

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Java defaults only when oauth2-server-uri is absent. This also defaults an explicitly empty value; please preserve that distinction or reject empty explicitly.

? AuthProperties::kOAuth2ServerUri.value()
: endpoint_it->second;

if (endpoint.starts_with("http://") || endpoint.starts_with("https://")) {
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 + "/" + std::string(TrimTrailingSlash(endpoint));

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Java preserves the relative endpoint suffix. Trimming oauth/token/ here changes the resolved URI; please keep the trailing slash and update the test.

}

} // namespace

std::unordered_map<std::string, std::string> AuthProperties::optional_oauth_params()
Expand All @@ -61,19 +84,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_.

Expand Down
14 changes: 14 additions & 0 deletions src/iceberg/catalog/rest/auth/auth_properties.h
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
55 changes: 31 additions & 24 deletions src/iceberg/catalog/rest/auth/auth_session.cc
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,18 @@ class OAuth2AuthSession : public AuthSession,
return request;
}

std::optional<OAuth2SessionInfo> OAuth2Info() const override {
std::shared_lock lock(mutex_);
return OAuth2SessionInfo{
.token = token_,
.issued_token_type = issued_token_type_,
.credential = Credential(config_),
.scope = config_.scope,
.oauth2_server_uri = config_.token_endpoint,
.optional_oauth_params = config_.optional_oauth_params,
};
}

Status Close() override { return CloseImpl(); }

~OAuth2AuthSession() override { std::ignore = CloseImpl(); }
Expand All @@ -107,12 +119,15 @@ class OAuth2AuthSession : 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()] =
config.client_id.empty() ? config.client_secret
: config.client_id + ":" + config.client_secret;
properties[AuthProperties::kCredential.key()] = Credential(config);
properties[AuthProperties::kScope.key()] = config.scope;
properties[AuthProperties::kOAuth2ServerUri.key()] = config.token_endpoint;

Expand Down Expand Up @@ -141,11 +156,14 @@ class OAuth2AuthSession : public AuthSession,
OAuth2AuthSession& session_;
};

void SetInitialToken(const OAuthTokenResponse& token_response) {
void UpdateTokenState(const OAuthTokenResponse& token_response) {
token_ = token_response.access_token;
headers_ = {{std::string(kAuthorizationHeader), std::string(kBearerPrefix) + token_}};
issued_token_type_ = token_response.issued_token_type.empty()
? AuthProperties::kAccessTokenType
: token_response.issued_token_type;
headers_ = AuthHeaders(token_);

// Determine expiration time
expires_at_ = std::chrono::steady_clock::time_point{};
if (token_response.expires_in_secs.has_value()) {
expires_at_ = std::chrono::steady_clock::now() +

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

expires_in is measured from the token request, but the session starts the clock later. A slow config request can consume part of the token lifetime; please carry the init fetch start time into the session.

std::chrono::seconds(*token_response.expires_in_secs);
Expand All @@ -157,6 +175,10 @@ class OAuth2AuthSession : public AuthSession,
std::chrono::system_clock::time_point(std::chrono::milliseconds(*exp_ms));
expires_at_ = now_steady + (exp_sys - now_sys);
}
}

void SetInitialToken(const OAuthTokenResponse& token_response) {
UpdateTokenState(token_response);

if (config_.keep_refreshed &&
expires_at_ != std::chrono::steady_clock::time_point{}) {
Expand Down Expand Up @@ -184,23 +206,7 @@ class OAuth2AuthSession : public AuthSession,
auto& response = result.value();
{
std::unique_lock lock(mutex_);
token_ = response.access_token;
headers_ = {
{std::string(kAuthorizationHeader), std::string(kBearerPrefix) + token_}};

// Reset before deriving new expiry
expires_at_ = std::chrono::steady_clock::time_point{};

if (response.expires_in_secs.has_value()) {
expires_at_ = std::chrono::steady_clock::now() +
std::chrono::seconds(*response.expires_in_secs);
} else if (auto exp_ms = ExpiresAtMillis(token_); exp_ms.has_value()) {
auto now_sys = std::chrono::system_clock::now();
auto now_steady = std::chrono::steady_clock::now();
auto exp_sys =
std::chrono::system_clock::time_point(std::chrono::milliseconds(*exp_ms));
expires_at_ = now_steady + (exp_sys - now_sys);
}
UpdateTokenState(response);
}
// Note: ScheduleRefresh must be called outside the lock.
ScheduleRefresh();
Expand Down Expand Up @@ -262,8 +268,9 @@ class OAuth2AuthSession : public AuthSession,
return std::max(wait_time, std::chrono::milliseconds(10));
}

mutable std::shared_mutex mutex_; // protects token_, headers_, expires_at_
mutable std::shared_mutex mutex_; // protects token state, headers, and expiration
std::string token_;
std::string issued_token_type_;
std::unordered_map<std::string, std::string> headers_;
std::chrono::steady_clock::time_point expires_at_{};

Expand Down
14 changes: 14 additions & 0 deletions src/iceberg/catalog/rest/auth/auth_session.h
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
#pragma once

#include <memory>
#include <optional>
#include <string>
#include <unordered_map>

Expand All @@ -33,6 +34,16 @@

namespace iceberg::rest::auth {

/// \brief OAuth2 metadata used to derive child authentication sessions.
struct ICEBERG_REST_EXPORT OAuth2SessionInfo {
std::string token;
std::string issued_token_type;
std::string credential;
std::string scope;
std::string oauth2_server_uri;
std::unordered_map<std::string, std::string> optional_oauth_params;
};

/// \brief An authentication session that can authenticate outgoing HTTP requests.
class ICEBERG_REST_EXPORT AuthSession {
public:
Expand All @@ -54,6 +65,9 @@ class ICEBERG_REST_EXPORT AuthSession {
/// - RestError: HTTP errors from authentication service
virtual Result<HttpRequest> Authenticate(HttpRequest request) = 0;

/// \brief Return OAuth2 metadata when this is an OAuth2 session.
virtual std::optional<OAuth2SessionInfo> OAuth2Info() const { return std::nullopt; }

/// \brief Close the session and release any resources.
///
/// This method is called when the session is no longer needed. For stateful
Expand Down
Loading
Loading