jnturton commented on code in PR #2544: URL: https://github.com/apache/drill/pull/2544#discussion_r873511567
########## 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: Nit, can we standardise teminolgy here to make our own lives easier and refer to the "query user" as already happens in some other places? I suppose there is no query running in this context, which makes things less obvious, but it is still eventually the query user that we're managing managing tokens for here. ########## exec/java-exec/src/main/java/org/apache/drill/exec/server/rest/PluginConfigWrapper.java: ########## @@ -114,4 +125,75 @@ public boolean isOauth() { return tokenCredentials.map(OAuthTokenCredentials::getClientID).orElse(null) != null; } + + @JsonIgnore + public String getClientID() { + CredentialedStoragePluginConfig securedStoragePluginConfig = (CredentialedStoragePluginConfig) config; + CredentialsProvider credentialsProvider = securedStoragePluginConfig.getCredentialsProvider(); + + return credentialsProvider.getCredentials().getOrDefault("clientID", ""); + } + + /** + * This function generates the authorization URI for use when a non-admin user is authorizing + * OAuth2.0 access for a storage plugin. This function is necessary as we do not wish to expose + * any plugin configuration information to the user. + * + * If the plugin is not OAuth, or is missing components, the function will return an empty string. + * @return The authorization URI for an OAuth enabled plugin. + */ + @JsonIgnore + public String getAuthorizationURIWithParams() { Review Comment: Just a note that I'll try to move this code, probably into OAuthTokenCredentials or the storage config, in DRILL-8200. ########## exec/java-exec/src/main/java/org/apache/drill/exec/server/rest/StorageResources.java: ########## @@ -212,61 +195,15 @@ public Response enablePlugin(@PathParam("name") String name, @PathParam("val") B @Consumes(MediaType.APPLICATION_JSON) @Produces(MediaType.APPLICATION_JSON) public Response updateRefreshToken(@PathParam("name") String name, OAuthTokenContainer tokens) { Review Comment: Are these endpoints different in any way from their counterparts in CredentialResources? Could we do everything we need to without the ones here? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
