transport/sasl/NonblockingSaslHandler.java

/*
 * 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.thrift.transport.sasl;

import java.nio.channels.SelectionKey;
import java.nio.charset.StandardCharsets;

import javax.security.sasl.SaslServer;

import org.apache.thrift.TByteArrayOutputStream;
import org.apache.thrift.TProcessor;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.protocol.TProtocolFactory;
import org.apache.thrift.server.ServerContext;
import org.apache.thrift.server.TServerEventHandler;
import org.apache.thrift.transport.TMemoryTransport;
import org.apache.thrift.transport.TNonblockingTransport;
import org.apache.thrift.transport.TTransportException;
import org.apache.thrift.transport.sasl.TSaslNegotiationException.ErrorType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import static org.apache.thrift.transport.sasl.NegotiationStatus.COMPLETE;
import static org.apache.thrift.transport.sasl.NegotiationStatus.OK;

/**
 * State machine managing one sasl connection in a nonblocking way.
 */
public class NonblockingSaslHandler {
  private static final Logger LOGGER = LoggerFactory.getLogger(NonblockingSaslHandler.class);

  private static final int INTEREST_NONE = 0;
  private static final int INTEREST_READ = SelectionKey.OP_READ;
  private static final int INTEREST_WRITE = SelectionKey.OP_WRITE;

  // Tracking the current running phase
  private Phase currentPhase = Phase.INITIIALIIZING;
  // Tracking the next phase on the next invocation of the state machine.
  // It should be the same as current phase if current phase is not yet finished.
  // Otherwise, if it is different from current phase, the statemachine is in a transition state:
  // current phase is done, and next phase is not yet started.
  private Phase nextPhase = currentPhase;

  // Underlying nonblocking transport
  private SelectionKey selectionKey;
  private TNonblockingTransport underlyingTransport;

  // APIs for intercepting event / customizing behaviors:
  // Factories (decorating the base implementations) & EventHandler (intercepting)
  private TSaslServerFactory saslServerFactory;
  private TSaslProcessorFactory processorFactory;
  private TProtocolFactory inputProtocolFactory;
  private TProtocolFactory outputProtocolFactory;
  private TServerEventHandler eventHandler;
  private ServerContext serverContext;
  // It turns out the event handler implementation in hive sometimes creates a null ServerContext.
  // In order to know whether TServerEventHandler#createContext is called we use such a flag.
  private boolean serverContextCreated = false;

  // Wrapper around sasl server
  private ServerSaslPeer saslPeer;

  // Sasl negotiation io
  private SaslNegotiationFrameReader saslResponse;
  private SaslNegotiationFrameWriter saslChallenge;
  // IO for request from and response to the socket
  private DataFrameReader requestReader;
  private DataFrameWriter responseWriter;
  // If sasl is negotiated for integrity/confidentiality protection
  private boolean dataProtected;

  public NonblockingSaslHandler(SelectionKey selectionKey, TNonblockingTransport underlyingTransport,
                                TSaslServerFactory saslServerFactory, TSaslProcessorFactory processorFactory,
                                TProtocolFactory inputProtocolFactory, TProtocolFactory outputProtocolFactory,
                                TServerEventHandler eventHandler) {
    this.selectionKey = selectionKey;
    this.underlyingTransport = underlyingTransport;
    this.saslServerFactory = saslServerFactory;
    this.processorFactory = processorFactory;
    this.inputProtocolFactory = inputProtocolFactory;
    this.outputProtocolFactory = outputProtocolFactory;
    this.eventHandler = eventHandler;

    saslResponse = new SaslNegotiationFrameReader();
    saslChallenge = new SaslNegotiationFrameWriter();
    requestReader = new DataFrameReader();
    responseWriter = new DataFrameWriter();
  }

  /**
   * Get current phase of the state machine.
   *
   * @return current phase.
   */
  public Phase getCurrentPhase() {
    return currentPhase;
  }

  /**
   * Get next phase of the state machine.
   * It is different from current phase iff current phase is done (and next phase not yet started).
   *
   * @return next phase.
   */
  public Phase getNextPhase() {
    return nextPhase;
  }

  /**
   *
   * @return underlying nonblocking socket
   */
  public TNonblockingTransport getUnderlyingTransport() {
    return underlyingTransport;
  }

