[SPARK-51156][CONNECT] Provide a basic authentication token when running Spark Connect server locally #49880

HyukjinKwon · 2025-02-11T06:09:39Z

What changes were proposed in this pull request?

This PR implements a simple authentication when running Spark Connect server locally.

Why are the changes needed?

To prevent security issues.

Does this PR introduce any user-facing change?

Yes. It requires the authentication token to access to the Spark Connect server.

How was this patch tested?

Enabled by default, and will be tested in CI.

Was this patch authored or co-authored using generative AI tooling?

No.

dongjoon-hyun

It's great to have the security feature. Are you going to add test cases?

HyukjinKwon · 2025-02-11T06:13:50Z

It's enabled by default .. so I think it's fine ...

HyukjinKwon · 2025-02-11T06:14:12Z

(and backward compat will be tested in the scheduled build )

pan3793 · 2025-02-11T08:48:15Z

Can AccessTokenCallCredentials added in SPARK-42533 protect the spark local connect server?

HyukjinKwon · 2025-02-11T10:26:25Z

This can secure individual connections themselves if I am not wrong .. but the problem is that any user can make a connection to the running server

HyukjinKwon · 2025-02-11T11:44:09Z

sql/connect/common/src/main/scala/org/apache/spark/sql/connect/SparkSession.scala

