sunchao commented on code in PR #6059: URL: https://github.com/apache/datafusion-comet/pull/6059#discussion_r4206712077
########## spark/src/main/scala/org/apache/comet/objectstore/AbfsAuthResolver.scala: ########## @@ -0,0 +1,495 @@ +/* + * 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.comet.objectstore + +import java.net.URI + +import scala.collection.mutable.ArrayBuffer +import scala.util.{Failure, Success, Try} +import scala.util.control.NonFatal + +import org.apache.commons.lang3.StringUtils +import org.apache.hadoop.conf.Configuration +import org.apache.spark.internal.Logging + +import org.apache.comet.objectstore.AbfsReflection.Handles + +/** + * Resolves the ABFS authentication for one `abfs[s]://container@account/...` URI on the driver, + * by asking the classpath's hadoop-azure the questions + * `AzureBlobFileSystemStore.initializeClient` asks, and encodes the answer as the options the + * native Azure store reads. + * + * Hadoop's `AbfsConfiguration` decides which mechanism applies and where each value comes from + * (container-, account- or globally-scoped keys, credential providers, `${...}` references). The + * resolver forwards the resolved values under the global key names plus `comet.azure.*` markers, + * never a key for another account, and never the environment. Anything Hadoop throws becomes an + * error marker: the native scan fails with it instead of trying a lookup of its own. + * + * `AbfsAuthParitySuite` pins the outcome of every configuration shape per hadoop-azure version + * (`spark/src/test/resources/abfs-auth-parity/expected-<version>.json`) and cross-checks it + * against the real `AzureBlobFileSystem`; its scaladoc says how to regenerate those files after a + * deliberate change here. + */ +private[comet] object AbfsAuthResolver extends Logging { + + /** Options the native Azure store reads. None of them starts with `fs.`. */ + object Markers { + val RESOLUTION = "comet.azure.resolution" + val AUTH_TYPE = "comet.azure.auth.type" + val OAUTH_PROVIDER_CLASS = "comet.azure.oauth.provider.class" + val SAS_PROVIDER_CLASS = "comet.azure.sas.provider.class" + val OAUTH_ASSERTION_PROVIDER_CLASS = "comet.azure.oauth.assertion.provider.class" + val ERROR_CLASS = "comet.azure.error.class" + val ERROR_MESSAGE = "comet.azure.error.message" + + val RESOLVED = "resolved" + val NONE = "none" + val DECLINED = "declined" + val ERROR = "error" + } + + /** Hadoop ABFS configuration keys, global form (`ConfigurationKeys`). */ + object Keys { + val AUTH_TYPE = "fs.azure.account.auth.type" + val ACCOUNT_KEY = "fs.azure.account.key" + val KEY_PROVIDER = "fs.azure.account.keyprovider" + val SHELL_KEY_PROVIDER_SCRIPT = "fs.azure.shellkeyprovider.script" + val OAUTH_PROVIDER_TYPE = "fs.azure.account.oauth.provider.type" + val CLIENT_ID = "fs.azure.account.oauth2.client.id" + val CLIENT_SECRET = "fs.azure.account.oauth2.client.secret" + val CLIENT_ENDPOINT = "fs.azure.account.oauth2.client.endpoint" + val MSI_TENANT = "fs.azure.account.oauth2.msi.tenant" + val MSI_ENDPOINT = "fs.azure.account.oauth2.msi.endpoint" + val MSI_AUTHORITY = "fs.azure.account.oauth2.msi.authority" + val USER_NAME = "fs.azure.account.oauth2.user.name" + val USER_PASSWORD = "fs.azure.account.oauth2.user.password" + val REFRESH_TOKEN = "fs.azure.account.oauth2.refresh.token" + val REFRESH_TOKEN_ENDPOINT = "fs.azure.account.oauth2.refresh.token.endpoint" + val TOKEN_FILE = "fs.azure.account.oauth2.token.file" + val CLIENT_ASSERTION_PROVIDER_TYPE = "fs.azure.account.oauth2.client.assertion.provider.type" + val SAS_TOKEN_PROVIDER_TYPE = "fs.azure.sas.token.provider.type" + val SAS_FIXED_TOKEN = "fs.azure.sas.fixed.token" + } + + /** The OAuth providers Hadoop accepts under `AuthType.OAuth`, matched by exact class name. */ + object Providers { + val CLIENT_CREDS = "org.apache.hadoop.fs.azurebfs.oauth2.ClientCredsTokenProvider" + val USER_PASSWORD = "org.apache.hadoop.fs.azurebfs.oauth2.UserPasswordTokenProvider" + val MSI = "org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider" + val REFRESH_TOKEN = "org.apache.hadoop.fs.azurebfs.oauth2.RefreshTokenBasedTokenProvider" + val WORKLOAD_IDENTITY = "org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider" + val BUILT_IN: Set[String] = + Set(CLIENT_CREDS, USER_PASSWORD, MSI, REFRESH_TOKEN, WORKLOAD_IDENTITY) + } + + private object AuthTypes { + val SHARED_KEY = "SharedKey" + val OAUTH = "OAuth" + val SAS = "SAS" + } + + private object Exceptions { + private val pkg = "org.apache.hadoop.fs.azurebfs.contracts.exceptions." + val CONFIGURATION_PROPERTY_NOT_FOUND: String = pkg + "ConfigurationPropertyNotFoundException" + val INVALID_CONFIGURATION_VALUE: String = pkg + "InvalidConfigurationValueException" + val INVALID_URI: String = pkg + "InvalidUriException" + val INVALID_URI_AUTHORITY: String = pkg + "InvalidUriAuthorityException" + } + + /** + * Exceptions whose messages name keys, classes or URIs and nothing else, so they are forwarded + * verbatim. Every other message (a `KeyProviderException` quotes the raw account key, custom + * providers say anything) is replaced by a pointer to the driver log. + */ + private val verbatimMessageClasses: Set[String] = Set( + Exceptions.CONFIGURATION_PROPERTY_NOT_FOUND, + Exceptions.INVALID_CONFIGURATION_VALUE, + Exceptions.INVALID_URI, + Exceptions.INVALID_URI_AUTHORITY, + classOf[ClassNotFoundException].getName, + classOf[NoSuchMethodException].getName, + classOf[NoSuchFieldException].getName) + + // An IllegalArgumentException is verbatim only when Hadoop or the JDK enum parser raised it + // ("No enum constant", "Failed to initialize", "Invalid account key."). One thrown + // by a user KeyProvider could carry anything, so it is reported by class only. + private val HADOOP_AZURE_PACKAGE = "org.apache.hadoop.fs.azurebfs." + private val JAVA_ENUM_CLASS = "java.lang.Enum" + + private def isVerbatim(t: Throwable): Boolean = t match { + case _: IllegalArgumentException => + t.getStackTrace.headOption.map(_.getClassName).exists { origin => + origin == JAVA_ENUM_CLASS || origin.startsWith(HADOOP_AZURE_PACKAGE) + } + case _ => verbatimMessageClasses.contains(t.getClass.getName) + } + + private val REDACTED = "<redacted>" + // Shorter values (`true`, a port) would mangle messages without protecting anything. + private val MIN_REDACTED_LENGTH = 8 + // `SimpleKeyProvider` (3.4.1+) puts the raw account key in its message as `key ="..."`. + private val QUOTED_KEY_PATTERN = "key =\"[^\"]*\"".r + private val MAX_CAUSE_DEPTH = 8 + + // The keys that name a mechanism or carry a credential in hadoop-azure. A configuration where + // none of them resolves for the account is `Unconfigured`; any one of them hands the decision to + // Hadoop, which defaults to SharedKey and fails in its own words when that does not fit. + private[objectstore] val plainAuthKeys: Seq[String] = Seq( + Keys.AUTH_TYPE, + Keys.KEY_PROVIDER, + Keys.SHELL_KEY_PROVIDER_SCRIPT, + Keys.OAUTH_PROVIDER_TYPE, + Keys.SAS_TOKEN_PROVIDER_TYPE) + private[objectstore] val passwordAuthKeys: Seq[String] = Seq( + Keys.ACCOUNT_KEY, + Keys.CLIENT_ID, + Keys.CLIENT_SECRET, + Keys.CLIENT_ENDPOINT, + Keys.MSI_TENANT, + Keys.MSI_ENDPOINT, + Keys.MSI_AUTHORITY, + Keys.USER_NAME, + Keys.USER_PASSWORD, + Keys.REFRESH_TOKEN, + Keys.REFRESH_TOKEN_ENDPOINT, + Keys.TOKEN_FILE, + Keys.CLIENT_ASSERTION_PROVIDER_TYPE, + Keys.SAS_FIXED_TOKEN) + + sealed trait Outcome + + /** Hadoop named a mechanism and every value it needs; `values` use the global key names. */ + final case class Resolved( + authType: String, + providerClass: Option[String], + values: Map[String, String]) + extends Outcome + + /** No auth-bearing key resolves for the account. */ + case object Unconfigured extends Outcome + + /** + * Hadoop named a mechanism the native scan cannot run; `providerClass` was never instantiated. + * `assertionProviderClass` is the trimmed value of + * `fs.azure.account.oauth2.client.assertion.provider.type` that a Workload Identity + * configuration names on hadoop-azure 3.5.0 and later. + */ + final case class Declined( + authType: String, + providerClass: Option[String], + assertionProviderClass: Option[String] = None) + extends Outcome + + /** Hadoop (or the reflection onto it) threw; `message` is already redacted. */ + final case class Failed(exceptionClass: String, message: String) extends Outcome + + def resolve(hadoopConf: Configuration, uri: URI): Outcome = + resolve(hadoopConf, uri, AbfsReflection.handles(AbfsReflection.loaderFor(hadoopConf))) + + private[objectstore] def resolve( + hadoopConf: Configuration, + uri: URI, + handles: Try[Handles]): Outcome = { + val secrets = ArrayBuffer[String]() + try { + handles match { + case Failure(e) => failed(uri, e, secrets) + case Success(h) => + authorityParts(uri) match { + case Left(failure) => failure + case Right((container, account)) => + val session = new Session( + h, + h.newAbfsConfiguration(hadoopConf, account, container, uri), + hadoopConf, + account, + secrets) + if (!session.isConfigured) { + Unconfigured + } else { + val authType: Enum[_] = h.getAuthType(session.abfsConf, account) + authType.name match { + case AuthTypes.SHARED_KEY => sharedKey(session, uri, account) + case AuthTypes.SAS => sas(session, authType) + case AuthTypes.OAUTH => oauth(session, authType) + // Custom and UserboundSASWithOAuth run Java providers per request; never + // instantiated here. + case other => Declined(other, None) + } + } + } + } + } catch { + // NonFatal alone misses NoSuchMethodError and friends, the reflection failures this must + // report as an error marker rather than let escape into planning. + case e @ (_: LinkageError | NonFatal(_)) => failed(uri, e, secrets) + } + } + + /** + * `(container, account)` exactly as `AzureBlobFileSystemStore.authorityParts`: the raw + * authority split once at the first `@`, so port, case and percent-encoding are kept. Failures + * carry Hadoop's exception class and message. + */ + private[objectstore] def authorityParts(uri: URI): Either[Failed, (String, String)] = { + Option(uri.getRawAuthority) match { + case None => + Left(Failed(Exceptions.INVALID_URI_AUTHORITY, s"$uri has invalid authority.")) + case Some(authority) if !authority.contains("@") => + Left(Failed(Exceptions.INVALID_URI_AUTHORITY, s"$uri has invalid authority.")) + case Some(authority) => + val parts = authority.split("@", 2) + if (parts.length < 2 || parts(0).isEmpty) { + Left( + Failed( + Exceptions.INVALID_URI, + s"Invalid URI '$uri' has a malformed authority, expected container name. " + + "Authority takes the form abfs://[<container name>@]<account name>")) + } else { + Right((parts(0), parts(1))) + } + } + } + + // One resolution's view of an AbfsConfiguration; records every credential it reads so a + // failure message can be redacted against them. + private final class Session( + val h: Handles, + val abfsConf: AnyRef, + hadoopConf: Configuration, + account: String, + secrets: ArrayBuffer[String]) { + + def password(key: String): Option[String] = { + val value = h.getPasswordString(abfsConf, key) + value.foreach(v => recordSecret(v)) + value + } + + // Both spellings, since Hadoop quotes the raw value and the providers use the trimmed one. + private def recordSecret(value: String): Unit = + Seq(value, value.trim).filter(_.length >= MIN_REDACTED_LENGTH).foreach(secrets += _) + + /** `getTrimmedPasswordString`: blank means the default, then trimmed. */ + def trimmedPassword(key: String, default: String): String = + password(key).filterNot(StringUtils.isBlank).getOrElse(default).trim + + def storageAccountKey(): String = { + val key = h.getStorageAccountKey(abfsConf) + recordSecret(key) + key + } + + def isConfigured: Boolean = { + // Every credential is read (no short circuit) so a later failure message is redacted + // against all of them. + val anyPassword = passwordAuthKeys.map(isSet(abfsConf, _, password = true)).contains(true) Review Comment: [P2] [P2] Avoid loading inactive credentials when the mechanism is already selected. With a local JCEKS credential provider, explicit `SharedKey` authentication and a valid account key in configuration, this eager `map` reads every OAuth/SAS password key before checking the auth type. Hadoop constructs and loads the keystore on every `getPassword` call. The exact-head resolver makes 48 such calls, although resolving the selected SharedKey requires only three. A warmed, seven-sample reproduction measured median option-extraction time of 449 ms at head versus 1.1 ms at base, with both calls succeeding. This adds roughly half a second of driver planning latency per native ABFS scan even when the keystore contains only unrelated credentials. The executor store cache cannot amortize it. Could presence detection short-circuit once the mechanism is configured, and defer additional redaction lookups until a failure occurs? Evidence: Offline reproduction: `/tmp/comet6059-review-final-ydfohbec/jvm/PlanningCost.scala`, results in `planning-cost-counted.log`. Configuration: `fs.azure.account.auth.type=SharedKey`, `fs.azure.account.key.myacct.dfs.core.windows.net=c2VjcmV0`, and `hadoop.security.credential.provider.path=jceks://file/<disposable>/test.jceks`. The synthetic keystore contains only an unrelated `fs.s3a.secret.key` entry. Invoke base/head `NativeConfig.extractObjectStoreOptions` for `abfss://[email protected]/path/file.parquet`. On Hadoop 3.4.2, a delegating `Configuration.getPassword` counter recorded base=0, head=48, direct Hadoop `getStorageAccountKey`=3. After three warmups, seven samples yielded medians 1.074 ms, 449.291 ms and 27.721 ms respectively. Source confirms each lookup calls `CredentialProviderFactory.getProviders`, constructs `JavaKeyStoreProvider` and loads the keystore. No storage or authentication requests were made. -- 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] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
