/*
 * 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.celeborn.common.network.sasl.registration;

import static org.apache.celeborn.common.network.sasl.SaslUtils.*;
import static org.apache.celeborn.common.protocol.MessageType.*;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.List;

import com.google.common.base.Throwables;
import com.google.common.collect.Lists;
import io.netty.channel.Channel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.celeborn.common.network.client.RpcResponseCallback;
import org.apache.celeborn.common.network.client.TransportClient;
import org.apache.celeborn.common.network.protocol.RequestMessage;
import org.apache.celeborn.common.network.protocol.RpcFailure;
import org.apache.celeborn.common.network.protocol.RpcRequest;
import org.apache.celeborn.common.network.protocol.TransportMessage;
import org.apache.celeborn.common.network.sasl.CelebornSaslServer;
import org.apache.celeborn.common.network.sasl.SaslRpcHandler;
import org.apache.celeborn.common.network.sasl.SecretRegistry;
import org.apache.celeborn.common.network.server.BaseMessageHandler;
import org.apache.celeborn.common.network.util.TransportConf;
import org.apache.celeborn.common.protocol.PbAuthType;
import org.apache.celeborn.common.protocol.PbAuthenticationInitiationRequest;
import org.apache.celeborn.common.protocol.PbAuthenticationInitiationResponse;
import org.apache.celeborn.common.protocol.PbRegisterApplicationRequest;
import org.apache.celeborn.common.protocol.PbRegisterApplicationResponse;
import org.apache.celeborn.common.protocol.PbSaslMechanism;
import org.apache.celeborn.common.protocol.PbSaslRequest;

/**
 * RPC Handler which registers an application. If an application is registered or the connection is
 * authenticated, subsequent messages are delegated to a child RPC handler.
 */
public class RegistrationRpcHandler extends BaseMessageHandler {
  private static final Logger LOG = LoggerFactory.getLogger(RegistrationRpcHandler.class);

  private static final String VERSION = "1.0";

  /**
   * TODO: This should be made configurable. For now, we only support ANONYMOUS for client-auth and
   * DIGEST-MD5 for connect-auth.
   */
  private static final List<PbSaslMechanism> SASL_MECHANISMS =
      Lists.newArrayList(
          PbSaslMechanism.newBuilder()
              .setMechanism(ANONYMOUS)
              .addAuthTypes(PbAuthType.CLIENT_AUTH)
              .build(),
          PbSaslMechanism.newBuilder()
              .setMechanism(DIGEST_MD5)
              .addAuthTypes(PbAuthType.CONNECTION_AUTH)
              .build());

  /** Transport configuration. */
  private final TransportConf conf;

  /** The client channel. */
  private final Channel channel;

  private final BaseMessageHandler delegate;

  private RegistrationState registrationState = RegistrationState.NONE;

  /** Class which provides secret keys which are shared by server and client on a per-app basis. */
  private final SecretRegistry secretRegistry;

  private SaslRpcHandler saslHandler;

  /** Used for client authentication. */
  private CelebornSaslServer saslServer = null;

  public RegistrationRpcHandler(
      TransportConf conf,
      Channel channel,
      BaseMessageHandler delegate,
      SecretRegistry secretRegistry) {
    this.conf = conf;
    this.channel = channel;
    this.secretRegistry = secretRegistry;
    this.delegate = delegate;
    this.saslHandler = new SaslRpcHandler(conf, channel, delegate, secretRegistry);
  }

  @Override
  public boolean checkRegistered() {
    return delegate.checkRegistered();
  }

  @Override
  public final void receive(
      TransportClient client, RequestMessage message, RpcResponseCallback callback) {
    // The message is delegated either if the client is already authenticated or if the connection
    // is authenticated.
    if (registrationState == RegistrationState.REGISTERED || saslHandler.isAuthenticated()) {
      LOG.trace("Already authenticated. Delegating {}", client.getClientId());
      delegate.receive(client, message, callback);
    } else {
      RpcRequest rpcRequest = (RpcRequest) message;
      try {
        processRpcMessage(client, rpcRequest, callback);
      } catch (Exception e) {
        LOG.error("Error while invoking RpcHandler#receive() on RPC id " + rpcRequest.requestId, e);
        registrationState = RegistrationState.FAILED;
        client
            .getChannel()
            .writeAndFlush(
                new RpcFailure(rpcRequest.requestId, Throwables.getStackTraceAsString(e)));
      }
    }
  }

  @Override
  public final void receive(TransportClient client, RequestMessage message) {
    if (registrationState == RegistrationState.REGISTERED || saslHandler.isAuthenticated()) {
      if (LOG.isTraceEnabled()) {
        LOG.trace("Already authenticated. Delegating {}", client.getClientId());
      }
      delegate.receive(client, message);
    } else {
      throw new SecurityException("Unauthenticated call to receive().");
    }
  }

