From bdb16c624cdffb9c2ea5aaa13d1b26e9fa82672d Mon Sep 17 00:00:00 2001
From: Sergey Lyubka <valenok@gmail.com>
Date: Thu, 21 Aug 2014 04:09:06 +0100
Subject: [PATCH] Added MG_RECV event and mg_send_file_data() func

---
 examples/big_upload.c | 84 +++++++++++++++++++++----------------------
 mongoose.c            | 41 +++++++++++++++------
 mongoose.h            |  4 +++
 3 files changed, 74 insertions(+), 55 deletions(-)

diff --git a/examples/big_upload.c b/examples/big_upload.c
index af41226b0..a756ddd3a 100644
--- a/examples/big_upload.c
+++ b/examples/big_upload.c
@@ -3,33 +3,27 @@
 #include <stdlib.h>
 #include "mongoose.h"
 
-struct file_data {
-  FILE *fp;
-  const char *data;
-  int data_len;
-  int written;
-};
-
 static int handle_request(struct mg_connection *conn) {
-  const char *data;
-  int data_len;
-  char var_name[100], file_name[100], path[100];
-
   if (strcmp(conn->uri, "/upload") == 0) {
-    if (mg_parse_multipart(conn->content, conn->content_len,
-                           var_name, sizeof(var_name),
-                           file_name, sizeof(file_name),
-                           &data, &data_len) > 0) {
-      struct file_data *p = (struct file_data *) malloc(sizeof(*p));
-      snprintf(path, sizeof(path), "UPLOAD_%s", file_name);
-      p->fp = fopen(path, "wb");
-      p->data = data;
-      p->data_len = data_len;
-      p->written = 0;
-      conn->connection_param = p;
-      mg_send_header(conn, "Content-Type", "text/html");
+    FILE *fp = (FILE *) conn->connection_param;
+    if (fp != NULL) {
+      fwrite(conn->content, 1, conn->content_len, fp); // Write last bits
+      mg_printf(conn, "HTTP/1.1 200 OK\r\n"
+                "Content-Type: text/plain\r\n"
+                "Connection: close\r\n\r\n"
+                "Written %ld of POST data to a temp file:\n\n",
+                (long) ftell(fp));
+
+      // Temp file will be destroyed after fclose(), do something with the
+      // data here -- for example, parse it and extract uploaded files.
+      // As an example, we just echo the whole POST buffer back to the client.
+      rewind(fp);
+      mg_send_file_data(conn, fileno(fp));
+      return MG_MORE;  // Tell Mongoose reply is not completed yet
+    } else {
+      mg_printf_data(conn, "%s", "Had no data to write...");
+      return MG_TRUE; // Tell Mongoose we're done with this request
     }
-    return MG_MORE; // Tell mongoose to keep this connection open
   } else {
     mg_printf_data(conn, "%s",
                    "<html><body>Upload example."
@@ -42,34 +36,36 @@ static int handle_request(struct mg_connection *conn) {
   }
 }
 
-static int handle_poll(struct mg_connection *conn) {
-  struct file_data *p = (struct file_data *) conn->connection_param;
-  if (p != NULL) {
-    // Write no more then 100 bytes in one go
-    int len = p->data_len - p->written;
-    int n = fwrite(p->data + p->written, 1, len > 100 ? 100 : len, p->fp);
-    if (n > 0) {
-      p->written += n;
-      mg_send_data(conn, " ", 1);  // Send something back to wake up select()
-    }
+// Mongoose sends MG_RECV for every received POST chunk.
+// When last POST chunk is received, Mongoose sends MG_REQUEST, then MG_CLOSE.
+static int handle_recv(struct mg_connection *conn) {
+  FILE *fp = (FILE *) conn->connection_param;
 
-    // If everything is written, close the connection
-    if (p->written >= p->data_len) {
-      mg_printf_data(conn, "Written %d bytes.", p->written);
-      fclose(p->fp);
-      free(p);
-      conn->connection_param = NULL;
-      return MG_TRUE; // Tell mongoose to close this connection
-    }
+  // Open temporary file where we going to write data
+  if (fp == NULL && ((conn->connection_param = fp = tmpfile())) == NULL) {
+    return -1;  // Close connection on error
+  }
+
+  // Return number of bytes written to a temporary file: that is how many
+  // bytes we want to discard from the receive buffer
+  return fwrite(conn->content, 1, conn->content_len, fp);
+}
+
+// Make sure we free all allocated resources
+static int handle_close(struct mg_connection *conn) {
+  if (conn->connection_param != NULL) {
+    fclose((FILE *) conn->connection_param);
+    conn->connection_param = NULL;
   }
-  return MG_FALSE;  // Tell mongoose to keep this connection open
+  return MG_TRUE;
 }
 
 static int ev_handler(struct mg_connection *conn, enum mg_event ev) {
   switch (ev) {
     case MG_AUTH:     return MG_TRUE;
     case MG_REQUEST:  return handle_request(conn);
-    case MG_POLL:     return handle_poll(conn);
+    case MG_RECV:     return handle_recv(conn);
+    case MG_CLOSE:    return handle_close(conn);
     default:          return MG_FALSE;
   }
 }
diff --git a/mongoose.c b/mongoose.c
index ff15bcd88..c5a72ffcd 100644
--- a/mongoose.c
+++ b/mongoose.c
@@ -1341,7 +1341,7 @@ struct connection {
   enum endpoint_type endpoint_type;
   char *path_info;
   char *request;
-  int64_t num_bytes_sent; // Total number of bytes sent
+  int64_t num_bytes_recv; // Total number of bytes received
   int64_t cl;             // Reply content length, for Range support
   int request_len;  // Request length, including last \r\n after last header
 };
@@ -3111,18 +3111,22 @@ static void open_file_endpoint(struct connection *conn, const char *path,
     conn->endpoint_type = EP_NONE;
   }
 }
+
+void mg_send_file_data(struct mg_connection *c, int fd) {
+  struct connection *conn = MG_CONN_2_CONN(c);
+  conn->endpoint_type = EP_FILE;
+  conn->endpoint.fd = fd;
+  ns_set_close_on_exec(conn->endpoint.fd);
+}
 #endif  // MONGOOSE_NO_FILESYSTEM
 
 static void call_request_handler_if_data_is_buffered(struct connection *conn) {
-  struct iobuf *loc = &conn->ns_conn->recv_iobuf;
-  struct mg_connection *c = &conn->mg_conn;
-
 #ifndef MONGOOSE_NO_WEBSOCKET
   if (conn->mg_conn.is_websocket) {
     do { } while (deliver_websocket_frame(conn));
   } else
 #endif
-  if ((size_t) loc->len >= c->content_len &&
+  if (conn->num_bytes_recv >= (conn->cl + conn->request_len) &&
       call_request_handler(conn) == MG_FALSE) {
     open_local_endpoint(conn, 1);
   }
@@ -4446,6 +4450,7 @@ static void do_proxy(struct connection *conn) {
 
 static void on_recv_data(struct connection *conn) {
   struct iobuf *io = &conn->ns_conn->recv_iobuf;
+  int n;
 
   if (conn->endpoint_type == EP_PROXY) {
     if (conn->endpoint.nc != NULL) do_proxy(conn);
@@ -4478,6 +4483,14 @@ static void on_recv_data(struct connection *conn) {
   }
 #endif
   if (conn->endpoint_type == EP_USER) {
+    conn->mg_conn.content = io->buf;
+    conn->mg_conn.content_len = io->len;
+    n = call_user(conn, MG_RECV);
+    if (n < 0) {
+      conn->ns_conn->flags |= NSF_FINISHED_SENDING_DATA;
+    } else if ((size_t) n <= io->len) {
+      iobuf_remove(io, n);
+    }
     call_request_handler_if_data_is_buffered(conn);
   }
 #ifndef MONGOOSE_NO_DAV
@@ -4499,7 +4512,7 @@ static void call_http_client_handler(struct connection *conn) {
   }
   iobuf_remove(&conn->ns_conn->recv_iobuf, conn->mg_conn.content_len);
   conn->mg_conn.status_code = 0;
-  conn->cl = conn->num_bytes_sent = conn->request_len = 0;
+  conn->cl = conn->num_bytes_recv = conn->request_len = 0;
   free(conn->request);
   conn->request = NULL;
 }
@@ -4568,12 +4581,12 @@ static void log_access(const struct connection *conn, const char *path) {
   flockfile(fp);
   mg_parse_header(mg_get_header(&conn->mg_conn, "Authorization"), "username",
                   user, sizeof(user));
-  fprintf(fp, "%s - %s [%s] \"%s %s%s%s HTTP/%s\" %d %" INT64_FMT,
+  fprintf(fp, "%s - %s [%s] \"%s %s%s%s HTTP/%s\" %d 0",
           c->remote_ip, user[0] == '\0' ? "-" : user, date,
           c->request_method ? c->request_method : "-",
           c->uri ? c->uri : "-", c->query_string ? "?" : "",
           c->query_string ? c->query_string : "",
-          c->http_version, c->status_code, conn->num_bytes_sent);
+          c->http_version, c->status_code);
   log_header(c, "Referer", fp);
   log_header(c, "User-Agent", fp);
   fputc('\n', fp);
@@ -4619,17 +4632,20 @@ static void close_local_endpoint(struct connection *conn) {
   iobuf_free(&conn->ns_conn->recv_iobuf);
   free(conn->request);
   free(conn->path_info);
+  conn->endpoint.nc = NULL;
+  conn->request = conn->path_info = NULL;
 
   conn->endpoint_type = EP_NONE;
-  conn->cl = conn->num_bytes_sent = conn->request_len = 0;
+  conn->cl = conn->num_bytes_recv = conn->request_len = 0;
   conn->ns_conn->flags &= ~(NSF_FINISHED_SENDING_DATA |
                             NSF_BUFFER_BUT_DONT_SEND | NSF_CLOSE_IMMEDIATELY |
                             MG_HEADERS_SENT | MG_LONG_RUNNING);
+  memset(c, 0, sizeof(*c));
+#if 0
   c->num_headers = c->status_code = c->is_websocket = c->content_len = 0;
-  conn->endpoint.nc = NULL;
   c->request_method = c->uri = c->http_version = c->query_string = NULL;
-  conn->request = conn->path_info = NULL;
   memset(c->http_headers, 0, sizeof(c->http_headers));
+#endif
 
   if (keep_alive) {
     on_recv_data(conn);  // Can call us recursively if pipelining is used
@@ -4981,6 +4997,9 @@ static void mg_ev_handler(struct ns_connection *nc, enum ns_event ev, void *p) {
       break;
 
     case NS_RECV:
+      if (conn != NULL) {
+        conn->num_bytes_recv += * (int *) p;
+      }
       if (nc->flags & NSF_ACCEPTED) {
         on_recv_data(conn);
 #ifndef MONGOOSE_NO_CGI
diff --git a/mongoose.h b/mongoose.h
index fbf30f1cc..7fc0bf0fa 100644
--- a/mongoose.h
+++ b/mongoose.h
@@ -66,6 +66,9 @@ enum mg_event {
   MG_AUTH,        // If callback returns MG_FALSE, authentication fails
   MG_REQUEST,     // If callback returns MG_FALSE, Mongoose continues with req
   MG_REPLY,       // If callback returns MG_FALSE, Mongoose closes connection
+  MG_RECV,        // Mongoose has received POST data chunk.
+                  // Callback should return a number of bytes to discard from
+                  // the receive buffer, or -1 to close the connection.
   MG_CLOSE,       // Connection is closed, callback return value is ignored
   MG_WS_HANDSHAKE,  // New websocket connection, handshake request
   MG_WS_CONNECT,  // New websocket connection established
@@ -112,6 +115,7 @@ size_t mg_websocket_printf(struct mg_connection* conn, int opcode,
                            const char *fmt, ...);
 
 void mg_send_file(struct mg_connection *, const char *path);
+void mg_send_file_data(struct mg_connection *, int fd);
 
 const char *mg_get_header(const struct mg_connection *, const char *name);
 const char *mg_get_mime_type(const char *name, const char *default_mime_type);
-- 
GitLab