From 4464c339f10bfad67077c3b39d938d939ff34d6b Mon Sep 17 00:00:00 2001 From: Yuming Wang Date: Sat, 25 Jul 2026 12:23:06 +0800 Subject: [PATCH] [SPARK-58322][CORE] Propagate driver instance token via environment variable Replace SparkConf-based token propagation with environment variable: - Remove EXECUTOR_DRIVER_INSTANCE_TOKEN config builder and isExecutorStartupConf entry - Add ENV_EXECUTOR_DRIVER_INSTANCE_TOKEN constant and getDriverInstanceToken helper reading from System.getenv - Use conf.setExecutorEnv + sc.executorEnvs for propagation - Revert sparkProperties filter (token no longer in SparkConf) Generated-by: GLM 5.2. fix [SPARK-58322][CORE] Simplify executor identity verification to token-only Remove redundant appId verification, readiness retry logic, and DriverAppIdNotReadyException. The per-driver instance token alone provides the security property: it is generated before the driver RPC endpoint is created, so there is no startup race window. - Rename config to spark.executor.driverInstanceToken for redaction coverage (matches default regex 'token') - Add .version("4.3.0").withBindingPolicy(NOT_APPLICABLE) - Simplify RetrieveSparkAppConfigWithIdentity to 2 fields - Remove realApplicationId() from SchedulerBackend and all overrides - Remove registration retry executor and fetchSparkAppConfig retry loop - Remove EXECUTOR_APP_ID_ATTR, DriverAppIdNotReadyException - Revert @volatile/visibility changes on Yarn/StandaloneSchedulerBackend Net: -212 lines of production code and tests [SPARK-58322][CORE] Bound identity readiness retries by RPC timeout Use one monotonic deadline derived from spark.rpc.askTimeout for bootstrap and registration retries, rather than multiplying the RPC timeout by a retry count. Retry delays are capped by the deadline's remaining time. Extend regressions to hold the driver identity unavailable beyond the prior 300 ms limit. [SPARK-58322][CORE] Bind executors to a driver instance token Generate a random token for every coarse-grained driver instance and pass it as executor startup configuration across Standalone, YARN, and Kubernetes. Require the token together with the application ID for bootstrap and registration, preventing driver swaps even when an app ID or RPC secret is reused. Reject the legacy credential-bearing request, derive startup retry limits from the RPC timeout, and add missing config binding policy and tests. [SPARK-58322][CORE] Retry executor identity checks while driver starts Reject disabled identity verification without RPC authentication, so legacy bootstrap requests cannot expose credentials unauthenticated. Introduce a transient driver-ID-not-ready error and retry both bootstrap config fetches and executor registration before failing. Share bootstrap retry logic with the Kubernetes executor backend and cover both retry paths with tests. [SPARK-58322][CORE] PMC review: volatile, config switch, SECURITY.md, code cleanup - Add @volatile to YarnSchedulerBackend.appId for cross-thread visibility - Fix realApplicationId() Scaladoc to say 'reject' not 'skip', add @since 4.3.0 - Fix error message: clarify cluster manager retries, not the executor - Add spark.executor.identityVerification.enabled (default true) for backward compatibility with external cluster managers - Extract EXECUTOR_APP_ID_ATTR constant in CoarseGrainedClusterMessages - Make StandaloneSchedulerBackend.appId private[scheduler] to replace fragile reflection in tests - Extract replySparkAppConfig helper to eliminate code duplication - Update SECURITY.md with executor identity verification threat model [SPARK-58322][CORE] Address review: fail-closed, preserve ABI, add production sender test - P1: Fix NPE by using this.realApplicationId() instead of scheduler.backend - P1: Fail closed when realApplicationId() returns None (reject both config fetch and registration instead of silently accepting) - P1: Require verified appId for credential-bearing RetrieveSparkAppConfig; gate the old one-field message on RPC auth being enabled - P2: Preserve RetrieveSparkAppConfig ABI by reverting to one-field and adding a new RetrieveSparkAppConfigWithIdentity message - P2: Add production sender test exercising CoarseGrainedExecutorBackend.onStart to verify RegisterExecutor carries spark.app.id in attributes [SPARK-58322][CORE] Validate executor app ID before returning credentials and at registration This patch implements a two-layer application-identity check to prevent an executor from registering with the wrong driver: Layer 1: RetrieveSparkAppConfig now carries executor appId; the driver validates it before returning I/O encryption key and Hadoop delegation tokens. Layer 2: RegisterExecutor carries appId in its attributes map (preserving the 8-field ABI for external backends); the driver validates it at registration time as defense-in-depth against driver swap after config fetch. A new SchedulerBackend.realApplicationId() method replaces the fragile "spark-application-*" string-prefix heuristic. StandaloneSchedulerBackend and YarnSchedulerBackend override it to return None before their respective cluster-assigned IDs are available (connected() / bindToYarn()), so the check is cleanly skipped during the startup race window. Added 11 tests: config-fetch rejection/acceptance/null-compat, registration rejection/acceptance/no-attr-compat, driver swap, standalone fallback, YARN client race, production CoarseGrainedExecutorBackend.run() sender, and Kubernetes executor backend sender. [SPARK-58322][CORE] Clean up SPARK-58322 test: fix indentation, add try/finally [SPARK-58322][CORE] Verify executor app ID matches driver at registration Add appId to RegisterExecutor message and validate it in DriverEndpoint against scheduler.sc.applicationId. This binds the identity check to the registration connection itself, preventing the config-fetch-vs-registration race where a wrong driver could bind the same port between config retrieval and registration. The check replaces the executor-side verifyAppId which ran on a transient RPC connection that was shut down before registration. The driver-side check is authoritative since scheduler.sc.applicationId is always set before any executor can register. Null appId is accepted for backward compatibility with custom/external cluster managers that construct RegisterExecutor without appId. verify appid --- .../CoarseGrainedExecutorBackend.scala | 19 +- .../cluster/CoarseGrainedClusterMessage.scala | 15 ++ .../CoarseGrainedSchedulerBackend.scala | 59 ++++-- .../cluster/StandaloneSchedulerBackend.scala | 5 +- .../apache/spark/HeartbeatReceiverSuite.scala | 5 +- .../StandaloneDynamicAllocationSuite.scala | 10 +- .../CoarseGrainedExecutorBackendSuite.scala | 70 ++++++- .../CoarseGrainedSchedulerBackendSuite.scala | 176 +++++++++++++++++- .../k8s/KubernetesExecutorBackend.scala | 5 +- ...bernetesClusterSchedulerBackendSuite.scala | 67 ++++++- 10 files changed, 394 insertions(+), 37 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala b/core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala index 206a6a0fe385c..8e139ee40acab 100644 --- a/core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala +++ b/core/src/main/scala/org/apache/spark/executor/CoarseGrainedExecutorBackend.scala @@ -104,7 +104,8 @@ private[spark] class CoarseGrainedExecutorBackend( driver = Some(ref) env.executorBackend = Option(this) ref.ask[Boolean](RegisterExecutor(executorId, self, hostname, cores, extractLogUrls, - extractAttributes, _resources, resourceProfile.id)) + extractAttributes ++ Map(EXECUTOR_DRIVER_INSTANCE_TOKEN -> getDriverInstanceToken), + _resources, resourceProfile.id)) }(ThreadUtils.sameThread).onComplete { case Success(_) => self.send(RegisteredExecutor) @@ -396,6 +397,19 @@ private[spark] object CoarseGrainedExecutorBackend extends Logging { // registration request. case object RegisteredExecutor + /** + * Read the driver instance token from the environment. In production the + * cluster manager sets it as an environment variable via + * `spark.executorEnv.SPARK_EXECUTOR_DRIVER_INSTANCE_TOKEN`. The system + * property fallback exists for unit tests, which run in-process and cannot + * set environment variables. + */ + private[spark] def getDriverInstanceToken: String = { + Option(System.getenv(EXECUTOR_DRIVER_INSTANCE_TOKEN)) + .orElse(Option(System.getProperty(EXECUTOR_DRIVER_INSTANCE_TOKEN))) + .getOrElse(throw new SparkException("Driver instance token is not set.")) + } + case class Arguments( driverUrl: String, executorId: String, @@ -454,7 +468,8 @@ private[spark] object CoarseGrainedExecutorBackend extends Logging { } } - val cfg = driver.askSync[SparkAppConfig](RetrieveSparkAppConfig(arguments.resourceProfileId)) + val cfg = driver.askSync[SparkAppConfig]( + RetrieveSparkAppConfigWithIdentity(arguments.resourceProfileId, getDriverInstanceToken)) val props = cfg.sparkProperties ++ Seq[(String, String)](("spark.app.id", arguments.appId)) fetcher.shutdown() diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala index 1f452ae7d109d..697a8ea2e914d 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedClusterMessage.scala @@ -29,8 +29,23 @@ private[spark] sealed trait CoarseGrainedClusterMessage extends Serializable private[spark] object CoarseGrainedClusterMessages { + /** + * Name of the driver instance token. Used as both the environment variable + * name for propagating the token to executor processes (set via + * [[org.apache.spark.SparkConf.setExecutorEnv]]) and the attribute key in + * [[RegisterExecutor]]. The driver validates this token to prevent an + * executor from connecting to the wrong driver after a port-reuse / + * driver-swap scenario. + */ + val EXECUTOR_DRIVER_INSTANCE_TOKEN = "SPARK_EXECUTOR_DRIVER_INSTANCE_TOKEN" + case class RetrieveSparkAppConfig(resourceProfileId: Int) extends CoarseGrainedClusterMessage + case class RetrieveSparkAppConfigWithIdentity( + resourceProfileId: Int, + driverInstanceToken: String) + extends CoarseGrainedClusterMessage + case class SparkAppConfig( sparkProperties: Seq[(String, String)], ioEncryptionKey: Option[Array[Byte]], diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala index ad92e22424c77..3dc86adedc373 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/CoarseGrainedSchedulerBackend.scala @@ -17,6 +17,7 @@ package org.apache.spark.scheduler.cluster +import java.util.UUID import java.util.concurrent.{ScheduledExecutorService, TimeUnit} import java.util.concurrent.atomic.{AtomicInteger, AtomicReference} import javax.annotation.concurrent.GuardedBy @@ -26,7 +27,7 @@ import scala.concurrent.Future import com.google.common.cache.CacheBuilder -import org.apache.spark.{ExecutorAllocationClient, SparkEnv, TaskState} +import org.apache.spark.{ExecutorAllocationClient, SparkEnv, SparkException, TaskState} import org.apache.spark.deploy.SparkHadoopUtil import org.apache.spark.errors.SparkCoreErrors import org.apache.spark.executor.ExecutorLogUrlHandler @@ -62,6 +63,10 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp // Total number of executors that are currently registered protected val totalRegisteredExecutors = new AtomicInteger(0) protected val conf = scheduler.sc.conf + private[spark] val driverInstanceToken = UUID.randomUUID().toString + // Propagate the token as an environment variable to executor processes (YARN, Kubernetes). + // Standalone mode adds it to Command.environment directly in StandaloneSchedulerBackend. + conf.setExecutorEnv(EXECUTOR_DRIVER_INSTANCE_TOKEN, driverInstanceToken) private val maxRpcMessageSize = RpcUtils.maxMessageSizeBytes(conf) private val defaultAskTimeout = RpcUtils.askRpcTimeout(conf) // Submit tasks only after (registered resources / total expected resources) @@ -149,7 +154,9 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp // Spark configuration sent to executors. This is a lazy val so that subclasses of the // scheduler can modify the SparkConf object before this view is created. private lazy val sparkProperties = scheduler.sc.conf.getAll - .filter { case (k, _) => k.startsWith("spark.") } + .filter { case (k, _) => + k.startsWith("spark.") && k != s"spark.executorEnv.$EXECUTOR_DRIVER_INSTANCE_TOKEN" + } .toImmutableArraySeq private val logUrlHandler: ExecutorLogUrlHandler = new ExecutorLogUrlHandler( @@ -244,11 +251,21 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp logError(log"Received unexpected message. ${MDC(ERROR, e)}") } + private def verifyDriverInstanceToken(token: Option[String]): Boolean = { + token.filter(_.nonEmpty).contains(driverInstanceToken) + } + override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = { case RegisterExecutor(executorId, executorRef, hostname, cores, logUrls, attributes, resources, resourceProfileId) => - if (executorDataMap.contains(executorId)) { + if (!verifyDriverInstanceToken(attributes.get(EXECUTOR_DRIVER_INSTANCE_TOKEN))) { + val msg = "Executor did not supply a matching driver instance token. " + + "This likely means the executor connected to the wrong driver." + logWarning(log"Rejecting executor ${MDC(LogKeys.EXECUTOR_ID, executorId)}: " + + log"${MDC(ERROR, msg)}") + context.sendFailure(new SparkException(msg)) + } else if (executorDataMap.contains(executorId)) { context.sendFailure(new IllegalStateException(s"Duplicate executor ID: $executorId")) } else if (scheduler.excludedNodes().contains(hostname) || isExecutorExcluded(executorId, hostname)) { @@ -276,6 +293,10 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp val resourcesInfo = resources.map { case (rName, info) => (info.name, new ExecutorResourceInfo(info.name, info.addresses.toIndexedSeq)) } + // Strip the driver instance token from attributes before storing or publishing. + // The token is only needed for verification; downstream listeners (event logging, + // REST API, status store) would otherwise expose it verbatim. + val redactedAttributes = attributes.removed(EXECUTOR_DRIVER_INSTANCE_TOKEN) // If we've requested the executor figure out when we did. val reqTs: Option[Long] = CoarseGrainedSchedulerBackend.this.synchronized { execRequestTimes.get(resourceProfileId).flatMap { @@ -294,7 +315,7 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp } val data = new ExecutorData(executorRef, executorAddress, hostname, - 0, cores, logUrlHandler.applyPattern(logUrls, attributes), attributes, + 0, cores, logUrlHandler.applyPattern(logUrls, redactedAttributes), redactedAttributes, resourcesInfo, resourceProfileId, registrationTs = System.currentTimeMillis(), requestTs = reqTs) // This must be synchronized because variables mutated @@ -353,14 +374,28 @@ class CoarseGrainedSchedulerBackend(scheduler: TaskSchedulerImpl, val rpcEnv: Rp triggeredByExecutor = true)) case RetrieveSparkAppConfig(resourceProfileId) => - val rp = scheduler.sc.resourceProfileManager.resourceProfileFromId(resourceProfileId) - val reply = SparkAppConfig( - sparkProperties, - SparkEnv.get.securityManager.getIOEncryptionKey(), - Option(delegationTokens.get()), - rp, - currentLogLevel) - context.reply(reply) + val msg = "Legacy RetrieveSparkAppConfig request rejected: it does not carry a " + + "driver instance token, so driver identity cannot be verified. " + + "Use RetrieveSparkAppConfigWithIdentity instead." + logWarning(msg) + context.sendFailure(new SparkException(msg)) + + case RetrieveSparkAppConfigWithIdentity(resourceProfileId, token) => + if (!verifyDriverInstanceToken(Option(token))) { + val msg = "RetrieveSparkAppConfigWithIdentity rejected: executor did not supply " + + "a matching driver instance token." + logWarning(log"${MDC(ERROR, msg)}") + context.sendFailure(new SparkException(msg)) + } else { + val rp = scheduler.sc.resourceProfileManager.resourceProfileFromId(resourceProfileId) + val reply = SparkAppConfig( + sparkProperties, + SparkEnv.get.securityManager.getIOEncryptionKey(), + Option(delegationTokens.get()), + rp, + currentLogLevel) + context.reply(reply) + } case IsExecutorAlive(executorId) => context.reply(isExecutorActive(executorId)) diff --git a/core/src/main/scala/org/apache/spark/scheduler/cluster/StandaloneSchedulerBackend.scala b/core/src/main/scala/org/apache/spark/scheduler/cluster/StandaloneSchedulerBackend.scala index 061b54914c839..d4f907c6cb363 100644 --- a/core/src/main/scala/org/apache/spark/scheduler/cluster/StandaloneSchedulerBackend.scala +++ b/core/src/main/scala/org/apache/spark/scheduler/cluster/StandaloneSchedulerBackend.scala @@ -34,7 +34,7 @@ import org.apache.spark.launcher.{LauncherBackend, SparkAppHandle} import org.apache.spark.resource.ResourceProfile import org.apache.spark.rpc.{RpcAddress, RpcEndpointAddress} import org.apache.spark.scheduler._ -import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages.RemoveExecutor +import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages.{EXECUTOR_DRIVER_INSTANCE_TOKEN, RemoveExecutor} import org.apache.spark.util.{ThreadUtils, Utils} import org.apache.spark.util.ArrayImplicits._ @@ -113,7 +113,8 @@ private[spark] class StandaloneSchedulerBackend( val sparkJavaOpts = Utils.sparkJavaOpts(conf, SparkConf.isExecutorStartupConf) val javaOpts = sparkJavaOpts ++ extraJavaOpts val command = Command("org.apache.spark.executor.CoarseGrainedExecutorBackend", - args, sc.executorEnvs, classPathEntries ++ testingClassPath, libraryPathEntries, javaOpts) + args, sc.executorEnvs.toMap + (EXECUTOR_DRIVER_INSTANCE_TOKEN -> driverInstanceToken), + classPathEntries ++ testingClassPath, libraryPathEntries, javaOpts) val webUrl = sc.ui.map(_.webUrl).getOrElse("") val coresPerExecutor = conf.getOption(config.EXECUTOR_CORES.key).map(_.toInt) // If we're using dynamic allocation, set our initial executor limit to 0 for now. diff --git a/core/src/test/scala/org/apache/spark/HeartbeatReceiverSuite.scala b/core/src/test/scala/org/apache/spark/HeartbeatReceiverSuite.scala index 7d1b8f8ea508b..3a520ff462ee0 100644 --- a/core/src/test/scala/org/apache/spark/HeartbeatReceiverSuite.scala +++ b/core/src/test/scala/org/apache/spark/HeartbeatReceiverSuite.scala @@ -173,15 +173,16 @@ class HeartbeatReceiverSuite // Register fake executors with our fake scheduler backend // This is necessary because the backend refuses to kill executors it does not know about fakeSchedulerBackend.start() + val tokenAttrs = Map(EXECUTOR_DRIVER_INSTANCE_TOKEN -> fakeSchedulerBackend.driverInstanceToken) val dummyExecutorEndpoint1 = new FakeExecutorEndpoint(rpcEnv) val dummyExecutorEndpoint2 = new FakeExecutorEndpoint(rpcEnv) val dummyExecutorEndpointRef1 = rpcEnv.setupEndpoint("fake-executor-1", dummyExecutorEndpoint1) val dummyExecutorEndpointRef2 = rpcEnv.setupEndpoint("fake-executor-2", dummyExecutorEndpoint2) fakeSchedulerBackend.driverEndpoint.askSync[Boolean]( - RegisterExecutor(executorId1, dummyExecutorEndpointRef1, "1.2.3.4", 0, Map.empty, Map.empty, + RegisterExecutor(executorId1, dummyExecutorEndpointRef1, "1.2.3.4", 0, Map.empty, tokenAttrs, Map.empty, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) fakeSchedulerBackend.driverEndpoint.askSync[Boolean]( - RegisterExecutor(executorId2, dummyExecutorEndpointRef2, "1.2.3.5", 0, Map.empty, Map.empty, + RegisterExecutor(executorId2, dummyExecutorEndpointRef2, "1.2.3.5", 0, Map.empty, tokenAttrs, Map.empty, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) heartbeatReceiverRef.askSync[Boolean](TaskSchedulerIsSet) addExecutorAndVerify(executorId1) diff --git a/core/src/test/scala/org/apache/spark/deploy/StandaloneDynamicAllocationSuite.scala b/core/src/test/scala/org/apache/spark/deploy/StandaloneDynamicAllocationSuite.scala index 90ef0aa510c24..5fd596356f4ed 100644 --- a/core/src/test/scala/org/apache/spark/deploy/StandaloneDynamicAllocationSuite.scala +++ b/core/src/test/scala/org/apache/spark/deploy/StandaloneDynamicAllocationSuite.scala @@ -35,7 +35,7 @@ import org.apache.spark.resource.ResourceProfile import org.apache.spark.rpc.{RpcAddress, RpcEndpointRef, RpcEnv} import org.apache.spark.scheduler.TaskSchedulerImpl import org.apache.spark.scheduler.cluster._ -import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages.{LaunchedExecutor, RegisterExecutor} +import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages.{EXECUTOR_DRIVER_INSTANCE_TOKEN, LaunchedExecutor, RegisterExecutor} /** * End-to-end tests for dynamic allocation in standalone mode. @@ -504,8 +504,6 @@ class StandaloneDynamicAllocationSuite val endpointRef = mock(classOf[RpcEndpointRef]) val mockAddress = mock(classOf[RpcAddress]) when(endpointRef.address).thenReturn(mockAddress) - val message = RegisterExecutor("one", endpointRef, "excluded-host", 10, Map.empty, - Map.empty, Map.empty, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID) val taskScheduler = mock(classOf[TaskSchedulerImpl]) when(taskScheduler.excludedNodes()).thenReturn(Set("excluded-host")) @@ -517,6 +515,9 @@ class StandaloneDynamicAllocationSuite val scheduler = new CoarseGrainedSchedulerBackend(taskScheduler, rpcEnv) try { scheduler.start() + val message = RegisterExecutor("one", endpointRef, "excluded-host", 10, Map.empty, + Map(EXECUTOR_DRIVER_INSTANCE_TOKEN -> scheduler.driverInstanceToken), + Map.empty, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID) val e = intercept[SparkException] { scheduler.driverEndpoint.askSync[Boolean](message) } @@ -627,7 +628,8 @@ class StandaloneDynamicAllocationSuite val endpointRef = mock(classOf[RpcEndpointRef]) val mockAddress = mock(classOf[RpcAddress]) when(endpointRef.address).thenReturn(mockAddress) - val message = RegisterExecutor(id, endpointRef, "localhost", 10, Map.empty, Map.empty, + val message = RegisterExecutor(id, endpointRef, "localhost", 10, Map.empty, + Map(EXECUTOR_DRIVER_INSTANCE_TOKEN -> backend.driverInstanceToken), Map.empty, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID) backend.driverEndpoint.askSync[Boolean](message) backend.driverEndpoint.send(LaunchedExecutor(id)) diff --git a/core/src/test/scala/org/apache/spark/executor/CoarseGrainedExecutorBackendSuite.scala b/core/src/test/scala/org/apache/spark/executor/CoarseGrainedExecutorBackendSuite.scala index a038d8e8613ef..06f2c4fa1c5a8 100644 --- a/core/src/test/scala/org/apache/spark/executor/CoarseGrainedExecutorBackendSuite.scala +++ b/core/src/test/scala/org/apache/spark/executor/CoarseGrainedExecutorBackendSuite.scala @@ -22,6 +22,7 @@ import java.nio.ByteBuffer import java.util.Properties import java.util.concurrent.ConcurrentHashMap import java.util.concurrent.atomic.AtomicInteger +import java.util.concurrent.atomic.AtomicReference import scala.collection.concurrent.TrieMap import scala.concurrent.duration._ @@ -41,9 +42,9 @@ import org.apache.spark.internal.config.{EXECUTOR_MEMORY, PLUGINS} import org.apache.spark.resource._ import org.apache.spark.resource.ResourceUtils._ import org.apache.spark.resource.TestResourceIDs._ -import org.apache.spark.rpc.RpcEnv +import org.apache.spark.rpc.{RpcCallContext, RpcEndpoint, RpcEnv} import org.apache.spark.scheduler.{SparkListener, SparkListenerExecutorAdded, SparkListenerExecutorRemoved, TaskDescription} -import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages.{KillTask, LaunchTask} +import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages.{EXECUTOR_DRIVER_INSTANCE_TOKEN, KillTask, LaunchTask, RegisterExecutor} import org.apache.spark.serializer.JavaSerializer import org.apache.spark.util.{SerializableBuffer, SslTestUtils, ThreadUtils, Utils} @@ -615,6 +616,71 @@ class CoarseGrainedExecutorBackendSuite extends SparkFunSuite } } + test("SPARK-58322: CoarseGrainedExecutorBackend sends token in RegisterExecutor") { + val testToken = "driver-instance-token-test" + + val driverRpcEnv = RpcEnv.create( + "test-driver", "localhost", 0, new SparkConf(), + new SecurityManager(new SparkConf()), clientMode = false) + + val capturedRegister = new AtomicReference[RegisterExecutor]() + val oldToken = System.getProperty(EXECUTOR_DRIVER_INSTANCE_TOKEN) + try { + driverRpcEnv.setupEndpoint("CoarseGrainedScheduler", new RpcEndpoint { + override val rpcEnv: RpcEnv = driverRpcEnv + override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = { + case msg: RegisterExecutor => + capturedRegister.set(msg) + context.sendFailure(new RuntimeException("test: captured RegisterExecutor")) + } + }) + + System.setProperty(EXECUTOR_DRIVER_INSTANCE_TOKEN, testToken) + val executorConf = new SparkConf() + .set(EXECUTOR_MEMORY.key, "512m") + val executorRpcEnv = RpcEnv.create( + "test-executor", "localhost", 0, executorConf, + new SecurityManager(executorConf), clientMode = true) + val serializer = new JavaSerializer(executorConf) + val env = createMockEnv(executorConf, serializer, Some(executorRpcEnv)) + val resourceProfile = ResourceProfile.getOrCreateDefaultProfile(executorConf) + + val backend = new CoarseGrainedExecutorBackend( + executorRpcEnv, + s"spark://CoarseGrainedScheduler@localhost:${driverRpcEnv.address.port}", + "1", "localhost", "localhost", 1, env, None, resourceProfile) { + override protected def exitExecutor(code: Int, reason: String, + throwable: Throwable = null, notifyDriver: Boolean = true): Unit = { + throw new RuntimeException(s"Test exit prevented: $reason") + } + } + executorRpcEnv.setupEndpoint("Executor", backend) + + try { + eventually(timeout(30.seconds)) { + val msg = capturedRegister.get() + assert(msg != null, "RegisterExecutor was not received by the driver") + assert(msg.attributes.contains(EXECUTOR_DRIVER_INSTANCE_TOKEN), + s"RegisterExecutor attributes should contain " + + s"$EXECUTOR_DRIVER_INSTANCE_TOKEN: ${msg.attributes}") + assert(msg.attributes(EXECUTOR_DRIVER_INSTANCE_TOKEN) == testToken, + s"Expected $EXECUTOR_DRIVER_INSTANCE_TOKEN=$testToken, " + + s"got ${msg.attributes(EXECUTOR_DRIVER_INSTANCE_TOKEN)}") + } + } finally { + executorRpcEnv.shutdown() + } + } finally { + driverRpcEnv.shutdown() + SparkEnv.set(null) + if (oldToken == null) { + System.clearProperty(EXECUTOR_DRIVER_INSTANCE_TOKEN) + } else { + System.setProperty(EXECUTOR_DRIVER_INSTANCE_TOKEN, oldToken) + } + } + } + private def createMockEnv(conf: SparkConf, serializer: JavaSerializer, rpcEnv: Option[RpcEnv] = None): SparkEnv = { val mockEnv = mock[SparkEnv] diff --git a/core/src/test/scala/org/apache/spark/scheduler/CoarseGrainedSchedulerBackendSuite.scala b/core/src/test/scala/org/apache/spark/scheduler/CoarseGrainedSchedulerBackendSuite.scala index 127e4b2e413ce..c1b866589c1b6 100644 --- a/core/src/test/scala/org/apache/spark/scheduler/CoarseGrainedSchedulerBackendSuite.scala +++ b/core/src/test/scala/org/apache/spark/scheduler/CoarseGrainedSchedulerBackendSuite.scala @@ -41,7 +41,7 @@ import org.apache.spark.resource.{ExecutorResourceRequests, ResourceInformation, import org.apache.spark.resource.ResourceAmountUtils.ONE_ENTIRE_RESOURCE import org.apache.spark.resource.ResourceUtils._ import org.apache.spark.resource.TestResourceIDs._ -import org.apache.spark.rpc.{RpcAddress, RpcEndpointRef, RpcEnv, RpcTimeout} +import org.apache.spark.rpc.{RpcAddress, RpcCallContext, RpcEndpoint, RpcEndpointRef, RpcEnv, RpcTimeout} import org.apache.spark.scheduler.cluster.{CoarseGrainedSchedulerBackend, ExecutorInfo} import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages._ import org.apache.spark.util.{RpcUtils, SerializableBuffer, Utils} @@ -249,7 +249,8 @@ class CoarseGrainedSchedulerBackendSuite extends SparkFunSuite with LocalSparkCo "CLUSTER_ID" -> "cl1", "USER" -> "dummy", "CONTAINER_ID" -> "container1", - "LOG_FILES" -> "stdout,stderr") + "LOG_FILES" -> "stdout,stderr", + EXECUTOR_DRIVER_INSTANCE_TOKEN -> backend.driverInstanceToken) val baseUrl = s"http://newhost:9999/logs/clusters/${attributes("CLUSTER_ID")}" + s"/users/${attributes("USER")}/containers/${attributes("CONTAINER_ID")}" @@ -325,13 +326,19 @@ class CoarseGrainedSchedulerBackendSuite extends SparkFunSuite with LocalSparkCo sc.addSparkListener(listener) backend.driverEndpoint.askSync[Boolean]( - RegisterExecutor("1", mockEndpointRef, mockAddress.host, 1, Map.empty, Map.empty, resources, + RegisterExecutor("1", mockEndpointRef, mockAddress.host, 1, Map.empty, + tokenAttrs(backend), + resources, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) backend.driverEndpoint.askSync[Boolean]( - RegisterExecutor("2", mockEndpointRef, mockAddress.host, 1, Map.empty, Map.empty, resources, + RegisterExecutor("2", mockEndpointRef, mockAddress.host, 1, Map.empty, + tokenAttrs(backend), + resources, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) backend.driverEndpoint.askSync[Boolean]( - RegisterExecutor("3", mockEndpointRef, mockAddress.host, 1, Map.empty, Map.empty, resources, + RegisterExecutor("3", mockEndpointRef, mockAddress.host, 1, Map.empty, + tokenAttrs(backend), + resources, rp.id)) val frameSize = RpcUtils.maxMessageSizeBytes(sc.conf) @@ -432,13 +439,19 @@ class CoarseGrainedSchedulerBackendSuite extends SparkFunSuite with LocalSparkCo sc.addSparkListener(listener) backend.driverEndpoint.askSync[Boolean]( - RegisterExecutor("1", mockEndpointRef, mockAddress.host, 1, Map.empty, Map.empty, resources, + RegisterExecutor("1", mockEndpointRef, mockAddress.host, 1, Map.empty, + tokenAttrs(backend), + resources, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) backend.driverEndpoint.askSync[Boolean]( - RegisterExecutor("2", mockEndpointRef, mockAddress.host, 1, Map.empty, Map.empty, resources, + RegisterExecutor("2", mockEndpointRef, mockAddress.host, 1, Map.empty, + tokenAttrs(backend), + resources, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) backend.driverEndpoint.askSync[Boolean]( - RegisterExecutor("3", mockEndpointRef, mockAddress.host, 1, Map.empty, Map.empty, resources, + RegisterExecutor("3", mockEndpointRef, mockAddress.host, 1, Map.empty, + tokenAttrs(backend), + resources, rp.id)) val frameSize = RpcUtils.maxMessageSizeBytes(sc.conf) @@ -531,7 +544,8 @@ class CoarseGrainedSchedulerBackendSuite extends SparkFunSuite with LocalSparkCo val ts = backend.getTaskSchedulerImpl() when(ts.resourceOffers(any[IndexedSeq[WorkerOffer]], any[Boolean])).thenReturn(Seq.empty) backend.driverEndpoint.askSync[Boolean]( - RegisterExecutor("1", mockEndpointRef, mockAddress.host, execCores, Map.empty, Map.empty, + RegisterExecutor("1", mockEndpointRef, mockAddress.host, execCores, Map.empty, + tokenAttrs(backend), Map.empty, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) backend.driverEndpoint.send(LaunchedExecutor("1")) eventually(timeout(5 seconds)) { @@ -599,13 +613,155 @@ class CoarseGrainedSchedulerBackendSuite extends SparkFunSuite with LocalSparkCo assert(!mockEndpointRef.decommissionReceived) backend.driverEndpoint.askSync[Boolean]( - RegisterExecutor("1", mockEndpointRef, mockAddress.host, 1, Map(), Map(), + RegisterExecutor("1", mockEndpointRef, mockAddress.host, 1, Map(), + tokenAttrs(backend), Map.empty, ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) sc.listenerBus.waitUntilEmpty(executorUpTimeout.toMillis) assert(mockEndpointRef.decommissionReceived) } + private def withBackend( + body: (SparkContext, CoarseGrainedSchedulerBackend) => Unit): Unit = { + val conf = new SparkConf() + .setMaster("local-cluster[0, 3, 1024]") + .setAppName("test") + sc = new SparkContext(conf) + body(sc, sc.schedulerBackend.asInstanceOf[CoarseGrainedSchedulerBackend]) + } + + private def tokenAttrs(backend: CoarseGrainedSchedulerBackend): Map[String, String] = { + Map(EXECUTOR_DRIVER_INSTANCE_TOKEN -> backend.driverInstanceToken) + } + + test("SPARK-58322: reject RetrieveSparkAppConfigWithIdentity with stale token") { + withBackend { (_, backend) => + val ex = intercept[SparkException] { + backend.driverEndpoint.askSync[SparkAppConfig]( + RetrieveSparkAppConfigWithIdentity( + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, + "stale-token")) + } + assert(ex.getCause.getMessage.contains( + "did not supply a matching driver instance token")) + } + } + + test("SPARK-58322: accept RetrieveSparkAppConfigWithIdentity with matching token") { + withBackend { (sc, backend) => + val cfg = backend.driverEndpoint.askSync[SparkAppConfig]( + RetrieveSparkAppConfigWithIdentity( + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, + backend.driverInstanceToken)) + assert(cfg.sparkProperties.find(_._1 == "spark.app.id").isDefined) + // The driver instance token must not be echoed back in the response. + assert(cfg.sparkProperties.forall { case (k, _) => + k != s"spark.executorEnv.$EXECUTOR_DRIVER_INSTANCE_TOKEN" }, + s"Token should not be present in sparkProperties: ${cfg.sparkProperties}") + } + } + + test("SPARK-58322: reject RetrieveSparkAppConfig without identity") { + withBackend { (_, backend) => + val ex = intercept[SparkException] { + backend.driverEndpoint.askSync[SparkAppConfig]( + RetrieveSparkAppConfig(ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) + } + assert(ex.getCause.getMessage.contains( + "Legacy RetrieveSparkAppConfig request rejected")) + } + } + + test("SPARK-58322: reject RegisterExecutor with stale token in attributes") { + withBackend { (_, backend) => + val mockEndpointRef = mock[RpcEndpointRef] + val mockAddress = mock[RpcAddress] + val ex = intercept[SparkException] { + backend.driverEndpoint.askSync[Boolean]( + RegisterExecutor("1111", mockEndpointRef, mockAddress.host, 1, Map.empty, + Map(EXECUTOR_DRIVER_INSTANCE_TOKEN -> "stale-token"), Map.empty, + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) + } + assert(ex.getCause.getMessage.contains( + "did not supply a matching driver instance token")) + } + } + + test("SPARK-58322: accept RegisterExecutor with matching token in attributes") { + withBackend { (sc, backend) => + val mockEndpointRef = mock[RpcEndpointRef] + val mockAddress = mock[RpcAddress] + val result = backend.driverEndpoint.askSync[Boolean]( + RegisterExecutor("1112", mockEndpointRef, mockAddress.host, 1, Map.empty, + tokenAttrs(backend), Map.empty, + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) + assert(result) + } + } + + test("SPARK-58322: reject RegisterExecutor without identity attributes") { + withBackend { (_, backend) => + val mockEndpointRef = mock[RpcEndpointRef] + val mockAddress = mock[RpcAddress] + val ex = intercept[SparkException] { + backend.driverEndpoint.askSync[Boolean]( + RegisterExecutor("1113", mockEndpointRef, mockAddress.host, 1, Map.empty, + Map.empty, Map.empty, + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) + } + assert(ex.getCause.getMessage.contains( + "did not supply a matching driver instance token")) + } + } + + test("SPARK-58322: reject RegisterExecutor with stale token after driver swap") { + // Port-reuse scenario: executor fetches config from driver A (token + // validated by fake driver A), then driver A dies and driver B binds + // the same address. The executor sends RegisterExecutor to driver B + // with driver A's token in attributes -- driver B rejects it. + val driverAToken = "driver-A-instance-token" + + val driverARpcEnv = RpcEnv.create("test-driverA", "localhost", 0, new SparkConf(), + new SecurityManager(new SparkConf()), clientMode = false) + try { + driverARpcEnv.setupEndpoint("fake-driverA", new RpcEndpoint { + override val rpcEnv: RpcEnv = driverARpcEnv + override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = { + case RetrieveSparkAppConfigWithIdentity(_, _) => + context.reply(SparkAppConfig( + Seq("spark.app.id" -> "app-driver-A"), + None, None, ResourceProfile.getOrCreateDefaultProfile(new SparkConf()), None)) + } + }) + + // Executor fetches config from driver A (token passes validation). + val driverARef = driverARpcEnv.setupEndpointRefByURI( + s"spark://fake-driverA@localhost:${driverARpcEnv.address.port}") + val cfg = driverARef.askSync[SparkAppConfig]( + RetrieveSparkAppConfigWithIdentity( + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, driverAToken)) + assert(cfg.sparkProperties.find(_._1 == "spark.app.id").isDefined) + + // Now create driver B (a real SparkContext with a different token). + withBackend { (_, backend) => + // Executor sends RegisterExecutor to driver B with driver A's token. + val mockEndpointRef = mock[RpcEndpointRef] + val mockAddress = mock[RpcAddress] + val ex = intercept[SparkException] { + backend.driverEndpoint.askSync[Boolean]( + RegisterExecutor("1114", mockEndpointRef, mockAddress.host, 1, Map.empty, + Map(EXECUTOR_DRIVER_INSTANCE_TOKEN -> driverAToken), + Map.empty, + ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID)) + } + assert(ex.getCause.getMessage.contains( + "did not supply a matching driver instance token")) + } + } finally { + driverARpcEnv.shutdown() + } + } + private def testSubmitJob(sc: SparkContext, rdd: RDD[Int]): Unit = { sc.submitJob( rdd, diff --git a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesExecutorBackend.scala b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesExecutorBackend.scala index e44d7e29ef606..10026f66a1d7f 100644 --- a/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesExecutorBackend.scala +++ b/resource-managers/kubernetes/core/src/main/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesExecutorBackend.scala @@ -94,7 +94,10 @@ private[spark] object KubernetesExecutorBackend extends Logging { } } - val cfg = driver.askSync[SparkAppConfig](RetrieveSparkAppConfig(arguments.resourceProfileId)) + val cfg = driver.askSync[SparkAppConfig]( + RetrieveSparkAppConfigWithIdentity( + arguments.resourceProfileId, + CoarseGrainedExecutorBackend.getDriverInstanceToken)) val props = cfg.sparkProperties ++ Seq[(String, String)](("spark.app.id", arguments.appId)) val execId: String = arguments.executorId match { case null | "EXECID" | "" => diff --git a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesClusterSchedulerBackendSuite.scala b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesClusterSchedulerBackendSuite.scala index cf172eb096d47..a5eff5ee9370f 100644 --- a/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesClusterSchedulerBackendSuite.scala +++ b/resource-managers/kubernetes/core/src/test/scala/org/apache/spark/scheduler/cluster/k8s/KubernetesClusterSchedulerBackendSuite.scala @@ -18,6 +18,7 @@ package org.apache.spark.scheduler.cluster.k8s import java.util.Arrays import java.util.concurrent.TimeUnit +import java.util.concurrent.atomic.AtomicReference import scala.jdk.CollectionConverters._ @@ -31,14 +32,14 @@ import org.mockito.ArgumentMatchers.{any, eq => mockitoEq} import org.mockito.Mockito.{atLeastOnce, mock, never, spy, verify, when} import org.scalatest.BeforeAndAfter -import org.apache.spark.{SparkConf, SparkContext, SparkEnv, SparkFunSuite} +import org.apache.spark.{SecurityManager, SparkConf, SparkContext, SparkEnv, SparkFunSuite} import org.apache.spark.deploy.k8s.Config._ import org.apache.spark.deploy.k8s.Constants._ import org.apache.spark.deploy.k8s.Fabric8Aliases._ import org.apache.spark.resource.{ResourceProfile, ResourceProfileManager} import org.apache.spark.rpc.{RpcCallContext, RpcEndpoint, RpcEndpointRef, RpcEnv} import org.apache.spark.scheduler.{ExecutorKilled, ExecutorLossReason, LiveListenerBus, TaskSchedulerImpl} -import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages.{RegisterExecutor, RemoveExecutor, StopDriver} +import org.apache.spark.scheduler.cluster.CoarseGrainedClusterMessages.{EXECUTOR_DRIVER_INSTANCE_TOKEN, RegisterExecutor, RemoveExecutor, RetrieveSparkAppConfigWithIdentity, SparkAppConfig, StopDriver} import org.apache.spark.scheduler.cluster.CoarseGrainedSchedulerBackend import org.apache.spark.scheduler.cluster.k8s.ExecutorLifecycleTestUtils.TEST_SPARK_APP_ID @@ -350,4 +351,66 @@ class KubernetesClusterSchedulerBackendSuite extends SparkFunSuite with BeforeAn assert(id1 === id2, "applicationId() must return the same value on repeated calls") assert(id1.startsWith("spark-"), "generated app ID should have the spark- prefix") } + + test("SPARK-58322: Kubernetes executor backend carries driver instance token in " + + "RetrieveSparkAppConfigWithIdentity") { + val testAppId = "app-k8s-production-test" + val testToken = "driver-instance-token-k8s-test" + + val driverRpcEnv = RpcEnv.create("test-k8s-driver", "localhost", 0, new SparkConf(), + new SecurityManager(new SparkConf()), clientMode = false) + val capturedMsg = new AtomicReference[RetrieveSparkAppConfigWithIdentity]() + try { + driverRpcEnv.setupEndpoint("CoarseGrainedScheduler", new RpcEndpoint { + override val rpcEnv: RpcEnv = driverRpcEnv + override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = { + case msg: RetrieveSparkAppConfigWithIdentity => + capturedMsg.set(msg) + context.reply(SparkAppConfig( + Seq("spark.app.id" -> testAppId), + None, None, ResourceProfile.getOrCreateDefaultProfile(new SparkConf()), None)) + } + }) + + val args = KubernetesExecutorBackend.Arguments( + driverUrl = s"spark://CoarseGrainedScheduler@localhost:${driverRpcEnv.address.port}", + executorId = "1", + bindAddress = "localhost", + hostname = "localhost", + cores = 1, + appId = testAppId, + workerUrl = None, + resourcesFileOpt = None, + resourceProfileId = ResourceProfile.DEFAULT_RESOURCE_PROFILE_ID, + podName = "test-pod") + + val prevEnv = SparkEnv.get + val oldToken = System.getProperty(EXECUTOR_DRIVER_INSTANCE_TOKEN) + try { + System.setProperty(EXECUTOR_DRIVER_INSTANCE_TOKEN, testToken) + intercept[Exception] { + KubernetesExecutorBackend.run(args, + (_, _, _, _, _) => throw new RuntimeException("expected")) + } + } finally { + if (oldToken == null) { + System.clearProperty(EXECUTOR_DRIVER_INSTANCE_TOKEN) + } else { + System.setProperty(EXECUTOR_DRIVER_INSTANCE_TOKEN, oldToken) + } + val currentEnv = SparkEnv.get + if (currentEnv != null && (currentEnv ne prevEnv)) { + currentEnv.stop() + } + SparkEnv.set(prevEnv) + } + + val msg = capturedMsg.get() + assert(msg != null, "RetrieveSparkAppConfigWithIdentity was not received by the driver") + assert(msg.driverInstanceToken == testToken, + s"Expected driverInstanceToken=$testToken, got ${msg.driverInstanceToken}") + } finally { + driverRpcEnv.shutdown() + } + } }