Merge pull request #208 from protegrity/av_add_auth_to_init_023

Adding JWT fetch during health check for server connection

Author: avalerio-tkd

src/client/dbps_api_client.cpp

@@ -173,7 +173,12 @@ std::string DBPSApiClient::HealthCheck() {
     if (!IsHttpSuccess(response.status_code)) {
         return "Health check failed with status: " + std::to_string(response.status_code);
     }
-    
+
+    auto auth_error = http_client_->PrefetchToken();
+    if (auth_error.has_value()) {
+        return "Health check succeeded but JWT fetch failed: " + auth_error.value();
+    }
+
     return response.result;
 }
 

src/client/http_client_base.cpp

@@ -88,6 +88,18 @@ HttpClientBase::HttpResponse HttpClientBase::Post(const std::string& endpoint,
     return result;
 }
 
+std::optional<std::string> HttpClientBase::PrefetchToken() {
+    std::string error;
+    auto token_opt = EnsureValidToken(error);
+    if (!token_opt.has_value()) {
+        if (error.empty()) {
+            return std::string("Failed to fetch auth token");
+        }
+        return error;
+    }
+    return std::nullopt;
+}
+
 std::string HttpClientBase::AddAuthorizationHeader(HeaderList& headers) {
     // Get the valid token or an error message.
     std::string token_or_error;

src/client/http_client_base.h

@@ -61,6 +61,9 @@ class HttpClientBase {
     HttpResponse Get(const std::string& endpoint, bool auth_required = true);
     HttpResponse Post(const std::string& endpoint, const std::string& json_body, bool auth_required = true);
 
+    // Fetches a JWT if missing. Returns nullopt on success, error message otherwise.
+    std::optional<std::string> PrefetchToken();
+
 protected:
     explicit HttpClientBase(std::string base_url,
                                  ClientCredentials credentials = {})

src/client/http_client_base_test.cpp

@@ -52,6 +52,7 @@ class FakeHttpClient final : public HttpClientBase {
     HeaderList last_get_headers;
     std::vector<HeaderList> get_headers_history;
     bool fail_first_get_with_401 = false;
+    bool fail_token_fetch = false;
 
 protected:
     HttpResponse DoGet(const std::string& endpoint, const HeaderList& headers) override {
@@ -72,6 +73,9 @@ class FakeHttpClient final : public HttpClientBase {
         ++post_calls;
         if (endpoint == "/token") {
             ++token_calls;
+            if (fail_token_fetch) {
+                return HttpResponse(401, "", "Unauthorized");
+            }
             // Use far-future expiry so the cached token is always considered valid.
             const TokenResp* tr = nullptr;
             if (!token_responses_.empty()) {
@@ -172,4 +176,28 @@ TEST(HttpClientBaseTest, RetryOnceOn401FetchesNewTokenAndRetries) {
     ASSERT_EQ(auth2->second, "Bearer t2");
 }
 
+TEST(HttpClientBaseTest, PrefetchTokenFetchesAndCachesToken) {
+    FakeHttpClient client({{"client_id", "clientA"}, {"api_key", "keyA"}});
+    client.SetTokenResponse("prefetch", "Bearer", 4102444800);
+
+    auto error = client.PrefetchToken();
+    ASSERT_FALSE(error.has_value());
+    ASSERT_EQ(client.token_calls.load(), 1);
+
+    auto r = client.Get("/statusz");
+    ASSERT_TRUE(r.error_message.empty());
+    ASSERT_EQ(r.status_code, 200);
+    ASSERT_EQ(client.token_calls.load(), 1);
+}
+
+TEST(HttpClientBaseTest, PrefetchTokenReturnsErrorOnFailure) {
+    FakeHttpClient client({{"client_id", "clientA"}, {"api_key", "keyA"}});
+    client.fail_token_fetch = true;
+
+    auto error = client.PrefetchToken();
+    ASSERT_TRUE(error.has_value());
+    ASSERT_NE(error->find("status code: 401"), std::string::npos);
+    ASSERT_EQ(client.token_calls.load(), 1);
+}
+