From fcbcbd53bd25a7707ac1ea6482987289f1e8a141 Mon Sep 17 00:00:00 2001
From: yhirose <yuji.hirose.bug@gmail.com>
Date: Thu, 26 Dec 2019 18:48:22 -0500
Subject: [PATCH] Fix #306

---
 httplib.h               | 23 ++++++++++++++++++++---
 test/test.cc            |  9 +++++++++
 test/www/dir/test.abcde |  1 +
 3 files changed, 30 insertions(+), 3 deletions(-)
 create mode 100644 test/www/dir/test.abcde

diff --git a/httplib.h b/httplib.h
index 6d3100f..a9d48be 100644
--- a/httplib.h
+++ b/httplib.h
@@ -528,6 +528,8 @@ public:
   Server &Options(const char *pattern, Handler handler);
 
   bool set_base_dir(const char *dir, const char *mount_point = nullptr);
+  void set_file_extension_and_mimetype_mapping(const char *ext,
+                                               const char *mime);
   void set_file_request_handler(Handler handler);
 
   void set_error_handler(Handler handler);
@@ -597,6 +599,7 @@ private:
   std::atomic<bool> is_running_;
   std::atomic<socket_t> svr_sock_;
   std::vector<std::pair<std::string, std::string>> base_dirs_;
+  std::map<std::string, std::string> file_extension_and_mimetype_map_;
   Handler file_request_handler_;
   Handlers get_handlers_;
   Handlers post_handlers_;
@@ -1526,8 +1529,14 @@ inline std::string get_remote_addr(socket_t sock) {
   return std::string();
 }
 
-inline const char *find_content_type(const std::string &path) {
+inline const char *
+find_content_type(const std::string &path,
+                  const std::map<std::string, std::string> &user_data) {
   auto ext = file_extension(path);
+
+  auto it = user_data.find(ext);
+  if (it != user_data.end()) { return it->second.c_str(); }
+
   if (ext == "txt") {
     return "text/plain";
   } else if (ext == "html" || ext == "htm") {
@@ -1550,6 +1559,8 @@ inline const char *find_content_type(const std::string &path) {
     return "application/pdf";
   } else if (ext == "js") {
     return "application/javascript";
+  } else if (ext == "wasm") {
+    return "application/wasm";
   } else if (ext == "xml") {
     return "application/xml";
   } else if (ext == "xhtml") {
@@ -2860,6 +2871,11 @@ inline bool Server::set_base_dir(const char *dir, const char *mount_point) {
   return false;
 }
 
+inline void Server::set_file_extension_and_mimetype_mapping(const char *ext,
+                                                            const char *mime) {
+  file_extension_and_mimetype_map_[ext] = mime;
+}
+
 inline void Server::set_file_request_handler(Handler handler) {
   file_request_handler_ = std::move(handler);
 }
@@ -3178,7 +3194,8 @@ inline bool Server::handle_file_request(Request &req, Response &res) {
 
         if (detail::is_file(path)) {
           detail::read_file(path, res.body);
-          auto type = detail::find_content_type(path);
+          auto type =
+              detail::find_content_type(path, file_extension_and_mimetype_map_);
           if (type) { res.set_header("Content-Type", type); }
           res.status = 200;
           if (file_request_handler_) { file_request_handler_(req, res); }
@@ -3666,7 +3683,7 @@ inline bool Client::write_request(Stream &strm, const Request &req,
   BufferStream bstrm;
 
   // Request line
-  const auto& path = detail::encode_url(req.path);
+  const auto &path = detail::encode_url(req.path);
 
   bstrm.write_format("%s %s HTTP/1.1\r\n", req.method.c_str(), path.c_str());
 
diff --git a/test/test.cc b/test/test.cc
index f8b3caf..0fb5ec7 100644
--- a/test/test.cc
+++ b/test/test.cc
@@ -647,6 +647,7 @@ protected:
   virtual void SetUp() {
     svr_.set_base_dir("./www");
     svr_.set_base_dir("./www2", "/mount");
+    svr_.set_file_extension_and_mimetype_mapping("abcde", "text/abcde");
 
     svr_.Get("/hi",
              [&](const Request & /*req*/, Response &res) {
@@ -1161,6 +1162,14 @@ TEST_F(ServerTest, GetMethodOutOfBaseDirMount2) {
   EXPECT_EQ(404, res->status);
 }
 
+TEST_F(ServerTest, UserDefinedMIMETypeMapping) {
+  auto res = cli_.Get("/dir/test.abcde");
+  ASSERT_TRUE(res != nullptr);
+  EXPECT_EQ(200, res->status);
+  EXPECT_EQ("text/abcde", res->get_header_value("Content-Type"));
+  EXPECT_EQ("abcde\n", res->body);
+}
+
 TEST_F(ServerTest, InvalidBaseDirMount) {
   EXPECT_EQ(false, svr_.set_base_dir("./www3", "invalid_mount_point"));
 }
diff --git a/test/www/dir/test.abcde b/test/www/dir/test.abcde
new file mode 100644
index 0000000..00dedf6
--- /dev/null
+++ b/test/www/dir/test.abcde
@@ -0,0 +1 @@
+abcde