  private void processRpcMessage(
      TransportClient client, RpcRequest message, RpcResponseCallback callback) throws IOException {
    TransportMessage pbMsg = TransportMessage.fromByteBuffer(message.body().nioByteBuffer());
    switch (pbMsg.getMessageTypeValue()) {
      case AUTHENTICATION_INITIATION_REQUEST_VALUE:
        // TODO: not validating the auth init request. Should we?
        PbAuthenticationInitiationRequest authInitRequest = pbMsg.getParsedPayload();
        checkRequestAllowed(RegistrationState.NONE);
        respondToAuthInitialization(callback);
        registrationState = RegistrationState.INIT;
        LOG.trace("Authentication initialization completed: rpcId {}", message.requestId);
        break;
      case SASL_REQUEST_VALUE:
        PbSaslRequest saslRequest = pbMsg.getParsedPayload();
        if (saslRequest.getAuthType().equals(PbAuthType.CLIENT_AUTH)) {
          LOG.trace("Received Sasl Message for client authentication");
          checkRequestAllowed(RegistrationState.INIT);
          authenticateClient(saslRequest, callback);
          if (saslServer.isComplete()) {
            LOG.debug("SASL authentication successful for channel {}", client);
            complete();
            registrationState = RegistrationState.AUTHENTICATED;
            LOG.trace("Client authenticated: rpcId {}", message.requestId);
          }
        } else {
          // It is a SASL message to authenticate the connection. If the application hasn't
          // registered, then
          // saslHandler will throw an exception that the app hasn't registered.
          LOG.trace("Delegating to sasl handler: rpcId {}", message.requestId);
          saslHandler.receive(client, message, callback);
        }
        break;
      case REGISTER_APPLICATION_REQUEST_VALUE:
        PbRegisterApplicationRequest registerApplicationRequest = pbMsg.getParsedPayload();
        checkRequestAllowed(RegistrationState.AUTHENTICATED);
        LOG.trace("Application registration started {}", registerApplicationRequest.getId());
        processRegisterApplicationRequest(registerApplicationRequest, callback);
        registrationState = RegistrationState.REGISTERED;
        client.setClientId(registerApplicationRequest.getId());
        LOG.info(
            "Application registered: appId {} rpcId {}",
            registerApplicationRequest.getId(),
            message.requestId);
        break;
      default:
        throw new SecurityException(
            "The app is not registered and the connection is not authenticated "
                + message.requestId);
    }
  }

  private void checkRequestAllowed(RegistrationState expectedState) {
    if (registrationState != expectedState) {
      throw new IllegalStateException(
          "Invalid registration state. Expected: "
              + expectedState
              + ", Actual: "
              + registrationState);
    }
  }

  private void respondToAuthInitialization(RpcResponseCallback callback) {
    PbAuthenticationInitiationResponse response =
        PbAuthenticationInitiationResponse.newBuilder()
            .setAuthEnabled(conf.authEnabled())
            .setVersion(VERSION)
            .addAllSaslMechanisms(SASL_MECHANISMS)
            .build();
    TransportMessage message =
        new TransportMessage(AUTHENTICATION_INITIATION_RESPONSE, response.toByteArray());
    callback.onSuccess(message.toByteBuffer());
  }

  private void authenticateClient(PbSaslRequest saslMessage, RpcResponseCallback callback) {
    if (saslServer == null || !saslServer.isComplete()) {
      if (saslServer == null) {
        saslServer = new CelebornSaslServer(ANONYMOUS, null, null);
      }
      byte[] response = saslServer.response(saslMessage.getPayload().toByteArray());
      callback.onSuccess(ByteBuffer.wrap(response));
    } else {
      throw new IllegalArgumentException("Unexpected message type " + saslMessage.toString());
    }
  }

  private void processRegisterApplicationRequest(
      PbRegisterApplicationRequest registerApplicationRequest, RpcResponseCallback callback) {
    if (secretRegistry.isRegistered(registerApplicationRequest.getId())) {
      // Re-registration is not allowed.
      throw new IllegalStateException(
          "Application is already registered " + registerApplicationRequest.getId());
    }
    secretRegistry.register(
        registerApplicationRequest.getId(), registerApplicationRequest.getSecret());
    PbRegisterApplicationResponse response =
        PbRegisterApplicationResponse.newBuilder().setStatus(true).build();
    TransportMessage message =
        new TransportMessage(REGISTER_APPLICATION_RESPONSE, response.toByteArray());
    callback.onSuccess(message.toByteBuffer());
  }

  @Override
  public void channelInactive(TransportClient client) {
    delegate.channelInactive(client);
    cleanup();
  }

  @Override
  public void exceptionCaught(Throwable cause, TransportClient client) {
    delegate.exceptionCaught(cause, client);
  }

  private void complete() {
    cleanup();
  }

  private void cleanup() {
    if (null != saslServer) {
      try {
        saslServer.dispose();
      } catch (RuntimeException e) {
        LOG.error("Error while disposing SASL server", e);
      } finally {
        saslServer = null;
      }
    }
    saslHandler.cleanup();
  }

  private enum RegistrationState {
    NONE,
    INIT,
    AUTHENTICATED,
    REGISTERED,
    FAILED
  }
}
