[SPARK-58322][CORE] Verify executor driver identity via per-instance token

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

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]],

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))
 

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.

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)

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))

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]