Browse Source

Fixed mg_fetch() by passing a buffer, placeholder for the reply headers

Sergey Lyubka 13 năm trước cách đây
mục cha
commit
a9bb7d78ab
3 tập tin đã thay đổi với 28 bổ sung20 xóa
  1. 15 10
      mongoose.c
  2. 2 1
      mongoose.h
  3. 11 9
      test/unit_test.c

+ 15 - 10
mongoose.c

@@ -3903,10 +3903,10 @@ struct mg_connection *mg_connect(struct mg_context *ctx,
 }
 
 FILE *mg_fetch(struct mg_context *ctx, const char *url, const char *path,
-               struct mg_request_info *ri) {
+               char *buf, size_t buf_len, struct mg_request_info *ri) {
   struct mg_connection *newconn;
   int n, req_length, data_length, port;
-  char host[1025], proto[10], buf[16384];
+  char host[1025], proto[10], buf2[BUFSIZ];
   FILE *fp = NULL;
 
   if (sscanf(url, "%9[htps]://%1024[^:]:%d/%n", proto, host, &port, &n) == 3) {
@@ -3924,7 +3924,7 @@ FILE *mg_fetch(struct mg_context *ctx, const char *url, const char *path,
     mg_printf(newconn, "GET /%s HTTP/1.0\r\n\r\n", url + n);
     data_length = 0;
     req_length = read_request(NULL, newconn->client.sock,
-                              newconn->ssl, buf, sizeof(buf), &data_length);
+                              newconn->ssl, buf, buf_len, &data_length);
     if (req_length <= 0) {
       cry(fc(ctx), "%s(%s): invalid HTTP reply", __func__, url);
     } else if (parse_http_response(buf, req_length, ri) <= 0) {
@@ -3933,15 +3933,22 @@ FILE *mg_fetch(struct mg_context *ctx, const char *url, const char *path,
       cry(fc(ctx), "%s: fopen(%s): %s", __func__, path, strerror(ERRNO));
     } else {
       data_length -= req_length;
-      memmove(buf, buf + req_length, data_length);
-      do {
-        if (fwrite(buf, 1, data_length, fp) != (size_t) data_length) {
+      // Write chunk of data that may be in the user's buffer
+      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) {
+        if (fwrite(buf2, 1, data_length, fp) != (size_t) data_length) {
+          cry(fc(ctx), "%s: fwrite(%s): %s", __func__, path, strerror(ERRNO));
           fclose(fp);
           fp = NULL;
           break;
         }
-        data_length = mg_read(newconn, buf, sizeof(buf));
-      } while (data_length > 0);
+      }
     }
     mg_close_connection(newconn);
   }
@@ -3949,8 +3956,6 @@ FILE *mg_fetch(struct mg_context *ctx, const char *url, const char *path,
   return fp;
 }
 
-
-
 static void discard_current_request_from_buffer(struct mg_connection *conn) {
   char *buffered;
   int buffered_len, body_len;

+ 2 - 1
mongoose.h

@@ -234,12 +234,13 @@ void mg_close_connection(struct mg_connection *conn);
 //   url: URL to download
 //   path: file name where to save the data
 //   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
 FILE *mg_fetch(struct mg_context *ctx, const char *url, const char *path,
-               struct mg_request_info *request_info);
+               char *buf, size_t buf_len, struct mg_request_info *request_info);
 
 
 // Convenience function -- create detached thread.

+ 11 - 9
test/unit_test.c

@@ -150,7 +150,7 @@ static void test_mg_fetch(void) {
     "listening_ports", "33796",
     NULL,
   };
-  char buf[1000];
+  char buf[2000], buf2[2000];
   int length;
   struct mg_context *ctx;
   struct mg_request_info ri;
@@ -160,26 +160,28 @@ static void test_mg_fetch(void) {
   ASSERT((ctx = mg_start(event_handler, NULL, options)) != NULL);
 
   // Failed fetch, pass invalid URL
-  ASSERT(mg_fetch(ctx, "localhost", tmp_file, &ri) == NULL);
-  ASSERT(mg_fetch(ctx, "localhost:33796", tmp_file, &ri) == NULL);
-  ASSERT(mg_fetch(ctx, "http://$$$.$$$", tmp_file, &ri) == NULL);
+  ASSERT(mg_fetch(ctx, "localhost", tmp_file, buf, sizeof(buf), &ri) == NULL);
+  ASSERT(mg_fetch(ctx, "localhost:33796", tmp_file,
+                  buf, sizeof(buf), &ri) == NULL);
+  ASSERT(mg_fetch(ctx, "http://$$$.$$$", tmp_file,
+                  buf, sizeof(buf), &ri) == NULL);
 
   // Failed fetch, pass invalid file name
   ASSERT(mg_fetch(ctx, "http://localhost:33796/data",
-                  "/this/file/must/not/exist/ever", &ri) == NULL);
+                  "/this/file/must/not/exist/ever",
+                  buf, sizeof(buf), &ri) == NULL);
 
   // Successful fetch
   ASSERT((fp = mg_fetch(ctx, "http://localhost:33796/data",
-                        tmp_file, &ri)) != NULL);
+                        tmp_file, buf, sizeof(buf), &ri)) != NULL);
   ASSERT(ri.num_headers == 2);
-  printf("%s: [%s]\n", __func__, ri.request_method);
   ASSERT(!strcmp(ri.request_method, "HTTP/1.1"));
   ASSERT(!strcmp(ri.uri, "200"));
   ASSERT(!strcmp(ri.http_version, "OK"));
   ASSERT((length = ftell(fp)) == (int) strlen(fetch_data));
   fseek(fp, 0, SEEK_SET);
-  ASSERT(fread(buf, 1, length, fp) == (size_t) length);
-  ASSERT(memcmp(buf, fetch_data, length) == 0);
+  ASSERT(fread(buf2, 1, length, fp) == (size_t) length);
+  ASSERT(memcmp(buf2, fetch_data, length) == 0);
 
   remove(tmp_file);
   mg_stop(ctx);