lib/d/src/thrift/transport/framed.d - thrift - 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.
  */

 module thrift.transport.framed;

 import core.bitop : bswap;
 import std.algorithm : min;
 import std.array : empty;
 import std.exception : enforce;
 import thrift.transport.base;

 /**
  * Framed transport.
  *
  * All writes go into an in-memory buffer until flush is called, at which point
  * the transport writes the length of the entire binary chunk followed by the
  * data payload. The receiver on the other end then performs a single
  * »fixed-length« read to get the whole message off the wire.
  */
 final class TFramedTransport : TBaseTransport {
   /**
    * Constructs a new framed transport.
    *
    * Params:
    *   transport = The underlying transport to wrap.
    */
   this(TTransport transport) {
     transport_ = transport;
   }

   /**
    * Returns the wrapped transport.
    */
   TTransport underlyingTransport() @property {
     return transport_;
   }

   override bool isOpen() @property {
     return transport_.isOpen;
   }

   override bool peek() {
     return rBuf_.length > 0 || transport_.peek();
   }

   override void open() {
     transport_.open();
   }

   override void close() {
     flush();
     transport_.close();
   }

   /**
    * Attempts to read data into the given buffer, stopping when the buffer is
    * exhausted or the frame end is reached.
    *
    * TODO: Contrary to the C++ implementation, this never does cross-frame
    * reads – is there actually a valid use case for that?
    *
    * Params:
    *   buf = Slice to use as buffer.
    *
    * Returns: How many bytes were actually read.
    *
    * Throws: TTransportException if an error occurs.
    */
   override size_t read(ubyte[] buf) {
     // If the buffer is empty, read a new frame off the wire.
     if (rBuf_.empty) {
       bool gotFrame = readFrame();
       if (!gotFrame) return 0;
     }

     auto size = min(rBuf_.length, buf.length);
     buf[0..size] = rBuf_[0..size];
     rBuf_ = rBuf_[size..$];
     return size;
   }

   override void write(in ubyte[] buf) {
     wBuf_ ~= buf;
   }

   override void flush() {
     if (wBuf_.empty) return;

     // Properly reset the write buffer even some of the protocol operations go
     // wrong.
     scope (exit) {
       wBuf_.length = 0;
       wBuf_.assumeSafeAppend();
     }

     int len = bswap(cast(int)wBuf_.length);
     transport_.write(cast(ubyte[])(&len)[0..1]);
     transport_.write(wBuf_);
     transport_.flush();
   }

   override const(ubyte)[] borrow(ubyte* buf, size_t len) {
     if (len <= rBuf_.length) {
       return rBuf_;
     } else {
       // Don't try attempting cross-frame borrows, trying that does not make
       // much sense anyway.
       return null;
     }
   }

   override void consume(size_t len) {
     enforce(len <= rBuf_.length, new TTransportException(
       "Invalid consume length", TTransportException.Type.BAD_ARGS));
     rBuf_ = rBuf_[len .. $];
   }

 private:
   bool readFrame() {
     // Read the size of the next frame. We can't use readAll() since that
     // always throws an exception on EOF, but want to throw an exception only
     // if EOF occurs after partial size data.
     int size;
     size_t size_read;
     while (size_read < size.sizeof) {
       auto data = (cast(ubyte*)&size)[size_read..size.sizeof];
       auto read = transport_.read(data);
       if (read == 0) {
         if (size_read == 0) {
           // EOF before any data was read.
           return false;
         } else {
           // EOF after a partial frame header – illegal.
           throw new TTransportException(
             "No more data to read after partial frame header",
             TTransportException.Type.END_OF_FILE
           );
         }
       }
       size_read += read;
     }

     size = bswap(size);
     enforce(size >= 0, new TTransportException("Frame size has negative value",
       TTransportException.Type.CORRUPTED_DATA));

     // TODO: Benchmark this.
     rBuf_.length = size;
     rBuf_.assumeSafeAppend();

     transport_.readAll(rBuf_);
     return true;
   }

   TTransport transport_;
   ubyte[] rBuf_;
   ubyte[] wBuf_;
 }

 /**
  * Wraps given transports into TFramedTransports.
  */
 alias TWrapperTransportFactory!TFramedTransport TFramedTransportFactory;

 version (unittest) {
   import std.random : Mt19937, uniform;
   import thrift.transport.memory;
 }

 // Some basic random testing, always starting with the same seed for
 // deterministic unit test results – more tests in transport_test.
 unittest {
   auto randGen = Mt19937(42);

   // 32 kiB of data to work with.
   auto data = new ubyte[1 << 15];
   foreach (ref b; data) {
     b = uniform!"[]"(cast(ubyte)0, cast(ubyte)255, randGen);
   }

   // Generate a list of chunk sizes to split the data into. A uniform
   // distribution is not quite realistic, but std.random doesn't have anything
   // else yet.
   enum MAX_FRAME_LENGTH = 512;
   auto chunkSizesList = new size_t[][2];
   foreach (ref chunkSizes; chunkSizesList) {
     size_t sum;
     while (true) {
       auto curLen = uniform(0, MAX_FRAME_LENGTH, randGen);
       sum += curLen;
       if (sum > data.length) break;
       chunkSizes ~= curLen;
     }
   }
   chunkSizesList ~= [data.length]; // Also test whole chunk at once.

   // Test writing data.
   {
     foreach (chunkSizes; chunkSizesList) {
       auto buf = new TMemoryBuffer;
       auto framed = new TFramedTransport(buf);

       auto remainingData = data;
       foreach (chunkSize; chunkSizes) {
         framed.write(remainingData[0..chunkSize]);
         remainingData = remainingData[chunkSize..$];
       }
       framed.flush();

       auto writtenData = data[0..($ - remainingData.length)];
       auto actualData = buf.getContents();

       // Check frame size.
       int frameSize = bswap((cast(int[])(actualData[0..int.sizeof]))[0]);
       enforce(frameSize == writtenData.length);

       // Check actual data.
       enforce(actualData[int.sizeof..$] == writtenData);
     }
   }

   // Test reading data.
   {
     foreach (chunkSizes; chunkSizesList) {
       auto buf = new TMemoryBuffer;

       auto size = bswap(cast(int)data.length);
       buf.write(cast(ubyte[])(&size)[0..1]);
       buf.write(data);

       auto framed = new TFramedTransport(buf);
       ubyte[] readData;
       readData.reserve(data.length);
       foreach (chunkSize; chunkSizes) {
         // This should work with read because we have one huge frame.
         auto oldReadLen = readData.length;
         readData.length += chunkSize;
         framed.read(readData[oldReadLen..$]);
       }

       enforce(readData == data[0..readData.length]);
     }
   }

   // Test combined reading/writing of multiple frames.
   foreach (flushProbability; [1, 2, 4, 8, 16, 32]) {
     foreach (chunkSizes; chunkSizesList) {
       auto buf = new TMemoryBuffer;
       auto framed = new TFramedTransport(buf);

       size_t[] frameSizes;

       // Write the data.
       size_t frameSize;
       auto remainingData = data;
       foreach (chunkSize; chunkSizes) {
         framed.write(remainingData[0..chunkSize]);
         remainingData = remainingData[chunkSize..$];

         frameSize += chunkSize;
         if (frameSize > 0 && uniform(0, flushProbability, randGen) == 0) {
           frameSizes ~= frameSize;
           frameSize = 0;
           framed.flush();
         }
       }
       if (frameSize > 0) {
         frameSizes ~= frameSize;
         frameSize = 0;
         framed.flush();
       }

       // Read it back.
       auto readData = new ubyte[data.length - remainingData.length];
       auto remainToRead = readData;
       foreach (fSize; frameSizes) {
         // We are exploiting an implementation detail of TFramedTransport:
         // The read buffer starts empty and it will never return more than one
         // frame per read, so by just requesting all of the data, we should
         // always get exactly one frame.
         auto got = framed.read(remainToRead);
         enforce(got == fSize);
         remainToRead = remainToRead[fSize..$];
       }

       enforce(remainToRead.empty);
       enforce(readData == data[0..readData.length]);
     }
   }
 }

 // Test flush()ing an empty buffer.
 unittest {
   auto buf = new TMemoryBuffer();
   auto framed = new TFramedTransport(buf);
   immutable out1 = [0, 0, 0, 1, 'a'];
   immutable out2 = [0, 0, 0, 1, 'a', 0, 0, 0, 2, 'b', 'c'];

   framed.flush();
   enforce(buf.getContents() == []);
   framed.flush();
   framed.flush();
   enforce(buf.getContents() == []);
   framed.write(cast(ubyte[])"a");
   enforce(buf.getContents() == []);
   framed.flush();
   enforce(buf.getContents() == out1);
   framed.flush();
   framed.flush();
   enforce(buf.getContents() == out1);
   framed.write(cast(ubyte[])"bc");
   enforce(buf.getContents() == out1);
   framed.flush();
   enforce(buf.getContents() == out2);
   framed.flush();
   framed.flush();
   enforce(buf.getContents() == out2);
 }
	/*
	* 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.
	*/

	module thrift.transport.framed;

	import core.bitop : bswap;
	import std.algorithm : min;
	import std.array : empty;
	import std.exception : enforce;
	import thrift.transport.base;

	/**
	* Framed transport.
	*
	* All writes go into an in-memory buffer until flush is called, at which point
	* the transport writes the length of the entire binary chunk followed by the
	* data payload. The receiver on the other end then performs a single
	* »fixed-length« read to get the whole message off the wire.
	*/
	final class TFramedTransport : TBaseTransport {
	/**
	* Constructs a new framed transport.
	*
	* Params:
	* transport = The underlying transport to wrap.
	*/
	this(TTransport transport) {
	transport_ = transport;
	}

	/**
	* Returns the wrapped transport.
	*/
	TTransport underlyingTransport() @property {
	return transport_;
	}

	override bool isOpen() @property {
	return transport_.isOpen;
	}

	override bool peek() {
	return rBuf_.length > 0 \|\| transport_.peek();
	}

	override void open() {
	transport_.open();
	}

	override void close() {
	flush();
	transport_.close();
	}

	/**
	* Attempts to read data into the given buffer, stopping when the buffer is
	* exhausted or the frame end is reached.
	*
	* TODO: Contrary to the C++ implementation, this never does cross-frame
	* reads – is there actually a valid use case for that?
	*
	* Params:
	* buf = Slice to use as buffer.
	*
	* Returns: How many bytes were actually read.
	*
	* Throws: TTransportException if an error occurs.
	*/
	override size_t read(ubyte[] buf) {
	// If the buffer is empty, read a new frame off the wire.
	if (rBuf_.empty) {
	bool gotFrame = readFrame();
	if (!gotFrame) return 0;
	}

	auto size = min(rBuf_.length, buf.length);
	buf[0..size] = rBuf_[0..size];
	rBuf_ = rBuf_[size..$];
	return size;
	}

	override void write(in ubyte[] buf) {
	wBuf_ ~= buf;
	}

	override void flush() {
	if (wBuf_.empty) return;

	// Properly reset the write buffer even some of the protocol operations go
	// wrong.
	scope (exit) {
	wBuf_.length = 0;
	wBuf_.assumeSafeAppend();
	}

	int len = bswap(cast(int)wBuf_.length);
	transport_.write(cast(ubyte[])(&len)[0..1]);
	transport_.write(wBuf_);
	transport_.flush();
	}

	override const(ubyte)[] borrow(ubyte* buf, size_t len) {
	if (len <= rBuf_.length) {
	return rBuf_;
	} else {
	// Don't try attempting cross-frame borrows, trying that does not make
	// much sense anyway.
	return null;
	}
	}

	override void consume(size_t len) {
	enforce(len <= rBuf_.length, new TTransportException(
	"Invalid consume length", TTransportException.Type.BAD_ARGS));
	rBuf_ = rBuf_[len .. $];
	}

	private:
	bool readFrame() {
	// Read the size of the next frame. We can't use readAll() since that
	// always throws an exception on EOF, but want to throw an exception only
	// if EOF occurs after partial size data.
	int size;
	size_t size_read;
	while (size_read < size.sizeof) {
	auto data = (cast(ubyte*)&size)[size_read..size.sizeof];
	auto read = transport_.read(data);
	if (read == 0) {
	if (size_read == 0) {
	// EOF before any data was read.
	return false;
	} else {
	// EOF after a partial frame header – illegal.
	throw new TTransportException(
	"No more data to read after partial frame header",
	TTransportException.Type.END_OF_FILE
	);
	}
	}
	size_read += read;
	}

	size = bswap(size);
	enforce(size >= 0, new TTransportException("Frame size has negative value",
	TTransportException.Type.CORRUPTED_DATA));

	// TODO: Benchmark this.
	rBuf_.length = size;
	rBuf_.assumeSafeAppend();

	transport_.readAll(rBuf_);
	return true;
	}

	TTransport transport_;
	ubyte[] rBuf_;
	ubyte[] wBuf_;
	}

	/**
	* Wraps given transports into TFramedTransports.
	*/
	alias TWrapperTransportFactory!TFramedTransport TFramedTransportFactory;

	version (unittest) {
	import std.random : Mt19937, uniform;
	import thrift.transport.memory;
	}

	// Some basic random testing, always starting with the same seed for
	// deterministic unit test results – more tests in transport_test.
	unittest {
	auto randGen = Mt19937(42);

	// 32 kiB of data to work with.
	auto data = new ubyte[1 << 15];
	foreach (ref b; data) {
	b = uniform!"[]"(cast(ubyte)0, cast(ubyte)255, randGen);
	}

	// Generate a list of chunk sizes to split the data into. A uniform
	// distribution is not quite realistic, but std.random doesn't have anything
	// else yet.
	enum MAX_FRAME_LENGTH = 512;
	auto chunkSizesList = new size_t[][2];
	foreach (ref chunkSizes; chunkSizesList) {
	size_t sum;
	while (true) {
	auto curLen = uniform(0, MAX_FRAME_LENGTH, randGen);
	sum += curLen;
	if (sum > data.length) break;
	chunkSizes ~= curLen;
	}
	}
	chunkSizesList ~= [data.length]; // Also test whole chunk at once.

	// Test writing data.
	{
	foreach (chunkSizes; chunkSizesList) {
	auto buf = new TMemoryBuffer;
	auto framed = new TFramedTransport(buf);

	auto remainingData = data;
	foreach (chunkSize; chunkSizes) {
	framed.write(remainingData[0..chunkSize]);
	remainingData = remainingData[chunkSize..$];
	}
	framed.flush();

	auto writtenData = data[0..($ - remainingData.length)];
	auto actualData = buf.getContents();

	// Check frame size.
	int frameSize = bswap((cast(int[])(actualData[0..int.sizeof]))[0]);
	enforce(frameSize == writtenData.length);

	// Check actual data.
	enforce(actualData[int.sizeof..$] == writtenData);
	}
	}

	// Test reading data.
	{
	foreach (chunkSizes; chunkSizesList) {
	auto buf = new TMemoryBuffer;

	auto size = bswap(cast(int)data.length);
	buf.write(cast(ubyte[])(&size)[0..1]);
	buf.write(data);

	auto framed = new TFramedTransport(buf);
	ubyte[] readData;
	readData.reserve(data.length);
	foreach (chunkSize; chunkSizes) {
	// This should work with read because we have one huge frame.
	auto oldReadLen = readData.length;
	readData.length += chunkSize;
	framed.read(readData[oldReadLen..$]);
	}

	enforce(readData == data[0..readData.length]);
	}
	}

	// Test combined reading/writing of multiple frames.
	foreach (flushProbability; [1, 2, 4, 8, 16, 32]) {
	foreach (chunkSizes; chunkSizesList) {
	auto buf = new TMemoryBuffer;
	auto framed = new TFramedTransport(buf);

	size_t[] frameSizes;

	// Write the data.
	size_t frameSize;
	auto remainingData = data;
	foreach (chunkSize; chunkSizes) {
	framed.write(remainingData[0..chunkSize]);
	remainingData = remainingData[chunkSize..$];

	frameSize += chunkSize;
	if (frameSize > 0 && uniform(0, flushProbability, randGen) == 0) {
	frameSizes ~= frameSize;
	frameSize = 0;
	framed.flush();
	}
	}
	if (frameSize > 0) {
	frameSizes ~= frameSize;
	frameSize = 0;
	framed.flush();
	}

	// Read it back.
	auto readData = new ubyte[data.length - remainingData.length];
	auto remainToRead = readData;
	foreach (fSize; frameSizes) {
	// We are exploiting an implementation detail of TFramedTransport:
	// The read buffer starts empty and it will never return more than one
	// frame per read, so by just requesting all of the data, we should
	// always get exactly one frame.
	auto got = framed.read(remainToRead);
	enforce(got == fSize);
	remainToRead = remainToRead[fSize..$];
	}

	enforce(remainToRead.empty);
	enforce(readData == data[0..readData.length]);
	}
	}
	}

	// Test flush()ing an empty buffer.
	unittest {
	auto buf = new TMemoryBuffer();
	auto framed = new TFramedTransport(buf);
	immutable out1 = [0, 0, 0, 1, 'a'];
	immutable out2 = [0, 0, 0, 1, 'a', 0, 0, 0, 2, 'b', 'c'];

	framed.flush();
	enforce(buf.getContents() == []);
	framed.flush();
	framed.flush();
	enforce(buf.getContents() == []);
	framed.write(cast(ubyte[])"a");
	enforce(buf.getContents() == []);
	framed.flush();
	enforce(buf.getContents() == out1);
	framed.flush();
	framed.flush();
	enforce(buf.getContents() == out1);
	framed.write(cast(ubyte[])"bc");
	enforce(buf.getContents() == out1);
	framed.flush();
	enforce(buf.getContents() == out2);
	framed.flush();
	framed.flush();
	enforce(buf.getContents() == out2);
	}