diff --git a/mongoose.c b/mongoose.c index 30eac158..1e174598 100644 --- a/mongoose.c +++ b/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; diff --git a/mongoose.h b/mongoose.h index c25753af..0eab0bf6 100644 --- a/mongoose.h +++ b/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. diff --git a/test/unit_test.c b/test/unit_test.c index 6f1e3631..7cbedacb 100644 --- a/test/unit_test.c +++ b/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);