You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@celeborn.apache.org by zh...@apache.org on 2023/12/14 14:53:00 UTC
(incubator-celeborn) branch main updated: [CELEBORN-1157] Add client-side support for Sasl Authentication in the transport layer
This is an automated email from the ASF dual-hosted git repository.
zhouky pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new a03ce6c16 [CELEBORN-1157] Add client-side support for Sasl Authentication in the transport layer
a03ce6c16 is described below
commit a03ce6c16597a485387e38491139b696455fada7
Author: Chandni Singh <si...@gmail.com>
AuthorDate: Thu Dec 14 22:52:49 2023 +0800
[CELEBORN-1157] Add client-side support for Sasl Authentication in the transport layer
### What changes were proposed in this pull request?
This adds the client side Sasl authentication support in the transport layer. Most of this code is taken from Apache Spark.
### Why are the changes needed?
The changes are needed for adding authentication to Celeborn. See [CELEBORN-1011](https://issues.apache.org/jira/browse/CELEBORN-1011).
### Does this PR introduce _any_ user-facing change?
Added a configuration for Sasl request timeout
### How was this patch tested?
Will be adding `CelebornSaslSuiteJ.java` (https://github.com/apache/incubator-celeborn/pull/2105) that tests the end-to-end Sasl flow.
Closes #2139 from otterc/CELEBORN-1157.
Authored-by: Chandni Singh <si...@gmail.com>
Signed-off-by: zky.zhoukeyong <zk...@alibaba-inc.com>
---
.../common/network/client/TransportClient.java | 23 +++
.../common/network/protocol/TransportMessage.java | 3 +
.../common/network/sasl/CelebornSaslClient.java | 165 +++++++++++++++++++++
.../common/network/sasl/SaslClientBootstrap.java | 120 +++++++++++++++
.../common/network/sasl/SaslCredentials.java | 42 ++++++
.../common/network/sasl/SaslTimeoutException.java | 33 +++++
.../celeborn/common/network/sasl/SaslUtils.java | 56 +++++++
.../common/network/util/TransportConf.java | 5 +
common/src/main/proto/TransportMessages.proto | 17 +++
.../org/apache/celeborn/common/CelebornConf.scala | 13 ++
docs/configuration/network.md | 1 +
11 files changed, 478 insertions(+)
diff --git a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java
index 482356b04..b5796fe98 100644
--- a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClient.java
@@ -24,6 +24,8 @@ import java.nio.ByteBuffer;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
+import javax.annotation.Nullable;
+
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.util.concurrent.SettableFuture;
@@ -79,6 +81,7 @@ public class TransportClient implements Closeable {
private final Channel channel;
private final TransportResponseHandler handler;
private volatile boolean timedOut;
+ @Nullable private String clientId;
public TransportClient(Channel channel, TransportResponseHandler handler) {
this.channel = Preconditions.checkNotNull(channel);
@@ -317,6 +320,7 @@ public class TransportClient implements Closeable {
public String toString() {
return new ToStringBuilder(this, ToStringStyle.SHORT_PREFIX_STYLE)
.append("remoteAdress", channel.remoteAddress())
+ .append("clientId", clientId)
.append("isActive", isActive())
.toString();
}
@@ -327,6 +331,25 @@ public class TransportClient implements Closeable {
return counter.getAndIncrement();
}
+ /**
+ * Returns the ID used by the client to authenticate itself when authentication is enabled.
+ *
+ * @return The client ID, or null if authentication is disabled.
+ */
+ public String getClientId() {
+ return clientId;
+ }
+
+ /**
+ * Sets the authenticated client ID. This is meant to be used by the authentication layer.
+ *
+ * <p>Trying to set a different client ID after it's been set will result in an exception.
+ */
+ public void setClientId(String id) {
+ Preconditions.checkState(clientId == null, "Client ID has already been set.");
+ this.clientId = id;
+ }
+
public class StdChannelListener implements GenericFutureListener<Future<? super Void>> {
final long startTime;
final Object requestId;
diff --git a/common/src/main/java/org/apache/celeborn/common/network/protocol/TransportMessage.java b/common/src/main/java/org/apache/celeborn/common/network/protocol/TransportMessage.java
index a716d7ded..0dbc93e66 100644
--- a/common/src/main/java/org/apache/celeborn/common/network/protocol/TransportMessage.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/protocol/TransportMessage.java
@@ -41,6 +41,7 @@ import org.apache.celeborn.common.protocol.PbRegionFinish;
import org.apache.celeborn.common.protocol.PbRegionStart;
import org.apache.celeborn.common.protocol.PbReportShuffleFetchFailure;
import org.apache.celeborn.common.protocol.PbReportShuffleFetchFailureResponse;
+import org.apache.celeborn.common.protocol.PbSaslRequest;
import org.apache.celeborn.common.protocol.PbStreamChunkSlice;
import org.apache.celeborn.common.protocol.PbStreamHandler;
import org.apache.celeborn.common.protocol.PbTransportableError;
@@ -102,6 +103,8 @@ public class TransportMessage implements Serializable {
return (T) PbReportShuffleFetchFailure.parseFrom(payload);
case REPORT_SHUFFLE_FETCH_FAILURE_RESPONSE_VALUE:
return (T) PbReportShuffleFetchFailureResponse.parseFrom(payload);
+ case SASL_REQUEST_VALUE:
+ return (T) PbSaslRequest.parseFrom(payload);
default:
logger.error("Unexpected type {}", type);
}
diff --git a/common/src/main/java/org/apache/celeborn/common/network/sasl/CelebornSaslClient.java b/common/src/main/java/org/apache/celeborn/common/network/sasl/CelebornSaslClient.java
new file mode 100644
index 000000000..96a3ee39a
--- /dev/null
+++ b/common/src/main/java/org/apache/celeborn/common/network/sasl/CelebornSaslClient.java
@@ -0,0 +1,165 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.common.network.sasl;
+
+import static org.apache.celeborn.common.network.sasl.SaslUtils.*;
+
+import java.util.Map;
+
+import javax.annotation.Nullable;
+import javax.annotation.concurrent.GuardedBy;
+import javax.security.auth.callback.Callback;
+import javax.security.auth.callback.CallbackHandler;
+import javax.security.auth.callback.NameCallback;
+import javax.security.auth.callback.PasswordCallback;
+import javax.security.auth.callback.UnsupportedCallbackException;
+import javax.security.sasl.RealmCallback;
+import javax.security.sasl.RealmChoiceCallback;
+import javax.security.sasl.Sasl;
+import javax.security.sasl.SaslClient;
+import javax.security.sasl.SaslException;
+
+import com.google.common.base.Preconditions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A SASL Client for Celeborn which simply keeps track of the state of a single SASL session, from
+ * the initial state to the "authenticated" state. This client initializes the protocol via a
+ * firstToken, which is then followed by a set of challenges and responses.
+ *
+ * <p>TODO: Currently, we have hardcoded Sasl mechanism to be DIGEST-MD5 for connection auth. We
+ * should make this configurable in the future (see:
+ * https://issues.apache.org/jira/browse/CELEBORN-1158).
+ */
+public class CelebornSaslClient {
+ private static final Logger logger = LoggerFactory.getLogger(CelebornSaslClient.class);
+
+ @GuardedBy("this")
+ private SaslClient saslClient;
+
+ public CelebornSaslClient(
+ String saslMechanism,
+ @Nullable Map<String, String> saslProps,
+ @Nullable CallbackHandler authCallbackHandler) {
+ Preconditions.checkNotNull(saslMechanism);
+ try {
+ this.saslClient =
+ Sasl.createSaslClient(
+ new String[] {saslMechanism},
+ null,
+ null,
+ DEFAULT_REALM,
+ saslProps,
+ authCallbackHandler);
+ } catch (SaslException e) {
+ throw new IllegalArgumentException(e);
+ }
+ }
+
+ /** Used to initiate SASL handshake with server. */
+ public synchronized byte[] firstToken() {
+ if (saslClient != null && saslClient.hasInitialResponse()) {
+ try {
+ return saslClient.evaluateChallenge(EMPTY_BYTE_ARRAY);
+ } catch (SaslException e) {
+ throw new IllegalArgumentException(e);
+ }
+ } else {
+ return EMPTY_BYTE_ARRAY;
+ }
+ }
+
+ /** Determines whether the authentication exchange has completed. */
+ public synchronized boolean isComplete() {
+ return saslClient != null && saslClient.isComplete();
+ }
+
+ /** Returns the value of a negotiated property. */
+ public synchronized Object getNegotiatedProperty(String name) {
+ return saslClient.getNegotiatedProperty(name);
+ }
+
+ /**
+ * Respond to server's SASL token.
+ *
+ * @param token contains server's SASL token
+ * @return client's response SASL token
+ */
+ public synchronized byte[] response(byte[] token) {
+ try {
+ return saslClient != null ? saslClient.evaluateChallenge(token) : EMPTY_BYTE_ARRAY;
+ } catch (SaslException e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ /**
+ * Disposes of any system resources or security-sensitive information the SaslClient might be
+ * using.
+ */
+ public synchronized void dispose() {
+ if (saslClient != null) {
+ try {
+ saslClient.dispose();
+ } catch (SaslException e) {
+ // ignore
+ } finally {
+ saslClient = null;
+ }
+ }
+ }
+
+ /**
+ * Implementation of javax.security.auth.callback.CallbackHandler that works with share secrets.
+ */
+ static class ClientCallbackHandler implements CallbackHandler {
+
+ private final String id;
+ private final String password;
+
+ public ClientCallbackHandler(String id, String password) {
+ this.id = Preconditions.checkNotNull(id, "id");
+ this.password = Preconditions.checkNotNull(password, "password");
+ }
+
+ @Override
+ public void handle(Callback[] callbacks) throws UnsupportedCallbackException {
+
+ for (Callback callback : callbacks) {
+ if (callback instanceof NameCallback) {
+ logger.trace("SASL client callback: setting username");
+ NameCallback nc = (NameCallback) callback;
+ nc.setName(encodeIdentifier(id));
+ } else if (callback instanceof PasswordCallback) {
+ logger.trace("SASL client callback: setting password");
+ PasswordCallback pc = (PasswordCallback) callback;
+ pc.setPassword(encodePassword(password));
+ } else if (callback instanceof RealmCallback) {
+ logger.trace("SASL client callback: setting realm");
+ RealmCallback rc = (RealmCallback) callback;
+ rc.setText(rc.getDefaultText());
+ } else if (callback instanceof RealmChoiceCallback) {
+ // ignore (?)
+ } else {
+ throw new UnsupportedCallbackException(callback, "Unrecognized callback: " + callback);
+ }
+ }
+ }
+ }
+}
diff --git a/common/src/main/java/org/apache/celeborn/common/network/sasl/SaslClientBootstrap.java b/common/src/main/java/org/apache/celeborn/common/network/sasl/SaslClientBootstrap.java
new file mode 100644
index 000000000..99a8054b2
--- /dev/null
+++ b/common/src/main/java/org/apache/celeborn/common/network/sasl/SaslClientBootstrap.java
@@ -0,0 +1,120 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.common.network.sasl;
+
+import static org.apache.celeborn.common.network.sasl.SaslUtils.*;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.concurrent.TimeoutException;
+
+import com.google.common.base.Preconditions;
+import com.google.protobuf.ByteString;
+import io.netty.channel.Channel;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.common.network.client.TransportClient;
+import org.apache.celeborn.common.network.client.TransportClientBootstrap;
+import org.apache.celeborn.common.network.protocol.TransportMessage;
+import org.apache.celeborn.common.network.util.TransportConf;
+import org.apache.celeborn.common.protocol.MessageType;
+import org.apache.celeborn.common.protocol.PbAuthType;
+import org.apache.celeborn.common.protocol.PbSaslRequest;
+import org.apache.celeborn.common.util.JavaUtils;
+
+/**
+ * Bootstraps a {@link TransportClient} by performing SASL authentication on the connection. The
+ * server should be setup with a {@code SaslRpcHandler} with matching keys for the given appId.
+ */
+public class SaslClientBootstrap implements TransportClientBootstrap {
+ private static final Logger logger = LoggerFactory.getLogger(SaslClientBootstrap.class);
+
+ private final TransportConf conf;
+ private final String appId;
+ private final SaslCredentials saslCredentials;
+
+ /**
+ * Creates a new SaslClientBootstrap. When this is used for authenticating a client connection,
+ * the user id in the saslCredentials should be the same as the appId.
+ *
+ * @param conf transport conf
+ * @param appId application id
+ * @param saslCredentials sasl userId and password.
+ */
+ public SaslClientBootstrap(TransportConf conf, String appId, SaslCredentials saslCredentials) {
+ this.conf = conf;
+ this.appId = appId;
+ this.saslCredentials = Preconditions.checkNotNull(saslCredentials);
+ }
+
+ /**
+ * Performs SASL authentication by sending a token, and then proceeding with the SASL
+ * challenge-response tokens until we either successfully authenticate or throw an exception due
+ * to mismatch.
+ */
+ @Override
+ public void doBootstrap(TransportClient client, Channel channel) {
+ // TODO: Hardcoding the SASL mechanism to DIGEST-MD5 for Connection Authentication. This
+ // should be configurable in the future.
+ CelebornSaslClient saslClient =
+ new CelebornSaslClient(
+ DIGEST_MD5,
+ DEFAULT_SASL_CLIENT_PROPS,
+ new CelebornSaslClient.ClientCallbackHandler(
+ saslCredentials.getUserId(), saslCredentials.getPassword()));
+ try {
+ byte[] payload = saslClient.firstToken();
+ boolean firstToken = true;
+ while (!saslClient.isComplete()) {
+ PbSaslRequest.Builder builder = PbSaslRequest.newBuilder();
+ if (firstToken) {
+ builder.setMethod(DIGEST_MD5).setAuthType(PbAuthType.CONNECTION_AUTH);
+ }
+ TransportMessage msg =
+ new TransportMessage(
+ MessageType.SASL_REQUEST,
+ builder.setPayload(ByteString.copyFrom(payload)).build().toByteArray());
+ ByteBuffer response;
+ try {
+ response = client.sendRpcSync(msg.toByteBuffer(), conf.saslTimeoutMs());
+ } catch (RuntimeException ex) {
+ // We know it is a Sasl timeout here if it is a TimeoutException.
+ if (ex.getCause() instanceof TimeoutException) {
+ throw new SaslTimeoutException(ex.getCause());
+ } else {
+ throw ex;
+ }
+ }
+ // The response of a SaslMessage is either a RpcResponse or a RpcFailure.
+ payload = saslClient.response(JavaUtils.bufferToArray(response));
+ firstToken = false;
+ }
+ client.setClientId(appId);
+ } catch (IOException ioe) {
+ throw new RuntimeException(ioe);
+ } finally {
+ try {
+ // Once authentication is complete, the server will trust all remaining communication.
+ saslClient.dispose();
+ } catch (RuntimeException e) {
+ logger.error("Error while disposing SASL client", e);
+ }
+ }
+ }
+}
diff --git a/common/src/main/java/org/apache/celeborn/common/network/sasl/SaslCredentials.java b/common/src/main/java/org/apache/celeborn/common/network/sasl/SaslCredentials.java
new file mode 100644
index 000000000..10bf6d7c7
--- /dev/null
+++ b/common/src/main/java/org/apache/celeborn/common/network/sasl/SaslCredentials.java
@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.common.network.sasl;
+
+import com.google.common.base.Preconditions;
+
+/** A class to encapsulate SASL credentials. */
+public class SaslCredentials {
+
+ /** For client to service authentication, this should be set to unique applicationId. */
+ private final String userId;
+
+ private final String password;
+
+ public SaslCredentials(String userId, String password) {
+ this.userId = Preconditions.checkNotNull(userId);
+ this.password = Preconditions.checkNotNull(password);
+ }
+
+ public String getUserId() {
+ return userId;
+ }
+
+ public String getPassword() {
+ return password;
+ }
+}
diff --git a/common/src/main/java/org/apache/celeborn/common/network/sasl/SaslTimeoutException.java b/common/src/main/java/org/apache/celeborn/common/network/sasl/SaslTimeoutException.java
new file mode 100644
index 000000000..c8b74f29d
--- /dev/null
+++ b/common/src/main/java/org/apache/celeborn/common/network/sasl/SaslTimeoutException.java
@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.common.network.sasl;
+
+/** An exception thrown if there is a SASL timeout. */
+public class SaslTimeoutException extends RuntimeException {
+ public SaslTimeoutException(Throwable cause) {
+ super(cause);
+ }
+
+ public SaslTimeoutException(String message) {
+ super(message);
+ }
+
+ public SaslTimeoutException(String message, Throwable cause) {
+ super(message, cause);
+ }
+}
diff --git a/common/src/main/java/org/apache/celeborn/common/network/sasl/SaslUtils.java b/common/src/main/java/org/apache/celeborn/common/network/sasl/SaslUtils.java
new file mode 100644
index 000000000..5d57057e4
--- /dev/null
+++ b/common/src/main/java/org/apache/celeborn/common/network/sasl/SaslUtils.java
@@ -0,0 +1,56 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.common.network.sasl;
+
+import java.nio.charset.StandardCharsets;
+import java.util.Base64;
+import java.util.Map;
+
+import javax.security.sasl.Sasl;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableMap;
+
+public class SaslUtils {
+ static final byte[] EMPTY_BYTE_ARRAY = new byte[0];
+
+ /** Sasl Mechanisms */
+ static final String DIGEST_MD5 = "DIGEST-MD5";
+
+ /** Quality of protection value that does not include encryption. */
+ static final String QOP_AUTH = "auth";
+
+ static final String DEFAULT_REALM = "default";
+
+ static final Map<String, String> DEFAULT_SASL_CLIENT_PROPS =
+ ImmutableMap.<String, String>builder().put(Sasl.QOP, QOP_AUTH).build();
+
+ /* Encode a byte[] identifier as a Base64-encoded string. */
+ static String encodeIdentifier(String identifier) {
+ Preconditions.checkNotNull(identifier, "User cannot be null if SASL is enabled");
+ return Base64.getEncoder().encodeToString(identifier.getBytes(StandardCharsets.UTF_8));
+ }
+
+ /** Encode a password as a base64-encoded char[] array. */
+ static char[] encodePassword(String password) {
+ Preconditions.checkNotNull(password, "Password cannot be null if SASL is enabled");
+ return Base64.getEncoder()
+ .encodeToString(password.getBytes(StandardCharsets.UTF_8))
+ .toCharArray();
+ }
+}
diff --git a/common/src/main/java/org/apache/celeborn/common/network/util/TransportConf.java b/common/src/main/java/org/apache/celeborn/common/network/util/TransportConf.java
index 468b70685..a93b545dc 100644
--- a/common/src/main/java/org/apache/celeborn/common/network/util/TransportConf.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/util/TransportConf.java
@@ -153,4 +153,9 @@ public class TransportConf {
public long clientHeartbeatInterval() {
return celebornConf.clientHeartbeatInterval(module);
}
+
+ /** Timeout for a single round trip of sasl message exchange, in milliseconds. */
+ public int saslTimeoutMs() {
+ return celebornConf.networkIoSaslTimoutMs(module);
+ }
}
diff --git a/common/src/main/proto/TransportMessages.proto b/common/src/main/proto/TransportMessages.proto
index 4fd2bbc6b..17fc93e3f 100644
--- a/common/src/main/proto/TransportMessages.proto
+++ b/common/src/main/proto/TransportMessages.proto
@@ -91,6 +91,7 @@ enum MessageType {
REPORT_SHUFFLE_FETCH_FAILURE_RESPONSE = 68;
GET_SHUFFLE_ID = 69;
GET_SHUFFLE_ID_RESPONSE = 70;
+ SASL_REQUEST = 71;
}
enum StreamType {
@@ -611,3 +612,19 @@ message PbTransportableError {
int64 streamId = 1;
string message = 2;
}
+
+enum PbAuthType {
+ CLIENT_AUTH = 0;
+ CONNECTION_AUTH = 1;
+}
+
+message PbSaslMechanism {
+ string mechanism = 1;
+ repeated PbAuthType authTypes = 2;
+}
+
+message PbSaslRequest {
+ string method = 1;
+ PbAuthType authType = 2;
+ bytes payload = 3;
+}
diff --git a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
index f5ed4745c..508171fe9 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -512,6 +512,11 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable with Logging with Se
def maxDefaultNettyThreads: Int = get(MAX_DEFAULT_NETTY_THREADS)
+ def networkIoSaslTimoutMs(module: String): Int = {
+ val key = NETWORK_IO_SASL_TIMEOUT.key.replace("<module>", module)
+ getTimeAsMs(key, s"${networkTimeout.duration.toMillis}ms").toInt
+ }
+
// //////////////////////////////////////////////////////
// Master //
// //////////////////////////////////////////////////////
@@ -4218,4 +4223,12 @@ object CelebornConf extends Logging {
.doc("Whether to filter excluded worker when register shuffle.")
.booleanConf
.createWithDefault(false)
+
+ val NETWORK_IO_SASL_TIMEOUT: ConfigEntry[Long] =
+ buildConf("celeborn.<module>.io.saslTimeout")
+ .categories("network")
+ .doc("Timeout for a single round trip of auth message exchange, in milliseconds.")
+ .version("0.4.0")
+ .timeConf(TimeUnit.MILLISECONDS)
+ .createWithDefaultString("30s")
}
diff --git a/docs/configuration/network.md b/docs/configuration/network.md
index 241f19836..677447b44 100644
--- a/docs/configuration/network.md
+++ b/docs/configuration/network.md
@@ -34,6 +34,7 @@ license: |
| celeborn.<module>.io.preferDirectBufs | true | If true, we will prefer allocating off-heap byte buffers within Netty. | |
| celeborn.<module>.io.receiveBuffer | 0b | Receive buffer size (SO_RCVBUF). Note: the optimal size for receive buffer and send buffer should be latency * network_bandwidth. Assuming latency = 1ms, network_bandwidth = 10Gbps buffer size should be ~ 1.25MB. | 0.2.0 |
| celeborn.<module>.io.retryWait | 5s | Time that we will wait in order to perform a retry after an IOException. Only relevant if maxIORetries > 0. | 0.2.0 |
+| celeborn.<module>.io.saslTimeout | 30s | Timeout for a single round trip of auth message exchange, in milliseconds. | 0.4.0 |
| celeborn.<module>.io.sendBuffer | 0b | Send buffer size (SO_SNDBUF). | 0.2.0 |
| celeborn.<module>.io.serverThreads | 0 | Number of threads used in the server thread pool. Default to 0, which is 2x#cores. | |
| celeborn.<module>.push.timeoutCheck.interval | 5s | Interval for checking push data timeout. If setting <module> to `data`, it works for shuffle client push data and should be configured on client side. If setting <module> to `replicate`, it works for worker replicate data to peer worker and should be configured on worker side. | 0.3.0 |