From b5ae7d0e7ad5593ac824d9610be9831c94791a6b Mon Sep 17 00:00:00 2001
From: yhirose <yuji.hirose.bug@gmail.com>
Date: Fri, 12 Oct 2012 17:41:58 -0400
Subject: [PATCH] Added HEAD method support.

---
 httplib.h    | 42 ++++++++++++++++++++++++++----------------
 test/test.cc | 19 ++++++++++++++++++-
 2 files changed, 44 insertions(+), 17 deletions(-)

diff --git a/httplib.h b/httplib.h
index 3c40b7a..4ac0dcb 100644
--- a/httplib.h
+++ b/httplib.h
@@ -81,13 +81,6 @@ struct Response {
     Response() : status(-1) {}
 };
 
-/*
-struct Connection {
-    Request  request;
-    Response response;
-};
-*/
-
 class Server {
 public:
     typedef std::function<void (const Request&, Response&)> Handler;
@@ -125,6 +118,7 @@ public:
     ~Client();
 
     std::shared_ptr<Response> get(const char* url);
+    std::shared_ptr<Response> head(const char* url);
     std::shared_ptr<Response> post(const char* url, const std::string& body, const char* content_type);
     std::shared_ptr<Response> post(const char* url, const Map& params);
 
@@ -249,11 +243,11 @@ inline socket_t create_client_socket(const char* host, int port)
 inline const char* status_message(int status)
 {
     switch (status) {
+    case 200: return "OK";
     case 400: return "Bad Request";
     case 404: return "Not Found";
     default:
-        status = 500;
-        return "Internal Server Error";
+        case 500: return "Internal Server Error";
     }
 }
 
@@ -333,13 +327,13 @@ inline void write_headers(FILE* fp, const T& x)
     fprintf(fp, "\r\n");
 }
 
-inline void write_response(FILE* fp, const Response& res)
+inline void write_response(FILE* fp, const Request& req, const Response& res)
 {
     fprintf(fp, "HTTP/1.0 %d %s\r\n", res.status, status_message(res.status));
 
     write_headers(fp, res);
 
-    if (!res.body.empty()) {
+    if (!res.body.empty() && req.method != "HEAD") {
         fprintf(fp, "%s", res.body.c_str());
     }
 }
@@ -643,7 +637,7 @@ inline bool Server::read_request_line(FILE* fp, Request& req)
         return false;
     }
 
-    static std::regex re("(GET|POST) ([^?]+)(?:\\?(.+?))? HTTP/1\\.[01]\r\n");
+    static std::regex re("(GET|HEAD|POST) ([^?]+)(?:\\?(.+?))? HTTP/1\\.[01]\r\n");
 
     std::cmatch m;
     if (std::regex_match(buf, m, re)) {
@@ -664,7 +658,7 @@ inline bool Server::read_request_line(FILE* fp, Request& req)
 
 inline bool Server::routing(Request& req, Response& res)
 {
-    if (req.method == "GET") {
+    if (req.method == "GET" || req.method == "HEAD") {
         return dispatch_request(req, res, get_handlers_);
     } else if (req.method == "POST") {
         return dispatch_request(req, res, post_handlers_);
@@ -722,7 +716,7 @@ inline void Server::process_request(socket_t sock)
         error_handler_(req, res);
     }
 
-    detail::write_response(fp_write, res);
+    detail::write_response(fp_write, req, res);
     fflush(fp_write);
 
     if (logger_) {
@@ -781,11 +775,16 @@ inline bool Client::send(const Request& req, Response& res)
     detail::write_request(fp_write, req);
     fflush(fp_write);
 
+    // Receive response
     if (!read_response_line(fp_read, res) ||
-        !detail::read_headers(fp_read, res.headers) ||
-        !detail::read_content(res, fp_read)) {
+        !detail::read_headers(fp_read, res.headers)) {
         return false;
     }
+    if (req.method != "HEAD") {
+        if (!detail::read_content(res, fp_read)) {
+            return false;
+        }
+    }
 
     detail::shutdown_socket(sock);
     detail::close_socket(sock);
@@ -804,6 +803,17 @@ inline std::shared_ptr<Response> Client::get(const char* url)
     return send(req, *res) ? res : nullptr;
 }
 
+inline std::shared_ptr<Response> Client::head(const char* url)
+{
+    Request req;
+    req.method = "HEAD";
+    req.url = url;
+
+    auto res = std::make_shared<Response>();
+
+    return send(req, *res) ? res : nullptr;
+}
+
 inline std::shared_ptr<Response> Client::post(
     const char* url, const std::string& body, const char* content_type)
 {
diff --git a/test/test.cc b/test/test.cc
index 531974f..d22dd17 100644
--- a/test/test.cc
+++ b/test/test.cc
@@ -15,7 +15,7 @@ using namespace std;
 using namespace httplib;
 
 const char* HOST = "localhost";
-const int   PORT = 8080;
+const int   PORT = 1234;
 
 class thread
 {
@@ -260,6 +260,23 @@ TEST_F(ServerTest, GetMethod404)
     EXPECT_EQ(404, res->status);
 }
 
+TEST_F(ServerTest, HeadMethod200)
+{
+    auto res = cli_.head("/hi");
+    ASSERT_TRUE(res != nullptr);
+    EXPECT_EQ(200, res->status);
+    EXPECT_EQ("text/plain", res->get_header_value("Content-Type"));
+    EXPECT_EQ("", res->body);
+}
+
+TEST_F(ServerTest, HeadMethod404)
+{
+    auto res = cli_.head("/invalid");
+    ASSERT_TRUE(res != nullptr);
+    EXPECT_EQ(404, res->status);
+    EXPECT_EQ("", res->body);
+}
+
 TEST_F(ServerTest, GetMethodPersonJohn)
 {
     auto res = cli_.get("/person/john");