path: root/common/network-shuffle/src/test/java/org/apache/spark/network/sasl/SaslIntegrationSuite.java

/*
 * 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.spark.network.sasl;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicReference;

import com.google.common.collect.Lists;
import org.junit.After;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.Test;

import static org.junit.Assert.*;
import static org.mockito.Mockito.*;

import org.apache.spark.network.TestUtils;
import org.apache.spark.network.TransportContext;
import org.apache.spark.network.buffer.ManagedBuffer;
import org.apache.spark.network.client.ChunkReceivedCallback;
import org.apache.spark.network.client.RpcResponseCallback;
import org.apache.spark.network.client.TransportClient;
import org.apache.spark.network.client.TransportClientBootstrap;
import org.apache.spark.network.client.TransportClientFactory;
import org.apache.spark.network.server.OneForOneStreamManager;
import org.apache.spark.network.server.RpcHandler;
import org.apache.spark.network.server.StreamManager;
import org.apache.spark.network.server.TransportServer;
import org.apache.spark.network.server.TransportServerBootstrap;
import org.apache.spark.network.shuffle.BlockFetchingListener;
import org.apache.spark.network.shuffle.ExternalShuffleBlockHandler;
import org.apache.spark.network.shuffle.ExternalShuffleBlockResolver;
import org.apache.spark.network.shuffle.OneForOneBlockFetcher;
import org.apache.spark.network.shuffle.protocol.BlockTransferMessage;
import org.apache.spark.network.shuffle.protocol.ExecutorShuffleInfo;
import org.apache.spark.network.shuffle.protocol.OpenBlocks;
import org.apache.spark.network.shuffle.protocol.RegisterExecutor;
import org.apache.spark.network.shuffle.protocol.StreamHandle;
import org.apache.spark.network.util.JavaUtils;
import org.apache.spark.network.util.MapConfigProvider;
import org.apache.spark.network.util.TransportConf;

public class SaslIntegrationSuite {

  // Use a long timeout to account for slow / overloaded build machines. In the normal case,
  // tests should finish way before the timeout expires.
  private static final long TIMEOUT_MS = 10_000;

  static TransportServer server;
  static TransportConf conf;
  static TransportContext context;
  static SecretKeyHolder secretKeyHolder;

  TransportClientFactory clientFactory;

  @BeforeClass
  public static void beforeAll() throws IOException {
    conf = new TransportConf("shuffle", MapConfigProvider.EMPTY);
    context = new TransportContext(conf, new TestRpcHandler());

    secretKeyHolder = mock(SecretKeyHolder.class);
    when(secretKeyHolder.getSaslUser(eq("app-1"))).thenReturn("app-1");
    when(secretKeyHolder.getSecretKey(eq("app-1"))).thenReturn("app-1");
    when(secretKeyHolder.getSaslUser(eq("app-2"))).thenReturn("app-2");
    when(secretKeyHolder.getSecretKey(eq("app-2"))).thenReturn("app-2");
    when(secretKeyHolder.getSaslUser(anyString())).thenReturn("other-app");
    when(secretKeyHolder.getSecretKey(anyString())).thenReturn("correct-password");

    TransportServerBootstrap bootstrap = new SaslServerBootstrap(conf, secretKeyHolder);
    server = context.createServer(Arrays.asList(bootstrap));
  }


  @AfterClass
  public static void afterAll() {
    server.close();
  }

  @After
  public void afterEach() {
    if (clientFactory != null) {
      clientFactory.close();
      clientFactory = null;
    }
  }

  @Test
  public void testGoodClient() throws IOException {
    clientFactory = context.createClientFactory(
      Lists.<TransportClientBootstrap>newArrayList(
        new SaslClientBootstrap(conf, "app-1", secretKeyHolder)));

    TransportClient client = clientFactory.createClient(TestUtils.getLocalHost(), server.getPort());
    String msg = "Hello, World!";
    ByteBuffer resp = client.sendRpcSync(JavaUtils.stringToBytes(msg), TIMEOUT_MS);
    assertEquals(msg, JavaUtils.bytesToString(resp));
  }

  @Test
  public void testBadClient() {
    SecretKeyHolder badKeyHolder = mock(SecretKeyHolder.class);
    when(badKeyHolder.getSaslUser(anyString())).thenReturn("other-app");
    when(badKeyHolder.getSecretKey(anyString())).thenReturn("wrong-password");
    clientFactory = context.createClientFactory(
      Lists.<TransportClientBootstrap>newArrayList(
        new SaslClientBootstrap(conf, "unknown-app", badKeyHolder)));

    try {
      // Bootstrap should fail on startup.
      clientFactory.createClient(TestUtils.getLocalHost(), server.getPort());
      fail("Connection should have failed.");
    } catch (Exception e) {
      assertTrue(e.getMessage(), e.getMessage().contains("Mismatched response"));
    }
  }

  @Test
  public void testNoSaslClient() throws IOException {
    clientFactory = context.createClientFactory(
      Lists.<TransportClientBootstrap>newArrayList());

    TransportClient client = clientFactory.createClient(TestUtils.getLocalHost(), server.getPort());
    try {
      client.sendRpcSync(ByteBuffer.allocate(13), TIMEOUT_MS);
      fail("Should have failed");
    } catch (Exception e) {
      assertTrue(e.getMessage(), e.getMessage().contains("Expected SaslMessage"));
    }

    try {
      // Guessing the right tag byte doesn't magically get you in...
      client.sendRpcSync(ByteBuffer.wrap(new byte[] { (byte) 0xEA }), TIMEOUT_MS);
      fail("Should have failed");
    } catch (Exception e) {
      assertTrue(e.getMessage(), e.getMessage().contains("java.lang.IndexOutOfBoundsException"));
    }
  }

  @Test
  public void testNoSaslServer() {
    RpcHandler handler = new TestRpcHandler();
    TransportContext context = new TransportContext(conf, handler);
    clientFactory = context.createClientFactory(
      Lists.<TransportClientBootstrap>newArrayList(
        new SaslClientBootstrap(conf, "app-1", secretKeyHolder)));
    TransportServer server = context.createServer();
    try {
      clientFactory.createClient(TestUtils.getLocalHost(), server.getPort());
    } catch (Exception e) {
      assertTrue(e.getMessage(), e.getMessage().contains("Digest-challenge format violation"));
    } finally {
      server.close();
    }
  }

  /**
   * This test is not actually testing SASL behavior, but testing that the shuffle service
   * performs correct authorization checks based on the SASL authentication data.
   */
  @Test
  public void testAppIsolation() throws Exception {
    // Start a new server with the correct RPC handler to serve block data.
    ExternalShuffleBlockResolver blockResolver = mock(ExternalShuffleBlockResolver.class);
    ExternalShuffleBlockHandler blockHandler = new ExternalShuffleBlockHandler(
      new OneForOneStreamManager(), blockResolver);
    TransportServerBootstrap bootstrap = new SaslServerBootstrap(conf, secretKeyHolder);
    TransportContext blockServerContext = new TransportContext(conf, blockHandler);
    TransportServer blockServer = blockServerContext.createServer(Arrays.asList(bootstrap));

    TransportClient client1 = null;
    TransportClient client2 = null;
    TransportClientFactory clientFactory2 = null;
    try {
      // Create a client, and make a request to fetch blocks from a different app.
      clientFactory = blockServerContext.createClientFactory(
        Lists.<TransportClientBootstrap>newArrayList(
          new SaslClientBootstrap(conf, "app-1", secretKeyHolder)));
      client1 = clientFactory.createClient(TestUtils.getLocalHost(),
        blockServer.getPort());

      final AtomicReference<Throwable> exception = new AtomicReference<>();

      final CountDownLatch blockFetchLatch = new CountDownLatch(1);
      BlockFetchingListener listener = new BlockFetchingListener() {
        @Override
        public void onBlockFetchSuccess(String blockId, ManagedBuffer data) {
          blockFetchLatch.countDown();
        }
        @Override
        public void onBlockFetchFailure(String blockId, Throwable t) {
          exception.set(t);
          blockFetchLatch.countDown();
        }
      };

      String[] blockIds = { "shuffle_2_3_4", "shuffle_6_7_8" };
      OneForOneBlockFetcher fetcher =
          new OneForOneBlockFetcher(client1, "app-2", "0", blockIds, listener);
      fetcher.start();
      blockFetchLatch.await();
      checkSecurityException(exception.get());

      // Register an executor so that the next steps work.
      ExecutorShuffleInfo executorInfo = new ExecutorShuffleInfo(
        new String[] { System.getProperty("java.io.tmpdir") }, 1,
          "org.apache.spark.shuffle.sort.SortShuffleManager");
      RegisterExecutor regmsg = new RegisterExecutor("app-1", "0", executorInfo);
      client1.sendRpcSync(regmsg.toByteBuffer(), TIMEOUT_MS);

      // Make a successful request to fetch blocks, which creates a new stream. But do not actually
      // fetch any blocks, to keep the stream open.
      OpenBlocks openMessage = new OpenBlocks("app-1", "0", blockIds);
      ByteBuffer response = client1.sendRpcSync(openMessage.toByteBuffer(), TIMEOUT_MS);
      StreamHandle stream = (StreamHandle) BlockTransferMessage.Decoder.fromByteBuffer(response);
      long streamId = stream.streamId;

      // Create a second client, authenticated with a different app ID, and try to read from
      // the stream created for the previous app.
      clientFactory2 = blockServerContext.createClientFactory(
        Lists.<TransportClientBootstrap>newArrayList(
          new SaslClientBootstrap(conf, "app-2", secretKeyHolder)));
      client2 = clientFactory2.createClient(TestUtils.getLocalHost(),
        blockServer.getPort());

      final CountDownLatch chunkReceivedLatch = new CountDownLatch(1);
      ChunkReceivedCallback callback = new ChunkReceivedCallback() {
        @Override
        public void onSuccess(int chunkIndex, ManagedBuffer buffer) {
          chunkReceivedLatch.countDown();
        }
        @Override
        public void onFailure(int chunkIndex, Throwable t) {
          exception.set(t);
          chunkReceivedLatch.countDown();
        }
      };

      exception.set(null);
      client2.fetchChunk(streamId, 0, callback);
      chunkReceivedLatch.await();
      checkSecurityException(exception.get());
    } finally {
      if (client1 != null) {
        client1.close();
      }
      if (client2 != null) {
        client2.close();
      }
      if (clientFactory2 != null) {
        clientFactory2.close();
      }
      blockServer.close();
    }
  }

  /** RPC handler which simply responds with the message it received. */
  public static class TestRpcHandler extends RpcHandler {
    @Override
    public void receive(TransportClient client, ByteBuffer message, RpcResponseCallback callback) {
      callback.onSuccess(message);
    }

    @Override
    public StreamManager getStreamManager() {
      return new OneForOneStreamManager();
    }
  }

  private void checkSecurityException(Throwable t) {
    assertNotNull("No exception was caught.", t);
    assertTrue("Expected SecurityException.",
      t.getMessage().contains(SecurityException.class.getName()));
  }
}