share the commmon http token fetching logic between oauth2 and jwt call credentials

Author: rockspore

src/core/BUILD

@@ -4887,6 +4887,7 @@ grpc_cc_library(
         "poll",
         "pollset_set",
         "ref_counted",
+        "status_conversion",
         "sync",
         "time",
         "useful",

src/core/credentials/call/gcp_service_account_identity/gcp_service_account_identity_credentials.cc

@@ -23,6 +23,7 @@
 #include "absl/strings/string_view.h"
 #include "src/core/call/metadata.h"
 #include "src/core/credentials/call/jwt_util.h"
+#include "src/core/credentials/call/token_fetcher/token_fetcher_credentials.h"
 #include "src/core/credentials/transport/transport_credentials.h"
 #include "src/core/lib/iomgr/error.h"
 #include "src/core/lib/transport/status_conversion.h"
@@ -36,78 +37,30 @@ namespace grpc_core {
 // JwtTokenFetcherCallCredentials
 //
 
-// State held for a pending HTTP request.
-class JwtTokenFetcherCallCredentials::HttpFetchRequest final
-    : public TokenFetcherCredentials::FetchRequest {
- public:
-  HttpFetchRequest(
-      JwtTokenFetcherCallCredentials* creds, Timestamp deadline,
-      absl::AnyInvocable<
-          void(absl::StatusOr<RefCountedPtr<TokenFetcherCredentials::Token>>)>
-          on_done)
-      : on_done_(std::move(on_done)) {
-    GRPC_CLOSURE_INIT(&on_http_response_, OnHttpResponse, this, nullptr);
-    Ref().release();  // Ref held by HTTP request callback.
-    http_request_ = creds->StartHttpRequest(creds->pollent(), deadline,
-                                            &response_, &on_http_response_);
-  }
-
-  ~HttpFetchRequest() override { grpc_http_response_destroy(&response_); }
-
-  void Orphan() override {
-    http_request_.reset();
-    Unref();
-  }
-
- private:
-  static void OnHttpResponse(void* arg, grpc_error_handle error) {
-    RefCountedPtr<HttpFetchRequest> self(static_cast<HttpFetchRequest*>(arg));
-    if (!error.ok()) {
-      // TODO(roth): It shouldn't be necessary to explicitly set the
-      // status to UNAVAILABLE here.  Once the HTTP client code is
-      // migrated to stop using legacy grpc_error APIs to create
-      // statuses, we should be able to just propagate the status as-is.
-      self->on_done_(absl::UnavailableError(StatusToString(error)));
-      return;
-    }
-    if (self->response_.status != 200) {
-      grpc_status_code status_code =
-          grpc_http2_status_to_grpc_status(self->response_.status);
-      if (status_code != GRPC_STATUS_UNAVAILABLE) {
-        status_code = GRPC_STATUS_UNAUTHENTICATED;
-      }
-      self->on_done_(absl::Status(static_cast<absl::StatusCode>(status_code),
-                                  absl::StrCat("JWT fetch failed with status ",
-                                               self->response_.status)));
-      return;
-    }
-    // Return token object.
-    absl::string_view body(self->response_.body, self->response_.body_length);
-    auto expiration_time = GetJwtExpirationTime(body);
-    if (!expiration_time.ok()) {
-      self->on_done_(expiration_time.status());
-      return;
-    }
-    self->on_done_(MakeRefCounted<Token>(
-        Slice::FromCopiedString(absl::StrCat("Bearer ", body)),
-        *expiration_time));
-  }
-
-  OrphanablePtr<HttpRequest> http_request_;
-  grpc_closure on_http_response_;
-  grpc_http_response response_;
-  absl::AnyInvocable<void(
-      absl::StatusOr<RefCountedPtr<TokenFetcherCredentials::Token>>)>
-      on_done_;
-};
-
 OrphanablePtr<TokenFetcherCredentials::FetchRequest>
 JwtTokenFetcherCallCredentials::FetchToken(
     Timestamp deadline,
     absl::AnyInvocable<
         void(absl::StatusOr<RefCountedPtr<TokenFetcherCredentials::Token>>)>
         on_done) {
-  return MakeOrphanable<HttpFetchRequest>(this, deadline, std::move(on_done));
+  return MakeOrphanable<HttpTokenFetcherCredentials::HttpFetchRequest>(
+      this, deadline,
+      [on_done = std::move(on_done)](
+          absl::StatusOr<grpc_http_response> response) mutable {
+        if (!response.ok()) {
+          on_done(response.status());
+          return;
+        }
+        absl::string_view body(response->body, response->body_length);
+        auto expiration_time = GetJwtExpirationTime(body);
+        if (!expiration_time.ok()) {
+          on_done(expiration_time.status());
+          return;
+        }
+        on_done(MakeRefCounted<Token>(
+            Slice::FromCopiedString(absl::StrCat("Bearer ", body)),
+            *expiration_time));
+      });
 }
 
 //

src/core/credentials/call/gcp_service_account_identity/gcp_service_account_identity_credentials.h

@@ -40,20 +40,13 @@ namespace grpc_core {
 
 // A base class for JWT token fetching credentials.
 // Subclasses must implement StartHttpRequest().
-class JwtTokenFetcherCallCredentials : public TokenFetcherCredentials {
+class JwtTokenFetcherCallCredentials : public HttpTokenFetcherCredentials {
  public:
   OrphanablePtr<FetchRequest> FetchToken(
       Timestamp deadline,
       absl::AnyInvocable<
           void(absl::StatusOr<RefCountedPtr<TokenFetcherCredentials::Token>>)>
           on_done) final;
-
- private:
-  class HttpFetchRequest;
-
-  virtual OrphanablePtr<HttpRequest> StartHttpRequest(
-      grpc_polling_entity* pollent, Timestamp deadline,
-      grpc_http_response* response, grpc_closure* on_complete) = 0;
 };
 
 // GCP service account identity call credentials.

src/core/credentials/call/oauth2/oauth2_credentials.cc

@@ -44,6 +44,7 @@
 #include "absl/strings/string_view.h"
 #include "src/core/call/metadata_batch.h"
 #include "src/core/credentials/call/json_util.h"
+#include "src/core/credentials/call/token_fetcher/token_fetcher_credentials.h"
 #include "src/core/credentials/transport/transport_credentials.h"
 #include "src/core/lib/debug/trace.h"
 #include "src/core/lib/iomgr/error.h"
@@ -211,58 +212,6 @@ grpc_oauth2_token_fetcher_credentials_parse_server_response(
 
 namespace grpc_core {
 
-// State held for a pending HTTP request.
-class Oauth2TokenFetcherCredentials::HttpFetchRequest final
-    : public TokenFetcherCredentials::FetchRequest {
- public:
-  HttpFetchRequest(
-      Oauth2TokenFetcherCredentials* creds, Timestamp deadline,
-      absl::AnyInvocable<
-          void(absl::StatusOr<RefCountedPtr<TokenFetcherCredentials::Token>>)>
-          on_done)
-      : on_done_(std::move(on_done)) {
-    GRPC_CLOSURE_INIT(&on_http_response_, OnHttpResponse, this, nullptr);
-    Ref().release();  // Ref held by HTTP request callback.
-    http_request_ = creds->StartHttpRequest(creds->pollent(), deadline,
-                                            &response_, &on_http_response_);
-  }
-
-  ~HttpFetchRequest() override { grpc_http_response_destroy(&response_); }
-
-  void Orphan() override {
-    http_request_.reset();
-    Unref();
-  }
-
- private:
-  static void OnHttpResponse(void* arg, grpc_error_handle error) {
-    RefCountedPtr<HttpFetchRequest> self(static_cast<HttpFetchRequest*>(arg));
-    if (!error.ok()) {
-      self->on_done_(std::move(error));
-      return;
-    }
-    // Parse oauth2 token.
-    std::optional<Slice> access_token_value;
-    Duration token_lifetime;
-    grpc_credentials_status status =
-        grpc_oauth2_token_fetcher_credentials_parse_server_response(
-            &self->response_, &access_token_value, &token_lifetime);
-    if (status != GRPC_CREDENTIALS_OK) {
-      self->on_done_(absl::UnavailableError("error parsing oauth2 token"));
-      return;
-    }
-    self->on_done_(MakeRefCounted<Token>(std::move(*access_token_value),
-                                         Timestamp::Now() + token_lifetime));
-  }
-
-  OrphanablePtr<HttpRequest> http_request_;
-  grpc_closure on_http_response_;
-  grpc_http_response response_;
-  absl::AnyInvocable<void(
-      absl::StatusOr<RefCountedPtr<TokenFetcherCredentials::Token>>)>
-      on_done_;
-};
-
 std::string Oauth2TokenFetcherCredentials::debug_string() {
   return "OAuth2TokenFetcherCredentials";
 }
@@ -278,7 +227,27 @@ Oauth2TokenFetcherCredentials::FetchToken(
     absl::AnyInvocable<
         void(absl::StatusOr<RefCountedPtr<TokenFetcherCredentials::Token>>)>
         on_done) {
-  return MakeOrphanable<HttpFetchRequest>(this, deadline, std::move(on_done));
+  return MakeOrphanable<HttpTokenFetcherCredentials::HttpFetchRequest>(
+      this, deadline,
+      [on_done = std::move(on_done)](
+          absl::StatusOr<grpc_http_response> response) mutable {
+        if (!response.ok()) {
+          on_done(response.status());
+          return;
+        }
+        // Parse oauth2 token.
+        std::optional<Slice> access_token_value;
+        Duration token_lifetime;
+        grpc_credentials_status status =
+            grpc_oauth2_token_fetcher_credentials_parse_server_response(
+                &(*response), &access_token_value, &token_lifetime);
+        if (status != GRPC_CREDENTIALS_OK) {
+          on_done(absl::UnavailableError("error parsing oauth2 token"));
+          return;
+        }
+        on_done(MakeRefCounted<Token>(std::move(*access_token_value),
+                                      Timestamp::Now() + token_lifetime));
+      });
 }
 
 }  // namespace grpc_core

src/core/credentials/call/oauth2/oauth2_credentials.h

@@ -90,7 +90,7 @@ namespace grpc_core {
 
 // A base class for oauth2 token fetching credentials.
 // Subclasses must implement StartHttpRequest().
-class Oauth2TokenFetcherCredentials : public TokenFetcherCredentials {
+class Oauth2TokenFetcherCredentials : public HttpTokenFetcherCredentials {
  public:
   std::string debug_string() override;
 
@@ -102,13 +102,7 @@ class Oauth2TokenFetcherCredentials : public TokenFetcherCredentials {
           void(absl::StatusOr<RefCountedPtr<TokenFetcherCredentials::Token>>)>
           on_done) final;
 
-  virtual OrphanablePtr<HttpRequest> StartHttpRequest(
-      grpc_polling_entity* pollent, Timestamp deadline,
-      grpc_http_response* response, grpc_closure* on_complete) = 0;
-
  private:
-  class HttpFetchRequest;
-
   int cmp_impl(const grpc_call_credentials* other) const override {
     // TODO(yashykt): Check if we can do something better here
     return QsortCompare(static_cast<const grpc_call_credentials*>(this), other);

src/core/credentials/call/token_fetcher/token_fetcher_credentials.cc

@@ -24,6 +24,7 @@
 #include "src/core/lib/promise/context.h"
 #include "src/core/lib/promise/poll.h"
 #include "src/core/lib/promise/promise.h"
+#include "src/core/lib/transport/status_conversion.h"
 
 namespace grpc_core {
 
@@ -301,4 +302,49 @@ TokenFetcherCredentials::GetRequestMetadata(
   };
 }
 
+//
+// HttpTokenFetcherCredentials
+//
+
+HttpTokenFetcherCredentials::HttpFetchRequest::HttpFetchRequest(
+    HttpTokenFetcherCredentials* creds, Timestamp deadline,
+    absl::AnyInvocable<void(absl::StatusOr<grpc_http_response>)> on_done)
+    : on_done_(std::move(on_done)) {
+  GRPC_CLOSURE_INIT(&on_http_response_, OnHttpResponse, this, nullptr);
+  Ref().release();  // Ref held by HTTP request callback.
+  http_request_ = creds->StartHttpRequest(creds->pollent(), deadline,
+                                          &response_, &on_http_response_);
+}
+
+void HttpTokenFetcherCredentials::HttpFetchRequest::Orphan() {
+  http_request_.reset();
+  Unref();
+}
+
+void HttpTokenFetcherCredentials::HttpFetchRequest::OnHttpResponse(
+    void* arg, grpc_error_handle error) {
+  RefCountedPtr<HttpFetchRequest> self(static_cast<HttpFetchRequest*>(arg));
+  if (!error.ok()) {
+    // TODO(roth): It shouldn't be necessary to explicitly set the
+    // status to UNAVAILABLE here.  Once the HTTP client code is
+    // migrated to stop using legacy grpc_error APIs to create
+    // statuses, we should be able to just propagate the status as-is.
+    self->on_done_(absl::UnavailableError(StatusToString(error)));
+    return;
+  }
+  if (self->response_.status != 200) {
+    grpc_status_code status_code =
+        grpc_http2_status_to_grpc_status(self->response_.status);
+    if (status_code != GRPC_STATUS_UNAVAILABLE) {
+      status_code = GRPC_STATUS_UNAUTHENTICATED;
+    }
+    self->on_done_(
+        absl::Status(static_cast<absl::StatusCode>(status_code),
+                     absl::StrCat("HTTP token fetch failed with status ",
+                                  self->response_.status)));
+    return;
+  }
+  self->on_done_(self->response_);
+}
+
 }  // namespace grpc_core

src/core/credentials/call/token_fetcher/token_fetcher_credentials.h

@@ -176,6 +176,35 @@ class TokenFetcherCredentials : public grpc_call_credentials {
   grpc_polling_entity pollent_ ABSL_GUARDED_BY(&mu_);
 };
 
+// A base class for fetching tokens via an HTTP request.
+class HttpTokenFetcherCredentials : public TokenFetcherCredentials {
+ public:
+  virtual OrphanablePtr<HttpRequest> StartHttpRequest(
+      grpc_polling_entity* pollent, Timestamp deadline,
+      grpc_http_response* response, grpc_closure* on_complete) = 0;
+
+ protected:
+  // State held for a pending HTTP request.
+  class HttpFetchRequest : public TokenFetcherCredentials::FetchRequest {
+   public:
+    // The given callback should assume the http response status has already
+    // been checked and handle the token parsing.
+    HttpFetchRequest(
+        HttpTokenFetcherCredentials* creds, Timestamp deadline,
+        absl::AnyInvocable<void(absl::StatusOr<grpc_http_response>)> on_done);
+    ~HttpFetchRequest() override { grpc_http_response_destroy(&response_); }
+
+    void Orphan() override;
+
+   private:
+    static void OnHttpResponse(void* arg, grpc_error_handle error);
+    OrphanablePtr<HttpRequest> http_request_;
+    grpc_closure on_http_response_;
+    grpc_http_response response_;
+    absl::AnyInvocable<void(absl::StatusOr<grpc_http_response>)> on_done_;
+  };
+};
+
 }  // namespace grpc_core
 
 #endif  // GRPC_SRC_CORE_CREDENTIALS_CALL_TOKEN_FETCHER_TOKEN_FETCHER_CREDENTIALS_H