-          override def run(): Unit = if (server.isDefined) {
-            new ProcessBuilder(maybeConnectScript.get.toString)
-              .start()
+      server.synchronized {


Just added a synchronized.

vicennial · 2025-02-11T13:25:46Z

...er/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectAddArtifactsHandler.scala

@@ -52,6 +53,8 @@ class SparkConnectAddArtifactsHandler(val responseObserver: StreamObserver[AddAr
  private var holder: SessionHolder = _

  override def onNext(req: AddArtifactsRequest): Unit = try {
+    ConnectCommon.CONNECT_LOCAL_AUTH_TOKEN.foreach(k =>


Nit: would prefer this check to be centralised in some companion object

We can aslo log any auth failure

vicennial · 2025-02-11T13:26:17Z

...er/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectAddArtifactsHandler.scala

@@ -52,6 +53,8 @@ class SparkConnectAddArtifactsHandler(val responseObserver: StreamObserver[AddAr
  private var holder: SessionHolder = _

  override def onNext(req: AddArtifactsRequest): Unit = try {
+    ConnectCommon.CONNECT_LOCAL_AUTH_TOKEN.foreach(k =>
+      assert(k == req.getUserContext.getLocalAuthToken))


Do we want to throw a more descriptive message in these cases of auth failures?

Further, assertions can technically be disabled in the JVM. Doing so might lead to accidental removal of auth

vicennial · 2025-02-11T13:35:30Z

...connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala

+    ConnectCommon.CONNECT_LOCAL_AUTH_TOKEN = Option(
+      System.getenv(ConnectCommon.CONNECT_LOCAL_AUTH_TOKEN_ENV_NAME))


Would we need to do any special handling when the env var is set to "" i.e empty string?

sql/connect/common/src/main/protobuf/spark/connect/base.proto

...connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala

hvanhovell · 2025-02-11T15:08:52Z

...connect/common/src/main/scala/org/apache/spark/sql/connect/common/config/ConnectCommon.scala

@@ -21,4 +21,10 @@ private[sql] object ConnectCommon {
  val CONNECT_GRPC_PORT_MAX_RETRIES: Int = 0
  val CONNECT_GRPC_MAX_MESSAGE_SIZE: Int = 128 * 1024 * 1024
  val CONNECT_GRPC_MARSHALLER_RECURSION_LIMIT: Int = 1024
+  // Set only when we locally run Spark Connect server.
+  val CONNECT_LOCAL_AUTH_TOKEN_ENV_NAME = "SPARK_CONNECT_LOCAL_AUTH_TOKEN"
+  var CONNECT_LOCAL_AUTH_TOKEN: Option[String] = None


Why use a var for storage? Use the conf system instead?

Using confs can be shown when other users ps 😢

Isn't that true for any security related config if you pass it on the command line? And for this local use case won't it just be in memory when creating the SparkContext through py4j?

hvanhovell · 2025-02-11T20:32:15Z

@HyukjinKwon can you make sure we are use SSL/TLS in this case. Otherwise it will be kind of easy to intercept the token.

HyukjinKwon · 2025-02-12T01:35:00Z

let me take a look

HyukjinKwon · 2025-02-12T11:58:11Z

Some tests might fail ... need to go sleep .. but should be able to review this. I addressed most of major comments.

hvanhovell · 2025-02-12T13:50:23Z

...connect/common/src/main/scala/org/apache/spark/sql/connect/common/config/ConnectCommon.scala

+  def setLocalAuthToken(token: String): Unit = CONNECT_LOCAL_AUTH_TOKEN = Option(token)
+  def getLocalAuthToken: Option[String] = CONNECT_LOCAL_AUTH_TOKEN
+  def assertLocalAuthToken(token: Option[String]): Unit = token.foreach { t =>
+    assert(CONNECT_LOCAL_AUTH_TOKEN.isDefined)


Again, you CANNOT use asserts for this. They will get elided if you disable assertions. Please throw a proper gRPC exception in the LocalAuthInterceptor.

For other places, I will fix it later. But here assert is correct because if token is set, CONNECT_LOCAL_AUTH_TOKEN must be set to for local usage.

hvanhovell · 2025-02-12T13:53:17Z

...connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala

-    val configuredInterceptors = SparkConnectInterceptorRegistry.createConfiguredInterceptors()
+    val configuredInterceptors =
+      SparkConnectInterceptorRegistry.createConfiguredInterceptors() ++
+        (if (localAuthToken != null) Seq(new LocalAuthInterceptor()) else Nil)


Just pass in the token as an argument to the LocalAuthInterceptor, there is absolute no reason for putting this in some global variable...

hvanhovell · 2025-02-12T13:56:06Z

...connect/common/src/main/scala/org/apache/spark/sql/connect/common/config/ConnectCommon.scala

+    Metadata.Key.of("Authentication", Metadata.ASCII_STRING_MARSHALLER)
+
+  val CONNECT_LOCAL_AUTH_TOKEN_ENV_NAME = "SPARK_CONNECT_LOCAL_AUTH_TOKEN"
+  private var CONNECT_LOCAL_AUTH_TOKEN: Option[String] = Option(


Don't store this in some random local variable. There is no need for this. On the client side the SparkConnectClient will store the token. On the server the LocalAuthInterceptor should just hold on to the token.

The problem is that local server can stop and start without turning off the JVM (for Python) unlike that we always stop/start JVM for Scala. So it has to be a variable.

hvanhovell · 2025-02-12T13:58:41Z

...onnect/server/src/main/scala/org/apache/spark/sql/connect/service/LocalAuthInterceptor.scala

+/**
+ * A gRPC interceptor to check if the header contains token for authentication.
+ */
+class LocalAuthInterceptor extends ServerInterceptor {


How about we name this PreSharedKeyAuthenticationInterceptor? It is not a Local interceptor.

+1 in general for making this more generic and usable beyond the local Spark Connect use case. Having a pre-shared secret capability built-in goes a long way in making Spark Connect more usable in shared computer clusters.

I will scope it down to local usage for now. Whole this is internal for now, and we don't need to generalize them at this moment.

hvanhovell · 2025-02-12T14:03:30Z

sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala

@@ -422,7 +422,13 @@ object SparkConnectClient {
   * port or a NameResolver-compliant URI connection string.
   */
  class Builder(private var _configuration: Configuration) {
-    def this() = this(Configuration())


Please do this in SparkConnectClient.loadFromEnvironment()

hvanhovell · 2025-02-12T14:11:43Z

python/pyspark/sql/connect/session.py

@@ -1072,6 +1083,13 @@ def _start_connect_server(master: str, opts: Dict[str, Any]) -> None:
                conf.setAll(list(overwrite_conf.items())).setAll(list(default_conf.items()))
                PySparkSession(SparkContext.getOrCreate(conf))

+                # In Python local mode, session.stop does not terminate JVM itself


How about you set the environment variable when we start spark?

Also how do you ensure the LocalAuthInterceptor will installed if there is no token yet?

hvanhovell · 2025-02-12T14:12:50Z

sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala

+    def this() = this {
+      ConnectCommon.getLocalAuthToken
+        .map { _ =>
+          Configuration(token = ConnectCommon.getLocalAuthToken, isSslEnabled = Some(true))


Setting SSL enabled does not mean the server enforces SSL only.

HyukjinKwon · 2025-02-13T10:47:10Z

I think I can't just enable SSL by default. We should expose the certificate or use insecure connection. The access token API cannot be used with SSL it seems so I can't reuse this existing token either.

HyukjinKwon · 2025-02-13T10:48:01Z

I think we will likely miss RC1 - I will have to be away from keyboard like 3 days. Since technically CVE isn't filed yet, and this is an optional distribution, I think we can go ahead with RC 1. I will try to target RC 2.

…r locally

Kimahriman · 2025-02-14T11:42:32Z

python/pyspark/sql/connect/client/core.py

@@ -125,6 +125,7 @@ class ChannelBuilder:
    PARAM_USER_ID = "user_id"
    PARAM_USER_AGENT = "user_agent"
    PARAM_SESSION_ID = "session_id"
+    CONNECT_LOCAL_AUTH_TOKEN_PARAM_NAME = "local_token"


It's unfortunate that the grpc client tries to force you to use TLS if you want to use call credentials when there's so many workarounds like simply using a different header. Though in this case you could theoretically use local_channel_credentials at least on the Python side to use the built-in token mechanism

github-actions bot added SQL PYTHON CONNECT labels Feb 11, 2025

HyukjinKwon marked this pull request as draft February 11, 2025 06:11

dongjoon-hyun reviewed Feb 11, 2025

View reviewed changes

HyukjinKwon force-pushed the localauth branch 5 times, most recently from 018c25a to bd46da2 Compare February 11, 2025 08:26

HyukjinKwon force-pushed the localauth branch from bd46da2 to 1edae37 Compare February 11, 2025 11:43

HyukjinKwon commented Feb 11, 2025

View reviewed changes

HyukjinKwon marked this pull request as ready for review February 11, 2025 11:44

vicennial reviewed Feb 11, 2025

View reviewed changes

hvanhovell reviewed Feb 11, 2025

View reviewed changes

sql/connect/common/src/main/protobuf/spark/connect/base.proto Outdated Show resolved Hide resolved

hvanhovell reviewed Feb 11, 2025

View reviewed changes

...connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala Outdated Show resolved Hide resolved

hvanhovell reviewed Feb 11, 2025

View reviewed changes

HyukjinKwon force-pushed the localauth branch 7 times, most recently from a935c0b to d368742 Compare February 12, 2025 11:47

HyukjinKwon force-pushed the localauth branch from d368742 to 66336ac Compare February 12, 2025 11:56

hvanhovell reviewed Feb 12, 2025

View reviewed changes

Provide a basic authentication token when running Spark Connect serve…

0558d14

…r locally

HyukjinKwon force-pushed the localauth branch from 66336ac to 0558d14 Compare February 13, 2025 10:49

Kimahriman reviewed Feb 14, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[SPARK-51156][CONNECT] Provide a basic authentication token when running Spark Connect server locally #49880

[SPARK-51156][CONNECT] Provide a basic authentication token when running Spark Connect server locally #49880

HyukjinKwon commented Feb 11, 2025

dongjoon-hyun left a comment

HyukjinKwon commented Feb 11, 2025

HyukjinKwon commented Feb 11, 2025

pan3793 commented Feb 11, 2025 •

edited

Loading

HyukjinKwon commented Feb 11, 2025

HyukjinKwon Feb 11, 2025

vicennial Feb 11, 2025

vicennial Feb 11, 2025

vicennial Feb 11, 2025

vicennial Feb 11, 2025

vicennial Feb 11, 2025

hvanhovell Feb 11, 2025

HyukjinKwon Feb 12, 2025

Kimahriman Feb 12, 2025

hvanhovell commented Feb 11, 2025

HyukjinKwon commented Feb 12, 2025

HyukjinKwon commented Feb 12, 2025

hvanhovell Feb 12, 2025

HyukjinKwon Feb 13, 2025

hvanhovell Feb 12, 2025

hvanhovell Feb 12, 2025

HyukjinKwon Feb 13, 2025

hvanhovell Feb 12, 2025

Kimahriman Feb 12, 2025

HyukjinKwon Feb 13, 2025

hvanhovell Feb 12, 2025

hvanhovell Feb 12, 2025

hvanhovell Feb 12, 2025

hvanhovell Feb 12, 2025

HyukjinKwon commented Feb 13, 2025

HyukjinKwon commented Feb 13, 2025

Kimahriman Feb 14, 2025

		ConnectCommon.CONNECT_LOCAL_AUTH_TOKEN = Option(
		System.getenv(ConnectCommon.CONNECT_LOCAL_AUTH_TOKEN_ENV_NAME))

[SPARK-51156][CONNECT] Provide a basic authentication token when running Spark Connect server locally #49880

Are you sure you want to change the base?

[SPARK-51156][CONNECT] Provide a basic authentication token when running Spark Connect server locally #49880

Conversation

HyukjinKwon commented Feb 11, 2025

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

Was this patch authored or co-authored using generative AI tooling?

dongjoon-hyun left a comment

Choose a reason for hiding this comment

HyukjinKwon commented Feb 11, 2025

HyukjinKwon commented Feb 11, 2025

pan3793 commented Feb 11, 2025 • edited Loading

HyukjinKwon commented Feb 11, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

hvanhovell commented Feb 11, 2025

HyukjinKwon commented Feb 12, 2025

HyukjinKwon commented Feb 12, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

HyukjinKwon commented Feb 13, 2025

HyukjinKwon commented Feb 13, 2025

Choose a reason for hiding this comment

pan3793 commented Feb 11, 2025 •

edited

Loading