support customized UserAgent in http client (#479)

Signed-off-by: liulanzheng <[email protected]>
alibaba · May 16, 2024 · 1738a58 · 1738a58
1 parent 32df410
commit 1738a58
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 2 deletions.
diff --git a/net/http/client.cpp b/net/http/client.cpp
@@ -224,9 +224,10 @@ class ClientImpl : public Client {
             LOG_ERROR_RETURN(EINVAL, ROUNDTRIP_FAILED,
                             "Content-Length and Transfer-Encoding conflicted");
         }
-        op->req.headers.insert("User-Agent", USERAGENT);
-        op->req.headers.insert("Connection", "keep-alive");
         op->req.headers.merge(m_common_headers);
+        op->req.headers.insert("User-Agent", m_user_agent.empty() ? std::string_view(USERAGENT)
+                                                                  : std::string_view(m_user_agent));
+        op->req.headers.insert("Connection", "keep-alive");
         if (m_cookie_jar && m_cookie_jar->set_cookies_to_headers(&op->req) != 0)
             LOG_ERROR_RETURN(0, -1, "set_cookies_to_headers failed");
         Timeout tmo(std::min(op->timeout.timeout(), m_timeout));

diff --git a/net/http/client.h b/net/http/client.h
@@ -109,6 +109,9 @@ class Client : public Object {
         m_proxy_url.from_string(proxy);
         m_proxy = true;
     }
+    void set_user_agent(std::string_view user_agent) {
+        m_user_agent = std::string(user_agent);
+    }
     StoredURL* get_proxy() {
         return &m_proxy_url;
     }
@@ -129,6 +132,7 @@ class Client : public Object {
                                     bool secure = false, uint64_t timeout = -1UL) = 0;
 protected:
     StoredURL m_proxy_url;
+    std::string m_user_agent;
     uint64_t m_timeout = -1UL;
     bool m_proxy = false;
 };

diff --git a/net/http/test/client_function_test.cpp b/net/http/test/client_function_test.cpp
@@ -512,6 +512,46 @@ TEST(http_client, partial_body) {
     EXPECT_EQ(true, buf == "http_clien");
 }
 
+int ua_check_handler(void*, Request &req, Response &resp, std::string_view) {
+    auto ua = req.headers["User-Agent"];
+    LOG_DEBUG(VALUE(ua));
+    EXPECT_EQ(ua, "TEST_UA");
+    resp.set_result(200);
+    std::string str = "success";
+    resp.headers.content_length(7);
+    resp.write((void*)str.data(), str.size());
+    return 0;
+}
+
+
+TEST(http_client, user_agent) {
+    auto tcpserver = new_tcp_socket_server();
+    DEFER(delete tcpserver);
+    tcpserver->bind(18731);
+    tcpserver->listen();
+    auto server = new_http_server();
+    DEFER(delete server);
+    server->add_handler({nullptr, &ua_check_handler});
+    tcpserver->set_handler(server->get_connection_handler());
+    tcpserver->start_loop();
+
+    std::string target_get = "http://localhost:18731/file";
+    auto client = new_http_client();
+    client->set_user_agent("TEST_UA");
+    DEFER(delete client);
+    auto op = client->new_operation(Verb::GET, target_get);
+    DEFER(delete op);
+    op->req.headers.content_length(0);
+    client->call(op);
+    EXPECT_EQ(op->status_code, 200);
+    std::string buf;
+    buf.resize(op->resp.headers.content_length());
+    op->resp.read((void*)buf.data(), op->resp.headers.content_length());
+    LOG_DEBUG(VALUE(buf));
+    EXPECT_EQ(true, buf == "success");
+}
+
+
 TEST(url, url_escape_unescape) {
     EXPECT_EQ(
         url_escape("?a=x:b&b=cd&c= feg&d=2/1[+]@alibaba.com&e='!bad';"),