  /**
   *
   * @return SaslServer instance
   */
  public SaslServer getSaslServer() {
    return saslPeer.getSaslServer();
  }

  /**
   *
   * @return true if current phase is done.
   */
  public boolean isCurrentPhaseDone() {
    return currentPhase != nextPhase;
  }

  /**
   * Run state machine.
   *
   * @throws IllegalStateException if current state is already done.
   */
  public void runCurrentPhase() {
    currentPhase.runStateMachine(this);
  }

  /**
   * When current phase is intrested in read selection, calling this will run the current phase and
   * its following phases if the following ones are interested to read, until there is nothing
   * available in the underlying transport.
   *
   * @throws IllegalStateException if is called in an irrelevant phase.
   */
  public void handleRead() {
    handleOps(INTEREST_READ);
  }

  /**
   * Similiar to handleRead. But it is for write ops.
   *
   * @throws IllegalStateException if it is called in an irrelevant phase.
   */
  public void handleWrite() {
    handleOps(INTEREST_WRITE);
  }

  private void handleOps(int interestOps) {
    if (currentPhase.selectionInterest != interestOps) {
      throw new IllegalStateException("Current phase " + currentPhase + " but got interest " +
          interestOps);
    }
    runCurrentPhase();
    if (isCurrentPhaseDone() && nextPhase.selectionInterest == interestOps) {
      stepToNextPhase();
      handleOps(interestOps);
    }
  }

  /**
   * When current phase is finished, it's expected to call this method first before running the
   * state machine again.
   * By calling this, "next phase" is marked as started (and not done), thus is ready to run.
   *
   * @throws IllegalArgumentException if current phase is not yet done.
   */
  public void stepToNextPhase() {
    if (!isCurrentPhaseDone()) {
      throw new IllegalArgumentException("Not yet done with current phase: " + currentPhase);
    }
    LOGGER.debug("Switch phase {} to {}", currentPhase, nextPhase);
    switch (nextPhase) {
      case INITIIALIIZING:
        throw new IllegalStateException("INITIALIZING cannot be the next phase of " + currentPhase);
      default:
    }
    // If next phase's interest is not the same as current,  nor the same as the selection key,
    // we need to change interest on the selector.
    if (!(nextPhase.selectionInterest == currentPhase.selectionInterest ||
        nextPhase.selectionInterest == selectionKey.interestOps())) {
      changeSelectionInterest(nextPhase.selectionInterest);
    }
    currentPhase = nextPhase;
  }

  private void changeSelectionInterest(int selectionInterest) {
    selectionKey.interestOps(selectionInterest);
  }

  // sasl negotiaion failure handling
  private void failSaslNegotiation(TSaslNegotiationException e) {
    LOGGER.error("Sasl negotiation failed", e);
    String errorMsg = e.getDetails();
    saslChallenge.withHeaderAndPayload(new byte[]{e.getErrorType().code.getValue()},
        errorMsg.getBytes(StandardCharsets.UTF_8));
    nextPhase = Phase.WRITING_FAILURE_MESSAGE;
  }

  private void fail(Exception e) {
    LOGGER.error("Failed io in " + currentPhase, e);
    nextPhase = Phase.CLOSING;
  }

  private void failIO(TTransportException e) {
    StringBuilder errorMsg = new StringBuilder("IO failure ")
        .append(e.getType())
        .append(" in ")
        .append(currentPhase);
    if (e.getMessage() != null) {
      errorMsg.append(": ").append(e.getMessage());
    }
    LOGGER.error(errorMsg.toString(), e);
    nextPhase = Phase.CLOSING;
  }

  // Read handlings

  private void handleInitializing() {
    try {
      saslResponse.read(underlyingTransport);
      if (saslResponse.isComplete()) {
        SaslNegotiationHeaderReader startHeader = saslResponse.getHeader();
        if (startHeader.getStatus() != NegotiationStatus.START) {
          throw new TInvalidSaslFrameException("Expecting START status but got " + startHeader.getStatus());
        }
        String mechanism = new String(saslResponse.getPayload(), StandardCharsets.UTF_8);
        saslPeer = saslServerFactory.getSaslPeer(mechanism);
        saslResponse.clear();
        nextPhase = Phase.READING_SASL_RESPONSE;
      }
    } catch (TSaslNegotiationException e) {
      failSaslNegotiation(e);
    } catch (TTransportException e) {
      failIO(e);
    }
  }

