/*
 * 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.hive.metastore;

import java.io.IOException;
import java.lang.reflect.InvocationTargetException;
import java.net.Socket;
import java.security.PrivilegedExceptionAction;
import java.util.List;
import java.util.Map;

import org.apache.hadoop.hive.metastore.api.ThriftHiveMetastore;
import org.apache.hadoop.hive.metastore.security.TUGIContainingTransport;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.hive.metastore.api.ThriftHiveMetastore.Iface;
import org.apache.hadoop.hive.metastore.api.ThriftHiveMetastore.set_ugi_args;
import org.apache.hadoop.hive.metastore.api.ThriftHiveMetastore.set_ugi_result;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.thrift.ProcessFunction;
import org.apache.thrift.TApplicationException;
import org.apache.thrift.TBase;
import org.apache.thrift.TException;
import org.apache.thrift.protocol.TMessage;
import org.apache.thrift.protocol.TMessageType;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.protocol.TProtocolException;
import org.apache.thrift.protocol.TProtocolUtil;
import org.apache.thrift.protocol.TType;
import org.apache.thrift.transport.TTransportException;

/** TUGIBasedProcessor is used in unsecure mode for thrift metastore client server communication.
 *  This processor checks whether the first rpc call after connection is set up is set_ugi()
 *  through which client sends ugi to server. Processor then perform all subsequent rpcs on the
 *  connection using ugi.doAs() so all actions are performed in client user context.
 *  Note that old clients will never call set_ugi() and thus ugi will never be received on server
 *  side, in which case server exhibits previous behavior and continues as usual.
 */
@SuppressWarnings("rawtypes")
public class TUGIBasedProcessor<I extends Iface> extends TSetIpAddressProcessor<Iface> {

  private final I iface;
  private final Map<String,  org.apache.thrift.ProcessFunction<Iface, ? extends  TBase>>
    functions;
  static final Logger LOG = LoggerFactory.getLogger(TUGIBasedProcessor.class);

  public TUGIBasedProcessor(I iface) throws SecurityException, NoSuchFieldException,
    IllegalArgumentException, IllegalAccessException, NoSuchMethodException,
    InvocationTargetException {
    super(iface);
    this.iface = iface;
    this.functions = getProcessMapView();
  }

  @SuppressWarnings("unchecked")
  @Override
  public void process(final TProtocol in, final TProtocol out) throws TException {
    setIpAddress(in);

    final TMessage msg;
    try {
      msg = in.readMessageBegin();
    } catch (TTransportException tte) {
      throw new RuntimeException(tte);
    }
    final ProcessFunction<Iface, ? extends  TBase> fn = functions.get(msg.name);
    if (fn == null) {
      TProtocolUtil.skip(in, TType.STRUCT);
      in.readMessageEnd();
      TApplicationException x = new TApplicationException(TApplicationException.UNKNOWN_METHOD,
          "Invalid method name: '"+msg.name+"'");
      out.writeMessageBegin(new TMessage(msg.name, TMessageType.EXCEPTION, msg.seqid));
      x.write(out);
      out.writeMessageEnd();
      out.getTransport().flush();
      return;
    }
    TUGIContainingTransport ugiTrans = (TUGIContainingTransport)in.getTransport();
    // Store ugi in transport if the rpc is set_ugi
    if (msg.name.equalsIgnoreCase("set_ugi")){
      try {
        handleSetUGI(ugiTrans, (ThriftHiveMetastore.Processor.set_ugi<Iface>)fn, msg, in, out);
      } catch (TException e) {
        throw e;
      } catch (Exception e) {
        throw new TException(e.getCause());
      }
      return;
    }
    UserGroupInformation clientUgi = ugiTrans.getClientUGI();
    if (null == clientUgi){
      // At this point, transport must contain client ugi, if it doesn't then its an old client.
      fn.process(msg.seqid, in, out, iface);
      return;
    } else { // Found ugi, perform doAs().
      PrivilegedExceptionAction<Void> pvea = new PrivilegedExceptionAction<Void>() {
        @Override
        public Void run() {
          try {
            fn.process(msg.seqid,in, out, iface);
            return null;
          } catch (TException te) {
            throw new RuntimeException(te);
          }
        }
      };
      try {
        clientUgi.doAs(pvea);
        return;
      } catch (RuntimeException rte) {
        if (rte.getCause() instanceof TException) {
          throw (TException)rte.getCause();
        }
        throw rte;
      } catch (InterruptedException ie) {
        throw new RuntimeException(ie); // unexpected!
      } catch (IOException ioe) {
        throw new RuntimeException(ioe); // unexpected!
      } finally {
          try {
            FileSystem.closeAllForUGI(clientUgi);
          } catch (IOException e) {
            LOG.error("Could not clean up file-system handles for UGI: " + clientUgi, e);
          }
      }
    }
  }

  private void handleSetUGI(TUGIContainingTransport ugiTrans,
                            ThriftHiveMetastore.Processor.set_ugi<Iface> fn, TMessage msg, TProtocol iprot, TProtocol oprot)
      throws TException, SecurityException, NoSuchMethodException, IllegalArgumentException,
      IllegalAccessException, InvocationTargetException{

    UserGroupInformation clientUgi = ugiTrans.getClientUGI();
    if( null != clientUgi){
      throw new TException(new IllegalStateException("UGI is already set. Resetting is not " +
      "allowed. Current ugi is: " + clientUgi.getUserName()));
    }

    set_ugi_args args = fn.getEmptyArgsInstance();
    try {
      args.read(iprot);
    } catch (TProtocolException e) {
      iprot.readMessageEnd();
      TApplicationException x = new TApplicationException(TApplicationException.PROTOCOL_ERROR,
          e.getMessage());
      oprot.writeMessageBegin(new TMessage(msg.name, TMessageType.EXCEPTION, msg.seqid));
      x.write(oprot);
      oprot.writeMessageEnd();
      oprot.getTransport().flush();
      return;
    }
    iprot.readMessageEnd();
    set_ugi_result result = fn.getResult(iface, args);
    List<String> principals = result.getSuccess();
    // Store the ugi in transport and then continue as usual.
    ugiTrans.setClientUGI(UserGroupInformation.createRemoteUser(principals.remove(principals.size()-1)));
    oprot.writeMessageBegin(new TMessage(msg.name, TMessageType.REPLY, msg.seqid));
    result.write(oprot);
    oprot.writeMessageEnd();
    oprot.getTransport().flush();
  }

  @Override
  protected void setIpAddress(final TProtocol in) {
    TUGIContainingTransport ugiTrans = (TUGIContainingTransport)in.getTransport();
    Socket socket = ugiTrans.getSocket();
    if (socket != null) {
      setIpAddress(socket);
    }
  }
}
