From 06e3d76dc93f3213bf480f9c70038e01effc9498 Mon Sep 17 00:00:00 2001
From: Jianhui Zhao <zhaojh329@gmail.com>
Date: Sat, 26 Dec 2020 17:30:28 +0800
Subject: [PATCH] handler: add new event: UH_EV_HEAD_COMPLETE

Signed-off-by: Jianhui Zhao <zhaojh329@gmail.com>
---
 example/example.c | 18 ++++++++++++++++--
 src/connection.c  |  5 +++++
 src/uhttpd.h      |  1 +
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/example/example.c b/example/example.c
index cefc201..2befc6e 100644
--- a/example/example.c
+++ b/example/example.c
@@ -64,7 +64,21 @@ static void upload_handler(struct uh_connection *conn, int event)
 {
     static int fd = -1;
 
-    if (event == UH_EV_BODY) {
+    if (event == UH_EV_HEAD_COMPLETE) {
+        struct uh_str str = conn->get_header(conn, "Content-Length");
+        int content_length;
+        char buf[128];
+
+        sprintf(buf, "%.*s\n", (int)str.len, str.p);
+
+        content_length = atoi(buf);
+
+        if (content_length > 1024 * 1024 * 1024) {
+            conn->error(conn, HTTP_STATUS_INTERNAL_SERVER_ERROR, "Too big");
+            return;
+        }
+
+    } if (event == UH_EV_BODY) {
         struct uh_str body = conn->extract_body(conn);
 
         if (fd < 0) {
@@ -80,7 +94,7 @@ static void upload_handler(struct uh_connection *conn, int event)
             close(fd);
             return;
         }
-    } else {
+    } else if (event == UH_EV_COMPLETE) {
         struct stat st;
         size_t size = 0;
 
diff --git a/src/connection.c b/src/connection.c
index 37d6f0d..b24afec 100644
--- a/src/connection.c
+++ b/src/connection.c
@@ -389,6 +389,11 @@ static int on_headers_complete(struct http_parser *parser)
         return -1;
     }
 
+    conn->handler(conn, UH_EV_HEAD_COMPLETE);
+
+    if (conn->flags & CONN_F_SEND_AND_CLOSE)
+        return -1;
+
     return 0;
 }
 
diff --git a/src/uhttpd.h b/src/uhttpd.h
index 12ca32d..db1659b 100644
--- a/src/uhttpd.h
+++ b/src/uhttpd.h
@@ -45,6 +45,7 @@ struct uh_plugin {
 };
 
 enum {
+    UH_EV_HEAD_COMPLETE,
     UH_EV_BODY,
     UH_EV_COMPLETE
 };