Improving CurlTask a bit.

4 files changed, 145 insertions, 9 deletions

diff --git a/includes/BSHA1.h b/includes/BSHA1.h
index 55710b2..b58a890 100644
--- a/includes/BSHA1.h
+++ b/includes/BSHA1.h
@@ -9,7 +9,7 @@ class SHA1 {
   public:
       SHA1() { reset(); }
     void reset() { sha1_init(&m_state); }
-    void update(const uint8_t* data, const size_t len) { sha1_process(&m_state, data, len); }
+    void update(const uint8_t * data, const size_t len) { sha1_process(&m_state, data, len); }
     void final(uint8_t * digest) { sha1_done(&m_state, digest); }
 
     enum { DIGEST_SIZE = 20 };
diff --git a/includes/BString.h b/includes/BString.h
index a321a30..0a5fceb 100644
--- a/includes/BString.h
+++ b/includes/BString.h
@@ -126,6 +126,16 @@ class String : private std::string {
     char & operator[](ssize_t i) { if (i < 0) i = strlen() + i; return at(i); }
 
     void reserve(size_t s) { std::string::reserve(s); }
+
+    size_t find_first_of(const String & set, size_t offset = 0) const { return std::string::find_first_of(set, offset); }
+    size_t find_last_of(const String & set, size_t offset = 0) const { return std::string::find_last_of(set, offset); }
+    size_t find_first_not_of(const String & set, size_t offset = 0) const { return std::string::find_first_not_of(set, offset); }
+    size_t find_last_not_of(const String & set, size_t offset = 0) const { return std::string::find_last_not_of(set, offset); }
+
+    bool isEmpty() { return std::string::empty(); }
+
+    using std::string::npos;
+    using std::string::size_type;
 };
 
 };
