/*
 * Licensed to the University of California, Berkeley 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 tachyon.security;

import java.io.IOException;
import java.lang.reflect.Field;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;

import tachyon.Constants;
import tachyon.LocalTachyonClusterResource;
import tachyon.security.MasterClientAuthenticationIntegrationTest.NameMatchAuthenticationProvider;
import tachyon.worker.ClientMetrics;
import tachyon.worker.WorkerClient;

/**
 * Test RPC authentication between worker and its client, in four modes: NOSASL, SIMPLE, CUSTOM,
 * KERBEROS.
 */
// TODO(bin): improve the way to set and isolate MasterContext/WorkerContext across test cases
public class WorkerClientAuthenticationIntegrationTest {
  @Rule
  public LocalTachyonClusterResource mLocalTachyonClusterResource =
      new LocalTachyonClusterResource(1000, 1000, Constants.GB);
  private ExecutorService mExecutorService;

  @Rule
  public ExpectedException mThrown = ExpectedException.none();

  @Before
  public void before() throws Exception {
    mExecutorService = Executors.newFixedThreadPool(2);
    clearLoginUser();
  }

  @After
  public void after() throws Exception {
    System.clearProperty(Constants.SECURITY_LOGIN_USERNAME);
  }

  @Test
  @LocalTachyonClusterResource.Config(
      tachyonConfParams = {Constants.SECURITY_AUTHENTICATION_TYPE, "NOSASL"})
  public void noAuthenticationOpenCloseTest() throws Exception {
    authenticationOperationTest();
  }

  @Test
  @LocalTachyonClusterResource.Config(
      tachyonConfParams = {Constants.SECURITY_AUTHENTICATION_TYPE, "SIMPLE"})
  public void simpleAuthenticationOpenCloseTest() throws Exception {
    authenticationOperationTest();
  }

  @Test
  @LocalTachyonClusterResource.Config(
      tachyonConfParams = {Constants.SECURITY_AUTHENTICATION_TYPE, "CUSTOM",
          Constants.SECURITY_AUTHENTICATION_CUSTOM_PROVIDER,
          NameMatchAuthenticationProvider.FULL_CLASS_NAME},
      startCluster = false)
  public void customAuthenticationOpenCloseTest() throws Exception {
    /**
     * Using tachyon as loginUser for unit testing, only tachyon user is allowed to connect to
     * Tachyon Worker.
     */
    System.setProperty(Constants.SECURITY_LOGIN_USERNAME, "tachyon");
    mLocalTachyonClusterResource.start();
    authenticationOperationTest();
  }

  @Test
  @LocalTachyonClusterResource.Config(tachyonConfParams = {Constants.SECURITY_AUTHENTICATION_TYPE,
      "CUSTOM", Constants.SECURITY_AUTHENTICATION_CUSTOM_PROVIDER,
      NameMatchAuthenticationProvider.FULL_CLASS_NAME}, startCluster = false)
  public void customAuthenticationDenyConnectTest() throws Exception {
    /**
     * Using tachyon as loginUser for unit testing, only tachyon user is allowed to connect to
     * Tachyon Master during starting cluster.
     */
    System.setProperty(Constants.SECURITY_LOGIN_USERNAME, "tachyon");
    mLocalTachyonClusterResource.start();

    // Using no-tachyon as loginUser to connect to Worker, the IOException will be thrown
    clearLoginUser();
    mThrown.expect(IOException.class);
    System.setProperty(Constants.SECURITY_LOGIN_USERNAME, "no-tachyon");

    WorkerClient workerClient =
        new WorkerClient(mLocalTachyonClusterResource.get().getWorkerAddress(), mExecutorService,
            mLocalTachyonClusterResource.get().getWorkerTachyonConf(), 1 /* fake session id */,
            true, new ClientMetrics());
    try {
      Assert.assertFalse(workerClient.isConnected());
      workerClient.mustConnect();
    } finally {
      workerClient.close();
    }
  }

  /**
   * Test Tachyon Worker client connects or disconnects to the Worker.
   *
   * @throws Exception
   */
  private void authenticationOperationTest() throws Exception {
    WorkerClient workerClient =
        new WorkerClient(mLocalTachyonClusterResource.get().getWorkerAddress(), mExecutorService,
            mLocalTachyonClusterResource.get().getWorkerTachyonConf(), 1 /* fake session id */,
            true, new ClientMetrics());

    Assert.assertFalse(workerClient.isConnected());
    workerClient.mustConnect();
    Assert.assertTrue(workerClient.isConnected());

    workerClient.close();
  }

  private void clearLoginUser() throws Exception {
    // User reflection to reset the private static member sLoginUser in LoginUser.
    Field field = LoginUser.class.getDeclaredField("sLoginUser");
    field.setAccessible(true);
    field.set(null, null);
  }
}

@@ -845,40 +845,8 @@
 on;%0A
-import java.lang.reflect.Field;%0A
 impo
@@ -1087,16 +1087,55 @@
 ception;
+%0Aimport org.powermock.reflect.Whitebox;
 %0A%0Aimport
@@ -3241,32 +3241,49 @@
 ();%0A  %7D%0A%0A  @Test
+(timeout = 10000)
 %0A  @LocalTachyon
@@ -3947,69 +3947,8 @@
 own%0A
-    clearLoginUser();%0A    mThrown.expect(IOException.class);%0A
     
@@ -4015,16 +4015,229 @@
 chyon%22);
+%0A    // Clear the login user so that it will be reloaded and pick up our no-tachyon change%0A    clearLoginUser();%0A    mThrown.expect(IOException.class);%0A    mThrown.expectMessage(%22Failed to connect to the worker%22);
 %0A%0A    Wo
@@ -5317,203 +5317,71 @@
     
-// User reflection to reset the private static member sLoginUser in LoginUser.%0A    Field field = LoginUser.class.getDeclaredField(%22sLoginUser%22);%0A    field.setAccessible(true);%0A    field.set(null,
+Whitebox.setInternalState(LoginUser.class, %22sLoginUser%22, (User)
  nul

Clean up WorkerClientAuthenticationIntegrationTest