/**
 * 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.hadoop.ipc;

import java.io.IOException;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.NetworkInterface;
import java.security.PrivilegedExceptionAction;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Enumeration;

import org.apache.hadoop.test.GenericTestUtils;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.ipc.protobuf.MiniRPCBenchmarkProtos.MiniDelegationTokenProto;
import org.apache.hadoop.ipc.protobuf.MiniRPCBenchmarkProtos.MiniGetDelegationTokenRequestProto;
import org.apache.hadoop.ipc.protobuf.MiniRPCBenchmarkProtos.MiniGetDelegationTokenResponseProto;
import org.apache.hadoop.ipc.protobuf.MiniRPCBenchmarkProtos.MiniProtocolService;
import org.apache.hadoop.net.NetUtils;
import org.apache.hadoop.security.KerberosInfo;
import org.apache.hadoop.security.SecurityUtil;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.security.authorize.DefaultImpersonationProvider;
import org.apache.hadoop.security.token.Token;
import org.apache.hadoop.security.token.TokenInfo;
import org.apache.hadoop.security.token.delegation.AbstractDelegationTokenSelector;
import org.apache.hadoop.security.token.delegation.TestDelegationToken.TestDelegationTokenIdentifier;
import org.apache.hadoop.security.token.delegation.TestDelegationToken.TestDelegationTokenSecretManager;
import org.apache.hadoop.thirdparty.protobuf.BlockingService;
import org.apache.hadoop.thirdparty.protobuf.ByteString;
import org.apache.hadoop.thirdparty.protobuf.RpcController;
import org.apache.hadoop.thirdparty.protobuf.ServiceException;
import org.apache.hadoop.util.Time;
import org.slf4j.event.Level;

import static org.junit.jupiter.api.Assertions.fail;

/**
 * MiniRPCBenchmark measures time to establish an RPC connection 
 * to a secure RPC server.
 * It sequentially establishes connections the specified number of times, 
 * and calculates the average time taken to connect.
 * The time to connect includes the server side authentication time.
 * The benchmark supports three authentication methods:
 * <ol>
 * <li>simple - no authentication. In order to enter this mode 
 * the configuration file <code>core-site.xml</code> should specify
 * <code>hadoop.security.authentication = simple</code>.
 * This is the default mode.</li>
 * <li>kerberos - kerberos authentication. In order to enter this mode 
 * the configuration file <code>core-site.xml</code> should specify
 * <code>hadoop.security.authentication = kerberos</code> and
 * the argument string should provide qualifying
 * <code>keytabFile</code> and <code>userName</code> parameters.
 * <li>delegation token - authentication using delegation token.
 * In order to enter this mode the benchmark should provide all the
 * mentioned parameters for kerberos authentication plus the
 * <code>useToken</code> argument option.
 * </ol>
 * Input arguments:
 * <ul>
 * <li>numIterations - number of connections to establish</li>
 * <li>keytabFile - keytab file for kerberos authentication</li>
 * <li>userName - principal name for kerberos authentication</li>
 * <li>useToken - should be specified for delegation token authentication</li>
 * <li>logLevel - logging level, see {@link Level}</li>
 * </ul>
 */
public class MiniRPCBenchmark {
  private static final String KEYTAB_FILE_KEY = "test.keytab.file";
  private static final String USER_NAME_KEY = "test.user.name";
  private static final String MINI_USER = "miniUser";
  private static final String RENEWER = "renewer";
  private static final String GROUP_NAME_1 = "MiniGroup1";
  private static final String GROUP_NAME_2 = "MiniGroup2";
  private static final String[] GROUP_NAMES = 
                                    new String[] {GROUP_NAME_1, GROUP_NAME_2};

  private UserGroupInformation currentUgi;
  private Level logLevel;

  MiniRPCBenchmark(Level l) {
    currentUgi = null;
    logLevel = l;
  }

  public static class TestDelegationTokenSelector extends 
  AbstractDelegationTokenSelector<TestDelegationTokenIdentifier>{

    protected TestDelegationTokenSelector() {
      super(new Text("MY KIND"));
    }    
  }

