/**
 * 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 com.google.protobuf.BlockingService;
import com.google.protobuf.RpcController;
import com.google.protobuf.ServiceException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.ipc.protobuf.TestProtos;
import org.apache.hadoop.ipc.protobuf.TestRpcServiceProtos;
import org.apache.hadoop.net.NetUtils;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.security.token.SecretManager;
import org.apache.hadoop.util.Time;
import org.junit.Assert;

import org.apache.hadoop.io.Text;
import org.apache.hadoop.io.retry.RetryPolicy;
import org.apache.hadoop.ipc.protobuf.TestProtos;
import org.apache.hadoop.ipc.protobuf.TestRpcServiceProtos;
import org.apache.hadoop.security.KerberosInfo;
import org.apache.hadoop.security.SaslRpcServer.AuthMethod;
import org.apache.hadoop.security.token.Token;
import org.apache.hadoop.security.token.TokenIdentifier;
import org.apache.hadoop.security.token.TokenInfo;
import org.apache.hadoop.security.token.TokenSelector;

import java.io.DataInput;
import java.io.DataOutput;
import java.io.IOException;
import java.lang.management.ManagementFactory;
import java.lang.management.ThreadInfo;
import java.lang.management.ThreadMXBean;
import java.net.InetSocketAddress;
import java.net.URISyntaxException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;

/** Test facilities for unit tests for RPC. */
public class TestRpcBase {

  protected final static String SERVER_PRINCIPAL_KEY =
      "test.ipc.server.principal";
  protected final static String ADDRESS = "0.0.0.0";
  protected final static int PORT = 0;
  protected static InetSocketAddress addr;
  protected static Configuration conf;

  protected void setupConf() {
    conf = new Configuration();
    // Set RPC engine to protobuf RPC engine
    RPC.setProtocolEngine(conf, TestRpcService.class, ProtobufRpcEngine.class);
    UserGroupInformation.setConfiguration(conf);
  }

  protected static RPC.Builder newServerBuilder(
      Configuration serverConf) throws IOException {
    // Create server side implementation
    PBServerImpl serverImpl = new PBServerImpl();
    BlockingService service = TestRpcServiceProtos.TestProtobufRpcProto
        .newReflectiveBlockingService(serverImpl);

    // Get RPC server for server side implementation
    RPC.Builder builder = new RPC.Builder(serverConf)
        .setProtocol(TestRpcService.class)
        .setInstance(service).setBindAddress(ADDRESS).setPort(PORT);

    return builder;
  }

  protected static RPC.Server setupTestServer(Configuration serverConf,
                                       int numHandlers) throws IOException {
    return setupTestServer(serverConf, numHandlers, null);
  }

  protected static RPC.Server setupTestServer(Configuration serverConf,
                                       int numHandlers,
                                       SecretManager<?> serverSm)
      throws IOException {
    RPC.Builder builder = newServerBuilder(serverConf);

    if (numHandlers > 0) {
      builder.setNumHandlers(numHandlers);
    }

    if (serverSm != null) {
      builder.setSecretManager(serverSm);
    }

    return setupTestServer(builder);
  }

  protected static RPC.Server setupTestServer(RPC.Builder builder) throws IOException {
    RPC.Server server = builder.build();

    server.start();

    addr = NetUtils.getConnectAddress(server);

    return server;
  }

  protected static TestRpcService getClient(InetSocketAddress serverAddr,
                                     Configuration clientConf)
      throws ServiceException {
    try {
      return RPC.getProxy(TestRpcService.class, 0, serverAddr, clientConf);
    } catch (IOException e) {
      throw new ServiceException(e);
    }
  }

  protected static TestRpcService getClient(InetSocketAddress serverAddr,
      Configuration clientConf, final RetryPolicy connectionRetryPolicy)
      throws ServiceException {
    try {
      return RPC.getProtocolProxy(
          TestRpcService.class,
          0,
          serverAddr,
          UserGroupInformation.getCurrentUser(),
          clientConf,
          NetUtils.getDefaultSocketFactory(clientConf),
          RPC.getRpcTimeout(clientConf),
          connectionRetryPolicy, null).getProxy();
    } catch (IOException e) {
      throw new ServiceException(e);
    }
  }

