/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio.security.authentication;

import alluxio.exception.status.AlluxioStatusException;
import alluxio.grpc.ChannelAuthenticationScheme;
import alluxio.grpc.SaslMessage;
import alluxio.grpc.SaslMessageType;

import io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.UUID;
import javax.security.sasl.SaslException;

/**
 * Responsible for driving authentication traffic from server-side.
 */
public class AuthenticatedChannelServerDriver implements StreamObserver<SaslMessage> {
  private static final Logger LOG = LoggerFactory.getLogger(AuthenticatedChannelServerDriver.class);

  /** Used to represent uninitialized channel Id. */
  private static final UUID EMPTY_UUID = new UUID(0L, 0L);

  /** Client's sasl stream. */
  private StreamObserver<SaslMessage> mRequestObserver = null;
  /** Authentication server. */
  private AuthenticationServer mAuthenticationServer;

  /** Id for client-side channel that is authenticating. */
  private UUID mChannelId = EMPTY_UUID;
  /** Reference for client owning the channel. */
  private String mChannelRef;
  /** Sasl server handler that will be used to build secure stream. */
  private SaslServerHandler mSaslServerHandler = null;
  /** Whether channel is authenticated. */
  private volatile boolean mChannelAuthenticated = false;

  /**
   * Creates {@link AuthenticatedChannelServerDriver} for given {@link AuthenticationServer}.
   *
   * @param authenticationServer authentication server
   */
  public AuthenticatedChannelServerDriver(AuthenticationServer authenticationServer) {
    mAuthenticationServer = authenticationServer;
  }

  /**
   * Sets the client's Sasl stream.
   *
   * @param requestObserver client Sasl stream
   */
  public void setClientObserver(StreamObserver<SaslMessage> requestObserver) {
    mRequestObserver = requestObserver;
  }

  private void initAuthenticatedChannel(ChannelAuthenticationScheme authScheme, UUID channelId,
      String channelRef) throws SaslException {
    LOG.debug("Initializing authentication for channel: {}. AuthType: {}", mChannelRef, authScheme);
    // Create sasl handler for the requested scheme.
    mSaslServerHandler = mAuthenticationServer.createSaslHandler(authScheme);
    // Unregister from registry if in case it was authenticated before.
    mAuthenticationServer.unregisterChannel(mChannelId);
    // Initialize channel id and ref
    mChannelId = channelId;
    mChannelRef = channelRef;
  }

  private void channelAuthenticated(AuthenticatedUserInfo authUserInfo) {
    // Register authorized user with the authentication server.
    mAuthenticationServer.registerChannel(mChannelId, authUserInfo, this);
    mChannelAuthenticated = true;
  }

  private void closeAuthenticatedChannel(boolean signalOwner) {
    if (mChannelAuthenticated) {
      mAuthenticationServer.unregisterChannel(mChannelId);
      mChannelAuthenticated = false;
    }

    if (mSaslServerHandler != null) {
      mSaslServerHandler.close();
    }

    if (signalOwner) {
      try {
        // Complete stream.
        mRequestObserver.onCompleted();
      } catch (Exception exc) {
        LOG.debug("Failed to close gRPC stream of channel: {}. Error: {}", mChannelRef, exc);
      }
    }
  }

  @Override
  public void onNext(SaslMessage saslMessage) {
    try {
      if (mChannelId.equals(EMPTY_UUID)) {
        initAuthenticatedChannel(
            saslMessage.getAuthenticationScheme(),
            UUID.fromString(saslMessage.getClientId()),
            saslMessage.getChannelRef());
      }

      LOG.debug("Responding to a message of channel: {}. Message: {}", mChannelRef, saslMessage);
      // Consult sasl server for handling the message.
      SaslMessage response = mSaslServerHandler.handleMessage(saslMessage);

      // Activate if sasl is secured.
      if (response.getMessageType() == SaslMessageType.SUCCESS) {
        channelAuthenticated(mSaslServerHandler.getAuthenticatedUserInfo());
      }
      // Push response to stream.
      mRequestObserver.onNext(response);
    } catch (Throwable t) {
      LOG.debug("Exception while handling message of channel: {}. Message: {}. Error: {}",
              mChannelRef, saslMessage, t);
      // Invalidate stream.
      mRequestObserver.onError(AlluxioStatusException.fromThrowable(t).toGrpcStatusException());
      closeAuthenticatedChannel(false);
    }
  }

  @Override
  public void onError(Throwable throwable) {
    LOG.debug("Authentication stream failed for server. Channel: {}. Error: {}",
        mChannelRef, throwable);
    closeAuthenticatedChannel(false);
  }

  @Override
  public void onCompleted() {
    LOG.debug("Authentication closed by client. Channel: {}.", mChannelRef);
    closeAuthenticatedChannel(true);
  }

  /**
   * Completes authenticated channel.
   */
  public void close() {
    LOG.debug("Authentication server-driver closing. Channel: {}", mChannelRef);
    closeAuthenticatedChannel(true);
  }
}