  /**
   * Protobuf based MiniProtocol used by {@link MiniRPCBenchmark}.
   * Replaces the legacy Writable based protocol now that
   * {@code WritableRpcEngine} has been removed.
   */
  @KerberosInfo(serverPrincipal = USER_NAME_KEY)
  @TokenInfo(TestDelegationTokenSelector.class)
  @ProtocolInfo(
      protocolName = "org.apache.hadoop.ipc.MiniRPCBenchmark$MiniProtocol",
      protocolVersion = 1)
  public interface MiniProtocol extends MiniProtocolService.BlockingInterface {
    long versionID = 1L;
  }

  /**
   * Primitive RPC server, which
   * allows clients to connect to it.
   */
  static class MiniServer implements MiniProtocol {
    private static final String DEFAULT_SERVER_ADDRESS = "0.0.0.0";

    private TestDelegationTokenSecretManager secretManager;
    private Server rpcServer;

    @Override // MiniProtocol
    public MiniGetDelegationTokenResponseProto getDelegationToken(
        RpcController controller,
        MiniGetDelegationTokenRequestProto request) throws ServiceException {
      try {
        Text renewer = new Text(request.getRenewer());
        String owner = UserGroupInformation.getCurrentUser().getUserName();
        String realUser =
            UserGroupInformation.getCurrentUser().getRealUser() == null ? "" :
                UserGroupInformation.getCurrentUser().getRealUser().getUserName();
        TestDelegationTokenIdentifier tokenId =
            new TestDelegationTokenIdentifier(
                new Text(owner), renewer, new Text(realUser));
        Token<TestDelegationTokenIdentifier> token =
            new Token<>(tokenId, secretManager);
        return MiniGetDelegationTokenResponseProto.newBuilder()
            .setToken(MiniDelegationTokenProto.newBuilder()
                .setIdentifier(ByteString.copyFrom(token.getIdentifier()))
                .setPassword(ByteString.copyFrom(token.getPassword()))
                .setKind(token.getKind().toString())
                .setService(token.getService().toString())
                .build())
            .build();
      } catch (IOException ioe) {
        throw new ServiceException(ioe);
      }
    }

    /** Start RPC server */
    MiniServer(Configuration conf, String user, String keytabFile)
    throws IOException {
      UserGroupInformation.setConfiguration(conf);
      UserGroupInformation.loginUserFromKeytab(user, keytabFile);
      secretManager = 
        new TestDelegationTokenSecretManager(24*60*60*1000,
            7*24*60*60*1000,24*60*60*1000,3600000);
      secretManager.startThreads();
      RPC.setProtocolEngine(conf, MiniProtocol.class, ProtobufRpcEngine2.class);
      BlockingService service =
          MiniProtocolService.newReflectiveBlockingService(this);
      rpcServer = new RPC.Builder(conf).setProtocol(MiniProtocol.class)
          .setInstance(service).setBindAddress(DEFAULT_SERVER_ADDRESS).setPort(0)
          .setNumHandlers(1).setVerbose(false).setSecretManager(secretManager)
          .build();
      rpcServer.start();
    }

    /** Stop RPC server */
    void stop() {
      if(rpcServer != null) rpcServer.stop();
      rpcServer = null;
    }

    /** Get RPC server address */
    InetSocketAddress getAddress() {
      if(rpcServer == null) return null;
      return NetUtils.getConnectAddress(rpcServer);
    }
  }

  long connectToServer(Configuration conf, InetSocketAddress addr)
  throws IOException {
    MiniProtocol client = null;
    try {
      long start = Time.now();
      RPC.setProtocolEngine(conf, MiniProtocol.class, ProtobufRpcEngine2.class);
      client = RPC.getProxy(MiniProtocol.class, MiniProtocol.versionID, addr, conf);
      long end = Time.now();
      return end - start;
    } finally {
      RPC.stopProxy(client);
    }
  }

