diff --git a/bin/interpreter.sh b/bin/interpreter.sh index 00ff0307313..07c18be4997 100755 --- a/bin/interpreter.sh +++ b/bin/interpreter.sh @@ -146,7 +146,7 @@ if [[ -z "$ZEPPELIN_IMPERSONATE_CMD" ]]; then ZEPPELIN_IMPERSONATE_RUN_CMD=("ssh" "${ZEPPELIN_IMPERSONATE_USER}@localhost") fi else - ZEPPELIN_IMPERSONATE_RUN_CMD=$(eval "echo ${ZEPPELIN_IMPERSONATE_CMD} ") + ZEPPELIN_IMPERSONATE_RUN_CMD=("${ZEPPELIN_IMPERSONATE_CMD[@]}") fi @@ -249,6 +249,10 @@ if [[ -n "$ZEPPELIN_IMPERSONATE_USER" ]]; then if [[ -f "${ZEPPELIN_CONF_DIR}/zeppelin-env.sh" ]]; then COMMAND_STRING+="source ${ZEPPELIN_CONF_DIR}/zeppelin-env.sh; " fi + + # Pass the callback credential over stdin so it is neither exposed in the + # impersonation command line nor lost when ssh/sudo changes the environment. + COMMAND_STRING+="IFS= read -r ZEPPELIN_INTERPRETER_EVENT_TOKEN; export ZEPPELIN_INTERPRETER_EVENT_TOKEN; " # Add interpreter command to the command string IFS=' ' read -r -a JAVA_INTP_OPTS_ARRAY <<< "${JAVA_INTP_OPTS}" @@ -256,7 +260,8 @@ if [[ -n "$ZEPPELIN_IMPERSONATE_USER" ]]; then COMMAND_STRING+="${ZEPPELIN_RUNNER} ${JAVA_INTP_OPTS_ARRAY[@]} ${ZEPPELIN_INTP_MEM_ARRAY[@]} -cp '${ZEPPELIN_INTP_CLASSPATH_OVERRIDES}:${ZEPPELIN_INTP_CLASSPATH}' ${ZEPPELIN_SERVER} ${CALLBACK_HOST} ${PORT} ${INTP_GROUP_ID} ${INTP_PORT}" # Set INTERPRETER_RUN_COMMAND with the impersonation command and command string - INTERPRETER_RUN_COMMAND=("${ZEPPELIN_IMPERSONATE_CMD[@]}" "${COMMAND_STRING}") + INTERPRETER_RUN_COMMAND=("${ZEPPELIN_IMPERSONATE_RUN_CMD[@]}" "${COMMAND_STRING}") + ZEPPELIN_INTERPRETER_EVENT_TOKEN_VIA_STDIN=true fi fi @@ -280,5 +285,13 @@ fi # Don't remove this echo, it is for diagnose, this line of output will be redirected to java log4j output. # Output that starts with `[INFO]` will be redirected to log4j INFO output. Other outputs from interpreter.sh # will be redirected to log4j DEBUG output. -echo "[INFO] Interpreter launch command: ${INTERPRETER_RUN_COMMAND[@]}" -exec "${INTERPRETER_RUN_COMMAND[@]}" +INTERPRETER_LOG_COMMAND="${INTERPRETER_RUN_COMMAND[*]}" +if [[ -n "${ZEPPELIN_INTERPRETER_EVENT_TOKEN}" ]]; then + INTERPRETER_LOG_COMMAND="${INTERPRETER_LOG_COMMAND//${ZEPPELIN_INTERPRETER_EVENT_TOKEN}/[REDACTED]}" +fi +echo "[INFO] Interpreter launch command: ${INTERPRETER_LOG_COMMAND}" +if [[ "${ZEPPELIN_INTERPRETER_EVENT_TOKEN_VIA_STDIN}" == "true" ]]; then + exec "${INTERPRETER_RUN_COMMAND[@]}" <<< "${ZEPPELIN_INTERPRETER_EVENT_TOKEN}" +else + exec "${INTERPRETER_RUN_COMMAND[@]}" +fi diff --git a/helium-dev/src/test/java/org/apache/zeppelin/helium/ZeppelinDevServerTest.java b/helium-dev/src/test/java/org/apache/zeppelin/helium/ZeppelinDevServerTest.java new file mode 100644 index 00000000000..25a0f95762c --- /dev/null +++ b/helium-dev/src/test/java/org/apache/zeppelin/helium/ZeppelinDevServerTest.java @@ -0,0 +1,52 @@ +/* + * 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.zeppelin.helium; + +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterUtils; +import org.junit.jupiter.api.Test; + +import java.util.HashMap; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class ZeppelinDevServerTest { + + @Test + void startsWithoutInterpreterCallbackCredentials() throws Exception { + int port = RemoteInterpreterUtils.findRandomAvailablePortOnAllLocalInterfaces(); + ZeppelinDevServer server = new ZeppelinDevServer(port); + try { + server.start(); + long deadline = System.currentTimeMillis() + 10_000; + while (!server.isRunning() && System.currentTimeMillis() < deadline) { + Thread.sleep(50); + } + + assertTrue(server.isRunning()); + assertTrue(RemoteInterpreterUtils.checkIfRemoteEndpointAccessible("localhost", port)); + server.init(new HashMap<>()); + } finally { + if (server.isRunning()) { + server.shutdown(); + } + server.join(10_000); + } + assertFalse(server.isRunning()); + } +} diff --git a/zeppelin-interpreter/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterDownloader.java b/zeppelin-interpreter/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterDownloader.java index f3854783f88..d7ebc0077e7 100644 --- a/zeppelin-interpreter/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterDownloader.java +++ b/zeppelin-interpreter/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterDownloader.java @@ -27,6 +27,7 @@ import org.apache.commons.io.FileUtils; +import org.apache.commons.lang3.StringUtils; import org.apache.zeppelin.interpreter.thrift.LibraryMetadata; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -55,8 +56,17 @@ public static void main(String[] args) { int port = Integer.parseInt(args[1]); String interpreter = args[2]; String localRepoPath = args[3]; + String interpreterGroupId = System.getenv("INTERPRETER_GROUP_ID"); + String callbackToken = System.getenv(RemoteInterpreterEventClient.CALLBACK_TOKEN_ENV); + if (StringUtils.isAnyBlank(interpreterGroupId, callbackToken)) { + // Externally managed interpreters may be started before Zeppelin can provision their + // callback credential through init(). They must use a pre-populated local repository. + LOGGER.warn("Skip remote interpreter library synchronization because callback " + + "credentials are not available"); + return; + } RemoteInterpreterEventClient intpEventClient = new RemoteInterpreterEventClient( - zeppelinServerHost, port, 3); + zeppelinServerHost, port, 3, interpreterGroupId, callbackToken); RemoteInterpreterDownloader downloader = new RemoteInterpreterDownloader(interpreter, intpEventClient, new File(localRepoPath)); diff --git a/zeppelin-interpreter/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterEventClient.java b/zeppelin-interpreter/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterEventClient.java index 174de3bc194..9f5ed19c42b 100644 --- a/zeppelin-interpreter/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterEventClient.java +++ b/zeppelin-interpreter/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterEventClient.java @@ -19,7 +19,9 @@ import com.google.gson.Gson; import org.apache.thrift.protocol.TBinaryProtocol; import org.apache.thrift.protocol.TProtocol; +import org.apache.thrift.transport.TSaslClientTransport; import org.apache.thrift.transport.TSocket; +import org.apache.thrift.transport.TTransport; import org.apache.thrift.transport.TTransportException; import org.apache.zeppelin.display.AngularObject; import org.apache.zeppelin.display.AngularObjectRegistryListener; @@ -45,9 +47,22 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +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.SaslException; import java.io.IOException; import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; import java.util.ArrayList; +import java.util.Base64; import java.util.List; import java.util.Map; @@ -60,15 +75,47 @@ public class RemoteInterpreterEventClient implements ResourcePoolConnector, private static final Logger LOGGER = LoggerFactory.getLogger(RemoteInterpreterEventClient.class); private static final Gson GSON = new Gson(); + public static final String CALLBACK_TOKEN_ENV = "ZEPPELIN_INTERPRETER_EVENT_TOKEN"; + public static final String CALLBACK_TOKEN_PROPERTY = "zeppelin.interpreter.event.token"; + public static final String INTERPRETER_GROUP_PROPERTY = "zeppelin.interpreter.group.id"; + + public static final String SASL_MECHANISM = "DIGEST-MD5"; + public static final String SASL_PROTOCOL = "zeppelin"; + public static final String SASL_SERVER_NAME = "interpreter-event"; + public static final Map SASL_PROPERTIES = Map.of( + Sasl.QOP, "auth-int", + Sasl.SERVER_AUTH, "true"); + private PooledRemoteClient remoteClient; private String intpGroupId; - public RemoteInterpreterEventClient(String intpEventHost, int intpEventPort, int connectionPoolSize) { + public RemoteInterpreterEventClient(String intpEventHost, + int intpEventPort, + int connectionPoolSize, + String intpGroupId, + String callbackToken) { + if (intpGroupId == null || intpGroupId.isEmpty()) { + throw new IllegalArgumentException("Interpreter group id must not be empty"); + } + if (callbackToken == null || callbackToken.isEmpty()) { + throw new IllegalArgumentException("Interpreter event callback token must not be empty"); + } + this.intpGroupId = intpGroupId; + String authenticationId = callbackAuthenticationId(intpGroupId, callbackToken); this.remoteClient = new PooledRemoteClient<>(() -> { - TSocket transport = new TSocket(intpEventHost, intpEventPort); + TSocket socket = new TSocket(intpEventHost, intpEventPort); + TTransport transport; try { + transport = new TSaslClientTransport( + SASL_MECHANISM, + authenticationId, + SASL_PROTOCOL, + SASL_SERVER_NAME, + SASL_PROPERTIES, + clientCallbackHandler(authenticationId, callbackToken), + socket); transport.open(); - } catch (TTransportException e) { + } catch (SaslException | TTransportException e) { throw new IOException(e); } TProtocol protocol = new TBinaryProtocol(transport); @@ -76,12 +123,41 @@ public RemoteInterpreterEventClient(String intpEventHost, int intpEventPort, int }, connectionPoolSize); } - public R callRemoteFunction(PooledRemoteClient.RemoteFunction func) { - return remoteClient.callRemoteFunction(func); + public static String callbackAuthenticationId(String intpGroupId, String callbackToken) { + try { + MessageDigest digest = MessageDigest.getInstance("SHA-256"); + digest.update(intpGroupId.getBytes(StandardCharsets.UTF_8)); + digest.update((byte) 0); + byte[] authenticationId = digest.digest(callbackToken.getBytes(StandardCharsets.UTF_8)); + return Base64.getUrlEncoder().withoutPadding().encodeToString(authenticationId); + } catch (NoSuchAlgorithmException e) { + throw new IllegalStateException("SHA-256 is unavailable", e); + } } - public void setIntpGroupId(String intpGroupId) { - this.intpGroupId = intpGroupId; + private static CallbackHandler clientCallbackHandler(String authenticationId, + String callbackToken) { + return callbacks -> { + for (Callback callback : callbacks) { + if (callback instanceof NameCallback) { + ((NameCallback) callback).setName(authenticationId); + } else if (callback instanceof PasswordCallback) { + ((PasswordCallback) callback).setPassword(callbackToken.toCharArray()); + } else if (callback instanceof RealmCallback) { + RealmCallback realmCallback = (RealmCallback) callback; + realmCallback.setText(realmCallback.getDefaultText()); + } else if (callback instanceof RealmChoiceCallback) { + ((RealmChoiceCallback) callback).setSelectedIndex(0); + } else { + throw new UnsupportedCallbackException(callback); + } + } + }; + } + + public R callRemoteFunction( + PooledRemoteClient.RemoteFunction func) { + return remoteClient.callRemoteFunction(func); } public void registerInterpreterProcess(RegisterInfo registerInfo) { diff --git a/zeppelin-interpreter/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterServer.java b/zeppelin-interpreter/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterServer.java index e733fd57f8c..356658dbc24 100644 --- a/zeppelin-interpreter/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterServer.java +++ b/zeppelin-interpreter/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterServer.java @@ -88,6 +88,7 @@ import java.util.Properties; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; @@ -101,6 +102,9 @@ public class RemoteInterpreterServer extends Thread implements RemoteInterpreterService.Iface { private static final Logger LOGGER = LoggerFactory.getLogger(RemoteInterpreterServer.class); + private static final String TEST_CALLBACK_CREDENTIAL = + "remote-interpreter-server-test-callback-credential"; + private static final long CALLBACK_REGISTRATION_TIMEOUT_SECONDS = 30; public static final int DEFAULT_SHUTDOWN_TIMEOUT = 2000; @@ -115,9 +119,14 @@ public class RemoteInterpreterServer extends Thread private String intpEventServerHost; private int intpEventServerPort; + private final String intpEventCallbackToken; + private final boolean callbackRegistrationRequired; + private final CountDownLatch callbackRegistrationReady; + private final CountDownLatch callbackRegistrationComplete; private String host; private int port; private TThreadPoolServer server; + private final Object intpEventClientLock = new Object(); RemoteInterpreterEventClient intpEventClient; private DependencyResolver depLoader; private LifecycleManager lifecycleManager; @@ -159,7 +168,34 @@ public RemoteInterpreterServer(String intpEventServerHost, String portRange, String interpreterGroupId, boolean isTest) throws Exception { + this(intpEventServerHost, intpEventServerPort, portRange, interpreterGroupId, isTest, + isTest ? TEST_CALLBACK_CREDENTIAL + : System.getenv(RemoteInterpreterEventClient.CALLBACK_TOKEN_ENV), + intpEventServerHost != null && !isTest); + } + + RemoteInterpreterServer(String intpEventServerHost, + int intpEventServerPort, + String portRange, + String interpreterGroupId, + boolean isTest, + String callbackCredential) throws Exception { + this(intpEventServerHost, intpEventServerPort, portRange, interpreterGroupId, isTest, + callbackCredential, intpEventServerHost != null && !isTest); + } + + RemoteInterpreterServer(String intpEventServerHost, + int intpEventServerPort, + String portRange, + String interpreterGroupId, + boolean isTest, + String callbackCredential, + boolean callbackRegistrationRequired) throws Exception { super("RemoteInterpreterServer-Thread"); + if (intpEventServerHost != null && StringUtils.isBlank(callbackCredential)) { + throw new IllegalArgumentException( + "Interpreter callback credential is required before startup"); + } if (null != intpEventServerHost) { this.intpEventServerHost = intpEventServerHost; this.intpEventServerPort = intpEventServerPort; @@ -171,6 +207,13 @@ public RemoteInterpreterServer(String intpEventServerHost, } this.isTest = isTest; this.interpreterGroupId = interpreterGroupId; + this.intpEventCallbackToken = callbackCredential; + this.callbackRegistrationRequired = callbackRegistrationRequired; + this.callbackRegistrationReady = new CountDownLatch( + interpreterGroupId != null + && interpreterGroupId.endsWith("-" + Constants.EXISTING_PROCESS) ? 1 : 0); + this.callbackRegistrationComplete = new CountDownLatch( + callbackRegistrationRequired ? 1 : 0); } @Override @@ -184,7 +227,7 @@ public void run() { .stopTimeoutUnit(TimeUnit.MILLISECONDS) .processor(processor)); - if (null != intpEventServerHost && !isTest) { + if (callbackRegistrationRequired) { Thread registerThread = new Thread(new RegisterRunnable()); registerThread.setName("RegisterThread"); registerThread.start(); @@ -201,7 +244,12 @@ public void run() { public void init(Map properties) throws InterpreterRPCException, TException { this.zProperties = new Properties(); this.zProperties.putAll(properties); - + this.zProperties.remove(RemoteInterpreterEventClient.CALLBACK_TOKEN_PROPERTY); + String configuredGroupId = zProperties.getProperty( + RemoteInterpreterEventClient.INTERPRETER_GROUP_PROPERTY); + if (configuredGroupId != null && !configuredGroupId.equals(interpreterGroupId)) { + throw new InterpreterRPCException("Interpreter callback group id does not match"); + } try { lifecycleManager = createLifecycleManager(); lifecycleManager.onInterpreterProcessStarted(interpreterGroupId); @@ -209,13 +257,37 @@ public void init(Map properties) throws InterpreterRPCException, throw new InterpreterRPCException("Fail to create LifecycleManager, cause: " + e.toString()); } - if (!isTest) { + if (intpEventServerHost != null && (!isTest || callbackRegistrationRequired)) { + if (StringUtils.isBlank(intpEventCallbackToken)) { + throw new InterpreterRPCException("Interpreter callback credential is required"); + } int connectionPoolSize = Integer.parseInt( zProperties.getProperty("zeppelin.interpreter.connection.poolsize", "100")); LOGGER.info("Creating RemoteInterpreterEventClient with connection pool size: {}", connectionPoolSize); - intpEventClient = new RemoteInterpreterEventClient(intpEventServerHost, intpEventServerPort, - connectionPoolSize); + synchronized (intpEventClientLock) { + // reconnect() has already installed and propagated the runtime client when this is a + // recovered process. Replacing it here would leave existing interpreter contexts with a + // closed client. + if (intpEventClient == null) { + intpEventClient = createInterpreterEventClient( + intpEventServerHost, intpEventServerPort, connectionPoolSize); + } + } + // An externally managed process can start before Zeppelin has installed its pre-shared + // credential. init() carries only this readiness signal, never the credential itself. + callbackRegistrationReady.countDown(); + try { + if (!callbackRegistrationComplete.await( + getCallbackRegistrationTimeoutSeconds(), TimeUnit.SECONDS)) { + throw new InterpreterRPCException( + "Timed out authenticating the interpreter callback registration"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new InterpreterRPCException( + "Interrupted while authenticating the interpreter callback registration"); + } } } @@ -276,6 +348,10 @@ private LifecycleManager createLifecycleManager() throws Exception { .newInstance(zProperties, this); } + long getCallbackRegistrationTimeoutSeconds() { + return CALLBACK_REGISTRATION_TIMEOUT_SECONDS; + } + public static void main(String[] args) throws Exception { String zeppelinServerHost = null; int port = Constants.ZEPPELIN_INTERPRETER_DEFAUlT_PORT; @@ -326,7 +402,6 @@ public void createInterpreter(String interpreterGroupId, String sessionId, Strin interpreterGroup.setInterpreterHookRegistry(hookRegistry); interpreterGroup.setAngularObjectRegistry(angularObjectRegistry); interpreterGroup.setResourcePool(resourcePool); - intpEventClient.setIntpGroupId(interpreterGroupId); String localRepoPath = properties.get("zeppelin.interpreter.localRepo"); if (properties.containsKey("zeppelin.interpreter.output.limit")) { @@ -481,28 +556,59 @@ public void close(String sessionId, String className) throws InterpreterRPCExcep @Override public void reconnect(String host, int port) throws InterpreterRPCException, TException { + RemoteInterpreterEventClient replacement = null; try { LOGGER.info("Reconnect to this interpreter process from {}:{}", host, port); - this.intpEventServerHost = host; - this.intpEventServerPort = port; - intpEventClient = new RemoteInterpreterEventClient(intpEventServerHost, intpEventServerPort, - Integer.parseInt(zProperties.getProperty("zeppelin.interpreter.connection.poolsize", "100"))); - intpEventClient.setIntpGroupId(interpreterGroupId); - - this.angularObjectRegistry = new AngularObjectRegistry(interpreterGroup.getId(), intpEventClient); - this.resourcePool = new DistributedResourcePool(interpreterGroup.getId(), intpEventClient); - - // reset all the available InterpreterContext's components that use intpEventClient. - for (InterpreterContext context : InterpreterContext.getAllContexts().values()) { - context.setIntpEventClient(intpEventClient); - context.setAngularObjectRegistry(angularObjectRegistry); - context.setResourcePool(resourcePool); + replacement = createInterpreterEventClient(host, port, + Integer.parseInt(zProperties.getProperty( + "zeppelin.interpreter.connection.poolsize", "100"))); + replacement.registerInterpreterProcess( + new RegisterInfo(this.host, this.port, interpreterGroupId)); + + RemoteInterpreterEventClient previous; + synchronized (intpEventClientLock) { + previous = intpEventClient; + this.intpEventServerHost = host; + this.intpEventServerPort = port; + intpEventClient = replacement; + + if (interpreterGroup != null) { + this.angularObjectRegistry = + new AngularObjectRegistry(interpreterGroup.getId(), intpEventClient); + this.resourcePool = + new DistributedResourcePool(interpreterGroup.getId(), intpEventClient); + interpreterGroup.setAngularObjectRegistry(angularObjectRegistry); + interpreterGroup.setResourcePool(resourcePool); + if (depLoader != null) { + appLoader = new ApplicationLoader(resourcePool, depLoader); + } + + // Reset every context in this interpreter process that retained the previous client. + for (InterpreterContext context : InterpreterContext.getAllContexts().values()) { + context.setIntpEventClient(intpEventClient); + context.setAngularObjectRegistry(angularObjectRegistry); + context.setResourcePool(resourcePool); + } + } + } + if (previous != null) { + previous.close(); } } catch (Exception e) { + if (replacement != null && replacement != intpEventClient) { + replacement.close(); + } throw new InterpreterRPCException(e.toString()); } } + RemoteInterpreterEventClient createInterpreterEventClient(String eventServerHost, + int eventServerPort, + int connectionPoolSize) { + return new RemoteInterpreterEventClient(eventServerHost, eventServerPort, + connectionPoolSize, interpreterGroupId, intpEventCallbackToken); + } + @Override public RemoteInterpreterResult interpret(String sessionId, String className, @@ -595,21 +701,53 @@ public void run() { if (!Thread.currentThread().isInterrupted()) { RegisterInfo registerInfo = new RegisterInfo(host, port, interpreterGroupId); try { - intpEventClient = new RemoteInterpreterEventClient(intpEventServerHost, intpEventServerPort, 10); - LOGGER.info("Registering interpreter process"); - intpEventClient.registerInterpreterProcess(registerInfo); - LOGGER.info("Registered interpreter process"); - } catch (Exception e) { - LOGGER.error("Error while registering interpreter: {}, cause: {}", registerInfo, e); - try { + callbackRegistrationReady.await(); + long registrationDeadline = System.nanoTime() + TimeUnit.SECONDS.toNanos( + getCallbackRegistrationTimeoutSeconds()); + while (!Thread.currentThread().isInterrupted() + && server != null && server.isServing() + && callbackRegistrationComplete.getCount() != 0 + && System.nanoTime() < registrationDeadline) { + RemoteInterpreterEventClient registrationClient = null; + try { + // Registration has its own short-lived transport. The server can call init() as + // soon as processStarted() is handled, before the registration reply is written; + // sharing the runtime client would let init() close that in-flight transport. + registrationClient = createInterpreterEventClient( + intpEventServerHost, intpEventServerPort, 1); + LOGGER.info("Registering interpreter process"); + registrationClient.registerInterpreterProcess(registerInfo); + callbackRegistrationComplete.countDown(); + LOGGER.info("Registered interpreter process"); + } catch (Exception e) { + LOGGER.warn("Interpreter callback registration failed; retrying: {}", e.getMessage()); + long remainingMillis = TimeUnit.NANOSECONDS.toMillis( + registrationDeadline - System.nanoTime()); + if (remainingMillis > 0) { + Thread.sleep(Math.min(1000, remainingMillis)); + } + } finally { + if (registrationClient != null) { + registrationClient.close(); + } + } + } + if (!Thread.currentThread().isInterrupted() + && server != null && server.isServing() + && callbackRegistrationComplete.getCount() != 0) { + LOGGER.error("Interpreter callback registration timed out; shutting down"); shutdown(); - } catch (Exception e1) { - LOGGER.warn("Exception occurs while shutting down", e1); } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + LOGGER.info("Interrupted while waiting to register the interpreter process"); + } catch (Exception e) { + LOGGER.error("Failed to shut down after callback registration failure", e); } } - if (launcherEnv != null && "yarn".endsWith(launcherEnv)) { + if (callbackRegistrationComplete.getCount() == 0 + && launcherEnv != null && "yarn".endsWith(launcherEnv)) { try { YarnUtils.register(host, port); ScheduledExecutorService yarnHeartbeat = ExecutorFactory.singleton() diff --git a/zeppelin-interpreter/src/test/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterServerTest.java b/zeppelin-interpreter/src/test/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterServerTest.java index 58ef866a744..5d026b3b889 100644 --- a/zeppelin-interpreter/src/test/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterServerTest.java +++ b/zeppelin-interpreter/src/test/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterServerTest.java @@ -23,6 +23,7 @@ import org.apache.zeppelin.interpreter.InterpreterException; import org.apache.zeppelin.interpreter.InterpreterResult; import org.apache.zeppelin.interpreter.LazyOpenInterpreter; +import org.apache.zeppelin.interpreter.thrift.InterpreterRPCException; import org.apache.zeppelin.interpreter.thrift.RemoteInterpreterContext; import org.apache.zeppelin.interpreter.thrift.RemoteInterpreterResult; import org.junit.jupiter.api.Test; @@ -35,12 +36,21 @@ import java.util.Properties; import java.util.Set; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.argThat; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; class RemoteInterpreterServerTest { @@ -66,6 +76,170 @@ void testStartStopWithQueuedEvents() throws Exception { stopRemoteInterpreterServer(server, 10 * 10000); } + @Test + void testInitDoesNotAcceptCallbackCredential() throws Exception { + RemoteInterpreterServer server = new RemoteInterpreterServer("localhost", + RemoteInterpreterUtils.findRandomAvailablePortOnAllLocalInterfaces(), ":", "groupId", + true, "launch-token"); + Map properties = new HashMap<>(); + properties.put(RemoteInterpreterEventClient.INTERPRETER_GROUP_PROPERTY, "groupId"); + properties.put(RemoteInterpreterEventClient.CALLBACK_TOKEN_PROPERTY, "untrusted-token"); + server.init(properties); + + assertFalse(server.getProperties().containsKey( + RemoteInterpreterEventClient.CALLBACK_TOKEN_PROPERTY)); + } + + @Test + void testExistingProcessRegistersOnlyAfterInitReadiness() throws Exception { + RemoteInterpreterEventClient runtimeClient = mock(RemoteInterpreterEventClient.class); + RemoteInterpreterEventClient registrationClient = mock(RemoteInterpreterEventClient.class); + AtomicInteger createdClients = new AtomicInteger(); + RemoteInterpreterServer server = new RemoteInterpreterServer("localhost", + RemoteInterpreterUtils.findRandomAvailablePortOnAllLocalInterfaces(), ":", + "setting-existing_process", true, "callback-token", true) { + @Override + RemoteInterpreterEventClient createInterpreterEventClient(String eventServerHost, + int eventServerPort, + int connectionPoolSize) { + return createdClients.getAndIncrement() == 0 ? runtimeClient : registrationClient; + } + }; + server.start(); + long deadline = System.currentTimeMillis() + 10_000; + while (!server.isRunning() && System.currentTimeMillis() < deadline) { + Thread.sleep(50); + } + assertTrue(server.isRunning()); + assertEquals(0, createdClients.get()); + + Map properties = new HashMap<>(); + properties.put(RemoteInterpreterEventClient.INTERPRETER_GROUP_PROPERTY, + "setting-existing_process"); + properties.put("zeppelin.interpreter.connection.poolsize", "1"); + server.init(properties); + + assertEquals(2, createdClients.get()); + assertSame(runtimeClient, server.intpEventClient); + verify(registrationClient).registerInterpreterProcess(argThat(info -> + info.getInterpreterGroupId().equals("setting-existing_process"))); + stopRemoteInterpreterServer(server, 10_000); + } + + @Test + void testRegistrationFailureStopsInterpreterServer() throws Exception { + RemoteInterpreterEventClient runtimeClient = mock(RemoteInterpreterEventClient.class); + RemoteInterpreterEventClient registrationClient = mock(RemoteInterpreterEventClient.class); + doThrow(new RuntimeException("callback registration failed")) + .when(registrationClient).registerInterpreterProcess( + org.mockito.ArgumentMatchers.any()); + AtomicInteger createdClients = new AtomicInteger(); + RemoteInterpreterServer server = new RemoteInterpreterServer("localhost", + RemoteInterpreterUtils.findRandomAvailablePortOnAllLocalInterfaces(), ":", + "setting-existing_process", true, "callback-token", true) { + @Override + RemoteInterpreterEventClient createInterpreterEventClient(String eventServerHost, + int eventServerPort, + int connectionPoolSize) { + return createdClients.getAndIncrement() == 0 ? runtimeClient : registrationClient; + } + + @Override + long getCallbackRegistrationTimeoutSeconds() { + return 1; + } + }; + server.start(); + long deadline = System.currentTimeMillis() + 10_000; + while (!server.isRunning() && System.currentTimeMillis() < deadline) { + Thread.sleep(50); + } + assertTrue(server.isRunning()); + + Map properties = new HashMap<>(); + properties.put(RemoteInterpreterEventClient.INTERPRETER_GROUP_PROPERTY, + "setting-existing_process"); + properties.put("zeppelin.interpreter.connection.poolsize", "1"); + assertThrows(InterpreterRPCException.class, () -> server.init(properties)); + + deadline = System.currentTimeMillis() + 10_000; + while (server.isRunning() && System.currentTimeMillis() < deadline) { + Thread.sleep(50); + } + assertFalse(server.isRunning()); + verify(registrationClient, atLeastOnce()).registerInterpreterProcess( + org.mockito.ArgumentMatchers.any()); + } + + @Test + void testInitKeepsReconnectedEventClient() throws Exception { + RemoteInterpreterEventClient initialClient = mock(RemoteInterpreterEventClient.class); + RemoteInterpreterEventClient reconnectedClient = mock(RemoteInterpreterEventClient.class); + AtomicInteger createdClients = new AtomicInteger(); + RemoteInterpreterServer server = new RemoteInterpreterServer("localhost", + RemoteInterpreterUtils.findRandomAvailablePortOnAllLocalInterfaces(), ":", "groupId", + false, "callback-token", false) { + @Override + RemoteInterpreterEventClient createInterpreterEventClient(String eventServerHost, + int eventServerPort, + int connectionPoolSize) { + return createdClients.getAndIncrement() == 0 ? initialClient : reconnectedClient; + } + }; + Map properties = new HashMap<>(); + properties.put(RemoteInterpreterEventClient.INTERPRETER_GROUP_PROPERTY, "groupId"); + properties.put(RemoteInterpreterEventClient.CALLBACK_TOKEN_PROPERTY, "untrusted-token"); + properties.put("zeppelin.interpreter.connection.poolsize", "1"); + server.init(properties); + assertSame(initialClient, server.intpEventClient); + + server.reconnect("localhost", + RemoteInterpreterUtils.findRandomAvailablePortOnAllLocalInterfaces()); + assertNotSame(initialClient, reconnectedClient); + assertSame(reconnectedClient, server.intpEventClient); + verify(reconnectedClient).registerInterpreterProcess(argThat(info -> + info.getInterpreterGroupId().equals("groupId") + && info.getHost() != null + && !info.getHost().isEmpty() + && info.getPort() == server.getPort())); + + server.init(properties); + assertSame(reconnectedClient, server.intpEventClient); + reconnectedClient.close(); + } + + @Test + void testFailedReconnectKeepsPreviousEventClient() throws Exception { + RemoteInterpreterEventClient initialClient = mock(RemoteInterpreterEventClient.class); + RemoteInterpreterEventClient failedReplacement = mock(RemoteInterpreterEventClient.class); + doThrow(new RuntimeException("callback proof failed")) + .when(failedReplacement).registerInterpreterProcess( + org.mockito.ArgumentMatchers.any()); + AtomicInteger createdClients = new AtomicInteger(); + RemoteInterpreterServer server = new RemoteInterpreterServer("localhost", + RemoteInterpreterUtils.findRandomAvailablePortOnAllLocalInterfaces(), ":", "groupId", + false, "callback-token", false) { + @Override + RemoteInterpreterEventClient createInterpreterEventClient(String eventServerHost, + int eventServerPort, + int connectionPoolSize) { + return createdClients.getAndIncrement() == 0 ? initialClient : failedReplacement; + } + }; + Map properties = new HashMap<>(); + properties.put(RemoteInterpreterEventClient.INTERPRETER_GROUP_PROPERTY, "groupId"); + properties.put(RemoteInterpreterEventClient.CALLBACK_TOKEN_PROPERTY, "untrusted-token"); + properties.put("zeppelin.interpreter.connection.poolsize", "1"); + server.init(properties); + + assertThrows(InterpreterRPCException.class, () -> server.reconnect("localhost", + RemoteInterpreterUtils.findRandomAvailablePortOnAllLocalInterfaces())); + + assertSame(initialClient, server.intpEventClient); + verify(failedReplacement).close(); + verify(initialClient, never()).close(); + } + private void startRemoteInterpreterServer(RemoteInterpreterServer server, int timeout) throws InterruptedException, TException { assertEquals(false, server.isRunning()); diff --git a/zeppelin-plugins/launcher/docker/src/main/java/org/apache/zeppelin/interpreter/launcher/DockerInterpreterProcess.java b/zeppelin-plugins/launcher/docker/src/main/java/org/apache/zeppelin/interpreter/launcher/DockerInterpreterProcess.java index 9c86a676083..c7f95e5a232 100644 --- a/zeppelin-plugins/launcher/docker/src/main/java/org/apache/zeppelin/interpreter/launcher/DockerInterpreterProcess.java +++ b/zeppelin-plugins/launcher/docker/src/main/java/org/apache/zeppelin/interpreter/launcher/DockerInterpreterProcess.java @@ -195,7 +195,7 @@ public void start(String userName) throws IOException { LOGGER.info("dockerCommand = {}", dockerCommand); List listEnv = getListEnvs(); - LOGGER.info("docker listEnv = {}", listEnv); + LOGGER.info("docker environment variables = {}", envs.keySet()); // check if the interpreter process exit script // if interpreter process exit, then container need exit diff --git a/zeppelin-plugins/launcher/docker/src/test/java/org/apache/zeppelin/interpreter/launcher/DockerInterpreterProcessTest.java b/zeppelin-plugins/launcher/docker/src/test/java/org/apache/zeppelin/interpreter/launcher/DockerInterpreterProcessTest.java index a4b0b4910bd..323e5842bd7 100644 --- a/zeppelin-plugins/launcher/docker/src/test/java/org/apache/zeppelin/interpreter/launcher/DockerInterpreterProcessTest.java +++ b/zeppelin-plugins/launcher/docker/src/test/java/org/apache/zeppelin/interpreter/launcher/DockerInterpreterProcessTest.java @@ -25,6 +25,7 @@ import org.apache.zeppelin.conf.ZeppelinConfiguration; import org.apache.zeppelin.conf.ZeppelinConfiguration.ConfVars; import org.apache.zeppelin.interpreter.InterpreterOption; +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterEventClient; import org.junit.jupiter.api.Test; import java.io.IOException; import java.util.Collections; @@ -186,11 +187,14 @@ void testCreateIntpProcess() throws IOException { InterpreterLaunchContext context = new InterpreterLaunchContext(properties, option, null, "user1", "intpGroupId", "groupId", "groupName", "name", 0, "host"); + context.setIntpEventCallbackToken("callback-token"); InterpreterClient client = launcher.launch(context); assertTrue(client instanceof DockerInterpreterProcess); DockerInterpreterProcess interpreterProcess = (DockerInterpreterProcess) client; assertEquals("name", interpreterProcess.getInterpreterSettingName()); + assertTrue(interpreterProcess.getListEnvs().contains( + RemoteInterpreterEventClient.CALLBACK_TOKEN_ENV + "=callback-token")); assertEquals("/opt/spark", interpreterProcess.containerSparkHome); assertTrue(interpreterProcess.uploadLocalLibToContainter); diff --git a/zeppelin-plugins/launcher/flink/src/main/java/org/apache/zeppelin/interpreter/launcher/FlinkInterpreterLauncher.java b/zeppelin-plugins/launcher/flink/src/main/java/org/apache/zeppelin/interpreter/launcher/FlinkInterpreterLauncher.java index 847cef2bff4..2247fa02792 100644 --- a/zeppelin-plugins/launcher/flink/src/main/java/org/apache/zeppelin/interpreter/launcher/FlinkInterpreterLauncher.java +++ b/zeppelin-plugins/launcher/flink/src/main/java/org/apache/zeppelin/interpreter/launcher/FlinkInterpreterLauncher.java @@ -22,6 +22,7 @@ import org.apache.commons.lang3.StringUtils; import org.apache.zeppelin.conf.ZeppelinConfiguration; import org.apache.zeppelin.interpreter.recovery.RecoveryStorage; +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterEventClient; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -231,6 +232,13 @@ private void updateEnvsForApplicationMode(String mode, yarnShipFiles.stream().collect(Collectors.joining(";"))); } + if (StringUtils.isNotBlank(context.getIntpEventCallbackToken())) { + flinkConfStringJoiner.add("-D"); + flinkConfStringJoiner.add("containerized.master.env." + + RemoteInterpreterEventClient.CALLBACK_TOKEN_ENV + "=" + + context.getIntpEventCallbackToken()); + } + // set yarn.application.name String yarnAppName = context.getProperties().getProperty("flink.yarn.appName"); if (StringUtils.isNotBlank(yarnAppName)) { diff --git a/zeppelin-plugins/launcher/flink/src/test/java/org/apache/zeppelin/interpreter/launcher/FlinkInterpreterLauncherTest.java b/zeppelin-plugins/launcher/flink/src/test/java/org/apache/zeppelin/interpreter/launcher/FlinkInterpreterLauncherTest.java new file mode 100644 index 00000000000..2cb03f90724 --- /dev/null +++ b/zeppelin-plugins/launcher/flink/src/test/java/org/apache/zeppelin/interpreter/launcher/FlinkInterpreterLauncherTest.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.zeppelin.interpreter.launcher; + +import org.apache.zeppelin.conf.ZeppelinConfiguration; +import org.apache.zeppelin.interpreter.InterpreterOption; +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterEventClient; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.nio.file.Path; +import java.util.Map; +import java.util.Properties; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class FlinkInterpreterLauncherTest { + + @Test + void passesCallbackCredentialToApplicationModeMaster(@TempDir Path flinkHome) throws Exception { + Properties properties = new Properties(); + properties.setProperty("FLINK_HOME", flinkHome.toString()); + properties.setProperty("flink.execution.mode", "kubernetes-application"); + properties.setProperty("flink.app.jar", "/tmp/flink-app.jar"); + InterpreterLaunchContext context = new InterpreterLaunchContext( + properties, new InterpreterOption(), null, "user", "group-id", + "setting-id", "flink", "flink", 0, "host"); + context.setIntpEventCallbackToken("callback-token"); + FlinkInterpreterLauncher launcher = new FlinkInterpreterLauncher( + ZeppelinConfiguration.load(), null); + + Map environment = launcher.buildEnvFromProperties(context); + + assertEquals("callback-token", + environment.get(RemoteInterpreterEventClient.CALLBACK_TOKEN_ENV)); + assertTrue(environment.get("ZEPPELIN_FLINK_APPLICATION_MODE_CONF").contains( + "containerized.master.env." + RemoteInterpreterEventClient.CALLBACK_TOKEN_ENV + + "=callback-token")); + } +} diff --git a/zeppelin-plugins/launcher/k8s-standard/src/main/java/org/apache/zeppelin/interpreter/launcher/K8sRemoteInterpreterProcess.java b/zeppelin-plugins/launcher/k8s-standard/src/main/java/org/apache/zeppelin/interpreter/launcher/K8sRemoteInterpreterProcess.java index 62fcc38df01..5ccdc547b77 100644 --- a/zeppelin-plugins/launcher/k8s-standard/src/main/java/org/apache/zeppelin/interpreter/launcher/K8sRemoteInterpreterProcess.java +++ b/zeppelin-plugins/launcher/k8s-standard/src/main/java/org/apache/zeppelin/interpreter/launcher/K8sRemoteInterpreterProcess.java @@ -33,6 +33,7 @@ import org.apache.commons.lang3.ArrayUtils; import org.apache.commons.lang3.StringUtils; import org.apache.zeppelin.interpreter.remote.RemoteInterpreterManagedProcess; +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterEventClient; import org.apache.zeppelin.interpreter.remote.RemoteInterpreterServer; import org.apache.zeppelin.interpreter.remote.RemoteInterpreterUtils; import org.slf4j.Logger; @@ -282,7 +283,10 @@ void apply(File path, boolean delete, Properties templateProperties) throws IOEx NamespaceListVisitFromServerGetDeleteRecreateWaitApplicable k8sObjects = client.load(IOUtils.toInputStream(template, StandardCharsets.UTF_8)); LOGGER.info("Apply {} with {} K8s Objects", path.getAbsolutePath(), k8sObjects.get().size()); - LOGGER.debug(template); + String callbackToken = getEnv().get(RemoteInterpreterEventClient.CALLBACK_TOKEN_ENV); + LOGGER.debug(StringUtils.isBlank(callbackToken) + ? template + : template.replace(callbackToken, "[REDACTED]")); if (delete) { k8sObjects.inNamespace(interpreterNamespace).delete(); } else { diff --git a/zeppelin-plugins/launcher/k8s-standard/src/main/java/org/apache/zeppelin/interpreter/launcher/K8sStandardInterpreterLauncher.java b/zeppelin-plugins/launcher/k8s-standard/src/main/java/org/apache/zeppelin/interpreter/launcher/K8sStandardInterpreterLauncher.java index 7a1ce482755..ef420bacccc 100644 --- a/zeppelin-plugins/launcher/k8s-standard/src/main/java/org/apache/zeppelin/interpreter/launcher/K8sStandardInterpreterLauncher.java +++ b/zeppelin-plugins/launcher/k8s-standard/src/main/java/org/apache/zeppelin/interpreter/launcher/K8sStandardInterpreterLauncher.java @@ -24,9 +24,11 @@ import java.util.HashMap; import java.util.Map; +import org.apache.commons.lang3.StringUtils; import org.apache.zeppelin.conf.ZeppelinConfiguration; import org.apache.zeppelin.interpreter.recovery.RecoveryStorage; import org.apache.zeppelin.interpreter.remote.RemoteInterpreterUtils; +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterEventClient; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -139,6 +141,10 @@ protected Map buildEnvFromProperties(InterpreterLaunchContext co } } env.put("INTERPRETER_GROUP_ID", context.getInterpreterGroupId()); + if (StringUtils.isNotBlank(context.getIntpEventCallbackToken())) { + env.put(RemoteInterpreterEventClient.CALLBACK_TOKEN_ENV, + context.getIntpEventCallbackToken()); + } return env; } diff --git a/zeppelin-plugins/launcher/k8s-standard/src/test/java/org/apache/zeppelin/interpreter/launcher/K8sStandardInterpreterLauncherTest.java b/zeppelin-plugins/launcher/k8s-standard/src/test/java/org/apache/zeppelin/interpreter/launcher/K8sStandardInterpreterLauncherTest.java index 651f4a9e6fc..5fee0a90381 100644 --- a/zeppelin-plugins/launcher/k8s-standard/src/test/java/org/apache/zeppelin/interpreter/launcher/K8sStandardInterpreterLauncherTest.java +++ b/zeppelin-plugins/launcher/k8s-standard/src/test/java/org/apache/zeppelin/interpreter/launcher/K8sStandardInterpreterLauncherTest.java @@ -17,6 +17,7 @@ package org.apache.zeppelin.interpreter.launcher; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -25,6 +26,7 @@ import org.apache.zeppelin.conf.ZeppelinConfiguration; import org.apache.zeppelin.interpreter.InterpreterOption; +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterEventClient; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -60,12 +62,17 @@ void testK8sLauncher() throws IOException { "name", 0, "host"); + context.setIntpEventCallbackToken("callback-token"); // when InterpreterClient client = launcher.launch(context); // then assertTrue(client instanceof K8sRemoteInterpreterProcess); + K8sRemoteInterpreterProcess process = (K8sRemoteInterpreterProcess) client; + assertEquals("callback-token", + process.getEnv().get(RemoteInterpreterEventClient.CALLBACK_TOKEN_ENV)); + process.close(); } @Test diff --git a/zeppelin-plugins/launcher/yarn/src/main/java/org/apache/zeppelin/interpreter/launcher/YarnInterpreterLauncher.java b/zeppelin-plugins/launcher/yarn/src/main/java/org/apache/zeppelin/interpreter/launcher/YarnInterpreterLauncher.java index 33ca438d007..e413c69af5e 100644 --- a/zeppelin-plugins/launcher/yarn/src/main/java/org/apache/zeppelin/interpreter/launcher/YarnInterpreterLauncher.java +++ b/zeppelin-plugins/launcher/yarn/src/main/java/org/apache/zeppelin/interpreter/launcher/YarnInterpreterLauncher.java @@ -17,8 +17,10 @@ package org.apache.zeppelin.interpreter.launcher; +import org.apache.commons.lang3.StringUtils; import org.apache.zeppelin.conf.ZeppelinConfiguration; import org.apache.zeppelin.interpreter.recovery.RecoveryStorage; +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterEventClient; import org.apache.zeppelin.interpreter.remote.RemoteInterpreterUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -59,6 +61,10 @@ protected Map buildEnvFromProperties(InterpreterLaunchContext co } } env.put("INTERPRETER_GROUP_ID", context.getInterpreterGroupId()); + if (StringUtils.isNotBlank(context.getIntpEventCallbackToken())) { + env.put(RemoteInterpreterEventClient.CALLBACK_TOKEN_ENV, + context.getIntpEventCallbackToken()); + } env.put("ZEPPELIN_INTERPRETER_LAUNCHER", "yarn"); return env; } diff --git a/zeppelin-plugins/launcher/yarn/src/test/java/org/apache/zeppelin/interpreter/launcher/YarnInterpreterLauncherTest.java b/zeppelin-plugins/launcher/yarn/src/test/java/org/apache/zeppelin/interpreter/launcher/YarnInterpreterLauncherTest.java new file mode 100644 index 00000000000..f1006023f89 --- /dev/null +++ b/zeppelin-plugins/launcher/yarn/src/test/java/org/apache/zeppelin/interpreter/launcher/YarnInterpreterLauncherTest.java @@ -0,0 +1,46 @@ +/* + * 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.zeppelin.interpreter.launcher; + +import org.apache.zeppelin.conf.ZeppelinConfiguration; +import org.apache.zeppelin.interpreter.InterpreterOption; +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterEventClient; +import org.junit.jupiter.api.Test; + +import java.util.Map; +import java.util.Properties; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +class YarnInterpreterLauncherTest { + + @Test + void passesCallbackCredentialToYarnLauncherEnvironment() { + YarnInterpreterLauncher launcher = new YarnInterpreterLauncher( + ZeppelinConfiguration.load(), null); + InterpreterLaunchContext context = new InterpreterLaunchContext( + new Properties(), new InterpreterOption(), null, "user", "group-id", + "setting-id", "group", "name", 0, "host"); + context.setIntpEventCallbackToken("callback-token"); + + Map environment = launcher.buildEnvFromProperties(context); + + assertEquals("callback-token", + environment.get(RemoteInterpreterEventClient.CALLBACK_TOKEN_ENV)); + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/InterpreterSetting.java b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/InterpreterSetting.java index f26bc54f0e7..c2da590eadd 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/InterpreterSetting.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/InterpreterSetting.java @@ -38,11 +38,13 @@ import org.apache.zeppelin.display.AngularObjectRegistryListener; import org.apache.zeppelin.helium.ApplicationEventListener; import org.apache.zeppelin.interpreter.launcher.InterpreterLaunchContext; +import org.apache.zeppelin.interpreter.launcher.InterpreterClient; import org.apache.zeppelin.interpreter.launcher.InterpreterLauncher; import org.apache.zeppelin.interpreter.recovery.NullRecoveryStorage; import org.apache.zeppelin.interpreter.recovery.RecoveryStorage; import org.apache.zeppelin.interpreter.remote.RemoteAngularObjectRegistry; import org.apache.zeppelin.interpreter.remote.RemoteInterpreter; +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterEventClient; import org.apache.zeppelin.interpreter.remote.RemoteInterpreterProcess; import org.apache.zeppelin.interpreter.remote.RemoteInterpreterProcessListener; import org.apache.zeppelin.plugin.PluginManager; @@ -861,12 +863,52 @@ synchronized RemoteInterpreterProcess createInterpreterProcess(String interprete Properties properties) throws IOException { InterpreterLauncher launcher = createLauncher(properties); + String callbackToken = interpreterEventServer.getCallbackToken(interpreterGroupId); + InterpreterClient recoveredClient = zConf.isRecoveryEnabled() + ? recoveryStorage.getInterpreterClient(interpreterGroupId) : null; + boolean reusableRecoveryCredential = recoveredClient != null && recoveredClient.isRunning(); + boolean installedLaunchCredential = false; + if (!reusableRecoveryCredential) { + if (option.isExistingProcess()) { + callbackToken = System.getenv(RemoteInterpreterEventClient.CALLBACK_TOKEN_ENV); + if (StringUtils.isBlank(callbackToken)) { + throw new IOException("Externally managed interpreters require " + + RemoteInterpreterEventClient.CALLBACK_TOKEN_ENV + + " to be configured in both processes"); + } + interpreterEventServer.registerCallbackToken(interpreterGroupId, callbackToken); + } else { + callbackToken = interpreterEventServer.issueCallbackToken(interpreterGroupId); + } + installedLaunchCredential = true; + } else if (StringUtils.isBlank(callbackToken)) { + throw new IOException("Recovered interpreter callback credential is unavailable"); + } InterpreterLaunchContext launchContext = new InterpreterLaunchContext(properties, option, interpreterRunner, userName, - interpreterGroupId, id, group, name, interpreterEventServer.getPort(), interpreterEventServer.getHost()); - RemoteInterpreterProcess process = (RemoteInterpreterProcess) launcher.launch(launchContext); - recoveryStorage.onInterpreterClientStart(process); - return process; + interpreterGroupId, id, group, name, interpreterEventServer.getPort(), + interpreterEventServer.getHost()); + launchContext.setIntpEventCallbackToken(callbackToken); + try { + RemoteInterpreterProcess process = (RemoteInterpreterProcess) launcher.launch(launchContext); + String launchedCallbackToken = callbackToken; + process.setTerminationListener(() -> { + interpreterEventServer.revokeCallbackToken( + interpreterGroupId, launchedCallbackToken); + try { + recoveryStorage.onInterpreterClientStop(process); + } catch (IOException e) { + LOGGER.warn("Fail to remove terminated interpreter process from recovery storage", e); + } + }); + recoveryStorage.onInterpreterClientStart(process); + return process; + } catch (IOException | RuntimeException e) { + if (installedLaunchCredential) { + interpreterEventServer.revokeCallbackToken(interpreterGroupId, callbackToken); + } + throw e; + } } List getOrCreateSession(String user, String noteId) { diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/ManagedInterpreterGroup.java b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/ManagedInterpreterGroup.java index 8f2c16c0743..716b6769f0b 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/ManagedInterpreterGroup.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/ManagedInterpreterGroup.java @@ -66,12 +66,24 @@ public RemoteInterpreterProcess getOrCreateInterpreterProcess(String userName, synchronized (interpreterProcessCreationLock) { if (remoteInterpreterProcess == null) { LOGGER.info("Create InterpreterProcess for InterpreterGroup: {}", getId()); - remoteInterpreterProcess = interpreterSetting.createInterpreterProcess(id, userName, - properties); - remoteInterpreterProcess.start(userName); - remoteInterpreterProcess.init(zConf); - getInterpreterSetting().getRecoveryStorage() - .onInterpreterClientStart(remoteInterpreterProcess); + RemoteInterpreterEventServer eventServer = interpreterSetting.getInterpreterSettingManager() + .getInterpreterEventServer(); + try { + remoteInterpreterProcess = interpreterSetting.createInterpreterProcess(id, userName, + properties); + remoteInterpreterProcess.start(userName); + remoteInterpreterProcess.init(zConf, id); + getInterpreterSetting().getRecoveryStorage() + .onInterpreterClientStart(remoteInterpreterProcess); + } catch (IOException | RuntimeException e) { + String callbackToken = eventServer.getCallbackToken(id); + eventServer.revokeCallbackToken(id, callbackToken); + if (remoteInterpreterProcess != null) { + remoteInterpreterProcess.stop(); + remoteInterpreterProcess = null; + } + throw e; + } } return remoteInterpreterProcess; } @@ -110,7 +122,14 @@ public synchronized void close(String sessionId) { interpreterSetting.removeInterpreterGroup(id); if (remoteInterpreterProcess != null) { LOGGER.info("Kill RemoteInterpreterProcess"); - remoteInterpreterProcess.stop(); + RemoteInterpreterEventServer eventServer = interpreterSetting + .getInterpreterSettingManager().getInterpreterEventServer(); + String callbackToken = eventServer.getCallbackToken(id); + try { + remoteInterpreterProcess.stop(); + } finally { + eventServer.revokeCallbackToken(id, callbackToken); + } try { interpreterSetting.getRecoveryStorage().onInterpreterClientStop(remoteInterpreterProcess); } catch (IOException e) { diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/RemoteInterpreterEventServer.java b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/RemoteInterpreterEventServer.java index bab3ee7b2ad..2eee85cd48d 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/RemoteInterpreterEventServer.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/RemoteInterpreterEventServer.java @@ -19,13 +19,15 @@ import com.google.gson.Gson; import com.google.gson.reflect.TypeToken; - -import java.util.OptionalInt; import org.apache.commons.io.FileUtils; import org.apache.commons.lang3.StringUtils; import org.apache.thrift.TException; +import org.apache.thrift.server.ServerContext; +import org.apache.thrift.server.TServerEventHandler; import org.apache.thrift.server.TThreadPoolServer; +import org.apache.thrift.transport.TSaslServerTransport; import org.apache.thrift.transport.TServerSocket; +import org.apache.thrift.transport.TTransport; import org.apache.thrift.transport.TTransportException; import org.apache.zeppelin.conf.ZeppelinConfiguration; import org.apache.zeppelin.display.AngularObject; @@ -33,6 +35,7 @@ import org.apache.zeppelin.interpreter.remote.AppendOutputRunner; import org.apache.zeppelin.interpreter.remote.InvokeResourceMethodEventMessage; import org.apache.zeppelin.interpreter.remote.RemoteAngularObject; +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterEventClient; import org.apache.zeppelin.interpreter.remote.RemoteInterpreterProcess; import org.apache.zeppelin.interpreter.remote.RemoteInterpreterProcessListener; import org.apache.zeppelin.interpreter.remote.RemoteInterpreterUtils; @@ -59,15 +62,26 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +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.AuthorizeCallback; +import javax.security.sasl.RealmCallback; +import javax.security.sasl.RealmChoiceCallback; import java.io.File; import java.io.IOException; import java.nio.ByteBuffer; +import java.security.SecureRandom; import java.util.ArrayList; +import java.util.Base64; import java.util.Collection; import java.util.Collections; import java.util.LinkedList; import java.util.List; import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; @@ -77,12 +91,20 @@ public class RemoteInterpreterEventServer implements RemoteInterpreterEventServi private static final Logger LOGGER = LoggerFactory.getLogger(RemoteInterpreterEventServer.class); private static final Gson GSON = new Gson(); + private static final int CALLBACK_TOKEN_BYTES = 32; + private static final SecureRandom SECURE_RANDOM = new SecureRandom(); + private static final ThreadLocal AUTHENTICATED_INTERPRETER_GROUP = new ThreadLocal<>(); + private static final ThreadLocal AUTHENTICATED_CREDENTIAL = + new ThreadLocal<>(); private int port; private String host; private ZeppelinConfiguration zConf; private TThreadPoolServer thriftServer; private InterpreterSettingManager interpreterSettingManager; + private final Map callbackCredentials = new ConcurrentHashMap<>(); + private final Map callbackCredentialsByAuthenticationId = + new ConcurrentHashMap<>(); private final ScheduledExecutorService appendService = Executors.newSingleThreadScheduledExecutor(); @@ -112,8 +134,17 @@ public void run() { LOGGER.info("InterpreterEventServer is starting at {}:{}", host, port); RemoteInterpreterEventService.Processor processor = new RemoteInterpreterEventService.Processor<>(RemoteInterpreterEventServer.this); + TSaslServerTransport.Factory transportFactory = new TSaslServerTransport.Factory( + RemoteInterpreterEventClient.SASL_MECHANISM, + RemoteInterpreterEventClient.SASL_PROTOCOL, + RemoteInterpreterEventClient.SASL_SERVER_NAME, + RemoteInterpreterEventClient.SASL_PROPERTIES, + createSaslCallbackHandler()); thriftServer = new TThreadPoolServer( - new TThreadPoolServer.Args(tSocket).processor(processor)); + new TThreadPoolServer.Args(tSocket) + .processor(processor) + .transportFactory(transportFactory)); + thriftServer.setServerEventHandler(new AuthenticationContextCleaner()); thriftServer.serve(); } catch (IOException | TTransportException e ) { throw new RuntimeException("Fail to create TServerSocket", e); @@ -164,46 +195,414 @@ public String getHost() { return host; } + public String issueCallbackToken(String interpreterGroupId) { + if (StringUtils.isBlank(interpreterGroupId)) { + throw new IllegalArgumentException("Interpreter group id is required"); + } + byte[] tokenBytes = new byte[CALLBACK_TOKEN_BYTES]; + SECURE_RANDOM.nextBytes(tokenBytes); + String token = Base64.getUrlEncoder().withoutPadding().encodeToString(tokenBytes); + installCallbackCredential( + interpreterGroupId, token, null, -1, CallbackCredentialState.BOOTSTRAP); + return token; + } + + public void registerCallbackToken(String interpreterGroupId, String callbackToken) { + if (StringUtils.isAnyBlank(interpreterGroupId, callbackToken)) { + throw new IllegalArgumentException("Interpreter group id and callback token are required"); + } + installCallbackCredential( + interpreterGroupId, callbackToken, null, -1, CallbackCredentialState.BOOTSTRAP); + } + + public synchronized CallbackCredentialRegistration registerCallbackToken( + String interpreterGroupId, + String callbackToken, + String registeredHost, + int registeredPort) { + if (StringUtils.isAnyBlank(interpreterGroupId, callbackToken, registeredHost) + || registeredPort < 1 || registeredPort > 65535) { + throw new IllegalArgumentException( + "Interpreter group, callback token, and registered endpoint are required"); + } + CallbackCredential current = callbackCredentials.get(interpreterGroupId); + if (current != null) { + boolean sameCredential = StringUtils.equals(current.token, callbackToken) + && StringUtils.equals(current.registeredHost, registeredHost) + && current.registeredPort == registeredPort + && (current.state == CallbackCredentialState.PENDING_RECOVERY + || current.state == CallbackCredentialState.ACTIVE); + if (sameCredential) { + return new CallbackCredentialRegistration(current); + } + throw new IllegalStateException( + "Recovery cannot replace an existing interpreter callback credential"); + } + return new CallbackCredentialRegistration(installCallbackCredential( + interpreterGroupId, callbackToken, registeredHost, registeredPort, + CallbackCredentialState.PENDING_RECOVERY)); + } + + public String getCallbackToken(String interpreterGroupId) { + CallbackCredential credential = callbackCredentials.get(interpreterGroupId); + return credential == null ? null : credential.token; + } + + public synchronized CallbackRecoveryCredential getCallbackRecoveryCredential( + String interpreterGroupId) { + CallbackCredential credential = callbackCredentials.get(interpreterGroupId); + if (credential == null + || credential.state != CallbackCredentialState.ACTIVE + || StringUtils.isBlank(credential.registeredHost) + || credential.registeredPort < 1 + || credential.registeredPort > 65535) { + return null; + } + return new CallbackRecoveryCredential( + credential.token, credential.registeredHost, credential.registeredPort); + } + + public void revokeCallbackToken(String interpreterGroupId, String callbackToken) { + CallbackCredential credential = callbackCredentials.get(interpreterGroupId); + if (credential != null && credential.token.equals(callbackToken)) { + removeCallbackCredential(interpreterGroupId, credential); + } + } + + public boolean isCallbackTokenActive(CallbackCredentialRegistration registration) { + if (registration == null) { + return false; + } + CallbackCredential credential = registration.credential; + return callbackCredentials.get(credential.interpreterGroupId) == credential + && credential.state == CallbackCredentialState.ACTIVE; + } + + public void revokeCallbackToken(CallbackCredentialRegistration registration) { + if (registration != null) { + CallbackCredential credential = registration.credential; + removeCallbackCredential(credential.interpreterGroupId, credential); + } + } + + private synchronized CallbackCredential installCallbackCredential( + String interpreterGroupId, + String callbackToken, + String registeredHost, + int registeredPort, + CallbackCredentialState state) { + CallbackCredential replacement = new CallbackCredential( + interpreterGroupId, callbackToken, registeredHost, registeredPort, state); + CallbackCredential previous = callbackCredentials.put(interpreterGroupId, replacement); + callbackCredentialsByAuthenticationId.put(replacement.authenticationId, replacement); + if (previous != null) { + previous.state = CallbackCredentialState.REVOKED; + callbackCredentialsByAuthenticationId.remove(previous.authenticationId, previous); + } + return replacement; + } + + private synchronized void removeCallbackCredential(String interpreterGroupId, + CallbackCredential credential) { + if (callbackCredentials.remove(interpreterGroupId, credential)) { + credential.state = CallbackCredentialState.REVOKED; + callbackCredentialsByAuthenticationId.remove(credential.authenticationId, credential); + } + } + + private CallbackHandler createSaslCallbackHandler() { + return callbacks -> { + String authenticationId = null; + for (Callback callback : callbacks) { + if (callback instanceof NameCallback) { + authenticationId = ((NameCallback) callback).getDefaultName(); + break; + } + } + + for (Callback callback : callbacks) { + if (callback instanceof NameCallback) { + // The DIGEST-MD5 provider supplies the requested authentication id as the default name. + } else if (callback instanceof PasswordCallback) { + CallbackCredential credential = + callbackCredentialsByAuthenticationId.get(authenticationId); + if (credential != null) { + ((PasswordCallback) callback).setPassword(credential.token.toCharArray()); + } + } else if (callback instanceof AuthorizeCallback) { + AuthorizeCallback authorizeCallback = (AuthorizeCallback) callback; + CallbackCredential credential = callbackCredentialsByAuthenticationId.get( + authorizeCallback.getAuthenticationID()); + boolean authorized = authorizeCallback.getAuthenticationID().equals( + authorizeCallback.getAuthorizationID()) + && credential != null + && callbackCredentials.get(credential.interpreterGroupId) == credential; + authorizeCallback.setAuthorized(authorized); + if (authorized) { + authorizeCallback.setAuthorizedID(authorizeCallback.getAuthorizationID()); + } + } else if (callback instanceof RealmCallback) { + RealmCallback realmCallback = (RealmCallback) callback; + realmCallback.setText(realmCallback.getDefaultText()); + } else if (callback instanceof RealmChoiceCallback) { + ((RealmChoiceCallback) callback).setSelectedIndex(0); + } else { + throw new UnsupportedCallbackException(callback); + } + } + }; + } + + private void requireAuthenticatedGroup(String interpreterGroupId) + throws InterpreterRPCException { + String authenticatedGroup = AUTHENTICATED_INTERPRETER_GROUP.get(); + if (!StringUtils.equals(authenticatedGroup, interpreterGroupId)) { + throw new InterpreterRPCException( + "Authenticated interpreter group does not match the requested interpreter group"); + } + } + + private void requireActiveCredential() throws InterpreterRPCException { + CallbackCredential credential = AUTHENTICATED_CREDENTIAL.get(); + if (credential == null + || callbackCredentials.get(credential.interpreterGroupId) != credential + || credential.state != CallbackCredentialState.ACTIVE) { + throw new InterpreterRPCException( + "Interpreter callback credential is not active"); + } + } + + private void requireLibraryCredential() throws InterpreterRPCException { + CallbackCredential credential = AUTHENTICATED_CREDENTIAL.get(); + if (credential == null + || callbackCredentials.get(credential.interpreterGroupId) != credential + || (credential.state != CallbackCredentialState.BOOTSTRAP + && credential.state != CallbackCredentialState.ACTIVE)) { + throw new InterpreterRPCException( + "Interpreter callback credential cannot access interpreter libraries"); + } + } + + private final class AuthenticationContextCleaner implements TServerEventHandler { + @Override + public void preServe() { + } + + @Override + public ServerContext createContext(org.apache.thrift.protocol.TProtocol input, + org.apache.thrift.protocol.TProtocol output) { + TTransport transport = input.getTransport(); + if (!(transport instanceof TSaslServerTransport) + || ((TSaslServerTransport) transport).getSaslServer() == null + || !((TSaslServerTransport) transport).getSaslServer().isComplete()) { + throw new SecurityException( + "Interpreter callback connection is not authenticated"); + } + String authenticationId = ((TSaslServerTransport) transport) + .getSaslServer().getAuthorizationID(); + CallbackCredential credential = + callbackCredentialsByAuthenticationId.get(authenticationId); + if (credential == null + || callbackCredentials.get(credential.interpreterGroupId) != credential) { + throw new SecurityException("Interpreter callback credential is no longer valid"); + } + return new AuthenticatedConnectionContext(credential.interpreterGroupId, credential); + } + + @Override + public void deleteContext(ServerContext serverContext, + org.apache.thrift.protocol.TProtocol input, + org.apache.thrift.protocol.TProtocol output) { + AUTHENTICATED_INTERPRETER_GROUP.remove(); + AUTHENTICATED_CREDENTIAL.remove(); + } + + @Override + public void processContext(ServerContext serverContext, TTransport input, TTransport output) { + if (!(serverContext instanceof AuthenticatedConnectionContext)) { + throw new SecurityException( + "Interpreter callback connection has no authentication context"); + } + AuthenticatedConnectionContext context = (AuthenticatedConnectionContext) serverContext; + if (callbackCredentials.get(context.interpreterGroupId) != context.credential) { + throw new SecurityException("Interpreter callback credential is no longer valid"); + } + AUTHENTICATED_INTERPRETER_GROUP.set(context.interpreterGroupId); + AUTHENTICATED_CREDENTIAL.set(context.credential); + } + } + + private enum CallbackCredentialState { + BOOTSTRAP, + PENDING_RECOVERY, + ACTIVE, + REVOKED + } + + public static final class CallbackRecoveryCredential { + private final String token; + private final String registeredHost; + private final int registeredPort; + + private CallbackRecoveryCredential(String token, String registeredHost, int registeredPort) { + this.token = token; + this.registeredHost = registeredHost; + this.registeredPort = registeredPort; + } + + public String getToken() { + return token; + } + + public String getRegisteredHost() { + return registeredHost; + } + + public int getRegisteredPort() { + return registeredPort; + } + } + + public static final class CallbackCredentialRegistration { + private final CallbackCredential credential; + + private CallbackCredentialRegistration(CallbackCredential credential) { + this.credential = credential; + } + } + + private static final class CallbackCredential { + private final String interpreterGroupId; + private final String token; + private final String authenticationId; + private String registeredHost; + private int registeredPort = -1; + private volatile CallbackCredentialState state; + + private CallbackCredential(String interpreterGroupId, + String token, + String registeredHost, + int registeredPort, + CallbackCredentialState state) { + this.interpreterGroupId = interpreterGroupId; + this.token = token; + this.authenticationId = RemoteInterpreterEventClient.callbackAuthenticationId( + interpreterGroupId, token); + this.registeredHost = registeredHost; + this.registeredPort = registeredPort; + this.state = state; + } + } + + private static final class AuthenticatedConnectionContext implements ServerContext { + private final String interpreterGroupId; + private final CallbackCredential credential; + + private AuthenticatedConnectionContext(String interpreterGroupId, + CallbackCredential credential) { + this.interpreterGroupId = interpreterGroupId; + this.credential = credential; + } + } + @Override public void registerInterpreterProcess(RegisterInfo registerInfo) throws InterpreterRPCException, TException { - InterpreterGroup interpreterGroup = - interpreterSettingManager.getInterpreterGroupById(registerInfo.getInterpreterGroupId()); - if (interpreterGroup == null) { - LOGGER.warn("Unable to register interpreter process, because no such interpreterGroup: {}", - registerInfo.getInterpreterGroupId()); - return; + requireAuthenticatedGroup(registerInfo.getInterpreterGroupId()); + if (StringUtils.isBlank(registerInfo.getHost()) + || registerInfo.getPort() < 1 || registerInfo.getPort() > 65535) { + throw new InterpreterRPCException("Interpreter process endpoint is invalid"); } - RemoteInterpreterProcess interpreterProcess = - ((ManagedInterpreterGroup) interpreterGroup).getInterpreterProcess(); - if (interpreterProcess == null) { - LOGGER.warn("Unable to register interpreter process, because no interpreter process associated with " + - "interpreterGroup: {}", registerInfo.getInterpreterGroupId()); - return; + CallbackCredential authenticatedCredential = AUTHENTICATED_CREDENTIAL.get(); + if (authenticatedCredential == null) { + throw new InterpreterRPCException("Interpreter callback credential is unavailable"); } - LOGGER.info("Register interpreter process: {}:{}, interpreterGroup: {}", + synchronized (this) { + if (callbackCredentials.get(registerInfo.getInterpreterGroupId()) + != authenticatedCredential + || authenticatedCredential.state == CallbackCredentialState.REVOKED) { + throw new InterpreterRPCException( + "Interpreter callback credential is no longer valid"); + } + boolean endpointMatches = authenticatedCredential.registeredPort == registerInfo.getPort() + && StringUtils.equals( + authenticatedCredential.registeredHost, registerInfo.getHost()); + if (authenticatedCredential.state == CallbackCredentialState.PENDING_RECOVERY) { + if (!endpointMatches) { + throw new InterpreterRPCException( + "Recovered interpreter endpoint does not match its persisted credential"); + } + authenticatedCredential.state = CallbackCredentialState.ACTIVE; + LOGGER.info("Authenticated recovered interpreter process at {}:{} for group {}", registerInfo.getHost(), registerInfo.getPort(), registerInfo.getInterpreterGroupId()); - interpreterProcess.processStarted(registerInfo.port, registerInfo.host); + return; + } + if (authenticatedCredential.state == CallbackCredentialState.ACTIVE) { + if (endpointMatches) { + LOGGER.debug("Interpreter process is already registered at {}:{} for group {}", + registerInfo.getHost(), registerInfo.getPort(), registerInfo.getInterpreterGroupId()); + return; + } + throw new InterpreterRPCException( + "Interpreter process endpoint is already registered for this launch credential"); + } + if (authenticatedCredential.state != CallbackCredentialState.BOOTSTRAP) { + throw new InterpreterRPCException("Interpreter callback credential cannot register"); + } + InterpreterGroup interpreterGroup = + interpreterSettingManager.getInterpreterGroupById(registerInfo.getInterpreterGroupId()); + if (interpreterGroup == null) { + throw new InterpreterRPCException( + "No interpreter group exists for callback registration"); + } + RemoteInterpreterProcess interpreterProcess = + ((ManagedInterpreterGroup) interpreterGroup).getInterpreterProcess(); + if (interpreterProcess == null) { + throw new InterpreterRPCException( + "No interpreter process exists for callback registration"); + } + LOGGER.info("Register interpreter process: {}:{}, interpreterGroup: {}", + registerInfo.getHost(), registerInfo.getPort(), registerInfo.getInterpreterGroupId()); + authenticatedCredential.registeredHost = registerInfo.getHost(); + authenticatedCredential.registeredPort = registerInfo.getPort(); + authenticatedCredential.state = CallbackCredentialState.ACTIVE; + try { + interpreterProcess.processStarted(registerInfo.port, registerInfo.host); + } catch (RuntimeException e) { + removeCallbackCredential(registerInfo.getInterpreterGroupId(), authenticatedCredential); + throw e; + } + } } @Override public void unRegisterInterpreterProcess(String intpGroupId) throws InterpreterRPCException, TException { - LOGGER.info("Unregister interpreter process: {}", intpGroupId); - InterpreterGroup interpreterGroup = - interpreterSettingManager.getInterpreterGroupById(intpGroupId); - if (interpreterGroup == null) { - LOGGER.warn("Unable to unregister interpreter process because no such interpreterGroup: {}", - intpGroupId); - return; + requireActiveCredential(); + requireAuthenticatedGroup(intpGroupId); + CallbackCredential authenticatedCredential = AUTHENTICATED_CREDENTIAL.get(); + try { + LOGGER.info("Unregister interpreter process: {}", intpGroupId); + InterpreterGroup interpreterGroup = + interpreterSettingManager.getInterpreterGroupById(intpGroupId); + if (interpreterGroup == null) { + LOGGER.warn("Unable to unregister interpreter process because no such interpreterGroup: {}", + intpGroupId); + return; + } + // Close RemoteInterpreter when RemoteInterpreterServer already timeout. + // Otherwise the ProgressBar will be missing when rerun after the + // RemoteInterpreterServer timeout + // and old RemoteInterpreterGroups will always alive after GC. + interpreterGroup.close(); + interpreterSettingManager.removeInterpreterGroup(intpGroupId); + } finally { + removeCallbackCredential(intpGroupId, authenticatedCredential); } - // Close RemoteInterpreter when RemoteInterpreterServer already timeout. - // Otherwise the ProgressBar will be missing when rerun after the RemoteInterpreterServer timeout - // and old RemoteInterpreterGroups will always alive after GC. - interpreterGroup.close(); - interpreterSettingManager.removeInterpreterGroup(intpGroupId); } @Override public void sendWebUrl(WebUrlInfo weburlInfo) throws InterpreterRPCException, TException { + requireActiveCredential(); + requireAuthenticatedGroup(weburlInfo.getInterpreterGroupId()); InterpreterGroup interpreterGroup = interpreterSettingManager.getInterpreterGroupById(weburlInfo.getInterpreterGroupId()); if (interpreterGroup == null) { @@ -216,6 +615,7 @@ public void sendWebUrl(WebUrlInfo weburlInfo) throws InterpreterRPCException, TE @Override public void appendOutput(OutputAppendEvent event) throws InterpreterRPCException, TException { + requireActiveCredential(); if (event.getAppId() == null) { runner.appendBuffer( event.getNoteId(), event.getParagraphId(), event.getIndex(), event.getData()); @@ -227,6 +627,7 @@ public void appendOutput(OutputAppendEvent event) throws InterpreterRPCException @Override public void updateOutput(OutputUpdateEvent event) throws InterpreterRPCException, TException { + requireActiveCredential(); if (event.getAppId() == null) { listener.onOutputUpdated(event.getNoteId(), event.getParagraphId(), event.getIndex(), InterpreterResult.Type.valueOf(event.getType()), event.getData()); @@ -238,6 +639,7 @@ public void updateOutput(OutputUpdateEvent event) throws InterpreterRPCException @Override public void updateAllOutput(OutputUpdateAllEvent event) throws InterpreterRPCException, TException { + requireActiveCredential(); listener.onOutputClear(event.getNoteId(), event.getParagraphId()); for (int i = 0; i < event.getMsg().size(); i++) { RemoteInterpreterResultMessage msg = event.getMsg().get(i); @@ -248,28 +650,33 @@ public void updateAllOutput(OutputUpdateAllEvent event) throws InterpreterRPCExc @Override public void appendAppOutput(AppOutputAppendEvent event) throws InterpreterRPCException, TException { + requireActiveCredential(); appListener.onOutputAppend(event.noteId, event.paragraphId, event.index, event.appId, event.data); } @Override public void updateAppOutput(AppOutputUpdateEvent event) throws InterpreterRPCException, TException { + requireActiveCredential(); appListener.onOutputUpdated(event.noteId, event.paragraphId, event.index, event.appId, InterpreterResult.Type.valueOf(event.type), event.data); } @Override public void updateAppStatus(AppStatusUpdateEvent event) throws InterpreterRPCException, TException { + requireActiveCredential(); appListener.onStatusChange(event.noteId, event.paragraphId, event.appId, event.status); } @Override public void checkpointOutput(String noteId, String paragraphId) throws InterpreterRPCException, TException { + requireActiveCredential(); listener.checkpointOutput(noteId, paragraphId); } @Override public void runParagraphs(RunParagraphsEvent event) throws InterpreterRPCException, TException { + requireActiveCredential(); try { listener.runParagraphs(event.getNoteId(), event.getParagraphIndices(), event.getParagraphIds(), event.getCurParagraphId()); @@ -285,6 +692,8 @@ public void runParagraphs(RunParagraphsEvent event) throws InterpreterRPCExcepti @Override public void addAngularObject(String intpGroupId, String json) throws InterpreterRPCException, TException { + requireActiveCredential(); + requireAuthenticatedGroup(intpGroupId); LOGGER.debug("Add AngularObject, interpreterGroupId: {}, json: {}", intpGroupId, json); AngularObject angularObject = AngularObject.fromJson(json); InterpreterGroup interpreterGroup = @@ -314,6 +723,8 @@ public void addAngularObject(String intpGroupId, String json) throws Interpreter @Override public void updateAngularObject(String intpGroupId, String json) throws InterpreterRPCException, TException { + requireActiveCredential(); + requireAuthenticatedGroup(intpGroupId); AngularObject angularObject = AngularObject.fromJson(json); InterpreterGroup interpreterGroup = interpreterSettingManager.getInterpreterGroupById(intpGroupId); @@ -351,6 +762,8 @@ public void removeAngularObject(String intpGroupId, String noteId, String paragraphId, String name) throws InterpreterRPCException, TException { + requireActiveCredential(); + requireAuthenticatedGroup(intpGroupId); InterpreterGroup interpreterGroup = interpreterSettingManager.getInterpreterGroupById(intpGroupId); if (interpreterGroup == null) { @@ -376,6 +789,8 @@ public void removeAngularObject(String intpGroupId, @Override public void sendParagraphInfo(String intpGroupId, String json) throws InterpreterRPCException, TException { + requireActiveCredential(); + requireAuthenticatedGroup(intpGroupId); InterpreterGroup interpreterGroup = interpreterSettingManager.getInterpreterGroupById(intpGroupId); if (interpreterGroup == null) { @@ -395,6 +810,8 @@ public void sendParagraphInfo(String intpGroupId, String json) throws Interprete @Override public List getAllResources(String intpGroupId) throws InterpreterRPCException, TException { + requireActiveCredential(); + requireAuthenticatedGroup(intpGroupId); ResourceSet resourceSet = getAllResourcePoolExcept(intpGroupId); List resourceList = new LinkedList<>(); for (Resource r : resourceSet) { @@ -405,6 +822,7 @@ public List getAllResources(String intpGroupId) throws InterpreterRPCExc @Override public ByteBuffer getResource(String resourceIdJson) throws InterpreterRPCException, TException { + requireActiveCredential(); ResourceId resourceId = ResourceId.fromJson(resourceIdJson); Object o = getResource(resourceId); ByteBuffer obj; @@ -430,6 +848,8 @@ public ByteBuffer getResource(String resourceIdJson) throws InterpreterRPCExcept @Override public ByteBuffer invokeMethod(String intpGroupId, String invokeMethodJson) throws InterpreterRPCException, TException { + requireActiveCredential(); + requireAuthenticatedGroup(intpGroupId); InvokeResourceMethodEventMessage invokeMethodMessage = InvokeResourceMethodEventMessage.fromJson(invokeMethodJson); Object ret = invokeResourceMethod(intpGroupId, invokeMethodMessage); @@ -449,6 +869,7 @@ public ByteBuffer invokeMethod(String intpGroupId, String invokeMethodJson) @Override public List getParagraphList(String user, String noteId) throws InterpreterRPCException, TException { + requireActiveCredential(); LOGGER.info("get paragraph list from remote interpreter noteId: {}, user = {}",noteId, user); if (user != null && noteId != null) { @@ -567,6 +988,7 @@ public void updateParagraphConfig(String noteId, String paragraphId, Map config) throws InterpreterRPCException, TException { + requireActiveCredential(); try { LOGGER.info("Update paragraph config"); interpreterSettingManager.getNotebook().processNote(noteId, @@ -583,6 +1005,7 @@ public void updateParagraphConfig(String noteId, @Override public List getAllLibraryMetadatas(String interpreter) throws TException { + requireLibraryCredential(); if (StringUtils.isBlank(interpreter)) { LOGGER.warn("Interpreter is blank"); return Collections.emptyList(); @@ -615,6 +1038,7 @@ public List getAllLibraryMetadatas(String interpreter) throws T @Override public ByteBuffer getLibrary(String interpreter, String libraryName) throws TException { + requireLibraryCredential(); if (StringUtils.isAnyBlank(interpreter, libraryName)) { LOGGER.warn("Interpreter \"{}\" or libraryName \"{}\" is blank", interpreter, libraryName); return null; diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/launcher/InterpreterLaunchContext.java b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/launcher/InterpreterLaunchContext.java index cc90a63ee0c..ad46ac8d365 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/launcher/InterpreterLaunchContext.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/launcher/InterpreterLaunchContext.java @@ -37,6 +37,7 @@ public class InterpreterLaunchContext { private String interpreterSettingName; private int intpEventServerPort; private String intpEventServerHost; + private String intpEventCallbackToken; public InterpreterLaunchContext(Properties properties, InterpreterOption option, @@ -99,4 +100,12 @@ public int getIntpEventServerPort() { public String getIntpEventServerHost() { return intpEventServerHost; } + + public String getIntpEventCallbackToken() { + return intpEventCallbackToken; + } + + public void setIntpEventCallbackToken(String intpEventCallbackToken) { + this.intpEventCallbackToken = intpEventCallbackToken; + } } diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/launcher/SparkInterpreterLauncher.java b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/launcher/SparkInterpreterLauncher.java index 98e0e5e0b82..7e7c2d8cbe9 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/launcher/SparkInterpreterLauncher.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/launcher/SparkInterpreterLauncher.java @@ -38,6 +38,7 @@ import org.apache.commons.lang3.StringUtils; import org.apache.zeppelin.conf.ZeppelinConfiguration; import org.apache.zeppelin.interpreter.recovery.RecoveryStorage; +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterEventClient; import org.apache.zeppelin.interpreter.remote.RemoteInterpreterUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -192,6 +193,23 @@ && getDeployMode(context).equals("cluster")) { } } + if (StringUtils.isNotBlank(context.getIntpEventCallbackToken()) + && "cluster".equals(getDeployMode(context))) { + if (isYarnMode(context)) { + sparkProperties.setProperty( + "spark.yarn.appMasterEnv." + RemoteInterpreterEventClient.CALLBACK_TOKEN_ENV, + context.getIntpEventCallbackToken()); + } else if (getSparkMaster(context).startsWith("k8s://")) { + sparkProperties.setProperty( + "spark.kubernetes.driverEnv." + RemoteInterpreterEventClient.CALLBACK_TOKEN_ENV, + context.getIntpEventCallbackToken()); + } else { + throw new IOException("Authenticated interpreter callbacks do not support Spark " + + "cluster deploy mode for master " + getSparkMaster(context) + + "; use client mode, Yarn, or Kubernetes"); + } + } + StringJoiner sparkConfSJ = new StringJoiner("|"); if (context.getOption().isUserImpersonate() && zConf.getZeppelinImpersonateSparkProxyUser()) { sparkConfSJ.add("--proxy-user"); @@ -261,7 +279,7 @@ && getDeployMode(context).equals("cluster")) { env.put("HADOOP_USER_NAME", userName); } } - LOGGER.info("buildEnvFromProperties: {}", env); + LOGGER.info("buildEnvFromProperties keys: {}", env.keySet()); return env; } diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/launcher/StandardInterpreterLauncher.java b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/launcher/StandardInterpreterLauncher.java index 77e6e7bddce..7487621058c 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/launcher/StandardInterpreterLauncher.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/launcher/StandardInterpreterLauncher.java @@ -25,6 +25,7 @@ import org.apache.zeppelin.interpreter.InterpreterRunner; import org.apache.zeppelin.interpreter.recovery.RecoveryStorage; import org.apache.zeppelin.interpreter.remote.ExecRemoteInterpreterProcess; +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterEventClient; import org.apache.zeppelin.interpreter.remote.RemoteInterpreterRunningProcess; import org.apache.zeppelin.interpreter.remote.RemoteInterpreterUtils; import org.slf4j.Logger; @@ -65,7 +66,9 @@ public InterpreterClient launchDirectly(InterpreterLaunchContext context) throws context.getIntpEventServerPort(), option.getHost(), option.getPort(), - false); + false, + null, + null); } else { // create new remote process String localRepoPath = zConf.getInterpreterLocalRepoPath() + File.separator @@ -89,6 +92,10 @@ public Map buildEnvFromProperties(InterpreterLaunchContext conte } } env.put("INTERPRETER_GROUP_ID", context.getInterpreterGroupId()); + if (StringUtils.isNotBlank(context.getIntpEventCallbackToken())) { + env.put(RemoteInterpreterEventClient.CALLBACK_TOKEN_ENV, + context.getIntpEventCallbackToken()); + } return env; } } diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/recovery/FileSystemRecoveryStorage.java b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/recovery/FileSystemRecoveryStorage.java index 26cecb256ed..64dd10f2c8a 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/recovery/FileSystemRecoveryStorage.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/recovery/FileSystemRecoveryStorage.java @@ -27,6 +27,8 @@ import org.slf4j.LoggerFactory; import java.io.IOException; +import java.nio.file.attribute.PosixFilePermission; +import java.util.EnumSet; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -76,9 +78,11 @@ private void save(String interpreterSettingName) throws IOException { InterpreterSetting interpreterSetting = interpreterSettingManager.getInterpreterSettingByName(interpreterSettingName); String recoveryData = RecoveryUtils.getRecoveryData(interpreterSetting); - LOGGER.debug("Updating recovery data of {}: {}", interpreterSettingName, recoveryData); + LOGGER.debug("Updating recovery data of {} with {} bytes", + interpreterSettingName, recoveryData.length()); Path recoveryFile = new Path(recoveryDir, interpreterSettingName + ".recovery"); - fs.writeFile(recoveryData, recoveryFile, true); + fs.writeFile(recoveryData, recoveryFile, true, + EnumSet.of(PosixFilePermission.OWNER_READ, PosixFilePermission.OWNER_WRITE)); } @Override diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/recovery/LocalRecoveryStorage.java b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/recovery/LocalRecoveryStorage.java index 078722e2a6f..c8d86b01386 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/recovery/LocalRecoveryStorage.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/recovery/LocalRecoveryStorage.java @@ -28,6 +28,8 @@ import java.io.File; import java.io.IOException; +import java.nio.file.attribute.PosixFilePermission; +import java.util.EnumSet; import java.util.HashMap; import java.util.Map; @@ -90,8 +92,10 @@ private void save(String interpreterSettingName) throws IOException { InterpreterSetting interpreterSetting = interpreterSettingManager.getInterpreterSettingByName(interpreterSettingName); String recoveryData = RecoveryUtils.getRecoveryData(interpreterSetting); - LOGGER.debug("Updating recovery data of {}: {}", interpreterSettingName, recoveryData); + LOGGER.debug("Updating recovery data of {} with {} bytes", + interpreterSettingName, recoveryData.length()); File recoveryFile = new File(recoveryDir, interpreterSettingName + ".recovery"); - org.apache.zeppelin.util.FileUtils.atomicWriteToFile(recoveryData, recoveryFile); + org.apache.zeppelin.util.FileUtils.atomicWriteToFile(recoveryData, recoveryFile, + EnumSet.of(PosixFilePermission.OWNER_READ, PosixFilePermission.OWNER_WRITE)); } } diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/recovery/RecoveryUtils.java b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/recovery/RecoveryUtils.java index b8d38a53459..cd2806a2a42 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/recovery/RecoveryUtils.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/recovery/RecoveryUtils.java @@ -23,6 +23,7 @@ import org.apache.zeppelin.interpreter.InterpreterSetting; import org.apache.zeppelin.interpreter.InterpreterSettingManager; import org.apache.zeppelin.interpreter.ManagedInterpreterGroup; +import org.apache.zeppelin.interpreter.RemoteInterpreterEventServer; import org.apache.zeppelin.interpreter.launcher.InterpreterClient; import org.apache.zeppelin.interpreter.remote.RemoteInterpreterProcess; import org.apache.zeppelin.interpreter.remote.RemoteInterpreterRunningProcess; @@ -54,8 +55,19 @@ public static String getRecoveryData(InterpreterSetting interpreterSetting) { for (ManagedInterpreterGroup interpreterGroup : interpreterSetting.getAllInterpreterGroups()) { RemoteInterpreterProcess interpreterProcess = interpreterGroup.getInterpreterProcess(); if (interpreterProcess != null && interpreterProcess.isRunning()) { - recoveryData.add(interpreterGroup.getId() + "\t" + interpreterProcess.getHost() + ":" + - interpreterProcess.getPort()); + RemoteInterpreterEventServer.CallbackRecoveryCredential callbackCredential = + interpreterSetting.getInterpreterSettingManager().getInterpreterEventServer() + .getCallbackRecoveryCredential(interpreterGroup.getId()); + if (callbackCredential == null) { + LOGGER.warn("Skip recovery data for interpreter group {} because its callback " + + "credential is unavailable", interpreterGroup.getId()); + continue; + } + recoveryData.add(interpreterGroup.getId() + + "\t" + interpreterProcess.getHost() + ":" + interpreterProcess.getPort() + + "\t" + callbackCredential.getToken() + + "\t" + callbackCredential.getRegisteredHost() + + "\t" + callbackCredential.getRegisteredPort()); } } } @@ -87,15 +99,47 @@ public static Map restoreFromRecoveryData(String reco if (!StringUtils.isBlank(recoveryData)) { for (String line : recoveryData.split(System.lineSeparator())) { - String[] tokens = line.split("\t"); + String[] tokens = line.split("\t", -1); + if (tokens.length < 2) { + LOGGER.warn("Ignore malformed interpreter recovery record"); + continue; + } String interpreterGroupId = tokens[0]; String[] hostPort = tokens[1].split(":"); + if (tokens.length < 5 + || StringUtils.isAnyBlank(tokens[2], tokens[3])) { + LOGGER.warn("Interpreter recovery record for {} has no callback identity; " + + "skipping it because it cannot be recovered securely", + interpreterGroupId); + continue; + } + + String callbackToken = tokens[2]; + String registeredHost = tokens[3]; + int registeredPort; + try { + registeredPort = Integer.parseInt(tokens[4]); + } catch (NumberFormatException e) { + LOGGER.warn("Interpreter recovery record for {} has an invalid callback port", + interpreterGroupId); + continue; + } + if (registeredPort < 1 || registeredPort > 65535) { + LOGGER.warn("Interpreter recovery record for {} has an invalid callback port", + interpreterGroupId); + continue; + } + RemoteInterpreterEventServer.CallbackCredentialRegistration callbackRegistration = + interpreterSettingManager.getInterpreterEventServer().registerCallbackToken( + interpreterGroupId, callbackToken, registeredHost, registeredPort); + RemoteInterpreterRunningProcess client = new RemoteInterpreterRunningProcess( interpreterSettingName, interpreterGroupId, connectTimeout, connectionPoolSize, interpreterSettingManager.getInterpreterEventServer().getHost(), interpreterSettingManager.getInterpreterEventServer().getPort(), - hostPort[0], Integer.parseInt(hostPort[1]), true); + hostPort[0], Integer.parseInt(hostPort[1]), true, + interpreterSettingManager.getInterpreterEventServer(), callbackRegistration); clients.put(interpreterGroupId, client); LOGGER.info("Recovering Interpreter Process: " + interpreterGroupId + ", " + hostPort[0] + ":" + hostPort[1]); diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/ExecRemoteInterpreterProcess.java b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/ExecRemoteInterpreterProcess.java index 6dd2793adf8..0cdb760f050 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/ExecRemoteInterpreterProcess.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/ExecRemoteInterpreterProcess.java @@ -247,6 +247,8 @@ public void onProcessComplete(int exitValue) { // is 0. if (exitValue != 0) { transition(State.TERMINATED); + ExecRemoteInterpreterProcess.this.processStopped( + "Interpreter process exited with status " + exitValue); } else { transition(State.COMPLETED); } diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterManagedProcess.java b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterManagedProcess.java index 02cedb322fe..d035523bedd 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterManagedProcess.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterManagedProcess.java @@ -100,6 +100,7 @@ public void processStarted(int port, String host) { // after detecting yarn app is killed/failed. public void processStopped(String errorMessage) { this.errorMessage = errorMessage; + notifyTermination(); } public Map getEnv() { diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterProcess.java b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterProcess.java index 95802a64fe7..4c450146a3f 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterProcess.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterProcess.java @@ -30,6 +30,8 @@ import java.io.IOException; import java.text.SimpleDateFormat; import java.util.Date; +import java.util.HashMap; +import java.util.Map; /** * Abstract class for interpreter process @@ -43,6 +45,8 @@ public abstract class RemoteInterpreterProcess implements InterpreterClient, Aut protected int intpEventServerPort; private PooledRemoteClient remoteClient; private String startTime; + private Runnable terminationListener = () -> { }; + private boolean terminated; public RemoteInterpreterProcess(int connectTimeout, int connectionPoolSize, @@ -72,6 +76,20 @@ public String getStartTime() { return startTime; } + public synchronized void setTerminationListener(Runnable terminationListener) { + this.terminationListener = terminationListener; + if (terminated) { + terminationListener.run(); + } + } + + protected synchronized void notifyTermination() { + if (!terminated) { + terminated = true; + terminationListener.run(); + } + } + @Override public void close() { if (remoteClient != null) { @@ -100,20 +118,30 @@ public R callRemoteFunction(PooledRemoteClient.RemoteFunction fun return remoteClient.callRemoteFunction(func); } - public void init(ZeppelinConfiguration zConf) { + public void init(ZeppelinConfiguration zConf, + String interpreterGroupId) { callRemoteFunction(client -> { - client.init(zConf.getCompleteConfiguration()); + client.init(createInitConfiguration(zConf, interpreterGroupId)); return null; }); } + static Map createInitConfiguration(ZeppelinConfiguration zConf, + String interpreterGroupId) { + Map configuration = new HashMap<>(zConf.getCompleteConfiguration()); + // The callback credential is bootstrap material supplied by the launcher. Never copy it + // into the unencrypted interpreter RPC payload, including when an operator accidentally + // configures the reserved property in zeppelin-site.xml. + configuration.remove(RemoteInterpreterEventClient.CALLBACK_TOKEN_PROPERTY); + configuration.put(RemoteInterpreterEventClient.INTERPRETER_GROUP_PROPERTY, + interpreterGroupId); + return configuration; + } + @Override public boolean recover() { try { - remoteClient.callRemoteFunction(client -> { - client.reconnect(intpEventServerHost, intpEventServerPort); - return null; - }); + reconnectToEventServer(); return true; } catch (Exception e) { LOGGER.error("Fail to recover remote interpreter process: {}" , e.getMessage()); @@ -121,6 +149,13 @@ public boolean recover() { } } + protected void reconnectToEventServer() { + remoteClient.callRemoteFunction(client -> { + client.reconnect(intpEventServerHost, intpEventServerPort); + return null; + }); + } + /** * called by RemoteInterpreterEventServer to notify that RemoteInterpreter Process is started diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterRunningProcess.java b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterRunningProcess.java index 71c3a69b3e9..eda2be2f03b 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterRunningProcess.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterRunningProcess.java @@ -16,6 +16,7 @@ */ package org.apache.zeppelin.interpreter.remote; +import org.apache.zeppelin.interpreter.RemoteInterpreterEventServer; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -30,6 +31,9 @@ public class RemoteInterpreterRunningProcess extends RemoteInterpreterProcess { private final String interpreterSettingName; private final String interpreterGroupId; private final boolean isRecovery; + private final RemoteInterpreterEventServer interpreterEventServer; + private final RemoteInterpreterEventServer.CallbackCredentialRegistration + callbackCredentialRegistration; public RemoteInterpreterRunningProcess( String interpreterSettingName, @@ -40,13 +44,18 @@ public RemoteInterpreterRunningProcess( int intpEventServerPort, String host, int port, - boolean isRecovery) { + boolean isRecovery, + RemoteInterpreterEventServer interpreterEventServer, + RemoteInterpreterEventServer.CallbackCredentialRegistration + callbackCredentialRegistration) { super(connectTimeout, connectionPoolSize, intpEventServerHost, intpEventServerPort); this.interpreterSettingName = interpreterSettingName; this.interpreterGroupId = interpreterGroupId; this.host = host; this.port = port; this.isRecovery = isRecovery; + this.interpreterEventServer = interpreterEventServer; + this.callbackCredentialRegistration = callbackCredentialRegistration; } @Override @@ -74,6 +83,19 @@ public void start(String userName) { // assume process is externally managed. nothing to do } + @Override + public boolean recover() { + boolean recovered = super.recover(); + if (recovered && interpreterEventServer != null) { + recovered = interpreterEventServer.isCallbackTokenActive( + callbackCredentialRegistration); + } + if (!recovered && interpreterEventServer != null) { + interpreterEventServer.revokeCallbackToken(callbackCredentialRegistration); + } + return recovered; + } + @Override public void stop() { // assume process is externally managed. nothing to do. But will kill it diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/notebook/FileSystemStorage.java b/zeppelin-server/src/main/java/org/apache/zeppelin/notebook/FileSystemStorage.java index 5fc60e74e55..3e222720a3b 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/notebook/FileSystemStorage.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/notebook/FileSystemStorage.java @@ -223,7 +223,10 @@ public void writeFile(final String content, final Path file, boolean writeTempFi writeFile(content, file, writeTempFileFirst, null); } - public void writeFile(final String content, final Path file, boolean writeTempFileFirst, Set permissions) + public void writeFile(final String content, + final Path file, + boolean writeTempFileFirst, + Set permissions) throws IOException { FsPermission fsPermission; if (permissions == null || permissions.isEmpty()) { @@ -240,8 +243,14 @@ public Void call() throws IOException { InputStream in = new ByteArrayInputStream(content.getBytes( zConf.getString(ZeppelinConfiguration.ConfVars.ZEPPELIN_ENCODING))); Path tmpFile = new Path(file.toString() + ".tmp"); - IOUtils.copyBytes(in, fs.create(tmpFile), hadoopConf); - fs.setPermission(tmpFile, fsPermission); + IOUtils.copyBytes(in, fs.create( + tmpFile, + fsPermission, + true, + hadoopConf.getInt("io.file.buffer.size", 4096), + fs.getDefaultReplication(tmpFile), + fs.getDefaultBlockSize(tmpFile), + null), hadoopConf); fs.delete(file, true); fs.rename(tmpFile, file); return null; diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/RemoteInterpreterEventServerTest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/RemoteInterpreterEventServerTest.java new file mode 100644 index 00000000000..532ea3e3260 --- /dev/null +++ b/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/RemoteInterpreterEventServerTest.java @@ -0,0 +1,331 @@ +/* + * 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.zeppelin.interpreter; + +import org.apache.thrift.protocol.TBinaryProtocol; +import org.apache.thrift.transport.TSocket; +import org.apache.zeppelin.conf.ZeppelinConfiguration; +import org.apache.zeppelin.interpreter.launcher.InterpreterClient; +import org.apache.zeppelin.interpreter.recovery.RecoveryUtils; +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterEventClient; +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterProcess; +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterRunningProcess; +import org.apache.zeppelin.interpreter.thrift.RegisterInfo; +import org.apache.zeppelin.interpreter.thrift.RemoteInterpreterEventService; +import org.apache.zeppelin.interpreter.thrift.RunParagraphsEvent; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.Collections; +import java.util.Map; +import java.util.OptionalInt; +import java.util.Properties; +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +class RemoteInterpreterEventServerTest { + + private RemoteInterpreterEventServer eventServer; + private InterpreterSettingManager interpreterSettingManager; + + @BeforeEach + void setUp() throws Exception { + ZeppelinConfiguration zConf = mock(ZeppelinConfiguration.class); + when(zConf.getZeppelinServerRpcPort()).thenReturn(OptionalInt.of(0)); + + interpreterSettingManager = mock(InterpreterSettingManager.class); + when(interpreterSettingManager.getAllInterpreterGroup()) + .thenReturn(Collections.emptyList()); + + eventServer = new RemoteInterpreterEventServer(zConf, interpreterSettingManager); + eventServer.start(); + } + + @AfterEach + void tearDown() { + eventServer.stop(); + } + + @Test + void authenticatesCallbackConnectionsAndInvalidatesRotatedCredentials() throws Exception { + String groupId = "group-a"; + String callbackToken = eventServer.issueCallbackToken(groupId); + String otherToken = eventServer.issueCallbackToken("group-b"); + assertNotEquals(callbackToken, otherToken); + allowRegistration(groupId); + + try (RemoteInterpreterEventClient client = new RemoteInterpreterEventClient( + eventServer.getHost(), eventServer.getPort(), 1, groupId, callbackToken)) { + assertEquals(Collections.emptyList(), client.getAllLibraryMetadatas("")); + assertThrows(RuntimeException.class, () -> client.callRemoteFunction( + remote -> remote.getAllResources(groupId))); + client.registerInterpreterProcess(new RegisterInfo("127.0.0.1", 12345, groupId)); + assertEquals(Collections.emptyList(), client.callRemoteFunction( + remote -> remote.getAllResources(groupId))); + + RuntimeException crossGroupFailure = assertThrows(RuntimeException.class, + () -> client.callRemoteFunction(remote -> remote.getAllResources("group-b"))); + assertEquals( + "Authenticated interpreter group does not match the requested interpreter group", + crossGroupFailure.getMessage()); + + String rotatedToken = eventServer.issueCallbackToken(groupId); + assertThrows(RuntimeException.class, () -> client.callRemoteFunction( + remote -> remote.getAllResources(groupId))); + + try (RemoteInterpreterEventClient rotatedClient = new RemoteInterpreterEventClient( + eventServer.getHost(), eventServer.getPort(), 1, groupId, rotatedToken)) { + assertThrows(RuntimeException.class, () -> rotatedClient.callRemoteFunction( + remote -> remote.getAllResources(groupId))); + rotatedClient.registerInterpreterProcess( + new RegisterInfo("127.0.0.1", 23456, groupId)); + assertEquals(Collections.emptyList(), rotatedClient.callRemoteFunction( + remote -> remote.getAllResources(groupId))); + } + } + + assertThrows(RuntimeException.class, () -> { + try (RemoteInterpreterEventClient client = new RemoteInterpreterEventClient( + eventServer.getHost(), eventServer.getPort(), 1, groupId, "wrong-token")) { + client.callRemoteFunction(remote -> remote.getAllResources(groupId)); + } + }); + + try (TSocket transport = new TSocket(eventServer.getHost(), eventServer.getPort(), 2_000)) { + transport.open(); + RemoteInterpreterEventService.Client rawClient = + new RemoteInterpreterEventService.Client(new TBinaryProtocol(transport)); + assertThrows(Exception.class, () -> rawClient.getAllResources(groupId)); + } + } + + @Test + void rejectsEndpointReplacementForTheSameLaunchCredential() { + String groupId = "group-a"; + String callbackToken = eventServer.issueCallbackToken(groupId); + RemoteInterpreterProcess interpreterProcess = allowRegistration(groupId); + AtomicBoolean activeDuringProcessStarted = new AtomicBoolean(); + doAnswer(invocation -> { + try (RemoteInterpreterEventClient activeClient = new RemoteInterpreterEventClient( + eventServer.getHost(), eventServer.getPort(), 1, groupId, callbackToken)) { + activeDuringProcessStarted.set(Collections.emptyList().equals( + activeClient.callRemoteFunction(remote -> remote.getAllResources(groupId)))); + } + return null; + }).when(interpreterProcess).processStarted(12345, "127.0.0.1"); + + try (RemoteInterpreterEventClient client = new RemoteInterpreterEventClient( + eventServer.getHost(), eventServer.getPort(), 1, groupId, callbackToken)) { + client.registerInterpreterProcess(new RegisterInfo("127.0.0.1", 12345, groupId)); + client.registerInterpreterProcess(new RegisterInfo("127.0.0.1", 12345, groupId)); + assertThrows(RuntimeException.class, () -> client.registerInterpreterProcess( + new RegisterInfo("127.0.0.1", 54321, groupId))); + } + + verify(interpreterProcess, times(1)).processStarted(12345, "127.0.0.1"); + assertTrue(activeDuringProcessStarted.get()); + } + + @Test + void revokesBootstrapCredentialWhenProcessStartupNotificationFails() { + String groupId = "group-a"; + String callbackToken = eventServer.issueCallbackToken(groupId); + RemoteInterpreterProcess interpreterProcess = allowRegistration(groupId); + doThrow(new IllegalStateException("failed to publish endpoint")) + .when(interpreterProcess).processStarted(12345, "127.0.0.1"); + + try (RemoteInterpreterEventClient client = new RemoteInterpreterEventClient( + eventServer.getHost(), eventServer.getPort(), 1, groupId, callbackToken)) { + assertThrows(RuntimeException.class, () -> client.registerInterpreterProcess( + new RegisterInfo("127.0.0.1", 12345, groupId))); + } + + assertNull(eventServer.getCallbackToken(groupId)); + } + + @Test + void recoveredCredentialRetainsItsRegisteredEndpoint() { + String groupId = "group-a"; + String callbackToken = "recovered-token"; + RemoteInterpreterEventServer.CallbackCredentialRegistration registration = + eventServer.registerCallbackToken( + groupId, callbackToken, "127.0.0.1", 12345); + + try (RemoteInterpreterEventClient client = new RemoteInterpreterEventClient( + eventServer.getHost(), eventServer.getPort(), 1, groupId, callbackToken)) { + RunParagraphsEvent runEvent = new RunParagraphsEvent( + "note", Collections.emptyList(), Collections.emptyList(), "paragraph"); + assertThrows(RuntimeException.class, () -> client.callRemoteFunction(remote -> { + remote.runParagraphs(runEvent); + return null; + })); + assertThrows(RuntimeException.class, () -> client.getAllLibraryMetadatas("")); + assertThrows(RuntimeException.class, () -> client.registerInterpreterProcess( + new RegisterInfo("127.0.0.1", 54321, groupId))); + client.registerInterpreterProcess(new RegisterInfo("127.0.0.1", 12345, groupId)); + assertEquals(Collections.emptyList(), client.callRemoteFunction( + remote -> remote.getAllResources(groupId))); + assertThrows(RuntimeException.class, () -> client.registerInterpreterProcess( + new RegisterInfo("127.0.0.1", 54321, groupId))); + } + + assertTrue(eventServer.isCallbackTokenActive(registration)); + } + + @Test + void staleRecoveryRegistrationCannotRevokeReplacementCredential() { + String groupId = "group-a"; + RemoteInterpreterEventServer.CallbackCredentialRegistration staleRegistration = + eventServer.registerCallbackToken( + groupId, "stale-token", "127.0.0.1", 12345); + String currentToken = eventServer.issueCallbackToken(groupId); + allowRegistration(groupId); + + try (RemoteInterpreterEventClient client = new RemoteInterpreterEventClient( + eventServer.getHost(), eventServer.getPort(), 1, groupId, currentToken)) { + client.registerInterpreterProcess(new RegisterInfo("127.0.0.1", 23456, groupId)); + } + + eventServer.revokeCallbackToken(staleRegistration); + + try (RemoteInterpreterEventClient client = new RemoteInterpreterEventClient( + eventServer.getHost(), eventServer.getPort(), 1, groupId, currentToken)) { + assertEquals(Collections.emptyList(), client.callRemoteFunction( + remote -> remote.getAllResources(groupId))); + } + } + + @Test + void recoveryRequiresCallbackProofAndRevokesUnprovenCredential() { + String groupId = "group-a"; + String callbackToken = "recovered-token"; + RemoteInterpreterEventServer.CallbackCredentialRegistration registration = + eventServer.registerCallbackToken( + groupId, callbackToken, "127.0.0.1", 12345); + RemoteInterpreterRunningProcess process = recoveryProcess( + groupId, registration, () -> { }); + + assertFalse(process.recover()); + assertFalse(eventServer.isCallbackTokenActive(registration)); + assertNull(eventServer.getCallbackToken(groupId)); + } + + @Test + void recoveryAcceptsProofFromInterpreterHoldingPersistedCredential() { + String groupId = "group-a"; + String callbackToken = "recovered-token"; + RemoteInterpreterEventServer.CallbackCredentialRegistration registration = + eventServer.registerCallbackToken( + groupId, callbackToken, "127.0.0.1", 12345); + RemoteInterpreterRunningProcess process = recoveryProcess(groupId, registration, () -> { + try (RemoteInterpreterEventClient client = new RemoteInterpreterEventClient( + eventServer.getHost(), eventServer.getPort(), 1, groupId, callbackToken)) { + client.registerInterpreterProcess(new RegisterInfo("127.0.0.1", 12345, groupId)); + } + }); + + assertTrue(process.recover()); + assertTrue(eventServer.isCallbackTokenActive(registration)); + } + + @Test + void persistsCallbackIdentitySeparatelyFromEffectiveCommandEndpoint() { + String groupId = "group-a"; + String callbackToken = eventServer.issueCallbackToken(groupId); + allowRegistration(groupId); + try (RemoteInterpreterEventClient client = new RemoteInterpreterEventClient( + eventServer.getHost(), eventServer.getPort(), 1, groupId, callbackToken)) { + client.registerInterpreterProcess( + new RegisterInfo("advertised-host", 2222, groupId)); + } + + when(interpreterSettingManager.getInterpreterEventServer()).thenReturn(eventServer); + InterpreterSetting interpreterSetting = mock(InterpreterSetting.class); + ManagedInterpreterGroup interpreterGroup = mock(ManagedInterpreterGroup.class); + RemoteInterpreterProcess interpreterProcess = mock(RemoteInterpreterProcess.class); + when(interpreterSetting.getAllInterpreterGroups()) + .thenReturn(Collections.singletonList(interpreterGroup)); + when(interpreterSetting.getInterpreterSettingManager()) + .thenReturn(interpreterSettingManager); + when(interpreterGroup.getId()).thenReturn(groupId); + when(interpreterGroup.getInterpreterProcess()).thenReturn(interpreterProcess); + when(interpreterProcess.isRunning()).thenReturn(true); + when(interpreterProcess.getHost()).thenReturn("command-host"); + when(interpreterProcess.getPort()).thenReturn(1111); + + String recoveryData = RecoveryUtils.getRecoveryData(interpreterSetting); + assertEquals("group-a\tcommand-host:1111\t" + callbackToken + + "\tadvertised-host\t2222", recoveryData); + + ZeppelinConfiguration recoveryConfig = mock(ZeppelinConfiguration.class); + when(recoveryConfig.getTime( + ZeppelinConfiguration.ConfVars.ZEPPELIN_INTERPRETER_CONNECT_TIMEOUT)) + .thenReturn(1_000L); + when(interpreterSettingManager.getByName("setting")).thenReturn(interpreterSetting); + when(interpreterSetting.getJavaProperties()).thenReturn(new Properties()); + Map recovered = RecoveryUtils.restoreFromRecoveryData( + recoveryData, "setting", interpreterSettingManager, recoveryConfig); + RemoteInterpreterRunningProcess process = + (RemoteInterpreterRunningProcess) recovered.get(groupId); + assertEquals("command-host", process.getHost()); + assertEquals(1111, process.getPort()); + + try (RemoteInterpreterEventClient client = new RemoteInterpreterEventClient( + eventServer.getHost(), eventServer.getPort(), 1, groupId, callbackToken)) { + assertThrows(RuntimeException.class, () -> client.registerInterpreterProcess( + new RegisterInfo("command-host", 1111, groupId))); + client.registerInterpreterProcess( + new RegisterInfo("advertised-host", 2222, groupId)); + } + } + + private RemoteInterpreterProcess allowRegistration(String groupId) { + ManagedInterpreterGroup interpreterGroup = mock(ManagedInterpreterGroup.class); + RemoteInterpreterProcess interpreterProcess = mock(RemoteInterpreterProcess.class); + when(interpreterSettingManager.getInterpreterGroupById(groupId)).thenReturn(interpreterGroup); + when(interpreterGroup.getInterpreterProcess()).thenReturn(interpreterProcess); + return interpreterProcess; + } + + private RemoteInterpreterRunningProcess recoveryProcess( + String groupId, + RemoteInterpreterEventServer.CallbackCredentialRegistration registration, + Runnable reconnectAction) { + return new RemoteInterpreterRunningProcess( + "setting", groupId, 1, 1, "localhost", 1, + "127.0.0.1", 12345, true, eventServer, registration) { + @Override + protected void reconnectToEventServer() { + reconnectAction.run(); + } + }; + } +} diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/launcher/InterpreterShellScriptTest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/launcher/InterpreterShellScriptTest.java new file mode 100644 index 00000000000..e5029bce3f1 --- /dev/null +++ b/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/launcher/InterpreterShellScriptTest.java @@ -0,0 +1,97 @@ +/* + * 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.zeppelin.interpreter.launcher; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Map; +import java.util.concurrent.TimeUnit; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class InterpreterShellScriptTest { + + @Test + void passesImpersonatedCallbackCredentialThroughStdinWithoutLoggingIt( + @TempDir Path temporaryDirectory) throws Exception { + Path zeppelinHome = Paths.get("..").toAbsolutePath().normalize(); + Path fakeJavaHome = Files.createDirectories(temporaryDirectory.resolve("java-home/bin")) + .getParent(); + Path captureFile = temporaryDirectory.resolve("java-capture"); + Path fakeJava = fakeJavaHome.resolve("bin/java"); + Files.writeString(fakeJava, + "#!/bin/bash\n" + + "if [[ \"$1\" == \"-version\" ]]; then\n" + + " echo 'openjdk version \"11.0.0\"' >&2\n" + + " exit 0\n" + + "fi\n" + + "if [[ \"$*\" == *RemoteInterpreterDownloader* ]]; then\n" + + " exit 0\n" + + "fi\n" + + "printf '%s\\n' \"${ZEPPELIN_INTERPRETER_EVENT_TOKEN}\" > \"${CAPTURE_FILE}\"\n" + + "printf '%s\\n' \"$*\" >> \"${CAPTURE_FILE}\"\n", + StandardCharsets.UTF_8); + fakeJava.toFile().setExecutable(true); + + Path confDirectory = Files.createDirectory(temporaryDirectory.resolve("conf")); + Files.writeString(confDirectory.resolve("zeppelin-env.sh"), + "ZEPPELIN_IMPERSONATE_CMD=(bash -c)\n", StandardCharsets.UTF_8); + Path interpreterDirectory = Files.createDirectory(temporaryDirectory.resolve("interpreter")); + Path localRepo = Files.createDirectory(temporaryDirectory.resolve("local-repo")); + Path logDirectory = Files.createDirectory(temporaryDirectory.resolve("logs")); + Path pidDirectory = Files.createDirectory(temporaryDirectory.resolve("run")); + + ProcessBuilder processBuilder = new ProcessBuilder( + zeppelinHome.resolve("bin/interpreter.sh").toString(), + "-p", "12345", + "-r", ":", + "-i", "group-id", + "-d", interpreterDirectory.toString(), + "-l", localRepo.toString(), + "-g", "test", + "-u", "impersonated-user"); + processBuilder.redirectErrorStream(true); + Map environment = processBuilder.environment(); + environment.put("JAVA_HOME", fakeJavaHome.toString()); + environment.put("ZEPPELIN_HOME", zeppelinHome.toString()); + environment.put("ZEPPELIN_CONF_DIR", confDirectory.toString()); + environment.put("ZEPPELIN_LOG_DIR", logDirectory.toString()); + environment.put("ZEPPELIN_PID_DIR", pidDirectory.toString()); + environment.put("INTERPRETER_GROUP_ID", "group-id"); + environment.put("ZEPPELIN_INTERPRETER_EVENT_TOKEN", "secret-callback-token"); + environment.put("CAPTURE_FILE", captureFile.toString()); + + Process process = processBuilder.start(); + assertTrue(process.waitFor(30, TimeUnit.SECONDS)); + String output = new String(process.getInputStream().readAllBytes(), StandardCharsets.UTF_8); + + assertEquals(0, process.exitValue(), output); + String capture = Files.readString(captureFile, StandardCharsets.UTF_8); + assertTrue(capture.startsWith("secret-callback-token\n")); + assertFalse(capture.substring(capture.indexOf('\n') + 1) + .contains("secret-callback-token")); + assertFalse(output.contains("secret-callback-token")); + } +} diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/launcher/SparkInterpreterCallbackCredentialTest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/launcher/SparkInterpreterCallbackCredentialTest.java new file mode 100644 index 00000000000..919d9783a90 --- /dev/null +++ b/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/launcher/SparkInterpreterCallbackCredentialTest.java @@ -0,0 +1,76 @@ +/* + * 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.zeppelin.interpreter.launcher; + +import org.apache.zeppelin.conf.ZeppelinConfiguration; +import org.apache.zeppelin.interpreter.InterpreterOption; +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterEventClient; +import org.apache.zeppelin.test.DownloadUtils; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.util.Map; +import java.util.Properties; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class SparkInterpreterCallbackCredentialTest { + + @Test + void passesCallbackCredentialToKubernetesClusterDriver() throws Exception { + Properties properties = new Properties(); + properties.setProperty("SPARK_HOME", DownloadUtils.downloadSpark()); + properties.setProperty("spark.master", "k8s://https://kubernetes.default.svc"); + properties.setProperty("spark.submit.deployMode", "cluster"); + InterpreterLaunchContext context = new InterpreterLaunchContext( + properties, new InterpreterOption(), null, "user", "group-id", + "setting-id", "spark", "spark", 0, "host"); + context.setIntpEventCallbackToken("callback-token"); + SparkInterpreterLauncher launcher = new SparkInterpreterLauncher( + ZeppelinConfiguration.load(), null); + + Map environment = launcher.buildEnvFromProperties(context); + + assertEquals("callback-token", + environment.get(RemoteInterpreterEventClient.CALLBACK_TOKEN_ENV)); + assertTrue(environment.get("ZEPPELIN_SPARK_CONF").contains( + "spark.kubernetes.driverEnv." + RemoteInterpreterEventClient.CALLBACK_TOKEN_ENV + + "=callback-token")); + } + + @Test + void rejectsClusterModeWithoutSecureCredentialPropagation() throws Exception { + Properties properties = new Properties(); + properties.setProperty("SPARK_HOME", DownloadUtils.downloadSpark()); + properties.setProperty("spark.master", "spark://standalone.example:7077"); + properties.setProperty("spark.submit.deployMode", "cluster"); + InterpreterLaunchContext context = new InterpreterLaunchContext( + properties, new InterpreterOption(), null, "user", "group-id", + "setting-id", "spark", "spark", 0, "host"); + context.setIntpEventCallbackToken("callback-token"); + SparkInterpreterLauncher launcher = new SparkInterpreterLauncher( + ZeppelinConfiguration.load(), null); + + IOException error = assertThrows(IOException.class, + () -> launcher.buildEnvFromProperties(context)); + + assertTrue(error.getMessage().contains("Spark cluster deploy mode")); + } +} diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/launcher/StandardInterpreterLauncherTest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/launcher/StandardInterpreterLauncherTest.java index 3cb2d9a7816..434841083ca 100644 --- a/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/launcher/StandardInterpreterLauncherTest.java +++ b/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/launcher/StandardInterpreterLauncherTest.java @@ -20,6 +20,7 @@ import org.apache.zeppelin.conf.ZeppelinConfiguration; import org.apache.zeppelin.interpreter.InterpreterOption; import org.apache.zeppelin.interpreter.remote.ExecRemoteInterpreterProcess; +import org.apache.zeppelin.interpreter.remote.RemoteInterpreterEventClient; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -47,6 +48,7 @@ void testLauncher() throws IOException { InterpreterOption option = new InterpreterOption(); option.setUserImpersonate(true); InterpreterLaunchContext context = new InterpreterLaunchContext(properties, option, null, "user1", "intpGroupId", "groupId", "groupName", "name", 0, "host"); + context.setIntpEventCallbackToken("callback-token"); InterpreterClient client = launcher.launch(context); assertTrue(client instanceof ExecRemoteInterpreterProcess); ExecRemoteInterpreterProcess interpreterProcess = (ExecRemoteInterpreterProcess) client; @@ -59,6 +61,8 @@ void testLauncher() throws IOException { assertTrue(interpreterProcess.getEnv().size() >= 2); assertEquals("VALUE_1", interpreterProcess.getEnv().get("ENV_1")); assertTrue(interpreterProcess.getEnv().containsKey("INTERPRETER_GROUP_ID")); + assertEquals("callback-token", + interpreterProcess.getEnv().get(RemoteInterpreterEventClient.CALLBACK_TOKEN_ENV)); assertEquals(true, interpreterProcess.isUserImpersonated()); interpreterProcess.close(); } diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterProcessTest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterProcessTest.java new file mode 100644 index 00000000000..46cecf2ef40 --- /dev/null +++ b/zeppelin-server/src/test/java/org/apache/zeppelin/interpreter/remote/RemoteInterpreterProcessTest.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.zeppelin.interpreter.remote; + +import org.apache.zeppelin.conf.ZeppelinConfiguration; +import org.junit.jupiter.api.Test; + +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +class RemoteInterpreterProcessTest { + + @Test + void initConfigurationNeverContainsCallbackCredential() { + ZeppelinConfiguration zConf = mock(ZeppelinConfiguration.class); + Map completeConfiguration = new HashMap<>(); + completeConfiguration.put(RemoteInterpreterEventClient.CALLBACK_TOKEN_PROPERTY, + "must-not-leak"); + when(zConf.getCompleteConfiguration()).thenReturn(completeConfiguration); + + Map configuration = + RemoteInterpreterProcess.createInitConfiguration(zConf, "group"); + + assertEquals("group", configuration.get( + RemoteInterpreterEventClient.INTERPRETER_GROUP_PROPERTY)); + assertFalse(configuration.containsKey( + RemoteInterpreterEventClient.CALLBACK_TOKEN_PROPERTY)); + } + + @Test + void terminationListenerRunsOnceWhenInstalledAfterTermination() { + TestRemoteInterpreterProcess process = new TestRemoteInterpreterProcess(); + AtomicInteger notifications = new AtomicInteger(); + + process.terminate(); + process.setTerminationListener(notifications::incrementAndGet); + process.terminate(); + + assertEquals(1, notifications.get()); + } + + private static final class TestRemoteInterpreterProcess extends RemoteInterpreterProcess { + + private TestRemoteInterpreterProcess() { + super(1, 1, "localhost", 1); + } + + private void terminate() { + notifyTermination(); + } + + @Override + public String getInterpreterGroupId() { + return "group"; + } + + @Override + public String getInterpreterSettingName() { + return "setting"; + } + + @Override + public void start(String userName) { + } + + @Override + public void stop() { + } + + @Override + public String getHost() { + return "localhost"; + } + + @Override + public int getPort() { + return 1; + } + + @Override + public boolean isAlive() { + return false; + } + + @Override + public boolean isRunning() { + return false; + } + + @Override + public void processStarted(int port, String host) { + } + + @Override + public String getErrorMessage() { + return null; + } + } +}