SimpleWeb::status_code(const std::string &) has been made more general and accepts custom status codes as well

1 files changed, 12 insertions, 3 deletions

diff --git a/status_code.hpp b/status_code.hpp
index 0792727..9f9ecc5 100644
--- a/status_code.hpp
+++ b/status_code.hpp
@@ -1,6 +1,7 @@
 #ifndef SIMPLE_WEB_STATUS_CODE_HPP
 #define SIMPLE_WEB_STATUS_CODE_HPP
 
+#include <cstdlib>
 #include <map>
 #include <string>
 #include <unordered_map>
@@ -140,17 +141,25 @@ namespace SimpleWeb {
   }
 
   inline StatusCode status_code(const std::string &status_code_string) noexcept {
+    if(status_code_string.size() < 3)
+      return StatusCode::unknown;
+
+    auto number = status_code_string.substr(0, 3);
+    if(number[0] < '0' || number[0] > '9' || number[1] < '0' || number[1] > '9' || number[2] < '0' || number[2] > '9')
+      return StatusCode::unknown;
+
     class StringToStatusCode : public std::unordered_map<std::string, SimpleWeb::StatusCode> {
     public:
       StringToStatusCode() {
         for(auto &status_code : status_code_strings())
-          emplace(status_code.second, status_code.first);
+          emplace(status_code.second.substr(0, 3), status_code.first);
       }
     };
     static StringToStatusCode string_to_status_code;
-    auto pos = string_to_status_code.find(status_code_string);
+
+    auto pos = string_to_status_code.find(number);
     if(pos == string_to_status_code.end())
-      return StatusCode::unknown;
+      return static_cast<StatusCode>(atoi(number.c_str()));
     return pos->second;
   }