From 4ec3eb2066a6c79b5dfb2f7ff45a4d96939feb76 Mon Sep 17 00:00:00 2001
From: Ian Gulliver <ian@flamingcow.tv>
Date: Sun, 28 Apr 2019 18:51:56 +0000
Subject: [PATCH] Clean Flush() implementation

---
 fastcgi_request.cpp | 77 ++++++++++++++++++++++++++++++---------------
 fastcgi_request.h   |  6 +++-
 2 files changed, 57 insertions(+), 26 deletions(-)

diff --git a/fastcgi_request.cpp b/fastcgi_request.cpp
index aefbbcc..7fe01c1 100644
--- a/fastcgi_request.cpp
+++ b/fastcgi_request.cpp
@@ -3,7 +3,17 @@
 #include "fastcgi_request.h"
 
 #include "fastcgi_conn.h"
-#include "fastcgi_parse.h"
+
+namespace {
+
+template<class T> void AppendVec(const T& obj, std::vector<iovec>* vec) {
+	vec->push_back(iovec{
+		.iov_base = (void*)(&obj),
+		.iov_len = sizeof(obj),
+	});
+}
+
+} // namespace
 
 FastCGIRequest::FastCGIRequest(uint16_t request_id, FastCGIConn* conn)
 		: request_id_(request_id),
@@ -23,7 +33,7 @@ const std::string& FastCGIRequest::GetParam(const std::string& key) {
 }
 
 void FastCGIRequest::WriteHeader(const std::string_view& name, const std::string_view& value) {
-	CHECK(!body_sent_);
+	CHECK(!body_written_);
 	CHECK(out_buf_.Write(name));
 	CHECK(out_buf_.Write(": "));
 	CHECK(out_buf_.Write(value));
@@ -31,38 +41,55 @@ void FastCGIRequest::WriteHeader(const std::string_view& name, const std::string
 }
 
 void FastCGIRequest::WriteBody(const std::string_view& body) {
-	if (!body_sent_) {
+	if (!body_written_) {
 		CHECK(out_buf_.Write("\n"));
-		body_sent_ = true;
+		body_written_ = true;
 	}
 	// TODO: make this able to span multiple packets
 	CHECK(out_buf_.Write(body));
 }
 
-void FastCGIRequest::End() {
-	const auto output_len = out_buf_.ReadMaxLen();
+void FastCGIRequest::Flush() {
+	std::vector<iovec> vecs;
+
+	auto header = OutputHeader();
+	AppendVec(header, &vecs);
+
+	vecs.push_back(OutputVec());
+
+	conn_->Write(vecs);
+	out_buf_.Commit();
+}
+
+void FastCGIRequest::End() {
+	// Fully empty response not allowed
+	WriteBody("");
+
+	std::vector<iovec> vecs;
+
+	// Must be outside if block, so it lives through Write() below
+	auto output_header = OutputHeader();
+	if (output_header.ContentLength()) {
+		AppendVec(output_header, &vecs);
+		vecs.push_back(OutputVec());
+	}
 
-	FastCGIHeader output_header(6, request_id_, output_len);
 	FastCGIEndRequest end;
 	FastCGIHeader end_header(3, request_id_, sizeof(end));
+	AppendVec(end_header, &vecs);
+	AppendVec(end, &vecs);
 
-	std::vector<iovec> vecs{
-		iovec{
-			.iov_base = &output_header,
-			.iov_len = sizeof(output_header),
-		},
-		iovec{
-			.iov_base = (void *)(CHECK_NOTNULL(out_buf_.Read(output_len))),
-			.iov_len = output_len,
-		},
-		{
-			.iov_base = &end_header,
-			.iov_len = sizeof(end_header),
-		},
-		{
-			.iov_base = &end,
-			.iov_len = sizeof(end),
-		},
-	};
 	conn_->Write(vecs);
 }
+
+iovec FastCGIRequest::OutputVec() {
+	const auto output_len = out_buf_.ReadMaxLen();
+	return iovec{
+		.iov_base = (void *)(CHECK_NOTNULL(out_buf_.Read(output_len))),
+		.iov_len = output_len,
+	};
+}
+
+FastCGIHeader FastCGIRequest::OutputHeader() {
+	return FastCGIHeader(6, request_id_, out_buf_.ReadMaxLen());
+}
diff --git a/fastcgi_request.h b/fastcgi_request.h
index 10bc12d..4391de8 100644
--- a/fastcgi_request.h
+++ b/fastcgi_request.h
@@ -3,6 +3,7 @@
 #include <unordered_map>
 
 #include "buffer.h"
+#include "fastcgi_parse.h"
 
 class FastCGIConn;
 
@@ -21,6 +22,9 @@ class FastCGIRequest {
 	void End();
 
   private:
+	FastCGIHeader OutputHeader();
+	iovec OutputVec();
+
 	const uint16_t request_id_;
 	FastCGIConn *conn_;
 
@@ -28,5 +32,5 @@ class FastCGIRequest {
 	std::string in_;
 
 	Buffer out_buf_;
-	bool body_sent_ = false;
+	bool body_written_ = false;
 };