main.c 4.5 KB

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