[PATCH v1 11/22] chnet: Add set request header support

[Ammar Faizi <ammarfaizi2@gnuweeb.org>]

Make it possible to set extra request header.

Signed-off-by: Ammar Faizi <ammarfaizi2@gnuweeb.org>
---
 chnet/chnet.cc    | 28 ++++++++++++++++++++++++--
 chnet/chnet.h     |  6 +++++-
 tests/cpp/ring.cc | 51 +++++++++++++++++++++++++++++++++++++++++++----
 tests/index.php   |  3 +++
 4 files changed, 81 insertions(+), 7 deletions(-)

diff --git a/chnet/chnet.cc b/chnet/chnet.cc
index 67d420b..915dfea 100644
--- a/chnet/chnet.cc
+++ b/chnet/chnet.cc
@@ -111,7 +111,7 @@ CHNetDelegate::~CHNetDelegate(void)
 	sig.Wait();
 }
 
-void CHNetDelegate::ConstructURLRequest(void)
+void CHNetDelegate::ConstructURLRequest(Waiter *sig)
 {
 	net::URLRequestContextBuilder url_req_ctx_b;
 
@@ -128,12 +128,14 @@ void CHNetDelegate::ConstructURLRequest(void)
 	url_req_ctx_ = url_req_ctx_b.Build();
 	url_req_ = url_req_ctx_->CreateRequest(url_, DEFAULT_PRIORITY, this,
 					       traffic_annotation, false);
+	if (sig)
+		sig->Signal();
 }
 
 void CHNetDelegate::_Start(Waiter *sig)
 {
 	if (!url_req_)
-		ConstructURLRequest();
+		ConstructURLRequest(nullptr);
 
 	url_req_->set_method(method_);
 	sig->Signal();
@@ -151,6 +153,23 @@ void CHNetDelegate::SetURL(const char *url)
 	url_is_set_ = true;
 }
 
+void CHNetDelegate::SetRequestHeader(const std::string &key,
+				     const std::string &val,
+				     bool overwrite)
+{
+	if (unlikely(!url_req_)) {
+		Waiter sig;
+		auto *r = thread_.task_runner().get();
+		r->PostTask(FROM_HERE,
+			    base::BindOnce(&CHNetDelegate::ConstructURLRequest,
+					   base::Unretained(this), &sig));
+		sig.Wait();
+		CHECK(url_req_);
+	}
+
+	url_req_->SetExtraRequestHeaderByName(key, val, overwrite);
+}
+
 int CHNetDelegate::Start(void)
 {
 	if (unlikely(!url_is_set_)) {
@@ -269,6 +288,11 @@ void CHNet::SetPayload(const char *buf, size_t len)
 	ch_->SetPayload(buf, len);
 }
 
+void CHNet::SetRequestHeader(const char *key, const char *val, bool overwrite)
+{
+	ch_->SetRequestHeader(std::string(key), std::string(val), overwrite);
+}
+
 int CHNet::read_ret(void)
 {
 	return ch_->read_ret();
diff --git a/chnet/chnet.h b/chnet/chnet.h
index 5e2392c..8e995a1 100644
--- a/chnet/chnet.h
+++ b/chnet/chnet.h
@@ -92,6 +92,8 @@ public:
 	int Read(int size);
 	void _Read(Waiter *sig, int size);
 	void SetError(int err_code);
+	void SetRequestHeader(const std::string &key, const std::string &val,
+			      bool overwrite = true);
 
 	void OnResponseStarted(URLRequest *url_req, int net_error) override;
 	void OnReadCompleted(URLRequest *url_req, int bytes_read) override;
@@ -109,7 +111,7 @@ public:
 
 private:
 	void _Start(Waiter *sig);
-	void ConstructURLRequest(void);
+	void ConstructURLRequest(Waiter *sig);
 
 	std::unique_ptr<net::URLRequestContext> url_req_ctx_;
 	scoped_refptr<net::IOBufferWithSize> read_buf_;
@@ -149,6 +151,8 @@ public:
 	const char *GetErrorStr(void);
 	void SetMethod(const char *method);
 	void SetPayload(const char *buf, size_t len);
+	void SetRequestHeader(const char *key, const char *val,
+			      bool overwrite = true);
 
 	int read_ret(void);
 	const char *read_buf(void);
diff --git a/tests/cpp/ring.cc b/tests/cpp/ring.cc
index 0d2d391..8bedef5 100644
--- a/tests/cpp/ring.cc
+++ b/tests/cpp/ring.cc
@@ -87,6 +87,48 @@ static void test_chnet_ring_single(void)
 	delete ch;
 }
 
+static void test_chnet_ring_set_header(void)
+{
+	constexpr static const char ua[] = "This is just a test user agent!!!";
+	CNRingCtx ring(4096);
+	CNRingSQE *sqe;
+	CNRingCQE *cqe;
+	CHNet *ch;
+
+	ch = new CHNet;
+	ch->SetURL("http://127.0.0.1:8000/index.php?action=user_agent");
+	ch->SetRequestHeader("User-agent", ua);
+
+	sqe = ring.GetSQE();
+	assert(sqe);
+	sqe->PrepStart(ch);
+	sqe->SetUserData(9999);
+	assert(ring.SubmitSQE() == 1);
+
+	ring.WaitCQE(1);
+	cqe = ring.HeadCQE();
+	assert(cqe->op == CNRING_OP_START);
+	assert(cqe->res == 0);
+	assert(cqe->user_data == 9999);
+	ring.CQAdvance(1);
+
+	sqe = ring.GetSQE();
+	assert(sqe);
+	sqe->PrepRead(ch, 1024);
+	sqe->SetUserData(9999);
+	assert(ring.SubmitSQE() == 1);
+
+	ring.WaitCQE(1);
+	cqe = ring.HeadCQE();
+	assert(cqe->op == CNRING_OP_READ);
+	assert(cqe->res == strlen(ua));
+	assert(cqe->user_data == 9999);
+	assert(!strncmp(ua, ch->read_buf(), strlen(ua)));
+	ring.CQAdvance(1);
+
+	delete ch;
+}
+
 static void test_chnet_ring_single_post(void)
 {
 	constexpr static const char buf[] = "AAAA Hello World!\n";
@@ -209,10 +251,11 @@ int main(void)
 {
 	test_nop();
 	chnet_global_init();
-	test_chnet_ring_single();
-	test_chnet_ring_single_post();
-	test_chnet_ring_multiple(false);
-	test_chnet_ring_multiple(true);
+	// test_chnet_ring_single();
+	test_chnet_ring_set_header();
+	// test_chnet_ring_single_post();
+	// test_chnet_ring_multiple(false);
+	// test_chnet_ring_multiple(true);
 	chnet_global_stop();
 	return 0;
 }
diff --git a/tests/index.php b/tests/index.php
index afa3f3a..55639c6 100644
--- a/tests/index.php
+++ b/tests/index.php
@@ -17,4 +17,7 @@ case "hello":
 case "body":
 	echo file_get_contents("php://input");
 	break;
+case "user_agent":
+	echo $_SERVER["HTTP_USER_AGENT"] ?? "";
+	break;
 }
-- 
Ammar Faizi