  void connectToServerAndGetDelegationToken(
      final Configuration conf, final InetSocketAddress addr) throws IOException {
    MiniProtocol client = null;
    try {
      UserGroupInformation current = UserGroupInformation.getCurrentUser();
      UserGroupInformation proxyUserUgi = 
        UserGroupInformation.createProxyUserForTesting(
            MINI_USER, current, GROUP_NAMES);
      
      try {
        client =  proxyUserUgi.doAs(new PrivilegedExceptionAction<MiniProtocol>() {
          @Override
          public MiniProtocol run() throws IOException {
            RPC.setProtocolEngine(conf, MiniProtocol.class,
                ProtobufRpcEngine2.class);
            MiniProtocol p = RPC.getProxy(MiniProtocol.class,
                MiniProtocol.versionID, addr, conf);
            try {
              MiniGetDelegationTokenResponseProto response =
                  p.getDelegationToken(null,
                      MiniGetDelegationTokenRequestProto.newBuilder()
                          .setRenewer(RENEWER).build());
              MiniDelegationTokenProto tokenProto = response.getToken();
              Token<TestDelegationTokenIdentifier> token = new Token<>(
                  tokenProto.getIdentifier().toByteArray(),
                  tokenProto.getPassword().toByteArray(),
                  new Text(tokenProto.getKind()),
                  new Text(tokenProto.getService()));
              currentUgi = UserGroupInformation.createUserForTesting(MINI_USER,
                  GROUP_NAMES);
              SecurityUtil.setTokenService(token, addr);
              currentUgi.addToken(token);
            } catch (ServiceException se) {
              throw new IOException(se);
            }
            return p;
          }
        });
      } catch (InterruptedException e) {
        fail(Arrays.toString(e.getStackTrace()));
      }
    } finally {
      RPC.stopProxy(client);
    }
  }

  long connectToServerUsingDelegationToken(
      final Configuration conf, final InetSocketAddress addr) throws IOException {
    MiniProtocol client = null;
    try {
      long start = Time.now();
      try {
        client = currentUgi.doAs(new PrivilegedExceptionAction<MiniProtocol>() {
          @Override
          public MiniProtocol run() throws IOException {
            RPC.setProtocolEngine(conf, MiniProtocol.class, ProtobufRpcEngine2.class);
            return RPC.getProxy(MiniProtocol.class,
                MiniProtocol.versionID, addr, conf);
          }
        });
      } catch (InterruptedException e) {
        e.printStackTrace();
      }
      long end = Time.now();
      return end - start;
    } finally {
      RPC.stopProxy(client);
    }
  }

  static void setLoggingLevel(Level level) {
    GenericTestUtils.setLogLevel(Server.LOG, level);
    GenericTestUtils.setLogLevel(Server.AUDITLOG, level);
    GenericTestUtils.setLogLevel(Client.LOG, level);
  }

  /**
   * Run MiniBenchmark with MiniServer as the RPC server.
   * 
   * @param conf - configuration
   * @param count - connect this many times
   * @param keytabKey - key for keytab file in the configuration
   * @param userNameKey - key for user name in the configuration
   * @return average time to connect
   * @throws IOException
   */
  long runMiniBenchmark(Configuration conf,
                        int count,
                        String keytabKey,
                        String userNameKey) throws IOException {
    // get login information
    String user = System.getProperty("user.name");
    if(userNameKey != null)
      user = conf.get(userNameKey, user);
    String keytabFile = null;
    if(keytabKey != null)
      keytabFile = conf.get(keytabKey, keytabFile);
    MiniServer miniServer = null;
    try {
      // start the server
      miniServer = new MiniServer(conf, user, keytabFile);
      InetSocketAddress addr = miniServer.getAddress();

      connectToServer(conf, addr);
      // connect to the server count times
      setLoggingLevel(logLevel);
      long elapsed = 0L;
      for(int idx = 0; idx < count; idx ++) {
        elapsed += connectToServer(conf, addr);
      }
      return elapsed;
    } finally {
      if(miniServer != null) miniServer.stop();
    }
  }

