Home | History | Annotate | Line # | Download | only in uvwget 
      1 #include <assert.h>
      2 #include <stdio.h>
      3 #include <stdlib.h>
      4 #include <uv.h>
      5 #include <curl/curl.h>
      6 
      7 uv_loop_t *loop;
      8 CURLM *curl_handle;
      9 uv_timer_t timeout;
     10 
     11 typedef struct curl_context_s {
     12     uv_poll_t poll_handle;
     13     curl_socket_t sockfd;
     14 } curl_context_t;
     15 
     16 curl_context_t *create_curl_context(curl_socket_t sockfd) {
     17     curl_context_t *context;
     18 
     19     context = (curl_context_t*) malloc(sizeof *context);
     20 
     21     context->sockfd = sockfd;
     22 
     23     int r = uv_poll_init_socket(loop, &context->poll_handle, sockfd);
     24     assert(r == 0);
     25     context->poll_handle.data = context;
     26 
     27     return context;
     28 }
     29 
     30 void curl_close_cb(uv_handle_t *handle) {
     31     curl_context_t *context = (curl_context_t*) handle->data;
     32     free(context);
     33 }
     34 
     35 void destroy_curl_context(curl_context_t *context) {
     36     uv_close((uv_handle_t*) &context->poll_handle, curl_close_cb);
     37 }
     38 
     39 
     40 void add_download(const char *url, int num) {
     41     char filename[50];
     42     sprintf(filename, "%d.download", num);
     43     FILE *file;
     44 
     45     file = fopen(filename, "w");
     46     if (file == NULL) {
     47         fprintf(stderr, "Error opening %s\n", filename);
     48         return;
     49     }
     50 
     51     CURL *handle = curl_easy_init();
     52     curl_easy_setopt(handle, CURLOPT_WRITEDATA, file);
     53     curl_easy_setopt(handle, CURLOPT_URL, url);
     54     curl_multi_add_handle(curl_handle, handle);
     55     fprintf(stderr, "Added download %s -> %s\n", url, filename);
     56 }
     57 
     58 void check_multi_info(void) {
     59     char *done_url;
     60     CURLMsg *message;
     61     int pending;
     62 
     63     while ((message = curl_multi_info_read(curl_handle, &pending))) {
     64         switch (message->msg) {
     65         case CURLMSG_DONE:
     66             curl_easy_getinfo(message->easy_handle, CURLINFO_EFFECTIVE_URL,
     67                             &done_url);
     68             printf("%s DONE\n", done_url);
     69 
     70             curl_multi_remove_handle(curl_handle, message->easy_handle);
     71             curl_easy_cleanup(message->easy_handle);
     72             break;
     73 
     74         default:
     75             fprintf(stderr, "CURLMSG default\n");
     76             abort();
     77         }
     78     }
     79 }
     80 
     81 void curl_perform(uv_poll_t *req, int status, int events) {
     82     uv_timer_stop(&timeout);
     83     int running_handles;
     84     int flags = 0;
     85     if (status < 0)                      flags = CURL_CSELECT_ERR;
     86     if (!status && events & UV_READABLE) flags |= CURL_CSELECT_IN;
     87     if (!status && events & UV_WRITABLE) flags |= CURL_CSELECT_OUT;
     88 
     89     curl_context_t *context;
     90 
     91     context = (curl_context_t*)req;
     92 
     93     curl_multi_socket_action(curl_handle, context->sockfd, flags, &running_handles);
     94     check_multi_info();
     95 }
     96 
     97 void on_timeout(uv_timer_t *req) {
     98     int running_handles;
     99     curl_multi_socket_action(curl_handle, CURL_SOCKET_TIMEOUT, 0, &running_handles);
    100     check_multi_info();
    101 }
    102 
    103 void start_timeout(CURLM *multi, long timeout_ms, void *userp) {
    104     if (timeout_ms <= 0)
    105         timeout_ms = 1; /* 0 means directly call socket_action, but we'll do it in a bit */
    106     uv_timer_start(&timeout, on_timeout, timeout_ms, 0);
    107 }
    108 
    109 int handle_socket(CURL *easy, curl_socket_t s, int action, void *userp, void *socketp) {
    110     curl_context_t *curl_context;
    111     if (action == CURL_POLL_IN || action == CURL_POLL_OUT) {
    112         if (socketp) {
    113             curl_context = (curl_context_t*) socketp;
    114         }
    115         else {
    116             curl_context = create_curl_context(s);
    117             curl_multi_assign(curl_handle, s, (void *) curl_context);
    118         }
    119     }
    120 
    121     switch (action) {
    122         case CURL_POLL_IN:
    123             uv_poll_start(&curl_context->poll_handle, UV_READABLE, curl_perform);
    124             break;
    125         case CURL_POLL_OUT:
    126             uv_poll_start(&curl_context->poll_handle, UV_WRITABLE, curl_perform);
    127             break;
    128         case CURL_POLL_REMOVE:
    129             if (socketp) {
    130                 uv_poll_stop(&((curl_context_t*)socketp)->poll_handle);
    131                 destroy_curl_context((curl_context_t*) socketp);
    132                 curl_multi_assign(curl_handle, s, NULL);
    133             }
    134             break;
    135         default:
    136             abort();
    137     }
    138 
    139     return 0;
    140 }
    141 
    142 int main(int argc, char **argv) {
    143     loop = uv_default_loop();
    144 
    145     if (argc <= 1)
    146         return 0;
    147 
    148     if (curl_global_init(CURL_GLOBAL_ALL)) {
    149         fprintf(stderr, "Could not init cURL\n");
    150         return 1;
    151     }
    152 
    153     uv_timer_init(loop, &timeout);
    154 
    155     curl_handle = curl_multi_init();
    156     curl_multi_setopt(curl_handle, CURLMOPT_SOCKETFUNCTION, handle_socket);
    157     curl_multi_setopt(curl_handle, CURLMOPT_TIMERFUNCTION, start_timeout);
    158 
    159     while (argc-- > 1) {
    160         add_download(argv[argc], argc);
    161     }
    162 
    163     uv_run(loop, UV_RUN_DEFAULT);
    164     curl_multi_cleanup(curl_handle);
    165     return 0;
    166 }
    167