diff --git a/src/iceberg/catalog/rest/auth/auth_manager.cc b/src/iceberg/catalog/rest/auth/auth_manager.cc index 0fad64eca..d0dd92753 100644 --- a/src/iceberg/catalog/rest/auth/auth_manager.cc +++ b/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 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 session_cache_; std::once_flag no_refresh_executor_warning_; diff --git a/src/iceberg/catalog/rest/auth/auth_properties.h b/src/iceberg/catalog/rest/auth/auth_properties.h index 0cf9be5e4..d65f2fc6c 100644 --- a/src/iceberg/catalog/rest/auth/auth_properties.h +++ b/src/iceberg/catalog/rest/auth/auth_properties.h @@ -79,8 +79,8 @@ class ICEBERG_REST_EXPORT AuthProperties : public ConfigBase { "v1/oauth/tokens"}; inline static Entry 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 kExchangeEnabled{"token-exchange-enabled", true}; inline static Entry kAudience{"audience", ""}; diff --git a/src/iceberg/catalog/rest/auth/auth_session.cc b/src/iceberg/catalog/rest/auth/auth_session.cc index bd74b5bb6..dd8f0d70c 100644 --- a/src/iceberg/catalog/rest/auth/auth_session.cc +++ b/src/iceberg/catalog/rest/auth/auth_session.cc @@ -64,28 +64,30 @@ Result> AuthSession::MakeOAuth2( const std::unordered_map& optional_oauth_params, std::shared_ptr 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> 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& optional_oauth_params, std::shared_ptr client, std::optional 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( diff --git a/src/iceberg/catalog/rest/auth/auth_session.h b/src/iceberg/catalog/rest/auth/auth_session.h index 393418e2a..b5dddcebd 100644 --- a/src/iceberg/catalog/rest/auth/auth_session.h +++ b/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(). diff --git a/src/iceberg/catalog/rest/auth/auth_session_internal.h b/src/iceberg/catalog/rest/auth/auth_session_internal.h index a4454aafc..a912a264d 100644 --- a/src/iceberg/catalog/rest/auth/auth_session_internal.h +++ b/src/iceberg/catalog/rest/auth/auth_session_internal.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -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 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 now = + std::chrono::steady_clock::now; }; static Result> 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 MakeRefreshProperties(const Config& config) { std::unordered_map 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 - token_request_started_at = std::nullopt) { + std::optional 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 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,21 +254,60 @@ class OAuth2Session final : public AuthSession, auto result = RetryRunner>({.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> RefreshToken() { + class ClientCredentialsSession final : public AuthSession { + public: + Result Authenticate(HttpRequest request) override { + // HttpClient merges its defaults before invoking the session. + request.headers.erase("Authorization"); + return request; + } + }; + + Result 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{{"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> RefreshToken(bool recover = false) { { std::lock_guard lock(refresh_mutex_); if (closed_) { @@ -276,21 +315,34 @@ class OAuth2Session final : public AuthSession, } } 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> 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& optional_oauth_params, std::shared_ptr client, std::optional token_request_started_at, - Executor* refresh_executor); + Executor* refresh_executor, bool exchange_enabled); } // namespace iceberg::rest::auth::internal diff --git a/src/iceberg/test/auth_manager_test.cc b/src/iceberg/test/auth_manager_test.cc index 076df633e..5d4746aad 100644 --- a/src/iceberg/test/auth_manager_test.cc +++ b/src/iceberg/test/auth_manager_test.cc @@ -28,6 +28,7 @@ #include #include #include +#include #include #include @@ -682,10 +683,22 @@ namespace { class RefreshHttpClient : public HttpClient { public: - Result PostForm(const std::string&, - const std::unordered_map&, - const std::unordered_map&, - const ErrorHandler&, AuthSession&) override { + struct Request { + std::string path; + std::unordered_map form; + HttpHeaders headers; + }; + + Result PostForm( + const std::string& path, const std::unordered_map& form, + const std::unordered_map& headers, const ErrorHandler&, + AuthSession& session) override { + ICEBERG_ASSIGN_OR_RAISE( + auto authenticated, + session.Authenticate({.method = HttpMethod::kPost, + .url = path, + .headers = {headers.begin(), headers.end()}})); + recorded_requests.push_back({path, form, std::move(authenticated.headers)}); ++requests; return response(); } @@ -693,6 +706,7 @@ class RefreshHttpClient : public HttpClient { return IOError("refresh failed"); }; std::atomic requests{0}; + std::vector recorded_requests; }; using test::MockExecutor; @@ -736,6 +750,760 @@ TEST(OAuth2SessionTest, RefreshesAndReschedulesThroughPublicFactory) { ICEBERG_UNWRAP_OR_FAIL(auto request, session->Authenticate({})); EXPECT_EQ(request.headers.at("Authorization"), "Bearer second"); EXPECT_EQ(client->requests, 2); + for (const auto& request : client->recorded_requests) { + EXPECT_EQ(request.form.at("grant_type"), "client_credentials"); + EXPECT_FALSE(request.headers.contains("Authorization")); + } +} + +class OAuth2ManagerRefreshTest : public ::testing::TestWithParam { + protected: + void SetUp() override { + if (!GetParam().empty()) { + properties_[AuthProperties::kExchangeEnabled.key()] = GetParam(); + } + ICEBERG_UNWRAP_OR_FAIL(manager_, AuthManagers::Load("test-catalog", properties_)); + manager_->SetAuthRefreshExecutor(&executor_); + client_->response = [this]() -> Result { + return HttpResponse( + 200, client_->requests == 1 + ? R"({"access_token":"initial","token_type":"bearer","expires_in":1})" + : R"({"access_token":"fresh","token_type":"bearer"})"); + }; + } + + void TearDown() override { + if (session_) { + EXPECT_THAT(session_->Close(), IsOk()); + } + if (manager_) { + EXPECT_THAT(manager_->Close(), IsOk()); + } + } + + void ExpectRefresh() { + EXPECT_CALL(executor_, Submit).WillOnce([this](ExecutorTask task) -> Status { + queued_.set_value(std::move(task)); + return {}; + }); + } + + void RunRefresh() { + auto task = queued_.get_future(); + ASSERT_EQ(task.wait_for(std::chrono::seconds(5)), std::future_status::ready); + std::move(task.get())(); + } + + void CheckRequests(const std::string& client_id, const std::string& client_secret, + const std::string& parent_token = "") { + ASSERT_EQ(client_->recorded_requests.size(), 2); + std::unordered_map credential_form{ + {"grant_type", "client_credentials"}, {"client_id", client_id}, + {"client_secret", client_secret}, {"scope", "catalog"}, + {"audience", "catalog-service"}, {"resource", "https://catalog.example"}, + }; + const auto& initial = client_->recorded_requests[0]; + EXPECT_THAT(initial.form, ::testing::ContainerEq(credential_form)); + if (parent_token.empty()) { + EXPECT_TRUE(initial.headers.empty()); + } else { + EXPECT_EQ(initial.headers.at("Authorization"), "Bearer " + parent_token); + } + + const auto& refresh = client_->recorded_requests[1]; + if (GetParam() == "false") { + EXPECT_THAT(refresh.form, ::testing::ContainerEq(credential_form)); + EXPECT_TRUE(refresh.headers.empty()); + } else { + const std::unordered_map exchange_form{ + {"grant_type", "urn:ietf:params:oauth:grant-type:token-exchange"}, + {"subject_token", "initial"}, + {"subject_token_type", AuthProperties::kAccessTokenType}, + {"scope", "catalog"}, + {"audience", "catalog-service"}, + {"resource", "https://catalog.example"}, + }; + EXPECT_THAT(refresh.form, ::testing::ContainerEq(exchange_form)); + EXPECT_THAT(refresh.headers, ::testing::ElementsAre(::testing::Pair( + "Authorization", "Bearer initial"))); + } + for (const auto& request : client_->recorded_requests) { + EXPECT_EQ(request.path, properties_.at(AuthProperties::kOAuth2ServerUri.key())); + } + ICEBERG_UNWRAP_OR_FAIL(auto request, session_->Authenticate({})); + EXPECT_EQ(request.headers.at("Authorization"), "Bearer fresh"); + EXPECT_EQ(session_->OAuth2Info()->credential, client_id + ":" + client_secret); + } + + MockExecutor executor_; + std::promise queued_; + std::shared_ptr client_ = std::make_shared(); + std::unordered_map properties_{ + {AuthProperties::kAuthType, "oauth2"}, + {AuthProperties::kCredential.key(), "catalog:catalog-secret"}, + {AuthProperties::kOAuth2ServerUri.key(), "https://auth.example/token"}, + {AuthProperties::kScope.key(), "catalog"}, + {AuthProperties::kAudience.key(), "catalog-service"}, + {AuthProperties::kResource.key(), "https://catalog.example"}, + }; + std::unique_ptr manager_; + std::shared_ptr session_; +}; + +TEST_P(OAuth2ManagerRefreshTest, CatalogRefreshUsesConfiguredGrant) { + ExpectRefresh(); + ICEBERG_UNWRAP_OR_FAIL(session_, manager_->CatalogSession(client_, properties_)); + ASSERT_NO_FATAL_FAILURE(RunRefresh()); + CheckRequests("catalog", "catalog-secret"); +} + +TEST_P(OAuth2ManagerRefreshTest, CredentialContextInheritsCatalogRefreshGrant) { + ExpectRefresh(); + properties_[AuthProperties::kToken.key()] = "parent-token"; + ICEBERG_UNWRAP_OR_FAIL(auto parent, manager_->CatalogSession(client_, properties_)); + SessionContext context{ + .session_id = "tenant-a", + .credentials = {{AuthProperties::kCredential.key(), "child:child-secret"}}, + }; + ICEBERG_UNWRAP_OR_FAIL(session_, manager_->ContextualSession(context, parent)); + ASSERT_NO_FATAL_FAILURE(RunRefresh()); + CheckRequests("child", "child-secret", "parent-token"); + ICEBERG_UNWRAP_OR_FAIL(auto cached, manager_->ContextualSession(context, parent)); + EXPECT_EQ(cached, session_); +} + +TEST_P(OAuth2ManagerRefreshTest, ExpiredCredentialContextRecoversWithOwnRawCredential) { + ExpectRefresh(); + properties_[AuthProperties::kToken.key()] = "parent-token"; + ICEBERG_UNWRAP_OR_FAIL(auto parent, manager_->CatalogSession(client_, properties_)); + client_->response = [this]() -> Result { + return HttpResponse( + 200, + client_->requests == 1 + ? R"({"access_token":"child-expired","token_type":"bearer","issued_token_type":"urn:ietf:params:oauth:token-type:jwt","expires_in":0})" + : R"({"access_token":"child-fresh","token_type":"bearer"})"); + }; + SessionContext context{ + .session_id = "tenant-recovery", + .credentials = {{AuthProperties::kCredential.key(), ":child-secret"}}, + }; + ICEBERG_UNWRAP_OR_FAIL(session_, manager_->ContextualSession(context, parent)); + EXPECT_EQ(session_->OAuth2Info()->credential, ":child-secret"); + ASSERT_NO_FATAL_FAILURE(RunRefresh()); + + ASSERT_EQ(client_->recorded_requests.size(), 2); + const auto& initial = client_->recorded_requests[0]; + const std::unordered_map credential_form{ + {"grant_type", "client_credentials"}, + {"client_secret", "child-secret"}, + {"scope", "catalog"}, + {"audience", "catalog-service"}, + {"resource", "https://catalog.example"}, + }; + EXPECT_THAT(initial.form, ::testing::ContainerEq(credential_form)); + EXPECT_EQ(initial.headers.at("Authorization"), "Bearer parent-token"); + const auto& recovery = client_->recorded_requests[1]; + if (GetParam() == "false") { + EXPECT_THAT(recovery.form, ::testing::ContainerEq(credential_form)); + EXPECT_TRUE(recovery.headers.empty()); + } else { + const std::unordered_map exchange_form{ + {"grant_type", "urn:ietf:params:oauth:grant-type:token-exchange"}, + {"subject_token", "child-expired"}, + {"subject_token_type", AuthProperties::kJwtTokenType}, + {"scope", "catalog"}, + {"audience", "catalog-service"}, + {"resource", "https://catalog.example"}, + }; + EXPECT_THAT(recovery.form, ::testing::ContainerEq(exchange_form)); + EXPECT_THAT(recovery.headers, ::testing::ElementsAre(::testing::Pair( + "Authorization", "Basic OmNoaWxkLXNlY3JldA=="))); + } + EXPECT_EQ(recovery.path, properties_.at(AuthProperties::kOAuth2ServerUri.key())); + EXPECT_EQ(session_->OAuth2Info()->credential, ":child-secret"); + ICEBERG_UNWRAP_OR_FAIL(auto request, session_->Authenticate({})); + EXPECT_EQ(request.headers.at("Authorization"), "Bearer child-fresh"); + ICEBERG_UNWRAP_OR_FAIL(auto parent_request, parent->Authenticate({})); + EXPECT_EQ(parent_request.headers.at("Authorization"), "Bearer parent-token"); + ICEBERG_UNWRAP_OR_FAIL(auto cached, manager_->ContextualSession(context, parent)); + EXPECT_EQ(cached, session_); + EXPECT_EQ(client_->requests, 2); +} + +TEST_P(OAuth2ManagerRefreshTest, RefreshDisabledDoesNotScheduleCatalogOrContext) { + EXPECT_CALL(executor_, Submit).Times(0); + properties_[AuthProperties::kKeepRefreshed.key()] = "false"; + client_->response = []() -> Result { + return HttpResponse( + 200, R"({"access_token":"initial","token_type":"bearer","expires_in":0})"); + }; + ICEBERG_UNWRAP_OR_FAIL(session_, manager_->CatalogSession(client_, properties_)); + ICEBERG_UNWRAP_OR_FAIL( + auto child, + manager_->ContextualSession( + {.session_id = "tenant-a", + .credentials = {{AuthProperties::kCredential.key(), "child:child-secret"}}}, + session_)); + ASSERT_TRUE(test::WaitForSharedTimer()); + EXPECT_EQ(client_->requests, 2); + ICEBERG_UNWRAP_OR_FAIL(auto request, child->Authenticate({})); + EXPECT_EQ(request.headers.at("Authorization"), "Bearer initial"); +} + +INSTANTIATE_TEST_SUITE_P(ExchangeModes, OAuth2ManagerRefreshTest, + ::testing::Values(std::string(""), std::string("true"), + std::string("false")), + [](const ::testing::TestParamInfo& info) { + return info.param.empty() ? "Default" + : info.param == "true" ? "Enabled" + : "Disabled"; + }); + +class OAuth2ExchangeRefreshTest : public ::testing::Test { + protected: + void SetUp() override { + config_.now = [this] { return now_; }; + } + void ExpectRefreshes(int count) { + if (count == 0) { + EXPECT_CALL(executor_, Submit).Times(0); + return; + } + EXPECT_CALL(executor_, Submit) + .Times(count) + .WillRepeatedly([this](ExecutorTask task) -> Status { + queued_.at(submitted_++).set_value(std::move(task)); + return {}; + }); + } + + Status Start() { + ICEBERG_ASSIGN_OR_RAISE( + session_, internal::OAuth2Session::Make(token_, config_, client_, std::nullopt)); + return {}; + } + + void RunRefresh(size_t index = 0) { + auto task = queued_.at(index).get_future(); + ASSERT_EQ(task.wait_for(std::chrono::seconds(5)), std::future_status::ready); + std::move(task.get())(); + } + + MockExecutor executor_; + std::array, 2> queued_; + int submitted_ = 0; + std::shared_ptr client_ = std::make_shared(); + std::chrono::steady_clock::time_point now_ = std::chrono::steady_clock::now(); + internal::OAuth2Session::Config config_{ + .token_endpoint = "https://auth.example/token", + .scope = "catalog", + .optional_oauth_params = {{"audience", "catalog-service"}, + {"resource", "https://catalog.example"}}, + .exchange_enabled = true, + .refresh_executor = &executor_, + }; + OAuthTokenResponse token_{ + .access_token = "initial", + .token_type = "bearer", + .expires_in_secs = 1, + .issued_token_type = AuthProperties::kAccessTokenType, + }; + // Destroy the session before its borrowed executor and task promises. + std::shared_ptr session_; +}; + +TEST_F(OAuth2ExchangeRefreshTest, ExchangesCurrentTokenAndTypeOnEachRefresh) { + ExpectRefreshes(2); + // Even when a credential is available, exchange must keep the current subject. + config_.credential = "client:secret"; + client_->response = [&]() -> Result { + if (client_->requests == 1) { + return HttpResponse( + 200, nlohmann::json({{"access_token", "first"}, + {"token_type", "bearer"}, + {"issued_token_type", AuthProperties::kJwtTokenType}, + {"expires_in", 1}}) + .dump()); + } + return HttpResponse(200, R"({"access_token":"second","token_type":"bearer"})"); + }; + ASSERT_THAT(Start(), IsOk()); + + ASSERT_NO_FATAL_FAILURE(RunRefresh()); + ICEBERG_UNWRAP_OR_FAIL(auto first_request, session_->Authenticate({})); + EXPECT_EQ(first_request.headers.at("Authorization"), "Bearer first"); + ASSERT_TRUE(session_->ExpiresAt().has_value()); + EXPECT_TRUE(session_->IsRefreshable()); + + ASSERT_NO_FATAL_FAILURE(RunRefresh(1)); + ASSERT_EQ(client_->recorded_requests.size(), 2); + for (size_t i = 0; i < client_->recorded_requests.size(); ++i) { + const auto& request = client_->recorded_requests[i]; + const std::string subject = i == 0 ? "initial" : "first"; + const std::unordered_map expected_form{ + {"grant_type", "urn:ietf:params:oauth:grant-type:token-exchange"}, + {"subject_token", subject}, + {"subject_token_type", + i == 0 ? AuthProperties::kAccessTokenType : AuthProperties::kJwtTokenType}, + {"scope", "catalog"}, + {"audience", "catalog-service"}, + {"resource", "https://catalog.example"}, + }; + EXPECT_EQ(request.path, config_.token_endpoint); + EXPECT_THAT(request.form, ::testing::ContainerEq(expected_form)); + EXPECT_EQ(request.headers.at("Authorization"), "Bearer " + subject); + } + ICEBERG_UNWRAP_OR_FAIL(auto second_request, session_->Authenticate({})); + EXPECT_EQ(second_request.headers.at("Authorization"), "Bearer second"); + EXPECT_EQ(session_->OAuth2Info()->issued_token_type, AuthProperties::kAccessTokenType); + EXPECT_FALSE(session_->ExpiresAt().has_value()); + EXPECT_FALSE(session_->IsRefreshable()); + EXPECT_THAT(session_->Close(), IsOk()); +} + +TEST_F(OAuth2ExchangeRefreshTest, RetriesExchangeWithoutChangingIdentity) { + ExpectRefreshes(1); + client_->response = [&]() -> Result { + if (client_->requests == 1) { + throw std::runtime_error("HTTP library failure"); + } + if (client_->requests == 2) { + return IOError("temporarily unavailable"); + } + return HttpResponse(200, R"({"access_token":"fresh","token_type":"bearer"})"); + }; + ASSERT_THAT(Start(), IsOk()); + test::FakeRetryEnvironment fake_retry; + ScopedRetryTestHooks hooks(fake_retry.hooks()); + ASSERT_NO_FATAL_FAILURE(RunRefresh()); + + EXPECT_EQ(client_->requests, 3); + EXPECT_THAT(fake_retry.observed_base_delays_ms(), ::testing::ElementsAre(100, 200)); + for (const auto& request : client_->recorded_requests) { + EXPECT_EQ(request.form.at("subject_token"), "initial"); + EXPECT_EQ(request.headers.at("Authorization"), "Bearer initial"); + EXPECT_FALSE(request.form.contains("client_secret")); + } + EXPECT_EQ(session_->OAuth2Info()->token, "fresh"); + EXPECT_THAT(session_->Close(), IsOk()); +} + +class OAuth2ExpiredExchangeRefreshTest : public OAuth2ExchangeRefreshTest, + public ::testing::WithParamInterface { + protected: + HttpResponse ExpiredResponse() const { + return HttpResponse( + 200, nlohmann::json( + {{"access_token", GetParam() ? MakeJwt(R"({"exp":0})") : "expired"}, + {"token_type", "bearer"}, + {"issued_token_type", GetParam() ? AuthProperties::kJwtTokenType + : AuthProperties::kAccessTokenType}, + {"expires_in", GetParam() ? 3600 : 0}}) + .dump()); + } + + void ExpectInitialState( + std::optional expiration) { + EXPECT_EQ(session_->OAuth2Info()->token, token_.access_token); + EXPECT_EQ(session_->OAuth2Info()->issued_token_type, token_.issued_token_type); + EXPECT_EQ(session_->ExpiresAt(), expiration); + ICEBERG_UNWRAP_OR_FAIL(auto request, session_->Authenticate({})); + EXPECT_EQ(request.headers.at("Authorization"), "Bearer " + token_.access_token); + } + + void ExpectInitialExchangeRequests(bool final_recovery = false) { + for (size_t i = 0; i < client_->recorded_requests.size(); ++i) { + const auto& request = client_->recorded_requests[i]; + EXPECT_EQ(request.form.at("subject_token"), token_.access_token); + EXPECT_EQ(request.form.at("subject_token_type"), token_.issued_token_type); + EXPECT_EQ(request.headers.at("Authorization"), + final_recovery && i == 6 ? "Basic Y2xpZW50OnNlY3JldA==" + : "Bearer " + token_.access_token); + } + } +}; + +TEST_P(OAuth2ExpiredExchangeRefreshTest, RetriesWithPreviousTokenUntilSuccess) { + ExpectRefreshes(1); + ASSERT_THAT(Start(), IsOk()); + auto initial_expiration = session_->ExpiresAt(); + client_->response = [&]() -> Result { + if (client_->requests == 1) { + return ExpiredResponse(); + } + ExpectInitialState(initial_expiration); + const auto& request = client_->recorded_requests.back(); + if (request.form.at("subject_token") != token_.access_token || + request.headers.at("Authorization") != "Bearer " + token_.access_token) { + return IOError("authorization server rejects the expired token"); + } + return HttpResponse(200, R"({"access_token":"fresh","token_type":"bearer"})"); + }; + test::FakeRetryEnvironment fake_retry; + ScopedRetryTestHooks hooks(fake_retry.hooks()); + ASSERT_NO_FATAL_FAILURE(RunRefresh()); + + EXPECT_EQ(client_->requests, 2); + EXPECT_THAT(fake_retry.observed_base_delays_ms(), ::testing::ElementsAre(100)); + ExpectInitialExchangeRequests(); + ICEBERG_UNWRAP_OR_FAIL(auto request, session_->Authenticate({})); + EXPECT_EQ(request.headers.at("Authorization"), "Bearer fresh"); + EXPECT_EQ(session_->OAuth2Info()->issued_token_type, AuthProperties::kAccessTokenType); + EXPECT_FALSE(session_->ExpiresAt().has_value()); + EXPECT_FALSE(session_->IsRefreshable()); + EXPECT_THAT(session_->Close(), IsOk()); +} + +TEST_P(OAuth2ExpiredExchangeRefreshTest, KeepsPreviousTokenAfterRetriesExhausted) { + ExpectRefreshes(1); + config_.credential = "client:secret"; + ASSERT_THAT(Start(), IsOk()); + auto initial_expiration = session_->ExpiresAt(); + client_->response = [&]() -> Result { return ExpiredResponse(); }; + test::FakeRetryEnvironment fake_retry; + ScopedRetryTestHooks hooks(fake_retry.hooks()); + ASSERT_NO_FATAL_FAILURE(RunRefresh()); + + EXPECT_EQ(client_->requests, 7); + ExpectInitialExchangeRequests(true); + ExpectInitialState(initial_expiration); + EXPECT_FALSE(session_->IsRefreshable()); + EXPECT_THAT(session_->Close(), IsOk()); +} + +INSTANTIATE_TEST_SUITE_P(ExpirationSources, OAuth2ExpiredExchangeRefreshTest, + ::testing::Bool(), + [](const ::testing::TestParamInfo& info) { + return info.param ? "Jwt" : "ExpiresIn"; + }); + +TEST_F(OAuth2ExchangeRefreshTest, StopsAfterExchangeRetryBudget) { + ExpectRefreshes(1); + config_.credential = "client:secret"; + ASSERT_THAT(Start(), IsOk()); + test::FakeRetryEnvironment fake_retry; + ScopedRetryTestHooks hooks(fake_retry.hooks()); + ASSERT_NO_FATAL_FAILURE(RunRefresh()); + + EXPECT_EQ(client_->requests, 7); + for (size_t i = 0; i < client_->recorded_requests.size(); ++i) { + const auto& request = client_->recorded_requests[i]; + EXPECT_EQ(request.headers.at("Authorization"), + i == 6 ? "Basic Y2xpZW50OnNlY3JldA==" : "Bearer initial"); + EXPECT_EQ(request.form.at("grant_type"), + "urn:ietf:params:oauth:grant-type:token-exchange"); + EXPECT_EQ(request.form.at("subject_token"), "initial"); + EXPECT_FALSE(request.form.contains("client_secret")); + } + EXPECT_EQ(session_->OAuth2Info()->token, "initial"); + EXPECT_FALSE(session_->IsRefreshable()); + EXPECT_THAT(session_->Close(), IsOk()); +} + +class OAuth2ExpiredCredentialTest + : public OAuth2ExchangeRefreshTest, + public ::testing::WithParamInterface> {}; + +TEST_P(OAuth2ExpiredCredentialTest, UsesOriginalCredentialForExpiredToken) { + ExpectRefreshes(1); + token_.expires_in_secs = 0; + config_.credential = GetParam().first; + client_->response = []() -> Result { + return HttpResponse(200, R"({"access_token":"fresh","token_type":"bearer"})"); + }; + ASSERT_THAT(Start(), IsOk()); + ASSERT_NO_FATAL_FAILURE(RunRefresh()); + + ASSERT_EQ(client_->recorded_requests.size(), 1); + const auto& request = client_->recorded_requests.front(); + EXPECT_EQ(request.headers.at("Authorization"), "Basic " + GetParam().second); + EXPECT_EQ(request.form.at("subject_token"), "initial"); + EXPECT_EQ(request.form.at("subject_token_type"), AuthProperties::kAccessTokenType); + EXPECT_FALSE(request.form.contains("actor_token")); + EXPECT_FALSE(request.form.contains("client_secret")); + EXPECT_EQ(session_->OAuth2Info()->credential, GetParam().first); + EXPECT_EQ(session_->OAuth2Info()->token, "fresh"); + EXPECT_FALSE(session_->IsRefreshable()); +} + +INSTANTIATE_TEST_SUITE_P( + CredentialForms, OAuth2ExpiredCredentialTest, + ::testing::Values(std::pair{"id:secret", "aWQ6c2VjcmV0"}, + std::pair{"secret", "c2VjcmV0"}, + std::pair{":secret", "OnNlY3JldA=="}, + std::pair{"id:secret:tail", + "aWQ6c2VjcmV0OnRhaWw="})); + +TEST_F(OAuth2ExchangeRefreshTest, SwitchesToBasicWhenTokenExpiresDuringRetries) { + ExpectRefreshes(1); + config_.credential = "client:secret"; + client_->response = [&]() -> Result { + if (client_->requests == 1) { + now_ += std::chrono::seconds(2); + return IOError("Bearer refresh failed"); + } + return HttpResponse(200, R"({"access_token":"fresh","token_type":"bearer"})"); + }; + ASSERT_THAT(Start(), IsOk()); + test::FakeRetryEnvironment fake_retry; + ScopedRetryTestHooks hooks(fake_retry.hooks()); + ASSERT_NO_FATAL_FAILURE(RunRefresh()); + + ASSERT_EQ(client_->recorded_requests.size(), 2); + EXPECT_EQ(client_->recorded_requests[0].headers.at("Authorization"), "Bearer initial"); + EXPECT_EQ(client_->recorded_requests[1].headers.at("Authorization"), + "Basic Y2xpZW50OnNlY3JldA=="); + for (const auto& request : client_->recorded_requests) { + EXPECT_EQ(request.form.at("subject_token"), "initial"); + EXPECT_EQ(request.form.at("subject_token_type"), AuthProperties::kAccessTokenType); + } + EXPECT_EQ(session_->OAuth2Info()->token, "fresh"); +} + +TEST_F(OAuth2ExchangeRefreshTest, FinalRecoveryUsesOwnRequestTimeAndReschedules) { + ExpectRefreshes(2); + config_.credential = "client:secret"; + client_->response = [&]() -> Result { + if (client_->requests <= 6) { + if (client_->requests == 6) { + now_ += std::chrono::milliseconds(250); + } + return IOError("Bearer refresh failed"); + } + if (client_->requests == 7) { + return HttpResponse( + 200, R"({"access_token":"recovered","token_type":"bearer","expires_in":1})"); + } + return HttpResponse(200, R"({"access_token":"next","token_type":"bearer"})"); + }; + ASSERT_THAT(Start(), IsOk()); + test::FakeRetryEnvironment fake_retry; + ScopedRetryTestHooks hooks(fake_retry.hooks()); + ASSERT_NO_FATAL_FAILURE(RunRefresh()); + + ASSERT_EQ(client_->recorded_requests.size(), 7); + for (size_t i = 0; i < 7; ++i) { + const auto& request = client_->recorded_requests[i]; + EXPECT_EQ(request.headers.at("Authorization"), + i == 6 ? "Basic Y2xpZW50OnNlY3JldA==" : "Bearer initial"); + EXPECT_EQ(request.form.at("subject_token"), "initial"); + EXPECT_FALSE(request.form.contains("actor_token")); + } + EXPECT_THAT(fake_retry.observed_base_delays_ms(), + ::testing::ElementsAre(100, 200, 400, 800, 1600)); + EXPECT_EQ(session_->ExpiresAt(), now_ + std::chrono::seconds(1)); + EXPECT_EQ(session_->OAuth2Info()->token, "recovered"); + EXPECT_TRUE(session_->IsRefreshable()); + ASSERT_NO_FATAL_FAILURE(RunRefresh(1)); + EXPECT_EQ(client_->recorded_requests.back().headers.at("Authorization"), + "Bearer recovered"); + EXPECT_EQ(client_->recorded_requests.back().form.at("subject_token"), "recovered"); + EXPECT_EQ(session_->OAuth2Info()->token, "next"); + EXPECT_FALSE(session_->IsRefreshable()); +} + +TEST_F(OAuth2ExchangeRefreshTest, ExpiredJwtSharesSevenAttemptBudget) { + ExpectRefreshes(1); + config_.credential = "client:secret"; + token_.access_token = MakeJwt(R"({"exp":0})"); + token_.issued_token_type = AuthProperties::kJwtTokenType; + token_.expires_in_secs = 3600; + ASSERT_THAT(Start(), IsOk()); + const auto expiration = session_->ExpiresAt(); + test::FakeRetryEnvironment fake_retry; + ScopedRetryTestHooks hooks(fake_retry.hooks()); + ASSERT_NO_FATAL_FAILURE(RunRefresh()); + + ASSERT_EQ(client_->recorded_requests.size(), 7); + for (const auto& request : client_->recorded_requests) { + EXPECT_EQ(request.headers.at("Authorization"), "Basic Y2xpZW50OnNlY3JldA=="); + EXPECT_EQ(request.form.at("subject_token"), token_.access_token); + EXPECT_EQ(request.form.at("subject_token_type"), AuthProperties::kJwtTokenType); + } + EXPECT_EQ(session_->OAuth2Info()->token, token_.access_token); + EXPECT_EQ(session_->ExpiresAt(), expiration); + EXPECT_FALSE(session_->IsRefreshable()); +} + +class OAuth2MissingCredentialTest : public OAuth2ExchangeRefreshTest, + public ::testing::WithParamInterface {}; + +TEST_P(OAuth2MissingCredentialTest, DoesNotSendUnrecoverableRequest) { + ExpectRefreshes(1); + config_.exchange_enabled = GetParam(); + token_.expires_in_secs = 0; + ASSERT_THAT(Start(), IsOk()); + test::FakeRetryEnvironment fake_retry; + ScopedRetryTestHooks hooks(fake_retry.hooks()); + ASSERT_NO_FATAL_FAILURE(RunRefresh()); + + EXPECT_EQ(client_->requests, 0); + EXPECT_EQ(session_->OAuth2Info()->token, "initial"); + EXPECT_FALSE(session_->IsRefreshable()); +} + +INSTANTIATE_TEST_SUITE_P(GrantModes, OAuth2MissingCredentialTest, ::testing::Bool()); + +TEST_F(OAuth2ExchangeRefreshTest, MissingCredentialSkipsFinalRecovery) { + ExpectRefreshes(1); + ASSERT_THAT(Start(), IsOk()); + test::FakeRetryEnvironment fake_retry; + ScopedRetryTestHooks hooks(fake_retry.hooks()); + ASSERT_NO_FATAL_FAILURE(RunRefresh()); + + ASSERT_EQ(client_->recorded_requests.size(), 6); + for (const auto& request : client_->recorded_requests) { + EXPECT_EQ(request.headers.at("Authorization"), "Bearer initial"); + } + EXPECT_EQ(session_->OAuth2Info()->token, "initial"); + EXPECT_FALSE(session_->IsRefreshable()); +} + +TEST_F(OAuth2ExchangeRefreshTest, MissingCredentialStopsWhenTokenExpiresDuringRetries) { + ExpectRefreshes(1); + client_->response = [&]() -> Result { + now_ += std::chrono::seconds(2); + return IOError("Bearer refresh failed"); + }; + ASSERT_THAT(Start(), IsOk()); + test::FakeRetryEnvironment fake_retry; + ScopedRetryTestHooks hooks(fake_retry.hooks()); + ASSERT_NO_FATAL_FAILURE(RunRefresh()); + + EXPECT_EQ(client_->requests, 1); + EXPECT_EQ(client_->recorded_requests.front().headers.at("Authorization"), + "Bearer initial"); + EXPECT_EQ(session_->OAuth2Info()->token, "initial"); + EXPECT_FALSE(session_->IsRefreshable()); +} + +TEST_F(OAuth2ExchangeRefreshTest, StopBeforeFinalRecoverySkipsRequest) { + ExpectRefreshes(1); + config_.credential = "client:secret"; + client_->response = [&]() -> Result { + if (client_->requests == 6) { + session_->StopRefreshing(); + } + return IOError("refresh failed"); + }; + ASSERT_THAT(Start(), IsOk()); + test::FakeRetryEnvironment fake_retry; + ScopedRetryTestHooks hooks(fake_retry.hooks()); + ASSERT_NO_FATAL_FAILURE(RunRefresh()); + + EXPECT_EQ(client_->requests, 6); + EXPECT_EQ(session_->OAuth2Info()->token, "initial"); + EXPECT_FALSE(session_->IsRefreshable()); +} + +TEST_F(OAuth2ExchangeRefreshTest, CloseWaitsForFinalRecovery) { + ExpectRefreshes(1); + config_.credential = "client:secret"; + std::promise entered; + std::promise release; + auto entered_future = entered.get_future(); + auto release_future = release.get_future(); + client_->response = [&]() -> Result { + if (client_->requests == 7) { + entered.set_value(); + release_future.wait(); + return HttpResponse( + 200, R"({"access_token":"recovered","token_type":"bearer","expires_in":1})"); + } + return IOError("refresh failed"); + }; + ASSERT_THAT(Start(), IsOk()); + auto task = queued_[0].get_future(); + ASSERT_EQ(task.wait_for(std::chrono::seconds(5)), std::future_status::ready); + auto refresh = std::async(std::launch::async, [&] { + test::FakeRetryEnvironment fake_retry; + ScopedRetryTestHooks hooks(fake_retry.hooks()); + std::move(task.get())(); + }); + auto entered_status = entered_future.wait_for(std::chrono::seconds(5)); + if (entered_status != std::future_status::ready) { + release.set_value(); + } + ASSERT_EQ(entered_status, std::future_status::ready); + auto closing = std::async(std::launch::async, [&] { return session_->Close(); }); + EXPECT_EQ(closing.wait_for(std::chrono::milliseconds(30)), std::future_status::timeout); + release.set_value(); + refresh.get(); + EXPECT_THAT(closing.get(), IsOk()); + + EXPECT_EQ(client_->requests, 7); + EXPECT_EQ(session_->OAuth2Info()->token, "recovered"); + EXPECT_FALSE(session_->IsRefreshable()); +} + +TEST_F(OAuth2ExchangeRefreshTest, CloseSkipsQueuedExchange) { + ExpectRefreshes(1); + ASSERT_THAT(Start(), IsOk()); + auto task = queued_[0].get_future(); + ASSERT_EQ(task.wait_for(std::chrono::seconds(5)), std::future_status::ready); + EXPECT_THAT(session_->Close(), IsOk()); + std::move(task.get())(); + EXPECT_EQ(client_->requests, 0); +} + +TEST_F(OAuth2ExchangeRefreshTest, StopDuringExchangePreventsRescheduling) { + ExpectRefreshes(1); + client_->response = [&]() -> Result { + session_->StopRefreshing(); + return HttpResponse( + 200, R"({"access_token":"fresh","token_type":"bearer","expires_in":1})"); + }; + ASSERT_THAT(Start(), IsOk()); + ASSERT_NO_FATAL_FAILURE(RunRefresh()); + EXPECT_EQ(session_->OAuth2Info()->token, "fresh"); + EXPECT_FALSE(session_->IsRefreshable()); + ASSERT_TRUE(test::WaitForSharedTimer(std::chrono::seconds(1))); + EXPECT_THAT(session_->Close(), IsOk()); +} + +TEST_F(OAuth2ExchangeRefreshTest, CloseWaitsForExchangeAndStopsRetries) { + ExpectRefreshes(1); + std::promise entered; + std::promise release; + auto entered_future = entered.get_future(); + auto release_future = release.get_future(); + client_->response = [&]() -> Result { + entered.set_value(); + release_future.wait(); + return IOError("exchange failed during Close"); + }; + ASSERT_THAT(Start(), IsOk()); + auto task = queued_[0].get_future(); + ASSERT_EQ(task.wait_for(std::chrono::seconds(5)), std::future_status::ready); + auto refresh = std::async(std::launch::async, [&] { std::move(task.get())(); }); + auto entered_status = entered_future.wait_for(std::chrono::seconds(5)); + if (entered_status != std::future_status::ready) { + release.set_value(); + } + ASSERT_EQ(entered_status, std::future_status::ready); + + auto closing = std::async(std::launch::async, [&] { return session_->Close(); }); + EXPECT_EQ(closing.wait_for(std::chrono::milliseconds(30)), std::future_status::timeout); + release.set_value(); + refresh.get(); + EXPECT_THAT(closing.get(), IsOk()); + EXPECT_EQ(client_->requests, 1); + EXPECT_FALSE(session_->IsRefreshable()); +} + +TEST_F(OAuth2ExchangeRefreshTest, DoesNotScheduleWithoutExpirationOrExecutor) { + ExpectRefreshes(0); + token_.expires_in_secs = std::nullopt; + ASSERT_THAT(Start(), IsOk()); + EXPECT_FALSE(session_->IsRefreshable()); + EXPECT_THAT(session_->Close(), IsOk()); + + token_.expires_in_secs = 0; + config_.refresh_executor = nullptr; + ASSERT_THAT(Start(), IsOk()); + EXPECT_FALSE(session_->IsRefreshable()); + ASSERT_TRUE(test::WaitForSharedTimer()); + EXPECT_EQ(client_->requests, 0); } TEST(OAuth2SessionTest, AlreadyExpiredRefreshedTokenBacksOff) { @@ -863,7 +1631,13 @@ TEST(OAuth2SessionTest, StopsAfterHttpRetryBudget) { test::FakeRetryEnvironment fake_retry; ScopedRetryTestHooks hooks(fake_retry.hooks()); std::move(task.get())(); - EXPECT_EQ(client->requests, 6); + EXPECT_EQ(client->requests, 7); + for (const auto& request : client->recorded_requests) { + EXPECT_EQ(request.form.at("grant_type"), "client_credentials"); + EXPECT_EQ(request.form.at("client_id"), "client"); + EXPECT_EQ(request.form.at("client_secret"), "secret"); + EXPECT_FALSE(request.headers.contains("Authorization")); + } EXPECT_THAT(fake_retry.observed_base_delays_ms(), ::testing::ElementsAre(100, 200, 400, 800, 1600)); // No further refresh is scheduled, so the session no longer counts as refreshable. @@ -885,8 +1659,7 @@ TEST(OAuth2SessionTest, StopRefreshingKeepsSessionUsable) { auto session, internal::OAuth2Session::Make(token_response, {.token_endpoint = "http://localhost/oauth/tokens", - .client_id = "client_id", - .client_secret = "client_secret", + .credential = "client_id:client_secret", .scope = "catalog", .refresh_executor = &executor}, std::make_shared(), std::nullopt)); diff --git a/src/iceberg/test/http_client_test.cc b/src/iceberg/test/http_client_test.cc index 6cc653131..dd5f2fca0 100644 --- a/src/iceberg/test/http_client_test.cc +++ b/src/iceberg/test/http_client_test.cc @@ -26,13 +26,18 @@ #include #include #include +#include +#include +#include #include #include #include +#include #include #include #include +#include #ifndef _WIN32 # include @@ -45,8 +50,11 @@ #endif #include "iceberg/catalog/rest/auth/auth_session.h" +#include "iceberg/catalog/rest/auth/auth_session_internal.h" #include "iceberg/catalog/rest/error_handlers.h" +#include "iceberg/test/executor.h" #include "iceberg/test/matchers.h" +#include "iceberg/util/url_encoder.h" namespace iceberg::rest { @@ -55,10 +63,16 @@ namespace iceberg::rest { namespace { /// \brief Minimal HTTP/1.1 server on 127.0.0.1 that answers every request with -/// a fixed response and closes the connection. +/// a response from the handler and closes the connection. class LoopbackHttpServer { public: - explicit LoopbackHttpServer(std::string response) : response_(std::move(response)) {} + using Handler = std::function; + + explicit LoopbackHttpServer(Handler handler) : handler_(std::move(handler)) {} + + explicit LoopbackHttpServer(std::string response) + : LoopbackHttpServer( + [response = std::move(response)](const std::string&) { return response; }) {} ~LoopbackHttpServer() { stop_ = true; @@ -116,16 +130,16 @@ class LoopbackHttpServer { if (conn < 0) { continue; } - ReadRequest(conn); + auto request = ReadRequest(conn); ++request_count_; - SendAll(conn, response_); + SendAll(conn, handler_(request)); close(conn); } } /// \brief Consume the request headers and any Content-Length body so the /// client does not see a connection reset before it reads the response. - static void ReadRequest(int conn) { + static std::string ReadRequest(int conn) { timeval timeout{.tv_sec = 10, .tv_usec = 0}; setsockopt(conn, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout)); setsockopt(conn, SOL_SOCKET, SO_SNDTIMEO, &timeout, sizeof(timeout)); @@ -137,7 +151,7 @@ class LoopbackHttpServer { while (true) { ssize_t n = recv(conn, buf.data(), buf.size(), 0); if (n <= 0) { - return; + return request; } request.append(buf.data(), static_cast(n)); if (header_end == std::string::npos) { @@ -148,7 +162,7 @@ class LoopbackHttpServer { content_length = ParseContentLength(request.substr(0, header_end)); } if (request.size() >= header_end + 4 + content_length) { - return; + return request; } } } @@ -183,7 +197,7 @@ class LoopbackHttpServer { } } - std::string response_; + Handler handler_; int listen_fd_ = -1; int port_ = 0; std::atomic stop_{false}; @@ -249,8 +263,228 @@ void ServeOnce(int listener, std::string_view response, bool end) { ::close(connection); } +HttpHeaders RequestHeaders(const std::string& request) { + HttpHeaders headers; + const auto header_end = request.find("\r\n\r\n"); + if (header_end == std::string::npos) return headers; + auto start = request.find("\r\n") + 2; + while (start < header_end) { + auto end = request.find("\r\n", start); + auto separator = request.find(':', start); + if (separator < end) { + auto value_start = request.find_first_not_of(" \t", separator + 1); + headers.emplace(request.substr(start, separator - start), + request.substr(value_start, end - value_start)); + } + start = end + 2; + } + return headers; +} + +std::unordered_map RequestForm(const std::string& request) { + std::unordered_map form; + auto header_end = request.find("\r\n\r\n"); + if (header_end == std::string::npos) return form; + auto body = request.substr(header_end + 4); + size_t start = 0; + while (start < body.size()) { + auto end = body.find('&', start); + auto field = body.substr(start, end == std::string::npos ? end : end - start); + auto separator = field.find('='); + if (separator != std::string::npos) { + form.emplace(UrlEncoder::Decode(field.substr(0, separator)), + UrlEncoder::Decode(field.substr(separator + 1))); + } + if (end == std::string::npos) break; + start = end + 1; + } + return form; +} + +std::string JsonResponse(int status, const nlohmann::json& json) { + auto body = json.dump(); + return std::format( + "HTTP/1.1 {} Response\r\nContent-Type: application/json\r\n" + "Content-Length: {}\r\nConnection: close\r\n\r\n{}", + status, body.size(), body); +} + } // namespace +TEST(HttpClientTest, OAuth2SessionRefreshesTokensOverHttp) { + using auth::AuthProperties; + using auth::internal::OAuth2Session; + const std::array tokens = {"Initial+/=", "First+/=", "Second+/="}; + const std::array types = {AuthProperties::kAccessTokenType, + AuthProperties::kJwtTokenType, + AuthProperties::kAccessTokenType}; + // The server alone owns generation; the client must follow the returned token/type. + size_t generation = 0; + LoopbackHttpServer server([&](const std::string& request) { + const auto headers = RequestHeaders(request); + const bool authenticated = + headers.contains("Authorization") && + headers.at("Authorization") == "Bearer " + tokens.at(generation); + if (request.starts_with("GET /protected ")) { + return authenticated ? JsonResponse(200, {{"ok", true}}) + : JsonResponse(401, {{"error", + {{"code", 401}, + {"type", "NotAuthorizedException"}, + {"message", "stale token"}}}}); + } + + const auto form = RequestForm(request); + const std::unordered_map expected_form{ + {"grant_type", "urn:ietf:params:oauth:grant-type:token-exchange"}, + {"subject_token", tokens.at(generation)}, + {"subject_token_type", types.at(generation)}, + {"scope", "catalog"}, + }; + if (!request.starts_with("POST /v1/oauth/tokens ") || !authenticated || + (headers.find("Content-Type") == headers.end() || + headers.at("Content-Type") != "application/x-www-form-urlencoded") || + form != expected_form || generation >= 2) { + return JsonResponse(400, {{"error", "invalid_request"}}); + } + ++generation; + nlohmann::json response{{"access_token", tokens[generation]}, + {"token_type", "bearer"}, + {"issued_token_type", types[generation]}}; + if (generation == 1) response["expires_in"] = 1; + return JsonResponse(200, response); + }); + ASSERT_TRUE(server.Start()); + + std::array, 2> completed; + test::ThreadExecutor workers; + test::MockExecutor executor; + size_t submitted = 0; + EXPECT_CALL(executor, Submit).Times(2).WillRepeatedly([&](ExecutorTask task) -> Status { + auto& done = completed.at(submitted++); + return workers.Submit(ExecutorTask([task = std::move(task), &done]() mutable { + std::move(task)(); + done.set_value(); + })); + }); + auto client = std::make_shared( + std::unordered_map{ + {"authorization", "Bearer stale-default"}}, + std::chrono::seconds(2), std::chrono::seconds(2)); + ICEBERG_UNWRAP_OR_FAIL( + auto session, OAuth2Session::Make({.access_token = tokens[0], + .token_type = "bearer", + .expires_in_secs = 1, + .issued_token_type = types[0]}, + {.token_endpoint = server.url("/v1/oauth/tokens"), + .scope = "catalog", + .exchange_enabled = true, + .refresh_executor = &executor}, + client, std::nullopt)); + for (auto& done : completed) { + ASSERT_EQ(done.get_future().wait_for(std::chrono::seconds(10)), + std::future_status::ready); + } + EXPECT_FALSE(session->IsRefreshable()); + + // A real resource request accepts the refreshed session and rejects the old token. + HttpClient resource_client({}, std::chrono::seconds(2), std::chrono::seconds(2)); + const auto& handler = *DefaultErrorHandler::Instance(); + EXPECT_THAT(resource_client.Get(server.url("/protected"), {}, {}, handler, *session), + IsOk()); + auto old_session = + auth::AuthSession::MakeDefault({{"Authorization", "Bearer " + tokens[0]}}); + EXPECT_THAT( + resource_client.Get(server.url("/protected"), {}, {}, handler, *old_session), + IsError(ErrorKind::kNotAuthorized)); + EXPECT_THAT(session->Close(), IsOk()); + EXPECT_EQ(server.request_count(), 4); +} + +class OAuth2HttpRecoveryTest : public ::testing::TestWithParam {}; + +TEST_P(OAuth2HttpRecoveryTest, RecoversExpiredTokenAndAuthenticatesResourceOverHttp) { + using auth::AuthProperties; + using auth::internal::OAuth2Session; + const bool exchange_enabled = GetParam(); + const std::string initial_token = "Expired+/="; + const std::string refreshed_token = "Recovered+/="; + LoopbackHttpServer server([&](const std::string& request) { + const auto headers = RequestHeaders(request); + if (request.starts_with("GET /protected ")) { + if (headers.contains("Authorization") && + headers.at("Authorization") == "Bearer " + refreshed_token) { + return JsonResponse(200, {{"ok", true}}); + } + return JsonResponse(401, {{"error", + {{"code", 401}, + {"type", "NotAuthorizedException"}, + {"message", "wrong token"}}}}); + } + std::unordered_map expected_form{{"scope", "catalog"}}; + bool valid_auth; + if (exchange_enabled) { + expected_form.insert({ + {"grant_type", "urn:ietf:params:oauth:grant-type:token-exchange"}, + {"subject_token", initial_token}, + {"subject_token_type", AuthProperties::kAccessTokenType}, + }); + valid_auth = headers.contains("Authorization") && + headers.at("Authorization") == "Basic OlNlQ3JFdDpUYWls"; + } else { + expected_form.insert( + {{"grant_type", "client_credentials"}, {"client_secret", "SeCrEt:Tail"}}); + valid_auth = !headers.contains("Authorization"); + } + if (!request.starts_with("POST /v1/oauth/tokens ") || !valid_auth || + RequestForm(request) != expected_form || !headers.contains("X-Trace") || + headers.at("X-Trace") != "KeepThisValue" || !headers.contains("Content-Type") || + headers.at("Content-Type") != "application/x-www-form-urlencoded") { + return JsonResponse(400, {{"error", "invalid_request"}}); + } + return JsonResponse(200, + {{"access_token", refreshed_token}, {"token_type", "bearer"}}); + }); + ASSERT_TRUE(server.Start()); + + test::MockExecutor executor; + std::promise queued; + auto ready = queued.get_future(); + EXPECT_CALL(executor, Submit).WillOnce([&](ExecutorTask task) -> Status { + queued.set_value(std::move(task)); + return {}; + }); + auto client = std::make_shared( + std::unordered_map{ + {"authorization", "Bearer StaleDefault"}, {"X-Trace", "KeepThisValue"}}, + std::chrono::seconds(2), std::chrono::seconds(2)); + ICEBERG_UNWRAP_OR_FAIL( + auto session, + OAuth2Session::Make( + {.access_token = initial_token, .token_type = "bearer", .expires_in_secs = 0}, + {.token_endpoint = server.url("/v1/oauth/tokens"), + .credential = ":SeCrEt:Tail", + .scope = "catalog", + .exchange_enabled = exchange_enabled, + .refresh_executor = &executor}, + client, std::nullopt)); + ASSERT_EQ(ready.wait_for(std::chrono::seconds(5)), std::future_status::ready); + std::move(ready.get())(); + EXPECT_EQ(session->OAuth2Info()->token, refreshed_token); + EXPECT_FALSE(session->IsRefreshable()); + + HttpClient resource_client({}, std::chrono::seconds(2), std::chrono::seconds(2)); + EXPECT_THAT(resource_client.Get(server.url("/protected"), {}, {}, + *DefaultErrorHandler::Instance(), *session), + IsOk()); + EXPECT_THAT(session->Close(), IsOk()); + EXPECT_EQ(server.request_count(), 2); +} + +INSTANTIATE_TEST_SUITE_P(ExchangeModes, OAuth2HttpRecoveryTest, ::testing::Bool(), + [](const ::testing::TestParamInfo& info) { + return info.param ? "Exchange" : "ClientCredentials"; + }); + TEST(HttpClientTest, GetFailsWhenItsTransferFails) { // A 200 arriving first does not make a truncated response a success. constexpr std::string_view kTruncated = diff --git a/src/iceberg/test/rest_catalog_unit_test.cc b/src/iceberg/test/rest_catalog_unit_test.cc index 1b162763b..640530128 100644 --- a/src/iceberg/test/rest_catalog_unit_test.cc +++ b/src/iceberg/test/rest_catalog_unit_test.cc @@ -54,6 +54,7 @@ #include "iceberg/table_update.h" #include "iceberg/test/executor.h" #include "iceberg/test/matchers.h" +#include "iceberg/test/retry.h" namespace iceberg::rest { @@ -669,6 +670,190 @@ class RestCatalogOAuth2UriTest : public RestCatalogLoadTableTest { int factory_calls_ = 0; }; +struct OAuth2RefreshModeCase { + std::string name; + std::string client; + std::string server_default; + std::string server_override; + bool exchange_enabled; +}; + +class RestCatalogOAuth2RefreshTest + : public RestCatalogOAuth2UriTest, + public ::testing::WithParamInterface {}; + +TEST_P(RestCatalogOAuth2RefreshTest, RefreshUsesMergedConfigAndUpdatesBusinessRequests) { + const auto& mode = GetParam(); + properties_["token-refresh-enabled"] = "true"; + if (!mode.client.empty()) { + properties_["token-exchange-enabled"] = mode.client; + } + CatalogConfig server_config; + if (!mode.server_default.empty()) { + server_config.defaults["token-exchange-enabled"] = mode.server_default; + } + if (!mode.server_override.empty()) { + server_config.overrides["token-exchange-enabled"] = mode.server_override; + } + ExpectInit(server_config, + R"({"access_token":"initial","token_type":"bearer","expires_in":1})"); + + std::unordered_map expected_form; + std::unordered_map expected_headers; + if (mode.exchange_enabled) { + expected_form = { + {"grant_type", "urn:ietf:params:oauth:grant-type:token-exchange"}, + {"subject_token", "initial"}, + {"subject_token_type", auth::AuthProperties::kAccessTokenType}, + {"scope", "catalog"}, + }; + expected_headers = {{"Authorization", "Bearer initial"}}; + } else { + expected_form = {{"grant_type", "client_credentials"}, + {"client_id", "client"}, + {"client_secret", "secret"}, + {"scope", "catalog"}}; + } + EXPECT_CALL(*mock_client_, PostForm("https://auth.example/token/", _, _, _, _)) + .WillOnce([&](const std::string&, const auto& form, const auto& headers, + const ErrorHandler&, auth::AuthSession&) -> Result { + EXPECT_THAT(form, ::testing::ContainerEq(expected_form)); + EXPECT_THAT(headers, ::testing::ContainerEq(expected_headers)); + return HttpResponse(200, R"({"access_token":"refreshed","token_type":"bearer"})"); + }); + + test::MockExecutor executor; + std::promise queued; + auto ready = queued.get_future(); + EXPECT_CALL(executor, Submit).WillOnce([&](ExecutorTask task) -> Status { + queued.set_value(std::move(task)); + return {}; + }); + ICEBERG_UNWRAP_OR_FAIL(auto root, MakeOAuth2Catalog(&executor)); + ASSERT_EQ(ready.wait_for(std::chrono::seconds(5)), std::future_status::ready); + std::move(ready.get())(); + + EXPECT_CALL(*mock_client_, Get("https://catalog.example/api/v1/namespaces", _, _, _, _)) + .WillOnce([](const std::string&, const auto&, const auto&, const ErrorHandler&, + auth::AuthSession& session) -> Result { + ICEBERG_ASSIGN_OR_RAISE(auto request, session.Authenticate({})); + EXPECT_EQ(request.headers.at("Authorization"), "Bearer refreshed"); + return HttpResponse(200, R"({"namespaces":[]})"); + }); + ICEBERG_UNWRAP_OR_FAIL(auto catalog, root->AsCatalog()); + EXPECT_THAT(catalog->ListNamespaces({}), IsOk()); + EXPECT_EQ(factory_calls_, 2); +} + +INSTANTIATE_TEST_SUITE_P( + ExchangeModes, RestCatalogOAuth2RefreshTest, + ::testing::Values(OAuth2RefreshModeCase{"Default", "", "", "", true}, + OAuth2RefreshModeCase{"Enabled", "true", "", "", true}, + OAuth2RefreshModeCase{"Disabled", "false", "", "", false}, + OAuth2RefreshModeCase{"ServerDisables", "true", "", "false", false}, + OAuth2RefreshModeCase{"ServerEnables", "false", "", "true", true}, + OAuth2RefreshModeCase{"ServerDefaultDisables", "", "false", "", + false}, + OAuth2RefreshModeCase{"ClientOverridesServerDefault", "false", + "true", "", false}), + [](const ::testing::TestParamInfo& info) { + return info.param.name; + }); + +struct OAuth2RecoveryCase { + std::string name; + bool expired; + bool exchange_enabled; +}; + +class RestCatalogOAuth2RecoveryTest + : public RestCatalogOAuth2UriTest, + public ::testing::WithParamInterface {}; + +TEST_P(RestCatalogOAuth2RecoveryTest, RecoversAndUpdatesBusinessRequests) { + const auto& mode = GetParam(); + properties_["token-refresh-enabled"] = "true"; + properties_["token-exchange-enabled"] = mode.exchange_enabled ? "true" : "false"; + properties_["credential"] = ":SeCrEt:Tail"; + ExpectInit({}, nlohmann::json({{"access_token", "Initial"}, + {"token_type", "bearer"}, + {"expires_in", mode.expired ? 0 : 1}}) + .dump()); + + const int expected_requests = mode.expired ? 1 : 7; + int requests = 0; + EXPECT_CALL(*mock_client_, PostForm("https://auth.example/token/", _, _, _, _)) + .Times(expected_requests) + .WillRepeatedly([&](const std::string&, const auto& form, const auto& headers, + const ErrorHandler&, + auth::AuthSession& session) -> Result { + ++requests; + HttpRequest request{.headers = {{"authorization", "Bearer StaleDefault"}, + {"X-Trace", "KeepThisValue"}}}; + for (const auto& [key, value] : headers) { + request.headers.insert_or_assign(key, value); + } + ICEBERG_ASSIGN_OR_RAISE(auto authenticated, + session.Authenticate(std::move(request))); + std::unordered_map expected_form{{"scope", "catalog"}}; + if (mode.exchange_enabled) { + expected_form.insert({ + {"grant_type", "urn:ietf:params:oauth:grant-type:token-exchange"}, + {"subject_token", "Initial"}, + {"subject_token_type", auth::AuthProperties::kAccessTokenType}, + }); + EXPECT_EQ(authenticated.headers.at("Authorization"), + mode.expired || requests == 7 ? "Basic OlNlQ3JFdDpUYWls" + : "Bearer Initial"); + } else { + expected_form.insert( + {{"grant_type", "client_credentials"}, {"client_secret", "SeCrEt:Tail"}}); + EXPECT_FALSE(authenticated.headers.contains("Authorization")); + } + EXPECT_THAT(form, ::testing::ContainerEq(expected_form)); + EXPECT_EQ(authenticated.headers.at("X-Trace"), "KeepThisValue"); + if (requests < expected_requests) { + return IOError("ordinary refresh failed"); + } + return HttpResponse(200, R"({"access_token":"Recovered","token_type":"bearer"})"); + }); + + test::MockExecutor executor; + std::promise queued; + auto ready = queued.get_future(); + EXPECT_CALL(executor, Submit).WillOnce([&](ExecutorTask task) -> Status { + queued.set_value(std::move(task)); + return {}; + }); + ICEBERG_UNWRAP_OR_FAIL(auto root, MakeOAuth2Catalog(&executor)); + ASSERT_EQ(ready.wait_for(std::chrono::seconds(5)), std::future_status::ready); + test::FakeRetryEnvironment retry_environment; + ScopedRetryTestHooks retry_hooks(retry_environment.hooks()); + std::move(ready.get())(); + EXPECT_EQ(requests, expected_requests); + + EXPECT_CALL(*mock_client_, Get("https://catalog.example/api/v1/namespaces", _, _, _, _)) + .WillOnce([](const std::string&, const auto&, const auto&, const ErrorHandler&, + auth::AuthSession& session) -> Result { + ICEBERG_ASSIGN_OR_RAISE(auto request, session.Authenticate({})); + EXPECT_EQ(request.headers.at("Authorization"), "Bearer Recovered"); + return HttpResponse(200, R"({"namespaces":[]})"); + }); + ICEBERG_UNWRAP_OR_FAIL(auto catalog, root->AsCatalog()); + EXPECT_THAT(catalog->ListNamespaces({}), IsOk()); + EXPECT_EQ(factory_calls_, 2); +} + +INSTANTIATE_TEST_SUITE_P( + RecoveryModes, RestCatalogOAuth2RecoveryTest, + ::testing::Values(OAuth2RecoveryCase{"ExpiredExchange", true, true}, + OAuth2RecoveryCase{"FinalExchange", false, true}, + OAuth2RecoveryCase{"ExpiredClientCredentials", true, false}, + OAuth2RecoveryCase{"FinalClientCredentials", false, false}), + [](const ::testing::TestParamInfo& info) { + return info.param.name; + }); + TEST_F(RestCatalogOAuth2UriTest, TrustsServerOverridesByDefault) { ExpectInit({.overrides = {{"uri", "HTTPS://other-catalog.example/base/"}, {"oauth2-server-uri", "oauth/token///"},