/*

  * 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.giraph.comm.netty.handler;

 import org.apache.giraph.comm.netty.NettyServer;

 import org.apache.giraph.comm.netty.SaslNettyServer;

 import org.apache.giraph.comm.requests.RequestType;

 import org.apache.giraph.comm.requests.SaslCompleteRequest;

 import org.apache.giraph.comm.requests.SaslTokenMessageRequest;

 import org.apache.giraph.comm.requests.WritableRequest;

 import org.apache.hadoop.conf.Configuration;

 import org.apache.hadoop.mapred.JobConf;

 import org.apache.hadoop.mapreduce.security.TokenCache;

 import org.apache.hadoop.mapreduce.security.token.JobTokenIdentifier;

 import org.apache.hadoop.mapreduce.security.token.JobTokenSecretManager;

 import org.apache.hadoop.security.Credentials;

 import org.apache.hadoop.security.UserGroupInformation;

 import org.apache.hadoop.security.token.Token;

 import org.apache.hadoop.security.token.TokenIdentifier;

 import org.apache.hadoop.util.ReflectionUtils;

 import org.apache.log4j.Logger;

 import org.apache.hadoop.security.SaslRpcServer.AuthMethod;

 import io.netty.channel.ChannelHandlerContext;

 import io.netty.channel.ChannelInboundHandlerAdapter;

 import java.io.ByteArrayInputStream;

 import java.io.DataInputStream;

 import java.io.IOException;

 import java.util.Collection;

 import static org.apache.giraph.conf.GiraphConstants.NETTY_SIMULATE_FIRST_REQUEST_CLOSED;

 /**

  * Generate SASL response tokens to client SASL tokens, allowing clients to

  * authenticate themselves with this server.

  */

 public class SaslServerHandler extends

     ChannelInboundHandlerAdapter {

     /** Class logger */

   private static final Logger LOG =

       Logger.getLogger(SaslServerHandler.class);

   // TODO: Move out into a separate, dedicated handler: ("FirstRequestHandler")

   // or similar.

   /** Already closed first request? */

   private static volatile boolean ALREADY_CLOSED_FIRST_REQUEST = false;

   /** Close connection on first request (used for simulating failure) */

   private final boolean closeFirstRequest;

   /** Used to store Hadoop Job Tokens to authenticate clients. */

   private JobTokenSecretManager secretManager;

   /**

    * Constructor

    *

    * @param conf Configuration

    */

   public SaslServerHandler(

       Configuration conf) throws IOException {

     SaslNettyServer.init(conf);

     setupSecretManager(conf);

     closeFirstRequest = NETTY_SIMULATE_FIRST_REQUEST_CLOSED.get(conf);

   }

   @Override

   public void channelRead(ChannelHandlerContext ctx, Object msg)

     throws Exception {

     if (LOG.isDebugEnabled()) {

       LOG.debug("messageReceived: Got " + msg.getClass());

     }

     WritableRequest writableRequest = (WritableRequest) msg;

     // Simulate a closed connection on the first request (if desired)

     // TODO: Move out into a separate, dedicated handler.

     if (closeFirstRequest && !ALREADY_CLOSED_FIRST_REQUEST) {

       LOG.info("messageReceived: Simulating closing channel on first " +

           "request " + writableRequest.getRequestId() + " from " +

           writableRequest.getClientId());

       setAlreadyClosedFirstRequest();

       ctx.close();

       return;

     }

     if (writableRequest.getType() == RequestType.SASL_TOKEN_MESSAGE_REQUEST) {

       // initialize server-side SASL functionality, if we haven't yet

       // (in which case we are looking at the first SASL message from the

       // client).

       SaslNettyServer saslNettyServer =

           ctx.attr(NettyServer.CHANNEL_SASL_NETTY_SERVERS).get();

       if (saslNettyServer == null) {

         if (LOG.isDebugEnabled()) {

           LOG.debug("No saslNettyServer for " + ctx.channel() +

               " yet; creating now, with secret manager: " + secretManager);

         }

         try {

           saslNettyServer = new SaslNettyServer(secretManager,

             AuthMethod.SIMPLE);

         } catch (IOException ioe) { //TODO:

           throw new RuntimeException(ioe);

         }

         ctx.attr(NettyServer.CHANNEL_SASL_NETTY_SERVERS).set(saslNettyServer);

       } else {

         if (LOG.isDebugEnabled()) {

           LOG.debug("Found existing saslNettyServer on server:" +

               ctx.channel().localAddress() + " for client " +

               ctx.channel().remoteAddress());

         }

       }

       ((SaslTokenMessageRequest) writableRequest).processToken(saslNettyServer);

       // Send response to client.

       ctx.write(writableRequest);

       if (saslNettyServer.isComplete()) {

         // If authentication of client is complete, we will also send a

         // SASL-Complete message to the client.

         if (LOG.isDebugEnabled()) {

           LOG.debug("SASL authentication is complete for client with " +

               "username: " + saslNettyServer.getUserName());

         }

         SaslCompleteRequest saslComplete = new SaslCompleteRequest();

         ctx.write(saslComplete);

         if (LOG.isDebugEnabled()) {

           LOG.debug("Removing SaslServerHandler from pipeline since SASL " +

               "authentication is complete.");

         }

         ctx.pipeline().remove(this);

       }

       ctx.flush();

       // do not send upstream to other handlers: no further action needs to be

       // done for SASL_TOKEN_MESSAGE_REQUEST requests.

       return;

     } else {

       // Client should not be sending other-than-SASL messages before

       // SaslServerHandler has removed itself from the pipeline. Such non-SASL

       // requests will be denied by the Authorize channel handler (the next

       // handler upstream in the server pipeline) if SASL authentication has

       // not completed.

       LOG.warn("Sending upstream an unexpected non-SASL message :  " +

           writableRequest);

       ctx.fireChannelRead(msg);

     }

   }

   /**

    * Set already closed first request flag

    */

   private static void setAlreadyClosedFirstRequest() {

     ALREADY_CLOSED_FIRST_REQUEST = true;

   }

   /**

    * Load Hadoop Job Token into secret manager.

    *

    * @param conf Configuration

    * @throws IOException

    */

   private void setupSecretManager(Configuration conf) throws IOException {

     secretManager = new JobTokenSecretManager();

     String localJobTokenFile = System.getenv().get(

         UserGroupInformation.HADOOP_TOKEN_FILE_LOCATION);

     if (localJobTokenFile == null) {

       throw new IOException("Could not find job credentials: environment " +

           "variable: " + UserGroupInformation.HADOOP_TOKEN_FILE_LOCATION +

           " was not defined.");

     }

     JobConf jobConf = new JobConf(conf);

     // Find the JobTokenIdentifiers among all the tokens available in the

     // jobTokenFile and store them in the secretManager.

     Credentials credentials =

         TokenCache.loadTokens(localJobTokenFile, jobConf);

     Collection<Token<? extends TokenIdentifier>> collection =

         credentials.getAllTokens();

     for (Token<? extends TokenIdentifier> token:  collection) {

       TokenIdentifier tokenIdentifier = decodeIdentifier(token,

           JobTokenIdentifier.class);

       if (tokenIdentifier instanceof JobTokenIdentifier) {

         Token<JobTokenIdentifier> theToken =

             (Token<JobTokenIdentifier>) token;

         JobTokenIdentifier jobTokenIdentifier =

             (JobTokenIdentifier) tokenIdentifier;

         secretManager.addTokenForJob(

             jobTokenIdentifier.getJobId().toString(), theToken);

       }

     }

     if (LOG.isDebugEnabled()) {

       LOG.debug("loaded JobToken credentials: " + credentials + " from " +

           "localJobTokenFile: " + localJobTokenFile);

     }

   }

   /**

    * Get the token identifier object, or null if it could not be constructed

    * (because the class could not be loaded, for example).

    * Hadoop 2.0.0 (and older Hadoop2 versions? (verify)) need this.

    * Hadoop 2.0.1 and newer have a Token.decodeIdentifier() method and do not

    * need this. Might want to create a munge flag to distinguish 2.0.0 vs newer.

    *

    * @param token the token to decode into a TokenIdentifier

    * @param cls the subclass of TokenIdentifier to decode the token into.

    * @return the token identifier.

    * @throws IOException

    */

   @SuppressWarnings("unchecked")

   private TokenIdentifier decodeIdentifier(

       Token<? extends TokenIdentifier> token,

       Class<? extends TokenIdentifier> cls) throws IOException {

     TokenIdentifier tokenIdentifier = ReflectionUtils.newInstance(cls, null);

     ByteArrayInputStream buf = new ByteArrayInputStream(token.getIdentifier());

     DataInputStream in = new DataInputStream(buf);

     tokenIdentifier.readFields(in);

     in.close();

     return tokenIdentifier;

   }

   /** Factory for {@link SaslServerHandler} */

   public static class Factory {

     /**

      * Constructor

      */

     public Factory() {

     }

     /**

      * Create new {@link SaslServerHandler}

      *

      * @param conf Configuration to use

      * @return New {@link SaslServerHandler}

      */

     public SaslServerHandler newHandler(

         Configuration conf) throws IOException {

       return new SaslServerHandler(conf);

     }

   }

 }