exec/rpc/src/main/java/org/apache/drill/exec/rpc/security/AuthenticationOutcomeListener.java - drill - Git at Google

 /*
  * 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.drill.exec.rpc.security;

 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.Internal.EnumLite;
 import com.google.protobuf.MessageLite;
 import io.netty.buffer.ByteBuf;
 import org.apache.drill.exec.proto.UserBitShared.SaslMessage;
 import org.apache.drill.exec.proto.UserBitShared.SaslStatus;
 import org.apache.drill.exec.rpc.BasicClient;
 import org.apache.drill.exec.rpc.ClientConnection;
 import org.apache.drill.exec.rpc.RpcException;
 import org.apache.drill.exec.rpc.RpcOutcomeListener;
 import org.apache.hadoop.security.UserGroupInformation;

 import javax.security.sasl.Sasl;
 import javax.security.sasl.SaslClient;
 import javax.security.sasl.SaslException;
 import java.io.IOException;
 import java.lang.reflect.UndeclaredThrowableException;
 import java.security.PrivilegedExceptionAction;
 import java.util.EnumMap;
 import java.util.Map;

 import static com.google.common.base.Preconditions.checkNotNull;

 /**
  * Handles SASL exchange, on the client-side.
  *
  * @param <T> handshake rpc type
  * @param <C> Client connection type
  * @param <HS> Handshake send type
  * @param <HR> Handshake receive type
  */
 public class AuthenticationOutcomeListener<T extends EnumLite, C extends ClientConnection,
     HS extends MessageLite, HR extends MessageLite>
     implements RpcOutcomeListener<SaslMessage> {
   private static final org.slf4j.Logger logger =
       org.slf4j.LoggerFactory.getLogger(AuthenticationOutcomeListener.class);

   private static final ImmutableMap<SaslStatus, SaslChallengeProcessor>
       CHALLENGE_PROCESSORS;
   static {
     final Map<SaslStatus, SaslChallengeProcessor> map = new EnumMap<>(SaslStatus.class);
     map.put(SaslStatus.SASL_IN_PROGRESS, new SaslInProgressProcessor());
     map.put(SaslStatus.SASL_SUCCESS, new SaslSuccessProcessor());
     map.put(SaslStatus.SASL_FAILED, new SaslFailedProcessor());
     CHALLENGE_PROCESSORS = Maps.immutableEnumMap(map);
   }

   private final BasicClient<T, C, HS, HR> client;
   private final C connection;
   private final T saslRpcType;
   private final UserGroupInformation ugi;
   private final RpcOutcomeListener<?> completionListener;

   public AuthenticationOutcomeListener(BasicClient<T, C, HS, HR> client,
                                        C connection, T saslRpcType, UserGroupInformation ugi,
                                        RpcOutcomeListener<?> completionListener) {
     this.client = client;
     this.connection = connection;
     this.saslRpcType = saslRpcType;
     this.ugi = ugi;
     this.completionListener = completionListener;
   }

   public void initiate(final String mechanismName) {
     logger.trace("Initiating SASL exchange.");
     try {
       final ByteString responseData;
       final SaslClient saslClient = connection.getSaslClient();
       if (saslClient.hasInitialResponse()) {
         responseData = ByteString.copyFrom(evaluateChallenge(ugi, saslClient, new byte[0]));
       } else {
         responseData = ByteString.EMPTY;
       }
       client.send(new AuthenticationOutcomeListener<>(client, connection, saslRpcType, ugi, completionListener),
           connection,
           saslRpcType,
           SaslMessage.newBuilder()
               .setMechanism(mechanismName)
               .setStatus(SaslStatus.SASL_START)
               .setData(responseData)
               .build(),
           SaslMessage.class,
           true /* the connection will not be backed up at this point */);
       logger.trace("Initiated SASL exchange.");
     } catch (final Exception e) {
       completionListener.failed(RpcException.mapException(e));
     }
   }

   @Override
   public void failed(RpcException ex) {
     completionListener.failed(RpcException.mapException(ex));
   }

   @Override
   public void success(SaslMessage value, ByteBuf buffer) {
     logger.trace("Server responded with message of type: {}", value.getStatus());
     final SaslChallengeProcessor processor = CHALLENGE_PROCESSORS.get(value.getStatus());
     if (processor == null) {
       completionListener.failed(RpcException.mapException(
           new SaslException("Server sent a corrupt message.")));
     } else {
       // SaslSuccessProcessor.process disposes saslClient so get mechanism here to use later in logging
       final String mechanism = connection.getSaslClient().getMechanismName();
       try {
         final SaslChallengeContext<C> context = new SaslChallengeContext<>(value, ugi, connection);
         final SaslMessage saslResponse = processor.process(context);

         if (saslResponse != null) {
           client.send(new AuthenticationOutcomeListener<>(client, connection, saslRpcType, ugi, completionListener),
               connection, saslRpcType, saslResponse, SaslMessage.class,
               true /* the connection will not be backed up at this point */);
         } else {
           // success
           completionListener.success(null, null);
           if (logger.isTraceEnabled()) {
             logger.trace("Successfully authenticated to server using {} mechanism and encryption context: {}",
                 mechanism, connection.getEncryptionCtxtString());
           }
         }
       } catch (final Exception e) {
         logger.error("Authentication with encryption context: {} using mechanism {} failed with {}",
             connection.getEncryptionCtxtString(), mechanism, e.getMessage());
         completionListener.failed(RpcException.mapException(e));
       }
     }
   }

   @Override
   public void interrupted(InterruptedException e) {
     completionListener.interrupted(e);
   }

   private static class SaslChallengeContext<C extends ClientConnection> {

     final SaslMessage challenge;
     final UserGroupInformation ugi;
     final C connection;

     SaslChallengeContext(SaslMessage challenge, UserGroupInformation ugi, C connection) {
       this.challenge = checkNotNull(challenge);
       this.ugi = checkNotNull(ugi);
       this.connection = checkNotNull(connection);
     }
   }

   private interface SaslChallengeProcessor {

     /**
      * Process challenge from server, and return a response.
      *
      * Returns null iff SASL exchange is complete and successful.
      *
      * @param context challenge context
      * @return response
      * @throws Exception in case of any failure
      */
     <CC extends ClientConnection>
     SaslMessage process(SaslChallengeContext<CC> context) throws Exception;

   }

   private static class SaslInProgressProcessor implements SaslChallengeProcessor {

     @Override
     public <CC extends ClientConnection> SaslMessage process(SaslChallengeContext<CC> context) throws Exception {
       final SaslMessage.Builder response = SaslMessage.newBuilder();
       final SaslClient saslClient = context.connection.getSaslClient();

       final byte[] responseBytes = evaluateChallenge(context.ugi, saslClient,
           context.challenge.getData().toByteArray());

       final boolean isComplete = saslClient.isComplete();
       logger.trace("Evaluated challenge. Completed? {}.", isComplete);
       response.setData(responseBytes != null ? ByteString.copyFrom(responseBytes) : ByteString.EMPTY);
       // if isComplete, the client will get one more response from server
       response.setStatus(isComplete ? SaslStatus.SASL_SUCCESS : SaslStatus.SASL_IN_PROGRESS);
       return response.build();
     }
   }

   private static class SaslSuccessProcessor implements SaslChallengeProcessor {

     @Override
     public <CC extends ClientConnection> SaslMessage process(SaslChallengeContext<CC> context) throws Exception {
       final SaslClient saslClient = context.connection.getSaslClient();

       if (saslClient.isComplete()) {
         handleSuccess(context);
         return null;
       } else {
         // server completed before client; so try once, fail otherwise
         evaluateChallenge(context.ugi, saslClient, context.challenge.getData().toByteArray()); // discard response

         if (saslClient.isComplete()) {
           handleSuccess(context);
           return null;
         } else {
           throw new SaslException("Server allegedly succeeded authentication, but client did not. Suspicious?");
         }
       }
     }
   }

   private static class SaslFailedProcessor implements SaslChallengeProcessor {

     @Override
     public <CC extends ClientConnection> SaslMessage process(SaslChallengeContext<CC> context) throws Exception {
       throw new SaslException(String.format("Authentication failed. Incorrect credentials? [Details: %s]",
           context.connection.getEncryptionCtxtString()));
     }
   }

   private static byte[] evaluateChallenge(final UserGroupInformation ugi, final SaslClient saslClient,
                                           final byte[] challengeBytes) throws SaslException {
     try {
       return ugi.doAs(new PrivilegedExceptionAction<byte[]>() {
         @Override
         public byte[] run() throws Exception {
           return saslClient.evaluateChallenge(challengeBytes);
         }
       });
     } catch (final UndeclaredThrowableException e) {
       throw new SaslException(
           String.format("Unexpected failure (%s)", saslClient.getMechanismName()), e.getCause());
     } catch (final IOException | InterruptedException e) {
       if (e instanceof SaslException) {
         throw (SaslException) e;
       } else {
         throw new SaslException(
             String.format("Unexpected failure (%s)", saslClient.getMechanismName()), e);
       }
     }
   }


   private static <CC extends ClientConnection> void handleSuccess(SaslChallengeContext<CC> context) throws
       SaslException {
     final CC connection = context.connection;
     final SaslClient saslClient = connection.getSaslClient();

     try {
       // Check if connection was marked for being secure then verify for negotiated QOP value for
       // correctness.
       final String negotiatedQOP = saslClient.getNegotiatedProperty(Sasl.QOP).toString();
       final String expectedQOP = connection.isEncryptionEnabled()
           ? SaslProperties.QualityOfProtection.PRIVACY.getSaslQop()
           : SaslProperties.QualityOfProtection.AUTHENTICATION.getSaslQop();

       if (!(negotiatedQOP.equals(expectedQOP))) {
         throw new SaslException(String.format("Mismatch in negotiated QOP value: %s and Expected QOP value: %s",
             negotiatedQOP, expectedQOP));
       }

       // Update the rawWrapChunkSize with the negotiated buffer size since we cannot call encode with more than
       // negotiated size of buffer.
       if (connection.isEncryptionEnabled()) {
         final int negotiatedRawSendSize = Integer.parseInt(
             saslClient.getNegotiatedProperty(Sasl.RAW_SEND_SIZE).toString());
         if (negotiatedRawSendSize <= 0) {
           throw new SaslException(String.format("Negotiated rawSendSize: %d is invalid. Please check the configured " +
               "value of encryption.sasl.max_wrapped_size. It might be configured to a very small value.",
               negotiatedRawSendSize));
         }
         connection.setWrapSizeLimit(negotiatedRawSendSize);
       }
     } catch (Exception e) {
       throw new SaslException(String.format("Unexpected failure while retrieving negotiated property values (%s)",
           e.getMessage()), e);
     }

     if (connection.isEncryptionEnabled()) {
       connection.addSecurityHandlers();
     } else {
       // Encryption is not required hence we don't need to hold on to saslClient object.
       connection.disposeSaslClient();
     }
   }
 }
	/*
	* 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.drill.exec.rpc.security;

	import com.google.common.collect.ImmutableMap;
	import com.google.common.collect.Maps;
	import com.google.protobuf.ByteString;
	import com.google.protobuf.Internal.EnumLite;
	import com.google.protobuf.MessageLite;
	import io.netty.buffer.ByteBuf;
	import org.apache.drill.exec.proto.UserBitShared.SaslMessage;
	import org.apache.drill.exec.proto.UserBitShared.SaslStatus;
	import org.apache.drill.exec.rpc.BasicClient;
	import org.apache.drill.exec.rpc.ClientConnection;
	import org.apache.drill.exec.rpc.RpcException;
	import org.apache.drill.exec.rpc.RpcOutcomeListener;
	import org.apache.hadoop.security.UserGroupInformation;

	import javax.security.sasl.Sasl;
	import javax.security.sasl.SaslClient;
	import javax.security.sasl.SaslException;
	import java.io.IOException;
	import java.lang.reflect.UndeclaredThrowableException;
	import java.security.PrivilegedExceptionAction;
	import java.util.EnumMap;
	import java.util.Map;

	import static com.google.common.base.Preconditions.checkNotNull;

	/**
	* Handles SASL exchange, on the client-side.
	*
	* @param <T> handshake rpc type
	* @param <C> Client connection type
	* @param <HS> Handshake send type
	* @param <HR> Handshake receive type
	*/
	public class AuthenticationOutcomeListener<T extends EnumLite, C extends ClientConnection,
	HS extends MessageLite, HR extends MessageLite>
	implements RpcOutcomeListener<SaslMessage> {
	private static final org.slf4j.Logger logger =
	org.slf4j.LoggerFactory.getLogger(AuthenticationOutcomeListener.class);

	private static final ImmutableMap<SaslStatus, SaslChallengeProcessor>
	CHALLENGE_PROCESSORS;
	static {
	final Map<SaslStatus, SaslChallengeProcessor> map = new EnumMap<>(SaslStatus.class);
	map.put(SaslStatus.SASL_IN_PROGRESS, new SaslInProgressProcessor());
	map.put(SaslStatus.SASL_SUCCESS, new SaslSuccessProcessor());
	map.put(SaslStatus.SASL_FAILED, new SaslFailedProcessor());
	CHALLENGE_PROCESSORS = Maps.immutableEnumMap(map);
	}

	private final BasicClient<T, C, HS, HR> client;
	private final C connection;
	private final T saslRpcType;
	private final UserGroupInformation ugi;
	private final RpcOutcomeListener<?> completionListener;

	public AuthenticationOutcomeListener(BasicClient<T, C, HS, HR> client,
	C connection, T saslRpcType, UserGroupInformation ugi,
	RpcOutcomeListener<?> completionListener) {
	this.client = client;
	this.connection = connection;
	this.saslRpcType = saslRpcType;
	this.ugi = ugi;
	this.completionListener = completionListener;
	}

	public void initiate(final String mechanismName) {
	logger.trace("Initiating SASL exchange.");
	try {
	final ByteString responseData;
	final SaslClient saslClient = connection.getSaslClient();
	if (saslClient.hasInitialResponse()) {
	responseData = ByteString.copyFrom(evaluateChallenge(ugi, saslClient, new byte[0]));
	} else {
	responseData = ByteString.EMPTY;
	}
	client.send(new AuthenticationOutcomeListener<>(client, connection, saslRpcType, ugi, completionListener),
	connection,
	saslRpcType,
	SaslMessage.newBuilder()
	.setMechanism(mechanismName)
	.setStatus(SaslStatus.SASL_START)
	.setData(responseData)
	.build(),
	SaslMessage.class,
	true /* the connection will not be backed up at this point */);
	logger.trace("Initiated SASL exchange.");
	} catch (final Exception e) {
	completionListener.failed(RpcException.mapException(e));
	}
	}

	@Override
	public void failed(RpcException ex) {
	completionListener.failed(RpcException.mapException(ex));
	}

	@Override
	public void success(SaslMessage value, ByteBuf buffer) {
	logger.trace("Server responded with message of type: {}", value.getStatus());
	final SaslChallengeProcessor processor = CHALLENGE_PROCESSORS.get(value.getStatus());
	if (processor == null) {
	completionListener.failed(RpcException.mapException(
	new SaslException("Server sent a corrupt message.")));
	} else {
	// SaslSuccessProcessor.process disposes saslClient so get mechanism here to use later in logging
	final String mechanism = connection.getSaslClient().getMechanismName();
	try {
	final SaslChallengeContext<C> context = new SaslChallengeContext<>(value, ugi, connection);
	final SaslMessage saslResponse = processor.process(context);

	if (saslResponse != null) {
	client.send(new AuthenticationOutcomeListener<>(client, connection, saslRpcType, ugi, completionListener),
	connection, saslRpcType, saslResponse, SaslMessage.class,
	true /* the connection will not be backed up at this point */);
	} else {
	// success
	completionListener.success(null, null);
	if (logger.isTraceEnabled()) {
	logger.trace("Successfully authenticated to server using {} mechanism and encryption context: {}",
	mechanism, connection.getEncryptionCtxtString());
	}
	}
	} catch (final Exception e) {
	logger.error("Authentication with encryption context: {} using mechanism {} failed with {}",
	connection.getEncryptionCtxtString(), mechanism, e.getMessage());
	completionListener.failed(RpcException.mapException(e));
	}
	}
	}

	@Override
	public void interrupted(InterruptedException e) {
	completionListener.interrupted(e);
	}

	private static class SaslChallengeContext<C extends ClientConnection> {

	final SaslMessage challenge;
	final UserGroupInformation ugi;
	final C connection;

	SaslChallengeContext(SaslMessage challenge, UserGroupInformation ugi, C connection) {
	this.challenge = checkNotNull(challenge);
	this.ugi = checkNotNull(ugi);
	this.connection = checkNotNull(connection);
	}
	}

	private interface SaslChallengeProcessor {

	/**
	* Process challenge from server, and return a response.
	*
	* Returns null iff SASL exchange is complete and successful.
	*
	* @param context challenge context
	* @return response
	* @throws Exception in case of any failure
	*/
	<CC extends ClientConnection>
	SaslMessage process(SaslChallengeContext<CC> context) throws Exception;

	}

	private static class SaslInProgressProcessor implements SaslChallengeProcessor {

	@Override
	public <CC extends ClientConnection> SaslMessage process(SaslChallengeContext<CC> context) throws Exception {
	final SaslMessage.Builder response = SaslMessage.newBuilder();
	final SaslClient saslClient = context.connection.getSaslClient();

	final byte[] responseBytes = evaluateChallenge(context.ugi, saslClient,
	context.challenge.getData().toByteArray());

	final boolean isComplete = saslClient.isComplete();
	logger.trace("Evaluated challenge. Completed? {}.", isComplete);
	response.setData(responseBytes != null ? ByteString.copyFrom(responseBytes) : ByteString.EMPTY);
	// if isComplete, the client will get one more response from server
	response.setStatus(isComplete ? SaslStatus.SASL_SUCCESS : SaslStatus.SASL_IN_PROGRESS);
	return response.build();
	}
	}

	private static class SaslSuccessProcessor implements SaslChallengeProcessor {

	@Override
	public <CC extends ClientConnection> SaslMessage process(SaslChallengeContext<CC> context) throws Exception {
	final SaslClient saslClient = context.connection.getSaslClient();

	if (saslClient.isComplete()) {
	handleSuccess(context);
	return null;
	} else {
	// server completed before client; so try once, fail otherwise
	evaluateChallenge(context.ugi, saslClient, context.challenge.getData().toByteArray()); // discard response

	if (saslClient.isComplete()) {
	handleSuccess(context);
	return null;
	} else {
	throw new SaslException("Server allegedly succeeded authentication, but client did not. Suspicious?");
	}
	}
	}
	}

	private static class SaslFailedProcessor implements SaslChallengeProcessor {

	@Override
	public <CC extends ClientConnection> SaslMessage process(SaslChallengeContext<CC> context) throws Exception {
	throw new SaslException(String.format("Authentication failed. Incorrect credentials? [Details: %s]",
	context.connection.getEncryptionCtxtString()));
	}
	}

	private static byte[] evaluateChallenge(final UserGroupInformation ugi, final SaslClient saslClient,
	final byte[] challengeBytes) throws SaslException {
	try {
	return ugi.doAs(new PrivilegedExceptionAction<byte[]>() {
	@Override
	public byte[] run() throws Exception {
	return saslClient.evaluateChallenge(challengeBytes);
	}
	});
	} catch (final UndeclaredThrowableException e) {
	throw new SaslException(
	String.format("Unexpected failure (%s)", saslClient.getMechanismName()), e.getCause());
	} catch (final IOException \| InterruptedException e) {
	if (e instanceof SaslException) {
	throw (SaslException) e;
	} else {
	throw new SaslException(
	String.format("Unexpected failure (%s)", saslClient.getMechanismName()), e);
	}
	}
	}


	private static <CC extends ClientConnection> void handleSuccess(SaslChallengeContext<CC> context) throws
	SaslException {
	final CC connection = context.connection;
	final SaslClient saslClient = connection.getSaslClient();

	try {
	// Check if connection was marked for being secure then verify for negotiated QOP value for
	// correctness.
	final String negotiatedQOP = saslClient.getNegotiatedProperty(Sasl.QOP).toString();
	final String expectedQOP = connection.isEncryptionEnabled()
	? SaslProperties.QualityOfProtection.PRIVACY.getSaslQop()
	: SaslProperties.QualityOfProtection.AUTHENTICATION.getSaslQop();

	if (!(negotiatedQOP.equals(expectedQOP))) {
	throw new SaslException(String.format("Mismatch in negotiated QOP value: %s and Expected QOP value: %s",
	negotiatedQOP, expectedQOP));
	}

	// Update the rawWrapChunkSize with the negotiated buffer size since we cannot call encode with more than
	// negotiated size of buffer.
	if (connection.isEncryptionEnabled()) {
	final int negotiatedRawSendSize = Integer.parseInt(
	saslClient.getNegotiatedProperty(Sasl.RAW_SEND_SIZE).toString());
	if (negotiatedRawSendSize <= 0) {
	throw new SaslException(String.format("Negotiated rawSendSize: %d is invalid. Please check the configured " +
	"value of encryption.sasl.max_wrapped_size. It might be configured to a very small value.",
	negotiatedRawSendSize));
	}
	connection.setWrapSizeLimit(negotiatedRawSendSize);
	}
	} catch (Exception e) {
	throw new SaslException(String.format("Unexpected failure while retrieving negotiated property values (%s)",
	e.getMessage()), e);
	}

	if (connection.isEncryptionEnabled()) {
	connection.addSecurityHandlers();
	} else {
	// Encryption is not required hence we don't need to hold on to saslClient object.
	connection.disposeSaslClient();
	}
	}
	}