Merge 5a4c3f6 into 2ce3bb0

jupp0r · Jun 21, 2022 · 1f7e816 · 1f7e816
2 parents 2ce3bb0 + 5a4c3f6
commit 1f7e816
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 1 deletion.
diff --git a/push/include/prometheus/gateway.h b/push/include/prometheus/gateway.h
@@ -58,6 +58,13 @@ class PROMETHEUS_CPP_PUSH_EXPORT Gateway {
   // Delete metrics from the given pushgateway (for configured instance labels).
   std::future<int> AsyncDeleteForInstance();
 
+  /**
+  * Add a custom http header, it must use after create gateway before push   
+  * `header` is the custom header, its content format is “key:value”, such 
+  *  as "user:gateway"
+  */
+  int AddHttpHeader(const std::string& header);
+
  private:
   std::string jobUri_;
   std::string labels_;

diff --git a/push/src/curl_wrapper.cc b/push/src/curl_wrapper.cc
@@ -22,12 +22,15 @@ CurlWrapper::CurlWrapper(const std::string& username,
     throw std::runtime_error("Cannot initialize easy curl!");
   }
 
+  optHttpHeader_ = nullptr;
+
   if (!username.empty()) {
     auth_ = username + ":" + password;
   }
 }
 
 CurlWrapper::~CurlWrapper() {
+  curl_slist_free_all(optHttpHeader_);
   curl_easy_cleanup(curl_);
   curl_global_cleanup();
 }
@@ -40,7 +43,26 @@ int CurlWrapper::performHttpRequest(HttpMethod method, const std::string& uri,
   curl_easy_setopt(curl_, CURLOPT_URL, uri.c_str());
 
   curl_slist* header_chunk = nullptr;
-  header_chunk = curl_slist_append(header_chunk, CONTENT_TYPE);
+  header_chunk = curl_slist_append(header_chunk, CONTENT_TYPE); 
+  if (nullptr == header_chunk)
+  {
+    return -1;
+  }
+
+  curl_slist* optHeader_tmp = optHttpHeader_;
+  while (optHeader_tmp)
+  {
+    curl_slist* header_tmp = curl_slist_append(header_chunk, optHeader_tmp->data);
+    if (nullptr == header_tmp)
+    {
+      curl_slist_free_all(header_chunk);
+      return -1;
+    }
+
+    header_chunk = header_tmp;
+    optHeader_tmp = optHeader_tmp->next;
+  }
+
   curl_easy_setopt(curl_, CURLOPT_HTTPHEADER, header_chunk);
 
   if (!body.empty()) {
@@ -86,5 +108,18 @@ int CurlWrapper::performHttpRequest(HttpMethod method, const std::string& uri,
   return response_code;
 }
 
+int CurlWrapper::addHttpHeader(const std::string& header)
+{
+  std::lock_guard<std::mutex> lock{mutex_};
+  auto header_tmp = curl_slist_append(optHttpHeader_, header.c_str());
+  if (nullptr == header_tmp)
+  {
+    return -1;
+  }
+
+  optHttpHeader_ = header_tmp;
+  return 0;
+}
+
 }  // namespace detail
 }  // namespace prometheus
diff --git a/push/src/curl_wrapper.h b/push/src/curl_wrapper.h
@@ -21,11 +21,13 @@ class CurlWrapper {
 
   int performHttpRequest(HttpMethod method, const std::string& uri,
                          const std::string& body);
+  int addHttpHeader(const std::string& header);
 
  private:
   CURL* curl_;
   std::string auth_;
   std::mutex mutex_;
+  curl_slist* optHttpHeader_;
 };
 
 }  // namespace detail

diff --git a/push/src/gateway.cc b/push/src/gateway.cc
@@ -164,4 +164,9 @@ void Gateway::CleanupStalePointers(
       std::end(collectables));
 }
 
+int Gateway::AddHttpHeader(const std::string& header)
+{
+  return curlWrapper_->addHttpHeader(header);
+}
+
 }  // namespace prometheus