From 5a04ea3cbd29e402688159d2a9d27b6cffc6a9bb Mon Sep 17 00:00:00 2001
From: Turiiya <34311583+tobealive@users.noreply.github.com>
Date: Sat, 1 Jul 2023 22:05:06 +0200
Subject: [PATCH] feat: add `custom_headers` (#11)

---
 src/_instruction_download.v |  4 ++--
 src/_instructions_get.v     |  4 ++--
 src/_instructions_head.v    |  2 +-
 src/_instructions_header.v  | 17 ++++++++---------
 src/_instructions_post.v    |  2 +-
 src/_state_common.v         |  2 +-
 src/_state_header.v         |  6 ++++++
 src/_tests_post_test.v      | 12 ++++++++++--
 8 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/src/_instruction_download.v b/src/_instruction_download.v
index cf923f0..e550e2e 100644
--- a/src/_instruction_download.v
+++ b/src/_instruction_download.v
@@ -5,7 +5,7 @@ import vibe.curl
 
 fn (req Request) download_file_(url string, file_path string) !Response {
 	h := curl.easy_init() or { return http_error(.easy_init, none) }
-	header := set_header(req.headers, h)
+	header := set_header(h, common: req.headers, custom: req.custom_headers)
 	defer {
 		curl.easy_cleanup(h)
 		curl.slist_free_all(header)
@@ -35,7 +35,7 @@ fn (req Request) download_file_(url string, file_path string) !Response {
 
 fn (req Request) download_file_with_progress_(url string, file_path string, mut dl Download) !Response {
 	h := curl.easy_init() or { return http_error(.easy_init, none) }
-	header := set_header(req.headers, h)
+	header := set_header(h, common: req.headers, custom: req.custom_headers)
 	defer {
 		curl.easy_cleanup(h)
 		curl.slist_free_all(header)
diff --git a/src/_instructions_get.v b/src/_instructions_get.v
index 0d5fd7e..919ca1e 100644
--- a/src/_instructions_get.v
+++ b/src/_instructions_get.v
@@ -5,7 +5,7 @@ import vibe.curl
 fn (req Request) get_(url string) !Response {
 	// Curl handle
 	h := curl.easy_init() or { return http_error(.easy_init, none) }
-	header := set_header(req.headers, h)
+	header := set_header(h, common: req.headers, custom: req.custom_headers)
 	defer {
 		curl.easy_cleanup(h)
 		curl.slist_free_all(header)
@@ -31,7 +31,7 @@ fn (req Request) get_(url string) !Response {
 
 fn (req Request) get_slice_(url string, start usize, max_size_ ?usize) !Response {
 	h := curl.easy_init() or { return http_error(.easy_init, none) }
-	header := set_header(req.headers, h)
+	header := set_header(h, common: req.headers, custom: req.custom_headers)
 	defer {
 		curl.easy_cleanup(h)
 		curl.slist_free_all(header)
diff --git a/src/_instructions_head.v b/src/_instructions_head.v
index e78f8fd..dfdfefc 100644
--- a/src/_instructions_head.v
+++ b/src/_instructions_head.v
@@ -4,7 +4,7 @@ import vibe.curl
 
 fn (req Request) head_(url string) !Response {
 	h := curl.easy_init() or { return http_error(.easy_init, none) }
-	header := set_header(req.headers, h)
+	header := set_header(h, common: req.headers, custom: req.custom_headers)
 	defer {
 		curl.easy_cleanup(h)
 		curl.slist_free_all(header)
diff --git a/src/_instructions_header.v b/src/_instructions_header.v
index c0db01c..ca0283c 100644
--- a/src/_instructions_header.v
+++ b/src/_instructions_header.v
@@ -2,27 +2,26 @@ module vibe
 
 import vibe.curl
 
-fn set_header(header_map map[HttpHeader]string, handle &C.CURL) &HeaderList {
+fn set_header(handle &C.CURL, headers HttpHeaders) &HeaderList {
 	mut list := &HeaderList(unsafe { nil })
 
-	// Default header
-	if header_map.len == 0 {
+	// Set default header and return if no headers were specified
+	if headers.common.len == 0 && headers.custom.len == 0 {
 		curl.easy_setopt(handle, .useragent, '${manifest.name}/${manifest.version}')
 		return list
 	}
 
-	mut arr := []string{}
 	mut has_user_agent := false
-	for k, v in header_map {
+	for k, v in headers.common {
 		if k == .user_agent {
 			has_user_agent = true
 		}
-		arr << '${k.str()}: ${v}'
+		list = curl.slist_append(list, '${k.str()}: ${v}')
 	}
-
-	for h in arr {
-		list = curl.slist_append(list, h)
+	for k, v in headers.custom {
+		list = curl.slist_append(list, '${k}: ${v}')
 	}
+
 	curl.easy_setopt(handle, .httpheader, list)
 
 	// Set default user agent if none was specified
diff --git a/src/_instructions_post.v b/src/_instructions_post.v
index 93d47b9..0aa0cfd 100644
--- a/src/_instructions_post.v
+++ b/src/_instructions_post.v
@@ -4,7 +4,7 @@ import vibe.curl
 
 fn (req Request) post_(url string, data string) !Response {
 	h := curl.easy_init() or { return http_error(.easy_init, none) }
-	header := set_header(req.headers, h)
+	header := set_header(h, common: req.headers, custom: req.custom_headers)
 	defer {
 		curl.easy_cleanup(h)
 		curl.slist_free_all(header)
diff --git a/src/_state_common.v b/src/_state_common.v
index 126112d..249bf37 100644
--- a/src/_state_common.v
+++ b/src/_state_common.v
@@ -6,7 +6,7 @@ import time
 pub struct Request {
 pub mut:
 	headers        map[HttpHeader]string
-	custom_headers map[string]string // TODO:
+	custom_headers map[string]string
 	cookie_jar     string
 	cookie_file    string
 	timeout        time.Duration
diff --git a/src/_state_header.v b/src/_state_header.v
index 38465b4..3dba42f 100644
--- a/src/_state_header.v
+++ b/src/_state_header.v
@@ -4,6 +4,12 @@ type HeaderList = C.curl_slist
 
 type Status = int
 
+[params]
+struct HttpHeaders {
+	common map[HttpHeader]string
+	custom map[string]string
+}
+
 // https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers
 [generated]
 pub enum HttpHeader {
diff --git a/src/_tests_post_test.v b/src/_tests_post_test.v
index 7fa4300..7a01f6c 100644
--- a/src/_tests_post_test.v
+++ b/src/_tests_post_test.v
@@ -1,6 +1,7 @@
 module vibe
 
 import time
+import x.json2 as json
 
 fn test_post() {
 	req := Request{
@@ -8,6 +9,9 @@ fn test_post() {
 			.user_agent:   'YourCustomUserAgent/v0.0.1'
 			.content_type: 'application/json; charset=utf-8'
 		}
+		custom_headers: {
+			'My-Custom-Header': 'FooBar'
+		}
 		timeout: time.second * 10
 	}
 	mut resp := req.post('https://httpbin.org/post', '{"msg":"hello from vibe"}')!
@@ -17,6 +21,10 @@ fn test_post() {
 	}
 
 	assert resp.status == 200
-	assert resp.body.contains('"User-Agent": "YourCustomUserAgent/v0.0.1"')
-	assert resp.body.contains('"msg": "hello from vibe"')
+	raw_json_resp := json.raw_decode(resp.body)!.as_map()
+	headers := raw_json_resp['headers']!.as_map()
+	assert headers['My-Custom-Header']!.str() == 'FooBar'
+	assert headers['User-Agent']!.str() == 'YourCustomUserAgent/v0.0.1'
+	json_data := raw_json_resp['json']!.as_map()
+	assert json_data['msg']!.str() == 'hello from vibe'
 }