3 files changed, 34 insertions, 7 deletions

diff --git a/src/extractors/haokan.c b/src/extractors/haokan.c
index 204e539..5659018 100644
--- a/src/extractors/haokan.c
+++ b/src/extractors/haokan.c
@@ -25,7 +25,7 @@ void haokan_extract(Options *options) {
       DEBUG_PRINT("videoURL: %s\n", videoURL);
 
       char *ct = NULL;
-      get_info(videoURL, NULL, &ct);
+      get_info(videoURL, NULL, &ct, NULL);
       if (ct == NULL) {
         goto end;
       }
diff --git a/src/process_url.c b/src/process_url.c
index 35abf46..77fd96d 100644
--- a/src/process_url.c
+++ b/src/process_url.c
@@ -244,7 +244,8 @@ static int parse_url(const char *URL, const char *outdir, char *fn) {
   return 0;
 }
 
-bool get_info(const char *URL, long *psize, char **p_content_type) {
+bool get_info(const char *URL, long *psize, char **p_content_type,
+              char **p_cookie) {
   CURL *curl;
   long resp_code;
   bool support_range = false;
@@ -268,13 +269,22 @@ bool get_info(const char *URL, long *psize, char **p_content_type) {
     *p_content_type = malloc(strlen(ct) + 1);
     strcpy(*p_content_type, ct);
   }
+  CURLHcode rh =
+      curl_easy_header(curl, "Set-Cookie", 0, CURLH_HEADER, -1, &pch);
+  if (logerr(rh)) {
+    goto end;
+  }
+  DEBUG_PRINT("Set-Cookie: %s\n", pch->value);
+  *p_cookie = malloc(strlen(pch->value) + 1);
+  strcpy(*p_cookie, pch->value);
+
   r = curl_easy_getinfo(curl, CURLINFO_CONTENT_LENGTH_DOWNLOAD_T,
                         (curl_off_t *)psize);
   if (logerr(r)) {
     goto end;
   }
-  CURLHcode rh =
-      curl_easy_header(curl, "Accept-Ranges", 0, CURLH_HEADER, -1, &pch);
+
+  rh = curl_easy_header(curl, "Accept-Ranges", 0, CURLH_HEADER, -1, &pch);
   if (logerr(rh) || strcmp(pch->value, "bytes")) {
     goto end;
   }
@@ -392,7 +402,7 @@ static int download(curl_conf_t *curl_c) {
 
   static thrd_info_t thrd_info[MAX_THREAD] = {0};
 
-  bool support_range = get_info(curl_c->URL, &cl, NULL);
+  bool support_range = get_info(curl_c->URL, &cl, NULL, NULL);
   DEBUG_PRINT("Size: %ld bytes.\n", cl);
   if (support_range && cl > 0L) {
     curl_c->dltotal = cl;
@@ -449,11 +459,24 @@ static char *callback_struct_convert_fullpath(char *filename) {
   return tmp;
 }
 
+void add_cookie(char *cookie) {
+  if (cookie_g) {
+    char *tmp = malloc(strlen(cookie_g) + strlen(cookie) + 3);
+    sprintf(tmp, "%s; %s", cookie_g, cookie);
+    free_and_nullify(cookie_g);
+    cookie_g = tmp;
+  } else {
+    cookie_g = cookie;
+  }
+}
+
+void set_referer(char *referer) { referer_g = referer; }
+
 void curl_init(char *cookie) {
   curl_global_init(CURL_GLOBAL_ALL);
   h = curl_url();
+  add_cookie(cookie);
   dl_queue = create_queue();
-  cookie_g = cookie;
   mtx_init(&mtx, mtx_plain);
   cnd_init(&cnd);
 }
diff --git a/src/process_url.h b/src/process_url.h
index 760ee49..00f7a74 100644
--- a/src/process_url.h
+++ b/src/process_url.h
@@ -58,9 +58,13 @@ void curl_cleanup(status_t *);
 
 void poll_status(status_t *);
 
+void add_cookie(char *cookie);
+
+void set_referer(char *referer);
+
 int get(const char *, char **);
 
-bool get_info(const char *URL, long *psize, char **p_content_type);
+bool get_info(const char *URL, long *psize, char **p_content_type, char **p_cookie);
 
 void add_url(const char *, const char *, const char *, const char *, callback_t,
              callback_struct_t *);