Merge pull request #69 from vibe-d/lazy_chunked_input_stream

Make ChunkedInputStream lazy.

Author: s-ludwig

source/vibe/http/common.d

@@ -389,19 +389,22 @@ final class ChunkedInputStream : InputStream
 	private {
 		InterfaceProxy!InputStream m_in;
 		ulong m_bytesInCurrentChunk = 0;
+		bool m_chunkHeaderRead = false;
 	}
 
 	/// private
 	this(InterfaceProxy!InputStream stream, bool dummy)
 	{
 		assert(!!stream);
 		m_in = stream;
-		readChunk();
 	}
 
-	@property bool empty() const { return m_bytesInCurrentChunk == 0; }
+	@property bool empty() { return leastSize == 0; }
 
-	@property ulong leastSize() const { return m_bytesInCurrentChunk; }
+	@property ulong leastSize() {
+		if (!m_chunkHeaderRead) readChunk();
+		return m_bytesInCurrentChunk;
+	}
 
 	@property bool dataAvailableForRead() { return m_bytesInCurrentChunk > 0 && m_in.dataAvailableForRead; }
 
@@ -417,21 +420,20 @@ final class ChunkedInputStream : InputStream
 		size_t nbytes = 0;
 
 		while (dst.length > 0) {
-			enforceBadRequest(m_bytesInCurrentChunk > 0, "Reading past end of chunked HTTP stream.");
+			enforceBadRequest(leastSize > 0, "Reading past end of chunked HTTP stream.");
 
 			auto sz = cast(size_t)min(m_bytesInCurrentChunk, dst.length);
 			m_in.read(dst[0 .. sz]);
 			dst = dst[sz .. $];
 			m_bytesInCurrentChunk -= sz;
 			nbytes += sz;
 
-			// FIXME: this blocks, but shouldn't for IOMode.once/immediat
 			if( m_bytesInCurrentChunk == 0 ){
 				// skip current chunk footer and read next chunk
 				ubyte[2] crlf;
 				m_in.read(crlf);
 				enforceBadRequest(crlf[0] == '\r' && crlf[1] == '\n');
-				readChunk();
+				m_chunkHeaderRead = false;
 			}
 
 			if (mode != IOMode.all) break;
@@ -444,12 +446,13 @@ final class ChunkedInputStream : InputStream
 
 	private void readChunk()
 	{
-		assert(m_bytesInCurrentChunk == 0);
+		assert(m_bytesInCurrentChunk == 0 && !m_chunkHeaderRead);
 		// read chunk header
 		logTrace("read next chunk header");
 		auto ln = () @trusted { return cast(string)m_in.readLine(); } ();
 		logTrace("got chunk header: %s", ln);
 		m_bytesInCurrentChunk = parse!ulong(ln, 16u);
+		m_chunkHeaderRead = true;
 
 		if( m_bytesInCurrentChunk == 0 ){
 			// empty chunk denotes the end
@@ -633,6 +636,34 @@ FreeListRef!ChunkedOutputStream createChunkedOutputStreamFL(OS, Allocator)(OS de
 	return FreeListRef!ChunkedOutputStream(interfaceProxy!OutputStream(destination_stream), allocator, true);
 }
 
+unittest {
+	import vibe.stream.memory : createMemoryStream;
+
+	ubyte[] bytes = new ubyte[](100*100);
+	foreach (i, ref bt; bytes) bt = (i + i / 100) % 256;
+	ubyte[] buf = new ubyte[](bytes.length);
+
+	auto dst = createMemoryStream(new ubyte[4*1024*1024], true, 0);
+
+	// write data with varying buffer and chunk sizes
+	auto outp = createChunkedOutputStreamFL(dst);
+	foreach (i; 0 .. 100) {
+		outp.write(bytes[0 .. i * 100]);
+		if (i % 5 == 0)
+			outp.flush();
+	}
+	outp.finalize();
+
+	// ensure decoding yields the same byte sequence
+	dst.seek(0);
+	auto inp = createChunkedInputStreamFL(dst);
+	foreach (i; 0 .. 100) {
+		inp.read(buf[0 .. i * 100]);
+		assert(buf[0 .. i * 100] == bytes[0 .. i * 100]);
+	}
+}
+
+
 /// Parses the cookie from a header field, returning the name of the cookie.
 /// Implements an algorithm equivalent to https://tools.ietf.org/html/rfc6265#section-5.2
 /// Returns: the cookie name as return value, populates the dst argument or allocates on the GC for the tuple overload.

source/vibe/http/websockets.d

@@ -906,6 +906,7 @@ scope:
 
 		try m_conn.close();
 		catch (Exception e) logException(e, "Failed to close WebSocket connection");
+		m_conn = ConnectionStream.init;
 		try m_readCondition.notifyAll();
 		catch (Exception e) assert(false, e.msg);
 	}
@@ -924,7 +925,7 @@ scope:
 			send((scope msg) { msg.write(nativeToLittleEndian(++m_lastPingIndex)); }, FrameOpcode.ping);
 			logDebugV("Ping sent");
 		} catch (Exception e) {
-			logError("Failed to acquire write mutex for sending a WebSocket ping frame: %s", e.msg);
+			logError("Failed to send WebSocket ping frame: %s", e.msg);
 		}
 	}
 }

tests/vibe.http.websocket.packetmaxlength/dub.sdl

@@ -1,3 +1,2 @@
 name "websocket_packetmaxlength"
 dependency "vibe-http" path="../../"
-versions "VibeDefaultMain"

tests/vibe.http.websocket.packetmaxlength/source/app.d

@@ -3,39 +3,39 @@ import vibe.core.log;
 import vibe.inet.url;
 import vibe.http.server;
 import vibe.http.websockets;
+import core.time : seconds;
 
-shared static this()
+
+void main()
 {
+	auto tm = setTimer(10.seconds, { assert(false, "Test timed out."); });
+
 	//test max payload
 	auto settings = new HTTPServerSettings;
 	settings.port = 0;
 	settings.bindAddresses = ["127.0.0.1"];
 	settings.webSocketPayloadMaxLength = 10;
-	immutable serverAddr = listenHTTP(settings, handleWebSockets((scope ws) {
+	auto listener = listenHTTP(settings, handleWebSockets((scope ws) {
 		assert(ws.connected); // issue #2104
 		assert(ws.receiveText() == "1234567890");
 		ws.send("ok");
 		try{
 			ws.receiveText();	//at this point the connection should close
 			assert(false);
-		}catch(Exception e)
-		{
-		}
+		} catch(Exception e) {} // expected
 		ws.close();
-	})).bindAddresses[0];
-
-	runTask({
-		scope(exit) exitEventLoop(true);
+	}));
 
-		try connectWebSocket(URL("http://" ~ serverAddr.toString), (scope ws) {
-			assert(ws.connected);
-			ws.send("1234567890");
-			assert(ws.receiveText() == "ok");
-			ws.send("1234567890a");		//should cause connection to close
-			assert(!ws.waitForData);
-			ws.close();
-			logInfo("WebSocket max payload test successful");
-		});
-		catch (Exception e) assert(false, e.msg);
+	try connectWebSocket(URL("http://" ~ listener.bindAddresses[0].toString), (scope ws) {
+		assert(ws.connected);
+		ws.send("1234567890");
+		assert(ws.receiveText() == "ok");
+		ws.send("1234567890a");		//should cause connection to close
+		assert(!ws.waitForData);
+		ws.close();
+		logInfo("WebSocket max payload test successful");
 	});
+	catch (Exception e) assert(false, e.msg);
+
+	listener.stopListening();
 }