  protected static void stop(Server server, TestRpcService proxy) {
    if (proxy != null) {
      try {
        RPC.stopProxy(proxy);
      } catch (Exception ignored) {}
    }

    if (server != null) {
      try {
        server.stop();
      } catch (Exception ignored) {}
    }
  }

  /**
   * Count the number of threads that have a stack frame containing
   * the given string
   */
  protected static int countThreads(String search) {
    ThreadMXBean threadBean = ManagementFactory.getThreadMXBean();

    int count = 0;
    ThreadInfo[] infos = threadBean.getThreadInfo(threadBean.getAllThreadIds(), 20);
    for (ThreadInfo info : infos) {
      if (info == null) continue;
      for (StackTraceElement elem : info.getStackTrace()) {
        if (elem.getClassName().contains(search)) {
          count++;
          break;
        }
      }
    }
    return count;
  }

  public static class TestTokenIdentifier extends TokenIdentifier {
    private Text tokenid;
    private Text realUser;
    final static Text KIND_NAME = new Text("test.token");

    public TestTokenIdentifier() {
      this(new Text(), new Text());
    }
    public TestTokenIdentifier(Text tokenid) {
      this(tokenid, new Text());
    }
    public TestTokenIdentifier(Text tokenid, Text realUser) {
      this.tokenid = tokenid == null ? new Text() : tokenid;
      this.realUser = realUser == null ? new Text() : realUser;
    }
    @Override
    public Text getKind() {
      return KIND_NAME;
    }
    @Override
    public UserGroupInformation getUser() {
      if (realUser.toString().isEmpty()) {
        return UserGroupInformation.createRemoteUser(tokenid.toString());
      } else {
        UserGroupInformation realUgi = UserGroupInformation
            .createRemoteUser(realUser.toString());
        return UserGroupInformation
            .createProxyUser(tokenid.toString(), realUgi);
      }
    }

    @Override
    public void readFields(DataInput in) throws IOException {
      tokenid.readFields(in);
      realUser.readFields(in);
    }
    @Override
    public void write(DataOutput out) throws IOException {
      tokenid.write(out);
      realUser.write(out);
    }
  }

  public static class TestTokenSecretManager extends
      SecretManager<TestTokenIdentifier> {
    @Override
    public byte[] createPassword(TestTokenIdentifier id) {
      return id.getBytes();
    }

    @Override
    public byte[] retrievePassword(TestTokenIdentifier id)
        throws InvalidToken {
      return id.getBytes();
    }

    @Override
    public TestTokenIdentifier createIdentifier() {
      return new TestTokenIdentifier();
    }
  }

  public static class TestTokenSelector implements
      TokenSelector<TestTokenIdentifier> {
    @SuppressWarnings("unchecked")
    @Override
    public Token<TestTokenIdentifier> selectToken(Text service,
                                                  Collection<Token<? extends TokenIdentifier>> tokens) {
      if (service == null) {
        return null;
      }
      for (Token<? extends TokenIdentifier> token : tokens) {
        if (TestTokenIdentifier.KIND_NAME.equals(token.getKind())
            && service.equals(token.getService())) {
          return (Token<TestTokenIdentifier>) token;
        }
      }
      return null;
    }
  }

  @KerberosInfo(serverPrincipal = SERVER_PRINCIPAL_KEY)
  @TokenInfo(TestTokenSelector.class)
  @ProtocolInfo(protocolName = "org.apache.hadoop.ipc.TestRpcBase$TestRpcService",
      protocolVersion = 1)
  public interface TestRpcService
      extends TestRpcServiceProtos.TestProtobufRpcProto.BlockingInterface {
  }

  public static class PBServerImpl implements TestRpcService {
    CountDownLatch fastPingCounter = new CountDownLatch(2);
    private List<Server.Call> postponedCalls = new ArrayList<>();
    private final Lock lock = new ReentrantLock();

