From 64c1da590ac3ec1ff28dd373ca964e78620c652c Mon Sep 17 00:00:00 2001 From: Eduard Tudenhoefner Date: Tue, 17 Jan 2023 12:41:39 +0100 Subject: [PATCH] Core: Allow customizing OAuth scope --- .../iceberg/rest/RESTSessionCatalog.java | 8 ++--- .../apache/iceberg/rest/auth/OAuth2Util.java | 32 +++++++++++++++---- .../apache/iceberg/rest/TestRESTCatalog.java | 15 +++++++-- 3 files changed, 40 insertions(+), 15 deletions(-) diff --git a/core/src/main/java/org/apache/iceberg/rest/RESTSessionCatalog.java b/core/src/main/java/org/apache/iceberg/rest/RESTSessionCatalog.java index f79dc9845cb7..5f861880a606 100644 --- a/core/src/main/java/org/apache/iceberg/rest/RESTSessionCatalog.java +++ b/core/src/main/java/org/apache/iceberg/rest/RESTSessionCatalog.java @@ -134,14 +134,12 @@ public void initialize(String name, Map unresolved) { ConfigResponse config; OAuthTokenResponse authResponse; String credential = props.get(OAuth2Properties.CREDENTIAL); - // TODO: if scope can be overridden, it should be done consistently + String scope = props.getOrDefault(OAuth2Properties.SCOPE, OAuth2Properties.CATALOG_SCOPE); try (RESTClient initClient = clientBuilder.apply(props)) { Map initHeaders = RESTUtil.merge(configHeaders(props), OAuth2Util.authHeaders(initToken)); if (credential != null && !credential.isEmpty()) { - authResponse = - OAuth2Util.fetchToken( - initClient, initHeaders, credential, OAuth2Properties.CATALOG_SCOPE); + authResponse = OAuth2Util.fetchToken(initClient, initHeaders, credential, scope); Map authHeaders = RESTUtil.merge(initHeaders, OAuth2Util.authHeaders(authResponse.token())); config = fetchConfig(initClient, authHeaders, props); @@ -166,7 +164,7 @@ public void initialize(String name, Map unresolved) { this.client = clientBuilder.apply(mergedProps); this.paths = ResourcePaths.forCatalogProperties(mergedProps); - this.catalogAuth = new AuthSession(baseHeaders, null, null, credential); + this.catalogAuth = new AuthSession(baseHeaders, null, null, credential, scope); if (authResponse != null) { this.catalogAuth = AuthSession.fromTokenResponse( diff --git a/core/src/main/java/org/apache/iceberg/rest/auth/OAuth2Util.java b/core/src/main/java/org/apache/iceberg/rest/auth/OAuth2Util.java index 4025be315ee6..259c6993f077 100644 --- a/core/src/main/java/org/apache/iceberg/rest/auth/OAuth2Util.java +++ b/core/src/main/java/org/apache/iceberg/rest/auth/OAuth2Util.java @@ -357,19 +357,25 @@ public static class AuthSession { private String tokenType; private Long expiresAtMillis; private final String credential; + private final String scope; private volatile boolean keepRefreshed = true; public AuthSession(Map baseHeaders, String token, String tokenType) { - this(baseHeaders, token, tokenType, null); + this(baseHeaders, token, tokenType, null, OAuth2Properties.CATALOG_SCOPE); } public AuthSession( - Map baseHeaders, String token, String tokenType, String credential) { + Map baseHeaders, + String token, + String tokenType, + String credential, + String scope) { this.headers = RESTUtil.merge(baseHeaders, authHeaders(token)); this.token = token; this.tokenType = tokenType; this.expiresAtMillis = OAuth2Util.expiresAtMillis(token); this.credential = credential; + this.scope = scope; } public Map headers() { @@ -388,6 +394,10 @@ public Long expiresAtMillis() { return expiresAtMillis; } + public String scope() { + return scope; + } + public void stopRefreshing() { this.keepRefreshed = false; } @@ -402,7 +412,7 @@ public String credential() { * @return A new {@link AuthSession} with empty headers. */ public static AuthSession empty() { - return new AuthSession(ImmutableMap.of(), null, null, null); + return new AuthSession(ImmutableMap.of(), null, null, null, OAuth2Properties.CATALOG_SCOPE); } /** @@ -457,7 +467,7 @@ private OAuthTokenResponse refreshCurrentToken(RESTClient client) { return refreshExpiredToken(client); } else { // attempt a normal refresh - return refreshToken(client, headers(), token, tokenType, OAuth2Properties.CATALOG_SCOPE); + return refreshToken(client, headers(), token, tokenType, scope); } } @@ -468,7 +478,7 @@ static boolean isExpired(Long expiresAtMillis, long now) { private OAuthTokenResponse refreshExpiredToken(RESTClient client) { if (credential != null) { Map basicHeaders = RESTUtil.merge(headers(), basicAuthHeaders(credential)); - return refreshToken(client, basicHeaders, token, tokenType, OAuth2Properties.CATALOG_SCOPE); + return refreshToken(client, basicHeaders, token, tokenType, scope); } return null; @@ -554,7 +564,11 @@ public static AuthSession fromAccessToken( AuthSession parent) { AuthSession session = new AuthSession( - parent.headers(), token, OAuth2Properties.ACCESS_TOKEN_TYPE, parent.credential()); + parent.headers(), + token, + OAuth2Properties.ACCESS_TOKEN_TYPE, + parent.credential(), + parent.scope()); long startTimeMillis = System.currentTimeMillis(); Long expiresAtMillis = session.expiresAtMillis(); @@ -610,7 +624,11 @@ private static AuthSession fromTokenResponse( String credential) { AuthSession session = new AuthSession( - parent.headers(), response.token(), response.issuedTokenType(), credential); + parent.headers(), + response.token(), + response.issuedTokenType(), + credential, + parent.scope()); if (response.expiresInSeconds() != null) { scheduleTokenRefresh( client, diff --git a/core/src/test/java/org/apache/iceberg/rest/TestRESTCatalog.java b/core/src/test/java/org/apache/iceberg/rest/TestRESTCatalog.java index 19422037e9f7..1199cb0fd6f9 100644 --- a/core/src/test/java/org/apache/iceberg/rest/TestRESTCatalog.java +++ b/core/src/test/java/org/apache/iceberg/rest/TestRESTCatalog.java @@ -50,6 +50,7 @@ import org.apache.iceberg.metrics.MetricsReporter; import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; import org.apache.iceberg.rest.RESTCatalogAdapter.HTTPMethod; +import org.apache.iceberg.rest.auth.OAuth2Properties; import org.apache.iceberg.rest.auth.OAuth2Util; import org.apache.iceberg.rest.responses.ConfigResponse; import org.apache.iceberg.rest.responses.ErrorResponse; @@ -918,8 +919,16 @@ public void testCatalogTokenRefresh() throws Exception { UUID.randomUUID().toString(), "user", contextCredentials, ImmutableMap.of()); RESTCatalog catalog = new RESTCatalog(context, (config) -> adapter); + String scope = "custom_catalog_scope"; catalog.initialize( - "prod", ImmutableMap.of(CatalogProperties.URI, "ignored", "credential", "catalog:secret")); + "prod", + ImmutableMap.of( + CatalogProperties.URI, + "ignored", + "credential", + "catalog:secret", + OAuth2Properties.SCOPE, + scope)); Thread.sleep(3_000); // sleep until after 2 refresh calls @@ -951,7 +960,7 @@ public void testCatalogTokenRefresh() throws Exception { "grant_type", "urn:ietf:params:oauth:grant-type:token-exchange", "subject_token", "client-credentials-token:sub=catalog", "subject_token_type", "urn:ietf:params:oauth:token-type:access_token", - "scope", "catalog"); + "scope", scope); Mockito.verify(adapter) .execute( eq(HTTPMethod.POST), @@ -968,7 +977,7 @@ public void testCatalogTokenRefresh() throws Exception { "grant_type", "urn:ietf:params:oauth:grant-type:token-exchange", "subject_token", "token-exchange-token:sub=client-credentials-token:sub=catalog", "subject_token_type", "urn:ietf:params:oauth:token-type:access_token", - "scope", "catalog"); + "scope", scope); Map secondRefreshHeaders = ImmutableMap.of( "Authorization",