  private void handleReadingSaslResponse() {
    try {
      saslResponse.read(underlyingTransport);
      if (saslResponse.isComplete()) {
        nextPhase = Phase.EVALUATING_SASL_RESPONSE;
      }
    } catch (TSaslNegotiationException e) {
      failSaslNegotiation(e);
    } catch (TTransportException e) {
      failIO(e);
    }
  }

  private void handleReadingRequest() {
    try {
      requestReader.read(underlyingTransport);
      if (requestReader.isComplete()) {
        nextPhase = Phase.PROCESSING;
      }
    } catch (TTransportException e) {
      failIO(e);
    }
  }

  // Computation executions

  private void executeEvaluatingSaslResponse() {
    if (!(saslResponse.getHeader().getStatus() == OK || saslResponse.getHeader().getStatus() == COMPLETE)) {
      String error = "Expect status OK or COMPLETE, but got " + saslResponse.getHeader().getStatus();
      failSaslNegotiation(new TSaslNegotiationException(ErrorType.PROTOCOL_ERROR, error));
      return;
    }
    try {
      byte[] response = saslResponse.getPayload();
      saslResponse.clear();
      byte[] newChallenge = saslPeer.evaluate(response);
      if (saslPeer.isAuthenticated()) {
        dataProtected = saslPeer.isDataProtected();
        saslChallenge.withHeaderAndPayload(new byte[]{COMPLETE.getValue()}, newChallenge);
        nextPhase = Phase.WRITING_SUCCESS_MESSAGE;
      } else {
        saslChallenge.withHeaderAndPayload(new byte[]{OK.getValue()}, newChallenge);
        nextPhase = Phase.WRITING_SASL_CHALLENGE;
      }
    } catch (TSaslNegotiationException e) {
      failSaslNegotiation(e);
    }
  }

  private void executeProcessing() {
    try {
      byte[] inputPayload = requestReader.getPayload();
      requestReader.clear();
      byte[] rawInput = dataProtected ? saslPeer.unwrap(inputPayload) : inputPayload;
      TMemoryTransport memoryTransport = new TMemoryTransport(rawInput);
      TProtocol requestProtocol = inputProtocolFactory.getProtocol(memoryTransport);
      TProtocol responseProtocol = outputProtocolFactory.getProtocol(memoryTransport);

      if (eventHandler != null) {
        if (!serverContextCreated) {
          serverContext = eventHandler.createContext(requestProtocol, responseProtocol);
          serverContextCreated = true;
        }
        eventHandler.processContext(serverContext, memoryTransport, memoryTransport);
      }

      TProcessor processor = processorFactory.getProcessor(this);
      processor.process(requestProtocol, responseProtocol);
      TByteArrayOutputStream rawOutput = memoryTransport.getOutput();
      if (rawOutput.len() == 0) {
        // This is a oneway request, no response to send back. Waiting for next incoming request.
        nextPhase = Phase.READING_REQUEST;
        return;
      }
      if (dataProtected) {
        byte[] outputPayload = saslPeer.wrap(rawOutput.get(), 0, rawOutput.len());
        responseWriter.withOnlyPayload(outputPayload);
      } else {
        responseWriter.withOnlyPayload(rawOutput.get(), 0 ,rawOutput.len());
      }
      nextPhase = Phase.WRITING_RESPONSE;
    } catch (TTransportException e) {
      failIO(e);
    } catch (Exception e) {
      fail(e);
    }
  }

  // Write handlings

  private void handleWritingSaslChallenge() {
    try {
      saslChallenge.write(underlyingTransport);
      if (saslChallenge.isComplete()) {
        saslChallenge.clear();
        nextPhase = Phase.READING_SASL_RESPONSE;
      }
    } catch (TTransportException e) {
      fail(e);
    }
  }

  private void handleWritingSuccessMessage() {
    try {
      saslChallenge.write(underlyingTransport);
      if (saslChallenge.isComplete()) {
        LOGGER.debug("Authentication is done.");
        saslChallenge = null;
        saslResponse = null;
        nextPhase = Phase.READING_REQUEST;
      }
    } catch (TTransportException e) {
      fail(e);
    }
  }

  private void handleWritingFailureMessage() {
    try {
      saslChallenge.write(underlyingTransport);
      if (saslChallenge.isComplete()) {
        nextPhase = Phase.CLOSING;
      }
    } catch (TTransportException e) {
      fail(e);
    }
  }