diff --git a/includes/CurlTask.h b/includes/CurlTask.h
index 41243ec..f1aba65 100644
--- a/includes/CurlTask.h
+++ b/includes/CurlTask.h
@@ -9,18 +9,28 @@ namespace Balau {
 class CurlTask : public StacklessTask {
   public:
       CurlTask();
-      ~CurlTask();
+      virtual ~CurlTask();
     friend class TaskMan;
+  
+    static String percentEncode(const String & src);
+    static String percentDecode(const String & src);
+    static std::vector<String> tokenize(const String & str, const String & delimiters = "&", bool trimEmpty = true);
+
+    virtual void prepareRequest() { }
+
   protected:
     CURL * m_curlHandle;
-    void registerCurlHandle() { getTaskMan()->registerCurlHandle(this); }
+    void registerCurlHandle() { prepareRequest(); getTaskMan()->registerCurlHandle(this); }
     void unregisterCurlHandle() { getTaskMan()->unregisterCurlHandle(this); }
+
   private:
-    virtual size_t writeFunction(char * ptr, size_t size, size_t nmemb) { return size * nmemb; }
-    virtual size_t readFunction(void * ptr, size_t size, size_t nmemb) { return CURL_READFUNC_ABORT; }
+    virtual size_t headerFunction(char * ptr, size_t size) { return size; }
+    virtual size_t writeFunction(char * ptr, size_t size) { return size; }
+    virtual size_t readFunction(void * ptr, size_t size) { return CURL_READFUNC_ABORT; }
     virtual int    debugFunction(curl_infotype info, char * str, size_t str_len) { return 0; }
     virtual void   curlDone(CURLcode result) { }
 
+    static  size_t headerFunctionStatic(char * ptr, size_t size, size_t nmemb, void * userdata);
     static  size_t writeFunctionStatic(char * ptr, size_t size, size_t nmemb, void * userdata);
     static  size_t readFunctionStatic(void * ptr, size_t size, size_t nmemb, void * userdata);
     static  int    debugFunctionStatic(CURL * easy, curl_infotype info, char * str, size_t str_len, void * userdata);
@@ -42,10 +52,22 @@ class DownloadTask : public CurlTask {
     virtual const char * getName() const override { return m_name.to_charp(); }
     virtual void Do() override;
     virtual void curlDone(CURLcode result) override;
-    virtual size_t writeFunction(char * ptr, size_t size, size_t nmemb) override { m_data += ptr; return size * nmemb; }
+    virtual void downloadDone() { }
+    virtual size_t writeFunction(char * ptr, size_t size) override { m_data += ptr; return size; }
     String m_name;
     Events::Custom m_evt;
     bool m_done = false;
 };
 
+class HttpDownloadTask : public DownloadTask {
+  public:
+      HttpDownloadTask(const String & url) : DownloadTask(url) { }
+      virtual ~HttpDownloadTask();
+    void addHeader(const String & header);
+    virtual void prepareRequest() override;
+
+  private:
+    struct curl_slist * m_headers = NULL;
+};
+
 };
diff --git a/src/CurlTask.cc b/src/CurlTask.cc
index c43210f..2a792f9 100644
--- a/src/CurlTask.cc
+++ b/src/CurlTask.cc
@@ -2,6 +2,8 @@
 
 Balau::CurlTask::CurlTask() {
     m_curlHandle = curl_easy_init();
+    curl_easy_setopt(m_curlHandle, CURLOPT_HEADERFUNCTION, reinterpret_cast<curl_write_callback>(headerFunctionStatic));
+    curl_easy_setopt(m_curlHandle, CURLOPT_HEADERDATA, this);
     curl_easy_setopt(m_curlHandle, CURLOPT_WRITEFUNCTION, reinterpret_cast<curl_write_callback>(writeFunctionStatic));
     curl_easy_setopt(m_curlHandle, CURLOPT_WRITEDATA, this);
     curl_easy_setopt(m_curlHandle, CURLOPT_READFUNCTION, reinterpret_cast<curl_read_callback>(readFunctionStatic));
@@ -15,14 +17,19 @@ Balau::CurlTask::~CurlTask() {
     curl_easy_cleanup(m_curlHandle);
 }
 
+size_t Balau::CurlTask::headerFunctionStatic(char * ptr, size_t size, size_t nmemb, void * userdata) {
+    CurlTask * curlTask = (CurlTask *) userdata;
+    return curlTask->headerFunction(ptr, size * nmemb);
+}
+
 size_t Balau::CurlTask::writeFunctionStatic(char * ptr, size_t size, size_t nmemb, void * userdata) {
     CurlTask * curlTask = (CurlTask *) userdata;
-    return curlTask->writeFunction(ptr, size, nmemb);
+    return curlTask->writeFunction(ptr, size * nmemb);
 }
 
 size_t Balau::CurlTask::readFunctionStatic(void * ptr, size_t size, size_t nmemb, void * userdata) {
     CurlTask * curlTask = (CurlTask *) userdata;
-    return curlTask->readFunction(ptr, size, nmemb);
+    return curlTask->readFunction(ptr, size * nmemb);
 }
 
 int Balau::CurlTask::debugFunctionStatic(CURL * easy, curl_infotype info, char * str, size_t str_len, void * userdata) {
@@ -31,14 +38,97 @@ int Balau::CurlTask::debugFunctionStatic(CURL * easy, curl_infotype info, char *
     return curlTask->debugFunction(info, str, str_len);
 }
 
+Balau::String Balau::CurlTask::percentEncode(const String & src) {
+    const size_t size = src.strlen();
+    String ret;
+    ret.reserve(size);
+
+    static char toHex[] = "0123456789ABCDEF";
+
+    for (size_t i = 0; i < size; i++) {
+        char c = src[i];
+        if (((c >= '0') && (c <= '9')) || ((c >= 'A') && (c <= 'Z')) || ((c >= 'a') && (c <= 'z')) || (c == '-') || (c == '.') || (c == '_') || (c == '~')) {
+            ret += c;
+        } else {
+            ret += '%';
+            ret += toHex[c >> 4];
+            ret += toHex[c & 15];
+        }
+    }
+
+    return ret;
+}
+
+Balau::String Balau::CurlTask::percentDecode(const String & src) {
+    const size_t size = src.strlen();
+    String ret;
+    ret.reserve(size);
+
+    for (size_t i = 0; i < size; i++) {
+        char c = src[i];
+        if (((c >= '0') && (c <= '9')) || ((c >= 'A') && (c <= 'Z')) || ((c >= 'a') && (c <= 'z')) || (c == '-') || (c == '.') || (c == '_') || (c == '~')) {
+            ret += c;
+        } else if ((c == '%') && ((i + 2) < size)) {
+            char h1 = src[i + 1];
+            char h2 = src[i + 2];
+            if ((h1 >= '0') && (h1 <= '9')) {
+                c = h1 - '0';
+            } else if ((h1 >= 'A') && (h1 <= 'F')) {
+                c = h1 - 'A' + 10;
+            } else {
+                // invalid
+                return ret;
+            }
+            c <<= 4;
+            if ((h2 >= '0') && (h2 <= '9')) {
+                c |= h2 - '0';
+            } else if ((h2 >= 'A') && (h2 <= 'F')) {
+                c |= h2 - 'A' + 10;
+            } else {
+                // invalid
+                return ret;
+            }
+            i += 2;
+        } else {
+            // invalid
+            return ret;
+        }
+    }
+
+    return ret;
+}
+
+std::vector<Balau::String> Balau::CurlTask::tokenize(const String & str, const String & delimiters, bool trimEmpty) {
+    std::vector<String> tokens;
+    size_t pos, lastPos = 0;
+    for (;;) {
+        pos = str.find_first_of(delimiters, lastPos);
+        if (pos == String::npos) {
+            pos = str.strlen();
+
+            if ((pos != lastPos) || !trimEmpty)
+                tokens.push_back(str.extract(lastPos, pos));
+
+            return tokens;
+        } else {
+            if ((pos != lastPos) || !trimEmpty)
+                tokens.push_back(str.extract(lastPos, pos));
+        }
+
+        lastPos = pos + 1;
+    }
+}
+
 Balau::DownloadTask::DownloadTask(const Balau::String & url) {
     curl_easy_setopt(m_curlHandle, CURLOPT_URL, url.to_charp());
     m_name.set("DownloadTask(%s)", url.to_charp());
 }
 
 void Balau::DownloadTask::Do() {
-    if (m_state)
+    if (m_state) {
+        downloadDone();
         return;
+    }
 
     m_state = 1;
     registerCurlHandle();
@@ -52,3 +142,17 @@ void Balau::DownloadTask::curlDone(CURLcode result) {
     m_evt.doSignal();
     m_done = true;
 }
+
+Balau::HttpDownloadTask::~HttpDownloadTask() {
+    curl_slist_free_all(m_headers);
+    m_headers = NULL;
+}
+
+void Balau::HttpDownloadTask::addHeader(const String & header) {
+    m_headers = curl_slist_append(m_headers, header.to_charp());
+}
+
+void Balau::HttpDownloadTask::prepareRequest() {
+    if (m_headers)
+        curl_easy_setopt(m_curlHandle, CURLOPT_HTTPHEADER, m_headers);
+}