Implemented mg_set_http_error_handler()

build/test/unit_test.c

 // Unit test for the mongoose web server.
 
 #define USE_WEBSOCKET
+#define USE_LUA
 
 #ifndef _WIN32
 #define USE_IPV6
@@ -406,6 +407,25 @@ static const char *test_server_param(void) {
   return NULL;
 }
 
+static int error_handler(struct mg_connection *conn) {
+  mg_printf(conn, "error: %d", conn->status_code);
+  return 1;
+}
+
+static const char *test_error_handler(void) {
+  int reply_len;
+  char *reply;
+
+  reply = wget("127.0.0.1", atoi(HTTP_PORT), &reply_len, "%s",
+               "GET /non_exist HTTP/1.0\r\n\r\n");
+  ASSERT(reply != NULL);
+  ASSERT(reply_len == 10);
+  ASSERT(memcmp(reply, "error: 404", 10) == 0);
+  free(reply);
+
+  return NULL;
+}
+
 static void *server_thread(void *param) {
   int i;
   for (i = 0; i < 10; i++) mg_poll_server((struct mg_server *) param, 1);
@@ -419,9 +439,11 @@ static const char *test_server(void) {
   ASSERT(mg_set_option(server, "listening_port", LISTENING_ADDR) == NULL);
   ASSERT(mg_set_option(server, "document_root", ".") == NULL);
   mg_add_uri_handler(server, "/cb1", cb1);
+  mg_set_http_error_handler(server, error_handler);
   mg_start_thread(server_thread, server);
   RUN_TEST(test_regular_file);
   RUN_TEST(test_server_param);
+  RUN_TEST(test_error_handler);
 
   // TODO(lsm): come up with a better way of thread sync
   sleep(1);

mongoose.c

@@ -285,6 +285,7 @@ struct mg_server {
   union socket_address lsa;   // Listening socket address
   struct ll active_connections;
   struct ll uri_handlers;
+  mg_handler_t error_handler;
   char *config_options[NUM_OPTIONS];
   void *server_data;
   void *ssl_ctx;    // SSL context
@@ -653,6 +654,15 @@ static void send_http_error(struct connection *conn, int code,
   va_list ap;
   int body_len, headers_len, match_code;
 
+  conn->mg_conn.status_code = code;
+
+  // Invoke error handler if it is set
+  if (conn->server->error_handler != NULL &&
+      conn->server->error_handler(&conn->mg_conn)) {
+    close_local_endpoint(conn);
+    return;
+  }
+
   // Handle error code rewrites
   while ((rewrites = next_option(rewrites, &a, &b)) != NULL) {
     if ((match_code = atoi(a.ptr)) > 0 && match_code == code) {
@@ -666,7 +676,6 @@ static void send_http_error(struct connection *conn, int code,
     }
   }
 
-  conn->mg_conn.status_code = code;
   body_len = mg_snprintf(body, sizeof(body), "%d %s\n", code, message);
   if (fmt != NULL) {
     body[body_len++] = '\n';
@@ -3876,6 +3885,11 @@ const char *mg_set_option(struct mg_server *server, const char *name,
   return error_msg;
 }
 
+
+void mg_set_http_error_handler(struct mg_server *server, mg_handler_t handler) {
+  server->error_handler = handler;
+}
+
 void mg_set_listening_socket(struct mg_server *server, int sock) {
   if (server->listening_sock != INVALID_SOCKET) {
     closesocket(server->listening_sock);