[
https://issues.apache.org/jira/browse/DRILL-8220?page=com.atlassian.jira.plugin.system.issuetabpanels:comment-tabpanel&focusedCommentId=17537465#comment-17537465
]
ASF GitHub Bot commented on DRILL-8220:
---------------------------------------
cgivre commented on code in PR #2544:
URL: https://github.com/apache/drill/pull/2544#discussion_r873578777
##########
exec/java-exec/src/main/java/org/apache/drill/exec/server/rest/OAuthRequests.java:
##########
@@ -0,0 +1,230 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.drill.exec.server.rest;
+
+import okhttp3.OkHttpClient;
+import okhttp3.OkHttpClient.Builder;
+import okhttp3.Request;
+import org.apache.drill.common.logical.CredentialedStoragePluginConfig;
+import org.apache.drill.common.logical.StoragePluginConfig;
+import org.apache.drill.common.logical.StoragePluginConfig.AuthMode;
+import org.apache.drill.common.logical.security.CredentialsProvider;
+import org.apache.drill.exec.oauth.OAuthTokenProvider;
+import org.apache.drill.exec.oauth.PersistentTokenTable;
+import org.apache.drill.exec.oauth.TokenRegistry;
+import org.apache.drill.exec.server.DrillbitContext;
+import org.apache.drill.exec.server.rest.DrillRestServer.UserAuthEnabled;
+import org.apache.drill.exec.server.rest.StorageResources.JsonResult;
+import org.apache.drill.exec.store.AbstractStoragePlugin;
+import org.apache.drill.exec.store.StoragePluginRegistry;
+import org.apache.drill.exec.store.StoragePluginRegistry.PluginException;
+import org.apache.drill.exec.store.http.oauth.OAuthUtils;
+import org.apache.drill.exec.store.security.oauth.OAuthTokenCredentials;
+import org.eclipse.jetty.util.resource.Resource;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.servlet.http.HttpServletRequest;
+import javax.ws.rs.core.Response;
+import javax.ws.rs.core.Response.Status;
+import javax.ws.rs.core.SecurityContext;
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.nio.charset.StandardCharsets;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+public class OAuthRequests {
+
+ private static final Logger logger =
LoggerFactory.getLogger(OAuthRequests.class);
+ private static final String OAUTH_SUCCESS_PAGE =
"/rest/storage/success.html";
+
+ public static Response updateAccessToken(String name,
+ OAuthTokenContainer tokens,
+ StoragePluginRegistry storage,
+ UserAuthEnabled authEnabled,
+ SecurityContext sc) {
+ try {
+ if (storage.getPlugin(name).getConfig() instanceof
CredentialedStoragePluginConfig) {
+ DrillbitContext context = ((AbstractStoragePlugin)
storage.getPlugin(name)).getContext();
+ OAuthTokenProvider tokenProvider = context.getoAuthTokenProvider();
+ PersistentTokenTable tokenTable =
tokenProvider.getOauthTokenRegistry(getActiveUser(storage.getPlugin(name).getConfig(),
authEnabled, sc)).getTokenTable(name);
+
+ // Set the access token
+ tokenTable.setAccessToken(tokens.getAccessToken());
+
+ return Response.status(Status.OK)
+ .entity("Access tokens have been updated.")
+ .build();
+ } else {
+ logger.error("{} does not support OAuth2.0. You can only add tokens
to OAuth enabled plugins.", name);
+ return Response.status(Status.INTERNAL_SERVER_ERROR)
+ .entity(message("Unable to add tokens: %s", name))
+ .build();
+ }
+ } catch (PluginException e) {
+ logger.error("Error when adding tokens to {}", name);
+ return Response.status(Status.INTERNAL_SERVER_ERROR)
+ .entity(message("Unable to add tokens: %s", e.getMessage()))
+ .build();
+ }
+ }
+
+ public static Response updateRefreshToken(String name, OAuthTokenContainer
tokens,
+ StoragePluginRegistry storage,
UserAuthEnabled authEnabled,
+ SecurityContext sc) {
+ try {
+ if (storage.getPlugin(name).getConfig() instanceof
CredentialedStoragePluginConfig) {
+ DrillbitContext context = ((AbstractStoragePlugin)
storage.getPlugin(name)).getContext();
+ OAuthTokenProvider tokenProvider = context.getoAuthTokenProvider();
+ PersistentTokenTable tokenTable = tokenProvider.getOauthTokenRegistry(
+ getActiveUser(storage.getPlugin(name).getConfig(), authEnabled,
sc)).getTokenTable(name);
+
+ // Set the access token
+ tokenTable.setRefreshToken(tokens.getRefreshToken());
+
+ return Response.status(Status.OK)
+ .entity("Refresh token have been updated.")
+ .build();
+ } else {
+ logger.error("{} is not a HTTP plugin. You can only add access tokens
to HTTP plugins.", name);
+ return Response.status(Status.INTERNAL_SERVER_ERROR)
+ .entity(message("Unable to add tokens: %s", name))
+ .build();
+ }
+ } catch (PluginException e) {
+ logger.error("Error when adding tokens to {}", name);
+ return Response.status(Status.INTERNAL_SERVER_ERROR)
+ .entity(message("Unable to add tokens: %s", e.getMessage()))
+ .build();
+ }
+ }
+
+ public static Response updateOAuthTokens(String name, OAuthTokenContainer
tokenContainer, StoragePluginRegistry storage,
+ UserAuthEnabled authEnabled,
SecurityContext sc) {
+ try {
+ if (storage.getPlugin(name).getConfig() instanceof
CredentialedStoragePluginConfig) {
+ DrillbitContext context = ((AbstractStoragePlugin)
storage.getPlugin(name)).getContext();
+ OAuthTokenProvider tokenProvider = context.getoAuthTokenProvider();
+ PersistentTokenTable tokenTable = tokenProvider
+
.getOauthTokenRegistry(getActiveUser(storage.getPlugin(name).getConfig(),
authEnabled, sc))
+ .getTokenTable(name);
+
+ // Set the access and refresh token
+ tokenTable.setAccessToken(tokenContainer.getAccessToken());
+ tokenTable.setRefreshToken(tokenContainer.getRefreshToken());
+
+ return Response.status(Status.OK)
+ .entity("Access tokens have been updated.")
+ .build();
+ } else {
+ logger.error("{} is not a HTTP plugin. You can only add access tokens
to HTTP plugins.", name);
+ return Response.status(Status.INTERNAL_SERVER_ERROR)
+ .entity(message("Unable to add tokens: %s", name))
+ .build();
+ }
+ } catch (PluginException e) {
+ logger.error("Error when adding tokens to {}", name);
+ return Response.status(Status.INTERNAL_SERVER_ERROR)
+ .entity(message("Unable to add tokens: %s", e.getMessage()))
+ .build();
+ }
+ }
+
+ public static Response updateAuthToken(String name, String code,
HttpServletRequest request,
+ StoragePluginRegistry storage,
UserAuthEnabled authEnabled,
+ SecurityContext sc) {
+ try {
+ if (storage.getPlugin(name).getConfig() instanceof
CredentialedStoragePluginConfig) {
+ CredentialedStoragePluginConfig securedStoragePluginConfig =
(CredentialedStoragePluginConfig) storage.getPlugin(name).getConfig();
+ CredentialsProvider credentialsProvider =
securedStoragePluginConfig.getCredentialsProvider();
+ String callbackURL = request.getRequestURL().toString();
+
+ // Now exchange the authorization token for an access token
+ Builder builder = new OkHttpClient.Builder();
+ OkHttpClient client = builder.build();
+
+ Request accessTokenRequest =
OAuthUtils.getAccessTokenRequest(credentialsProvider, code, callbackURL);
+ Map<String, String> updatedTokens = OAuthUtils.getOAuthTokens(client,
accessTokenRequest);
+
+ // Add to token registry
+ // If USER_TRANSLATION is enabled, Drill will create a token table for
each user.
+ TokenRegistry tokenRegistry = ((AbstractStoragePlugin)
storage.getPlugin(name))
+ .getContext()
+ .getoAuthTokenProvider()
+
.getOauthTokenRegistry(getActiveUser(storage.getPlugin(name).getConfig(),
authEnabled, sc));
+
+ // Add a token registry table if none exists
+ tokenRegistry.createTokenTable(name);
+ PersistentTokenTable tokenTable = tokenRegistry.getTokenTable(name);
+
+ // Add tokens to persistent storage
+
tokenTable.setAccessToken(updatedTokens.get(OAuthTokenCredentials.ACCESS_TOKEN));
+
tokenTable.setRefreshToken(updatedTokens.get(OAuthTokenCredentials.REFRESH_TOKEN));
+
+ // Get success page
+ String successPage = null;
+ try (InputStream inputStream =
Resource.newClassPathResource(OAUTH_SUCCESS_PAGE).getInputStream()) {
+ InputStreamReader reader = new InputStreamReader(inputStream,
StandardCharsets.UTF_8);
+ BufferedReader bufferedReader = new BufferedReader(reader);
+ successPage = bufferedReader.lines()
+ .collect(Collectors.joining("\n"));
+ bufferedReader.close();
+ reader.close();
+ } catch (IOException e) {
+ return Response.status(Status.OK).entity("You may close this
window.").build();
+ }
+
+ return Response.status(Status.OK).entity(successPage).build();
+ } else {
+ logger.error("{} is not a HTTP plugin. You can only add auth code to
HTTP plugins.", name);
+ return Response.status(Status.INTERNAL_SERVER_ERROR)
+ .entity(message("Unable to add authorization code: %s", name))
+ .build();
+ }
+ } catch (PluginException e) {
+ logger.error("Error when adding auth token to {}", name);
+ return Response.status(Status.INTERNAL_SERVER_ERROR)
+ .entity(message("Unable to add authorization code: %s",
e.getMessage()))
+ .build();
+ }
+ }
+
+ private static JsonResult message(String message, Object... args) {
+ return new JsonResult(String.format(message, args)); // lgtm
[java/tainted-format-string]
+ }
+
+ /**
+ * This function checks to see if a given storage plugin is using
USER_TRANSLATION mode and if user
+ * authentication is enabled. If so, it will return the active user name.
If not it will return null.
+ * @param config {@link StoragePluginConfig} The current plugin configuration
+ * @return If USER_TRANSLATION is enabled, returns the active user. If not,
returns null.
+ */
+ private static String getActiveUser(StoragePluginConfig config,
Review Comment:
Renamed function to `getQueryUser()`.
> Add User Translation Support for OAuth Enabled Plugins
> ------------------------------------------------------
>
> Key: DRILL-8220
> URL: https://issues.apache.org/jira/browse/DRILL-8220
> Project: Apache Drill
> Issue Type: Improvement
> Components: Storage - Other
> Affects Versions: 1.20.0
> Reporter: Charles Givre
> Assignee: Charles Givre
> Priority: Major
> Fix For: 2.0.0
>
>
> This PR adds support for individual users to provide their own credentials
> for plugins that use OAuth 2.0 as a means of authorization and
> authentication. Currently, only the HTTP storage plugin supports OAuth,
> however, this PR moves some of the core features out of the HTTP plugin so
> that other plugins can access this.
--
This message was sent by Atlassian Jira
(v8.20.7#820007)
