[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.

No, you can't use an assert. Read the comment please. The only way you should hit this code path is through GRPC. You can throw an unauthenticated grpc exception.

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.

How can the local server be restarted? What does that mean?

In Python, you call SparkSession.stop, it will only stop SparkContext but not the JVM itself in the test. So, later, when you restart SparkSession, you can't change the environment variable set to the JVM process.

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.

@HyukjinKwon is not about local usage or not. In theory this thing can be used for remote connections as well. Please give it a name that describes what it does, don't give it a name that describes how you want to use it.

In theory yes but we don't, and this is internal for Spark right?

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?

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

Yeah I think that's a good idea.

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

It is ensured by being blocked SparkConext with Spark Connect server is running. And then, here we create a Spark Connect client. So the token will always be there.

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.

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

@Kimahriman are you suggesting to use a different header than Authentication? If so, my issue with that is that folks will invariably use this to 'secure' remote connections as well, and in that case having TLS might not be such a bad thing.

We are planning to add UDS support in a follow-up. That should make this less of an issue.

No more the opposite, I assume the custom local_token header was used here to avoid the TLS requirement, but it's probably ok to allow non-TLS uses of the authorization bearer token header for a local connections, which the Python gRPC client even has support for. Adding the UDS supported would just improve that further. The local_channel_credentials even supports both of those cases, so you could use that with the authorization header to avoid the TLS requirement for this use case, while still requiring it for remote connections.

The "workarounds" are simply that you can still use the authorization header without TLS by using a custom interceptor that injects it in the metadata after the fact. Figured that out while trying to build a custom dynamic proxy for launching cluster deploy mode connect sessions to replace something like Livy.

HyukjinKwon · 2025-02-18T23:24:54Z

I am back. I am going back to work on this.

pan3793 · 2025-02-19T08:52:16Z

python/pyspark/sql/connect/session.py

@@ -1060,6 +1068,9 @@ def _start_connect_server(master: str, opts: Dict[str, Any]) -> None:
                overwrite_conf["spark.connect.grpc.binding.port"] = "0"

            origin_remote = os.environ.get("SPARK_REMOTE", None)
+            local_auth_token = str(uuid.uuid4())
+            SparkConnectClient._local_auth_token = local_auth_token
+            os.environ["SPARK_CONNECT_LOCAL_AUTH_TOKEN"] = local_auth_token


hmm... should it be enabled unconditionally?

currently, there are many spark services have no authN by default, this looks like an exception

Actually Py4J gateway server enables it by default :-).

…r locally

HyukjinKwon · 2025-02-20T00:21:50Z

#50006 replaces this. Thanks @Kimahriman

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

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

Kimahriman reviewed Feb 14, 2025

View reviewed changes

HyukjinKwon force-pushed the localauth branch 2 times, most recently from b7d7809 to d751d10 Compare February 19, 2025 08:30

pan3793 reviewed Feb 19, 2025

View reviewed changes

HyukjinKwon force-pushed the localauth branch from d751d10 to 49bdde7 Compare February 19, 2025 09:54

Provide a basic authentication token when running Spark Connect serve…

d14401a

…r locally

HyukjinKwon force-pushed the localauth branch 2 times, most recently from f76d420 to b86c198 Compare February 19, 2025 10:28

fixup

a3a349e

HyukjinKwon force-pushed the localauth branch from b86c198 to a3a349e Compare February 19, 2025 10:33

HyukjinKwon marked this pull request as draft February 19, 2025 12:46

HyukjinKwon mentioned this pull request Feb 19, 2025

[SPARK-51156][CONNECT] Static token authentication support in Spark Connect #50006

Closed

HyukjinKwon closed this Feb 20, 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

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

Uh oh!

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?

Uh oh!

dongjoon-hyun left a comment

Choose a reason for hiding this comment

Uh oh!

HyukjinKwon commented Feb 11, 2025

Uh oh!

HyukjinKwon commented Feb 11, 2025

Uh oh!

pan3793 commented Feb 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

HyukjinKwon commented Feb 11, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hvanhovell commented Feb 11, 2025

Uh oh!

HyukjinKwon commented Feb 12, 2025

Uh oh!

HyukjinKwon commented Feb 12, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

pan3793 commented Feb 11, 2025 •

edited

Loading