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() + } + } }