-
Notifications
You must be signed in to change notification settings - Fork 121
feat(rest): support OAuth token exchange sessions #867
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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,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() | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Java defaults only when |
||
| ? 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)); | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Java preserves the relative endpoint suffix. Trimming |
||
| } | ||
|
|
||
| } // namespace | ||
|
|
||
| std::unordered_map<std::string, std::string> AuthProperties::optional_oauth_params() | ||
|
|
@@ -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_. | ||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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(); } | ||
|
|
@@ -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; | ||
|
|
||
|
|
@@ -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() + | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
| std::chrono::seconds(*token_response.expires_in_secs); | ||
|
|
@@ -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{}) { | ||
|
|
@@ -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(); | ||
|
|
@@ -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_{}; | ||
|
|
||
|
|
||
There was a problem hiding this comment.
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 astd::shared_ptr<HttpClient>through the manager/session API so the client lifetime is explicit.