    @Override
    public TestProtos.EmptyResponseProto ping(RpcController unused,
                TestProtos.EmptyRequestProto request) throws ServiceException {
      // Ensure clientId is received
      byte[] clientId = Server.getClientId();
      Assert.assertNotNull(clientId);
      Assert.assertEquals(ClientId.BYTE_LENGTH, clientId.length);
      return TestProtos.EmptyResponseProto.newBuilder().build();
    }

    @Override
    public TestProtos.EchoResponseProto echo(
        RpcController unused, TestProtos.EchoRequestProto request)
        throws ServiceException {
      return TestProtos.EchoResponseProto.newBuilder().setMessage(
          request.getMessage())
          .build();
    }

    @Override
    public TestProtos.EmptyResponseProto error(
        RpcController unused, TestProtos.EmptyRequestProto request)
        throws ServiceException {
      throw new ServiceException("error", new RpcServerException("error"));
    }

    @Override
    public TestProtos.EmptyResponseProto error2(
        RpcController unused, TestProtos.EmptyRequestProto request)
        throws ServiceException {
      throw new ServiceException("error", new URISyntaxException("",
          "testException"));
    }

    @Override
    public TestProtos.EmptyResponseProto slowPing(
        RpcController unused, TestProtos.SlowPingRequestProto request)
        throws ServiceException {
      boolean shouldSlow = request.getShouldSlow();
      if (shouldSlow) {
        try {
          fastPingCounter.await(); //slow response until two fast pings happened
        } catch (InterruptedException ignored) {}
      } else {
        fastPingCounter.countDown();
      }

      return TestProtos.EmptyResponseProto.newBuilder().build();
    }

    @Override
    public TestProtos.EchoResponseProto2 echo2(
        RpcController controller, TestProtos.EchoRequestProto2 request)
        throws ServiceException {
      return TestProtos.EchoResponseProto2.newBuilder().addAllMessage(
          request.getMessageList()).build();
    }

    @Override
    public TestProtos.AddResponseProto add(
        RpcController controller, TestProtos.AddRequestProto request)
        throws ServiceException {
      return TestProtos.AddResponseProto.newBuilder().setResult(
          request.getParam1() + request.getParam2()).build();
    }

    @Override
    public TestProtos.AddResponseProto add2(
        RpcController controller, TestProtos.AddRequestProto2 request)
        throws ServiceException {
      int sum = 0;
      for (Integer num : request.getParamsList()) {
        sum += num;
      }
      return TestProtos.AddResponseProto.newBuilder().setResult(sum).build();
    }

    @Override
    public TestProtos.EmptyResponseProto testServerGet(
        RpcController controller, TestProtos.EmptyRequestProto request)
        throws ServiceException {
      if (!(Server.get() instanceof RPC.Server)) {
        throw new ServiceException("Server.get() failed");
      }
      return TestProtos.EmptyResponseProto.newBuilder().build();
    }

    @Override
    public TestProtos.ExchangeResponseProto exchange(
        RpcController controller, TestProtos.ExchangeRequestProto request)
        throws ServiceException {
      Integer[] values = new Integer[request.getValuesCount()];
      for (int i = 0; i < values.length; i++) {
        values[i] = i;
      }
      return TestProtos.ExchangeResponseProto.newBuilder()
          .addAllValues(Arrays.asList(values)).build();
    }

    @Override
    public TestProtos.EmptyResponseProto sleep(
        RpcController controller, TestProtos.SleepRequestProto request)
        throws ServiceException {
      try {
        Thread.sleep(request.getMilliSeconds());
      } catch (InterruptedException ignore) {}
      return  TestProtos.EmptyResponseProto.newBuilder().build();
    }

    @Override
    public TestProtos.EmptyResponseProto lockAndSleep(
        RpcController controller, TestProtos.SleepRequestProto request)
        throws ServiceException {
      ProcessingDetails details =
          Server.getCurCall().get().getProcessingDetails();
      lock.lock();
      long startNanos = Time.monotonicNowNanos();
      try {
        Thread.sleep(request.getMilliSeconds());
      } catch (InterruptedException ignore) {
        // ignore
      } finally {
        lock.unlock();
      }
      // Add some arbitrary large lock wait time since in any test scenario
      // the lock wait time will probably actually be too small to notice
      details.add(ProcessingDetails.Timing.LOCKWAIT, 10, TimeUnit.SECONDS);
      details.add(ProcessingDetails.Timing.LOCKEXCLUSIVE,
          Time.monotonicNowNanos() - startNanos, TimeUnit.NANOSECONDS);
      return  TestProtos.EmptyResponseProto.newBuilder().build();
    }