  private void handleWritingResponse() {
    try {
      responseWriter.write(underlyingTransport);
      if (responseWriter.isComplete()) {
        responseWriter.clear();
        nextPhase = Phase.READING_REQUEST;
      }
    } catch (TTransportException e) {
      fail(e);
    }
  }

  /**
   * Release all the resources managed by this state machine (connection, selection and sasl server).
   * To avoid being blocked, this should be invoked in the network thread that manages the selector.
   */
  public void close() {
    underlyingTransport.close();
    selectionKey.cancel();
    if (saslPeer != null) {
      saslPeer.dispose();
    }
    if (serverContextCreated) {
      eventHandler.deleteContext(serverContext,
          inputProtocolFactory.getProtocol(underlyingTransport),
          outputProtocolFactory.getProtocol(underlyingTransport));
    }
    nextPhase = Phase.CLOSED;
    currentPhase = Phase.CLOSED;
    LOGGER.trace("Connection closed: {}", underlyingTransport);
  }

  public enum Phase {
    INITIIALIIZING(INTEREST_READ) {
      @Override
      void unsafeRun(NonblockingSaslHandler statemachine) {
        statemachine.handleInitializing();
      }
    },
    READING_SASL_RESPONSE(INTEREST_READ) {
      @Override
      void unsafeRun(NonblockingSaslHandler statemachine) {
        statemachine.handleReadingSaslResponse();
      }
    },
    EVALUATING_SASL_RESPONSE(INTEREST_NONE) {
      @Override
      void unsafeRun(NonblockingSaslHandler statemachine) {
        statemachine.executeEvaluatingSaslResponse();
      }
    },
    WRITING_SASL_CHALLENGE(INTEREST_WRITE) {
      @Override
      void unsafeRun(NonblockingSaslHandler statemachine) {
        statemachine.handleWritingSaslChallenge();
      }
    },
    WRITING_SUCCESS_MESSAGE(INTEREST_WRITE) {
      @Override
      void unsafeRun(NonblockingSaslHandler statemachine) {
        statemachine.handleWritingSuccessMessage();
      }
    },
    WRITING_FAILURE_MESSAGE(INTEREST_WRITE) {
      @Override
      void unsafeRun(NonblockingSaslHandler statemachine) {
        statemachine.handleWritingFailureMessage();
      }
    },
    READING_REQUEST(INTEREST_READ) {
      @Override
      void unsafeRun(NonblockingSaslHandler statemachine) {
        statemachine.handleReadingRequest();
      }
    },
    PROCESSING(INTEREST_NONE) {
      @Override
      void unsafeRun(NonblockingSaslHandler statemachine) {
        statemachine.executeProcessing();
      }
    },
    WRITING_RESPONSE(INTEREST_WRITE) {
      @Override
      void unsafeRun(NonblockingSaslHandler statemachine) {
        statemachine.handleWritingResponse();
      }
    },
    CLOSING(INTEREST_NONE) {
      @Override
      void unsafeRun(NonblockingSaslHandler statemachine) {
        statemachine.close();
      }
    },
    CLOSED(INTEREST_NONE) {
      @Override
      void unsafeRun(NonblockingSaslHandler statemachine) {
        // Do nothing.
      }
    }
    ;

    // The interest on the selection key during the phase
    private int selectionInterest;

    Phase(int selectionInterest) {
      this.selectionInterest = selectionInterest;
    }

    /**
     * Provide the execution to run for the state machine in current phase. The execution should
     * return the next phase after running on the state machine.
     *
     * @param statemachine The state machine to run.
     * @throws IllegalArgumentException if the state machine's current phase is different.
     * @throws IllegalStateException if the state machine' current phase is already done.
     */
    void runStateMachine(NonblockingSaslHandler statemachine) {
      if (statemachine.currentPhase != this) {
        throw new IllegalArgumentException("State machine is " + statemachine.currentPhase +
            " but is expected to be " + this);
      }
      if (statemachine.isCurrentPhaseDone()) {
        throw new IllegalStateException("State machine should step into " + statemachine.nextPhase);
      }
      unsafeRun(statemachine);
    }

    // Run the state machine without checkiing its own phase
    // It should not be called direcly by users.
    abstract void unsafeRun(NonblockingSaslHandler statemachine);
  }
}