be/src/transport/TSaslTransport.cpp - impala - Git at Google

 // This file will be removed when the code is accepted into the Thrift library.
 /*
  * 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.
  */

 #include "config.h"

 #ifdef HAVE_SASL_SASL_H
 #include <stdint.h>
 #include <sstream>
 #include <boost/shared_ptr.hpp>
 #include <boost/scoped_ptr.hpp>

 #include <thrift/transport/TBufferTransports.h>
 #include "transport/TSaslTransport.h"

 #include "common/names.h"

 // Default size, in bytes, for the memory buffer used to stage reads.
 const int32_t DEFAULT_MEM_BUF_SIZE = 32 * 1024;

 namespace apache { namespace thrift { namespace transport {

   TSaslTransport::TSaslTransport(boost::shared_ptr<TTransport> transport)
       : transport_(transport),
         memBuf_(new TMemoryBuffer(DEFAULT_MEM_BUF_SIZE)),
         sasl_(NULL),
         shouldWrap_(false),
         isClient_(false) {
   }

   TSaslTransport::TSaslTransport(boost::shared_ptr<sasl::TSasl> saslClient,
                                  boost::shared_ptr<TTransport> transport)
       : transport_(transport),
         memBuf_(new TMemoryBuffer()),
         sasl_(saslClient),
         shouldWrap_(false),
         isClient_(true) {
   }

   TSaslTransport::~TSaslTransport() {
     delete memBuf_;
   }

   bool TSaslTransport::isOpen() {
     return transport_->isOpen();
   }

   bool TSaslTransport::peek(){
     return (transport_->peek());
   }

   string TSaslTransport::getUsername() {
     return sasl_->getUsername();
   }

   void TSaslTransport::doSaslNegotiation() {
     NegotiationStatus status = TSASL_INVALID;
     uint32_t resLength;

     try {
       // Setup Sasl context.
       setupSaslNegotiationState();

       // Initiate SASL message.
       handleSaslStartMessage();

       // SASL connection handshake
       while (!sasl_->isComplete()) {
         uint8_t* message = receiveSaslMessage(&status, &resLength);
         if (status == TSASL_COMPLETE) {
           if (isClient_) {
             if (!sasl_->isComplete()) {
               // Server sent COMPLETE out of order.
               throw TTransportException("Received COMPLETE but no handshake occurred");
             }
             break; // handshake complete
           }
         } else if (status != TSASL_OK) {
           stringstream ss;
           ss << "Expected COMPLETE or OK, got " << status;
           throw TTransportException(ss.str());
         }
         uint32_t challengeLength;
         uint8_t* challenge = sasl_->evaluateChallengeOrResponse(
             message, resLength, &challengeLength);
         sendSaslMessage(sasl_->isComplete() ? TSASL_COMPLETE : TSASL_OK,
                         challenge, challengeLength);
       }

       // If the server isn't complete yet, we need to wait for its response.
       // This will occur with ANONYMOUS auth, for example, where we send an
       // initial response and are immediately complete.
       if (isClient_ && (status == TSASL_INVALID || status == TSASL_OK)) {
         receiveSaslMessage(&status, &resLength);
         if (status != TSASL_COMPLETE) {
           stringstream ss;
           ss << "Expected COMPLETE or OK, got " << status;
           throw TTransportException(ss.str());
         }
       }
       // TODO : need to set the shouldWrap_ based on QOP
       /*
       String qop = (String) sasl.getNegotiatedProperty(Sasl.QOP);
       if (qop != null && !qop.equalsIgnoreCase("auth"))
         shouldWrap_ = true;
       */
     } catch (const TException& e) {
       // If we hit an exception, that means the Sasl negotiation failed. We explicitly
       // reset the negotiation state here since the caller may retry an open() which would
       // start a new connection negotiation.
       resetSaslNegotiationState();
       throw e;
     }
   }

   void TSaslTransport::open() {
     // Only client should open the underlying transport.
     if (isClient_ && !transport_->isOpen()) {
       transport_->open();
     }

     // Start the SASL negotiation protocol.
     doSaslNegotiation();
   }

   void TSaslTransport::close() {
     transport_->close();
   }

   uint32_t TSaslTransport::readLength() {
     uint8_t lenBuf[PAYLOAD_LENGTH_BYTES];

     transport_->readAll(lenBuf, PAYLOAD_LENGTH_BYTES);
     int32_t len = decodeInt(lenBuf, 0);
     if (len < 0) {
       throw TTransportException("Frame size has negative value");
     }
     return static_cast<uint32_t>(len);
   }

   void TSaslTransport::shrinkBuffer() {
     // readEnd() returns the number of bytes already read, i.e. the number of 'junk' bytes
     // taking up space at the front of the memory buffer.
     uint32_t read_end = memBuf_->readEnd();

     // If the size of the junk space at the beginning of the buffer is too large, and
     // there's no data left in the buffer to read (number of bytes read == number of bytes
     // written), then shrink the buffer back to the default. We don't want to do this on
     // every read that exhausts the buffer, since the layer above often reads in small
     // chunks, which is why we only resize if there's too much junk. The write and read
     // pointers will eventually catch up after every RPC, so we will always take this path
     // eventually once the buffer becomes sufficiently full.
     //
     // readEnd() may reset the write / read pointers (but only once if there's no
     // intervening read or write between calls), so needs to be called a second time to
     // get their current position.
     if (read_end > DEFAULT_MEM_BUF_SIZE && memBuf_->writeEnd() == memBuf_->readEnd()) {
       memBuf_->resetBuffer(DEFAULT_MEM_BUF_SIZE);
     }
   }

   uint32_t TSaslTransport::read(uint8_t* buf, uint32_t len) {
     uint32_t read_bytes = memBuf_->read(buf, len);

     if (read_bytes > 0) {
       shrinkBuffer();
       return read_bytes;
     }

     // if there's not enough data in cache, read from underlying transport
     uint32_t dataLength = readLength();

     // Fast path
     if (len == dataLength && !shouldWrap_) {
       transport_->readAll(buf, len);
       return len;
     }

     uint8_t* tmpBuf = new uint8_t[dataLength];
     transport_->readAll(tmpBuf, dataLength);
     if (shouldWrap_) {
       tmpBuf = sasl_->unwrap(tmpBuf, 0, dataLength, &dataLength);
     }

     // We will consume all the data, no need to put it in the memory buffer.
     if (len == dataLength) {
       memcpy(buf, tmpBuf, len);
       delete[] tmpBuf;
       return len;
     }

     memBuf_->write(tmpBuf, dataLength);
     memBuf_->flush();
     delete[] tmpBuf;

     uint32_t ret = memBuf_->read(buf, len);
     shrinkBuffer();
     return ret;
   }

   void TSaslTransport::writeLength(uint32_t length) {
     uint8_t lenBuf[PAYLOAD_LENGTH_BYTES];

     encodeInt(length, lenBuf, 0);
     transport_->write(lenBuf, PAYLOAD_LENGTH_BYTES);
   }

   void TSaslTransport::write(const uint8_t* buf, uint32_t len) {
     const uint8_t* newBuf;

     if (shouldWrap_) {
       newBuf = sasl_->wrap(const_cast<uint8_t*>(buf), 0, len, &len);
     } else {
       newBuf = buf;
     }
     writeLength(len);
     transport_->write(newBuf, len);
   }

   void TSaslTransport::flush() {
     transport_->flush();
   }

   void TSaslTransport::sendSaslMessage(const NegotiationStatus status,
       const uint8_t* payload, const uint32_t length, bool flush) {
     uint8_t messageHeader[STATUS_BYTES + PAYLOAD_LENGTH_BYTES];
     uint8_t dummy = 0;
     if (payload == NULL) {
       payload = &dummy;
     }
     messageHeader[0] = (uint8_t)status;
     encodeInt(length, messageHeader, STATUS_BYTES);
     transport_->write(messageHeader, HEADER_LENGTH);
     transport_->write(payload, length);
     if (flush) transport_->flush();
   }

   uint8_t* TSaslTransport::receiveSaslMessage(NegotiationStatus* status,
                                               uint32_t* length) {
     uint8_t messageHeader[HEADER_LENGTH];

     // read header
     transport_->readAll(messageHeader, HEADER_LENGTH);

     // get payload status
     *status = (NegotiationStatus)messageHeader[0];
     if ((*status < TSASL_START) || (*status > TSASL_COMPLETE)) {
       throw TTransportException("invalid sasl status");
     } else if (*status == TSASL_BAD || *status == TSASL_ERROR) {
         throw TTransportException("sasl Peer indicated failure: ");
     }

     // get the length
     *length = decodeInt(messageHeader, STATUS_BYTES);

     // get payload
     protoBuf_.reset(new uint8_t[*length]);
     transport_->readAll(protoBuf_.get(), *length);

     return protoBuf_.get();
   }
 }
 }
 }

 #endif
	// This file will be removed when the code is accepted into the Thrift library.
	/*
	* 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.
	*/

	#include "config.h"

	#ifdef HAVE_SASL_SASL_H
	#include <stdint.h>
	#include <sstream>
	#include <boost/shared_ptr.hpp>
	#include <boost/scoped_ptr.hpp>

	#include <thrift/transport/TBufferTransports.h>
	#include "transport/TSaslTransport.h"

	#include "common/names.h"

	// Default size, in bytes, for the memory buffer used to stage reads.
	const int32_t DEFAULT_MEM_BUF_SIZE = 32 * 1024;

	namespace apache { namespace thrift { namespace transport {

	TSaslTransport::TSaslTransport(boost::shared_ptr<TTransport> transport)
	: transport_(transport),
	memBuf_(new TMemoryBuffer(DEFAULT_MEM_BUF_SIZE)),
	sasl_(NULL),
	shouldWrap_(false),
	isClient_(false) {
	}

	TSaslTransport::TSaslTransport(boost::shared_ptr<sasl::TSasl> saslClient,
	boost::shared_ptr<TTransport> transport)
	: transport_(transport),
	memBuf_(new TMemoryBuffer()),
	sasl_(saslClient),
	shouldWrap_(false),
	isClient_(true) {
	}

	TSaslTransport::~TSaslTransport() {
	delete memBuf_;
	}

	bool TSaslTransport::isOpen() {
	return transport_->isOpen();
	}

	bool TSaslTransport::peek(){
	return (transport_->peek());
	}

	string TSaslTransport::getUsername() {
	return sasl_->getUsername();
	}

	void TSaslTransport::doSaslNegotiation() {
	NegotiationStatus status = TSASL_INVALID;
	uint32_t resLength;

	try {
	// Setup Sasl context.
	setupSaslNegotiationState();

	// Initiate SASL message.
	handleSaslStartMessage();

	// SASL connection handshake
	while (!sasl_->isComplete()) {
	uint8_t* message = receiveSaslMessage(&status, &resLength);
	if (status == TSASL_COMPLETE) {
	if (isClient_) {
	if (!sasl_->isComplete()) {
	// Server sent COMPLETE out of order.
	throw TTransportException("Received COMPLETE but no handshake occurred");
	}
	break; // handshake complete
	}
	} else if (status != TSASL_OK) {
	stringstream ss;
	ss << "Expected COMPLETE or OK, got " << status;
	throw TTransportException(ss.str());
	}
	uint32_t challengeLength;
	uint8_t* challenge = sasl_->evaluateChallengeOrResponse(
	message, resLength, &challengeLength);
	sendSaslMessage(sasl_->isComplete() ? TSASL_COMPLETE : TSASL_OK,
	challenge, challengeLength);
	}

	// If the server isn't complete yet, we need to wait for its response.
	// This will occur with ANONYMOUS auth, for example, where we send an
	// initial response and are immediately complete.
	if (isClient_ && (status == TSASL_INVALID \|\| status == TSASL_OK)) {
	receiveSaslMessage(&status, &resLength);
	if (status != TSASL_COMPLETE) {
	stringstream ss;
	ss << "Expected COMPLETE or OK, got " << status;
	throw TTransportException(ss.str());
	}
	}
	// TODO : need to set the shouldWrap_ based on QOP
	/*
	String qop = (String) sasl.getNegotiatedProperty(Sasl.QOP);
	if (qop != null && !qop.equalsIgnoreCase("auth"))
	shouldWrap_ = true;
	*/
	} catch (const TException& e) {
	// If we hit an exception, that means the Sasl negotiation failed. We explicitly
	// reset the negotiation state here since the caller may retry an open() which would
	// start a new connection negotiation.
	resetSaslNegotiationState();
	throw e;
	}
	}

	void TSaslTransport::open() {
	// Only client should open the underlying transport.
	if (isClient_ && !transport_->isOpen()) {
	transport_->open();
	}

	// Start the SASL negotiation protocol.
	doSaslNegotiation();
	}

	void TSaslTransport::close() {
	transport_->close();
	}

	uint32_t TSaslTransport::readLength() {
	uint8_t lenBuf[PAYLOAD_LENGTH_BYTES];

	transport_->readAll(lenBuf, PAYLOAD_LENGTH_BYTES);
	int32_t len = decodeInt(lenBuf, 0);
	if (len < 0) {
	throw TTransportException("Frame size has negative value");
	}
	return static_cast<uint32_t>(len);
	}

	void TSaslTransport::shrinkBuffer() {
	// readEnd() returns the number of bytes already read, i.e. the number of 'junk' bytes
	// taking up space at the front of the memory buffer.
	uint32_t read_end = memBuf_->readEnd();

	// If the size of the junk space at the beginning of the buffer is too large, and
	// there's no data left in the buffer to read (number of bytes read == number of bytes
	// written), then shrink the buffer back to the default. We don't want to do this on
	// every read that exhausts the buffer, since the layer above often reads in small
	// chunks, which is why we only resize if there's too much junk. The write and read
	// pointers will eventually catch up after every RPC, so we will always take this path
	// eventually once the buffer becomes sufficiently full.
	//
	// readEnd() may reset the write / read pointers (but only once if there's no
	// intervening read or write between calls), so needs to be called a second time to
	// get their current position.
	if (read_end > DEFAULT_MEM_BUF_SIZE && memBuf_->writeEnd() == memBuf_->readEnd()) {
	memBuf_->resetBuffer(DEFAULT_MEM_BUF_SIZE);
	}
	}

	uint32_t TSaslTransport::read(uint8_t* buf, uint32_t len) {
	uint32_t read_bytes = memBuf_->read(buf, len);

	if (read_bytes > 0) {
	shrinkBuffer();
	return read_bytes;
	}

	// if there's not enough data in cache, read from underlying transport
	uint32_t dataLength = readLength();

	// Fast path
	if (len == dataLength && !shouldWrap_) {
	transport_->readAll(buf, len);
	return len;
	}

	uint8_t* tmpBuf = new uint8_t[dataLength];
	transport_->readAll(tmpBuf, dataLength);
	if (shouldWrap_) {
	tmpBuf = sasl_->unwrap(tmpBuf, 0, dataLength, &dataLength);
	}

	// We will consume all the data, no need to put it in the memory buffer.
	if (len == dataLength) {
	memcpy(buf, tmpBuf, len);
	delete[] tmpBuf;
	return len;
	}

	memBuf_->write(tmpBuf, dataLength);
	memBuf_->flush();
	delete[] tmpBuf;

	uint32_t ret = memBuf_->read(buf, len);
	shrinkBuffer();
	return ret;
	}

	void TSaslTransport::writeLength(uint32_t length) {
	uint8_t lenBuf[PAYLOAD_LENGTH_BYTES];

	encodeInt(length, lenBuf, 0);
	transport_->write(lenBuf, PAYLOAD_LENGTH_BYTES);
	}

	void TSaslTransport::write(const uint8_t* buf, uint32_t len) {
	const uint8_t* newBuf;

	if (shouldWrap_) {
	newBuf = sasl_->wrap(const_cast<uint8_t*>(buf), 0, len, &len);
	} else {
	newBuf = buf;
	}
	writeLength(len);
	transport_->write(newBuf, len);
	}

	void TSaslTransport::flush() {
	transport_->flush();
	}

	void TSaslTransport::sendSaslMessage(const NegotiationStatus status,
	const uint8_t* payload, const uint32_t length, bool flush) {
	uint8_t messageHeader[STATUS_BYTES + PAYLOAD_LENGTH_BYTES];
	uint8_t dummy = 0;
	if (payload == NULL) {
	payload = &dummy;
	}
	messageHeader[0] = (uint8_t)status;
	encodeInt(length, messageHeader, STATUS_BYTES);
	transport_->write(messageHeader, HEADER_LENGTH);
	transport_->write(payload, length);
	if (flush) transport_->flush();
	}

	uint8_t* TSaslTransport::receiveSaslMessage(NegotiationStatus* status,
	uint32_t* length) {
	uint8_t messageHeader[HEADER_LENGTH];

	// read header
	transport_->readAll(messageHeader, HEADER_LENGTH);

	// get payload status
	*status = (NegotiationStatus)messageHeader[0];
	if ((status < TSASL_START) \|\| (status > TSASL_COMPLETE)) {
	throw TTransportException("invalid sasl status");
	} else if (status == TSASL_BAD \|\| status == TSASL_ERROR) {
	throw TTransportException("sasl Peer indicated failure: ");
	}

	// get the length
	*length = decodeInt(messageHeader, STATUS_BYTES);

	// get payload
	protoBuf_.reset(new uint8_t[*length]);
	transport_->readAll(protoBuf_.get(), *length);

	return protoBuf_.get();
	}
	}
	}
	}

	#endif