瀏覽代碼

Setting Host header in mg_fetch. Added test for fetching large file.

Sergey Lyubka 13 年之前
父節點
當前提交
df7da95cb6
共有 3 個文件被更改,包括 18 次插入6 次删除
  1. 6 4
      mongoose.c
  2. 3 2
      mongoose.h
  3. 9 0
      test/unit_test.c

+ 6 - 4
mongoose.c

@@ -3921,7 +3921,7 @@ FILE *mg_fetch(struct mg_context *ctx, const char *url, const char *path,
                             !strcmp(proto, "https"))) == NULL) {
     cry(fc(ctx), "%s: mg_connect(%s): %s", __func__, url, strerror(ERRNO));
   } else {
-    mg_printf(newconn, "GET /%s HTTP/1.0\r\n\r\n", url + n);
+    mg_printf(newconn, "GET /%s HTTP/1.0\r\nHost: %s\r\n\r\n", url + n, host);
     data_length = 0;
     req_length = read_request(NULL, newconn->client.sock,
                               newconn->ssl, buf, buf_len, &data_length);
@@ -3932,16 +3932,18 @@ FILE *mg_fetch(struct mg_context *ctx, const char *url, const char *path,
     } else if ((fp = fopen(path, "w+b")) == NULL) {
       cry(fc(ctx), "%s: fopen(%s): %s", __func__, path, strerror(ERRNO));
     } else {
-      data_length -= req_length;
       // Write chunk of data that may be in the user's buffer
+      data_length -= req_length;
       if (data_length > 0 &&
         fwrite(buf + req_length, 1, data_length, fp) != (size_t) data_length) {
         cry(fc(ctx), "%s: fwrite(%s): %s", __func__, path, strerror(ERRNO));
         fclose(fp);
         fp = NULL;
       }
-      // Read the rest of the response and write it to the file
-      while (fp && (data_length = mg_read(newconn, buf2, sizeof(buf2))) > 0) {
+      // Read the rest of the response and write it to the file. Do not use
+      // mg_read() cause we didn't set newconn->content_len properly.
+      while (fp && (data_length = pull(NULL, newconn->client.sock, newconn->ssl,
+                                       buf2, sizeof(buf2))) > 0) {
         if (fwrite(buf2, 1, data_length, fp) != (size_t) data_length) {
           cry(fc(ctx), "%s: fwrite(%s): %s", __func__, path, strerror(ERRNO));
           fclose(fp);

+ 3 - 2
mongoose.h

@@ -236,9 +236,10 @@ void mg_close_connection(struct mg_connection *conn);
 //   request_info: pointer to a structure that will hold parsed reply headers
 //   buf, bul_len: a buffer for the reply headers
 // Return:
-//   On success, opened file stream to the downloaded contents. The stream
-//   is positioned to the end of the file.
 //   On error, NULL
+//   On success, opened file stream to the downloaded contents. The stream
+//   is positioned to the end of the file. It is a user responsibility
+//   to fclose() opened file stream.
 FILE *mg_fetch(struct mg_context *ctx, const char *url, const char *path,
                char *buf, size_t buf_len, struct mg_request_info *request_info);
 

+ 9 - 0
test/unit_test.c

@@ -155,6 +155,7 @@ static void test_mg_fetch(void) {
   struct mg_context *ctx;
   struct mg_request_info ri;
   const char *tmp_file = "temporary_file_name_for_unit_test.txt";
+  struct mgstat st;
   FILE *fp;
 
   ASSERT((ctx = mg_start(event_handler, NULL, options)) != NULL);
@@ -182,6 +183,14 @@ static void test_mg_fetch(void) {
   fseek(fp, 0, SEEK_SET);
   ASSERT(fread(buf2, 1, length, fp) == (size_t) length);
   ASSERT(memcmp(buf2, fetch_data, length) == 0);
+  fclose(fp);
+
+  // Fetch big file, mongoose.c
+  ASSERT((fp = mg_fetch(ctx, "http://localhost:33796/mongoose.c",
+                        tmp_file, buf, sizeof(buf), &ri)) != NULL);
+  ASSERT(mg_stat("mongoose.c", &st) == 0);
+  ASSERT(st.size == ftell(fp));
+  ASSERT(!strcmp(ri.request_method, "HTTP/1.1"));
 
   remove(tmp_file);
   mg_stop(ctx);