  /**
   * Run MiniBenchmark using delegation token authentication.
   * 
   * @param conf - configuration
   * @param count - connect this many times
   * @param keytabKey - key for keytab file in the configuration
   * @param userNameKey - key for user name in the configuration
   * @return average time to connect
   * @throws IOException
   */
  long runMiniBenchmarkWithDelegationToken(Configuration conf,
                                           int count,
                                           String keytabKey,
                                           String userNameKey)
  throws IOException {
    // get login information
    String user = System.getProperty("user.name");
    if(userNameKey != null)
      user = conf.get(userNameKey, user);
    String keytabFile = null;
    if(keytabKey != null)
      keytabFile = conf.get(keytabKey, keytabFile);
    MiniServer miniServer = null;
    UserGroupInformation.setConfiguration(conf);
    String shortUserName =
      UserGroupInformation.createRemoteUser(user).getShortUserName();
    try {
      conf.setStrings(DefaultImpersonationProvider.getTestProvider().
              getProxySuperuserGroupConfKey(shortUserName), GROUP_NAME_1);
      configureSuperUserIPAddresses(conf, shortUserName);
      // start the server
      miniServer = new MiniServer(conf, user, keytabFile);
      InetSocketAddress addr = miniServer.getAddress();

      connectToServerAndGetDelegationToken(conf, addr);
      // connect to the server count times
      setLoggingLevel(logLevel);
      long elapsed = 0L;
      for(int idx = 0; idx < count; idx ++) {
        elapsed += connectToServerUsingDelegationToken(conf, addr);
      }
      return elapsed;
    } finally {
      if(miniServer != null) miniServer.stop();
    }
  }

  static void printUsage() {
    System.err.println(
        "Usage: MiniRPCBenchmark <numIterations> [<keytabFile> [<userName> " +
        "[useToken|useKerberos [<logLevel>]]]]");
    System.exit(-1);
  }

  public static void main(String[] args) throws Exception {
    System.out.println("Benchmark: RPC session establishment.");
    if(args.length < 1)
      printUsage();

    Configuration conf = new Configuration();
    int count = Integer.parseInt(args[0]);
    if(args.length > 1)
      conf.set(KEYTAB_FILE_KEY, args[1]);
    if(args.length > 2)
      conf.set(USER_NAME_KEY, args[2]);
    boolean useDelegationToken = false;
    if(args.length > 3)
      useDelegationToken = args[3].equalsIgnoreCase("useToken");
    Level l = Level.ERROR;
    if(args.length > 4)
      l = GenericTestUtils.toLevel(args[4]);

    MiniRPCBenchmark mb = new MiniRPCBenchmark(l);
    long elapsedTime = 0;
    if(useDelegationToken) {
      System.out.println(
          "Running MiniRPCBenchmark with delegation token authentication.");
      elapsedTime = mb.runMiniBenchmarkWithDelegationToken(
                              conf, count, KEYTAB_FILE_KEY, USER_NAME_KEY);
    } else {
      String auth = SecurityUtil.getAuthenticationMethod(conf).toString();
      System.out.println(
          "Running MiniRPCBenchmark with " + auth + " authentication.");
      elapsedTime = mb.runMiniBenchmark(
                              conf, count, KEYTAB_FILE_KEY, USER_NAME_KEY);
    }
    System.out.println(org.apache.hadoop.util.VersionInfo.getVersion());
    System.out.println("Number  of  connects: " + count);
    System.out.println("Average connect time: " + ((double)elapsedTime/count));
  }
  
  private void configureSuperUserIPAddresses(Configuration conf,
      String superUserShortName) throws IOException {
    ArrayList<String> ipList = new ArrayList<String>();
    Enumeration<NetworkInterface> netInterfaceList = NetworkInterface
        .getNetworkInterfaces();
    while (netInterfaceList.hasMoreElements()) {
      NetworkInterface inf = netInterfaceList.nextElement();
      Enumeration<InetAddress> addrList = inf.getInetAddresses();
      while (addrList.hasMoreElements()) {
        InetAddress addr = addrList.nextElement();
        ipList.add(addr.getHostAddress());
      }
    }
    StringBuilder builder = new StringBuilder();
    for (String ip : ipList) {
      builder.append(ip);
      builder.append(',');
    }
    builder.append("127.0.1.1,");
    builder.append(InetAddress.getLocalHost().getCanonicalHostName());
    conf.setStrings(DefaultImpersonationProvider.getTestProvider().
            getProxySuperuserIpConfKey(superUserShortName), builder.toString());
  }
}
