5 files changed, 41 insertions, 26 deletions

diff --git a/src/ClientConnection.cpp b/src/ClientConnection.cpp
index 3bf93c3..da0cf9f 100644
--- a/src/ClientConnection.cpp
+++ b/src/ClientConnection.cpp
@@ -26,24 +26,19 @@ void ClientConnection::dump_request_data()
 	}
 }
 
-bool ClientConnection::send(const HttpResponse& response, const char* message)
+bool ClientConnection::send(const HttpResponse& response)
 {
 	if (!m_is_open)
 		return false;
 
-	std::stringstream ss;
-	ss << response.to_string()
-	   << "\r\n\r\n"
-	   << message;
-
-	std::string result = ss.str();
+	std::string result = response.to_string();
 	write(m_socket_fd, result.c_str(), result.length());
 
 	return true;
 }
 
-void ClientConnection::close()
+void ClientConnection::close_connection()
 {
 	m_is_open = false;
-	::close(m_socket_fd);
+	close(m_socket_fd);
 }
diff --git a/src/ClientConnection.h b/src/ClientConnection.h
index 6fdf99e..55be620 100644
--- a/src/ClientConnection.h
+++ b/src/ClientConnection.h
@@ -8,8 +8,8 @@ public:
 
 	void dump_request_data();
 
-	bool send(const HttpResponse&, const char*);
-	void close();
+	bool send(const HttpResponse&);
+	void close_connection();
 private:
 	int m_socket_fd;
 	bool m_is_open { true };
diff --git a/src/HttpResponse.cpp b/src/HttpResponse.cpp
index c9c5ff4..d696e03 100644
--- a/src/HttpResponse.cpp
+++ b/src/HttpResponse.cpp
@@ -1,16 +1,32 @@
 #include "HttpResponse.h"
 
-HttpResponse::HttpResponse(ResponseCode response_code)
+#include <sstream>
+
+static std::string status_code_string(HttpStatusCode status_code)
 {
-	m_string_stream << "HTTP/1.0 200 OK";
+	switch (status_code) {
+	case HttpStatusCode::OK:
+		return "200 OK";
+	}
 }
 
-void HttpResponse::add_header(std::string header)
+HttpResponse::HttpResponse(HttpStatusCode status_code)
+	: m_status_code(status_code)
 {
-	m_string_stream << '\n' << header;
+}
+
+void HttpResponse::add_header(const std::string& header, const std::string& value)
+{
+	m_headers[header] = value;
 }
 
 std::string HttpResponse::to_string() const
 {
-	return m_string_stream.str();
+	std::stringstream string_stream;
+	string_stream << "HTTP/1.0 " << status_code_string(m_status_code) << "\r\n";
+	for (const auto& header : m_headers)
+		string_stream << header.first << ": " << header.second << "\r\n";
+	string_stream << "\r\n" << m_content;
+
+	return string_stream.str();
 }
diff --git a/src/HttpResponse.h b/src/HttpResponse.h
index 33148da..8a1d2d4 100644
--- a/src/HttpResponse.h
+++ b/src/HttpResponse.h
@@ -1,18 +1,22 @@
 #pragma once
 
-#include <sstream>
+#include <map>
 #include <string>
 
-enum class ResponseCode {
+enum class HttpStatusCode {
 	OK = 200
 };
 
 class HttpResponse {
 public:
-	HttpResponse(ResponseCode);
+	HttpResponse(HttpStatusCode status_code);
+
+	void add_header(const std::string& header, const std::string& value);
+	void set_content(const std::string& content) { m_content = content; }
 
-	void add_header(std::string header);
 	std::string to_string() const;
 private:
-	std::stringstream m_string_stream;
+	HttpStatusCode m_status_code;
+	std::map<std::string, std::string> m_headers;
+	std::string m_content;
 };
diff --git a/src/main.cpp b/src/main.cpp
index 60c7b4f..3521079 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -15,11 +15,11 @@ int main(int argc, char** argv)
 		ClientConnection client = server.accept_client_connection();
 		client.dump_request_data();
 
-		HttpResponse http_response(ResponseCode::OK);
-		http_response.add_header("Server: cfws");
+		HttpResponse response(HttpStatusCode::OK);
+		response.add_header("Server", "cfws");
+		response.set_content("Welcome to the page.");
 
-		const char* message = "Welcome to the page.";
-		client.send(http_response, message);
-		client.close();
+		client.send(response);
+		client.close_connection();
 	}
 }