    @Override
    public TestProtos.AuthMethodResponseProto getAuthMethod(
        RpcController controller, TestProtos.EmptyRequestProto request)
        throws ServiceException {
      AuthMethod authMethod = null;
      try {
        authMethod = UserGroupInformation.getCurrentUser()
            .getAuthenticationMethod().getAuthMethod();
      } catch (IOException e) {
        throw new ServiceException(e);
      }

      return TestProtos.AuthMethodResponseProto.newBuilder()
          .setCode(authMethod.code)
          .setMechanismName(authMethod.getMechanismName())
          .build();
    }

    @Override
    public TestProtos.AuthUserResponseProto getAuthUser(
        RpcController controller, TestProtos.EmptyRequestProto request)
        throws ServiceException {
      UserGroupInformation authUser = null;
      try {
        authUser = UserGroupInformation.getCurrentUser();
      } catch (IOException e) {
        throw new ServiceException(e);
      }

      return TestProtos.AuthUserResponseProto.newBuilder()
          .setAuthUser(authUser.getUserName())
          .build();
    }

    @Override
    public TestProtos.EchoResponseProto echoPostponed(
        RpcController controller, TestProtos.EchoRequestProto request)
        throws ServiceException {
      Server.Call call = Server.getCurCall().get();
      call.postponeResponse();
      postponedCalls.add(call);

      return TestProtos.EchoResponseProto.newBuilder().setMessage(
          request.getMessage())
          .build();
    }

    @Override
    public TestProtos.EmptyResponseProto sendPostponed(
        RpcController controller, TestProtos.EmptyRequestProto request)
        throws ServiceException {
      Collections.shuffle(postponedCalls);
      try {
        for (Server.Call call : postponedCalls) {
          call.sendResponse();
        }
      } catch (IOException e) {
        throw new ServiceException(e);
      }
      postponedCalls.clear();

      return TestProtos.EmptyResponseProto.newBuilder().build();
    }
  }

  protected static TestProtos.EmptyRequestProto newEmptyRequest() {
    return TestProtos.EmptyRequestProto.newBuilder().build();
  }

  protected static TestProtos.EmptyResponseProto newEmptyResponse() {
    return TestProtos.EmptyResponseProto.newBuilder().build();
  }

  protected static TestProtos.EchoRequestProto newEchoRequest(String msg) {
    return TestProtos.EchoRequestProto.newBuilder().setMessage(msg).build();
  }

  protected static String convert(TestProtos.EchoResponseProto response) {
    return response.getMessage();
  }

  protected static TestProtos.SlowPingRequestProto newSlowPingRequest(
      boolean shouldSlow) throws ServiceException {
    return TestProtos.SlowPingRequestProto.newBuilder().
        setShouldSlow(shouldSlow).build();
  }

  protected static TestProtos.SleepRequestProto newSleepRequest(
      int milliSeconds) {
    return TestProtos.SleepRequestProto.newBuilder()
        .setMilliSeconds(milliSeconds).build();
  }

  protected static TestProtos.EchoResponseProto newEchoResponse(String msg) {
    return TestProtos.EchoResponseProto.newBuilder().setMessage(msg).build();
  }

  protected static AuthMethod convert(
      TestProtos.AuthMethodResponseProto authMethodResponse) {
    String mechanism = authMethodResponse.getMechanismName();
    if (mechanism.equals(AuthMethod.SIMPLE.getMechanismName())) {
      return AuthMethod.SIMPLE;
    } else if (mechanism.equals(AuthMethod.KERBEROS.getMechanismName())) {
      return AuthMethod.KERBEROS;
    } else if (mechanism.equals(AuthMethod.TOKEN.getMechanismName())) {
      return AuthMethod.TOKEN;
    }
    return null;
  }

  protected static String convert(TestProtos.AuthUserResponseProto response) {
    return response.getAuthUser();
  }
}