1 // 2 // Copyright 2016 gRPC authors. 3 // 4 // Licensed under the Apache License, Version 2.0 (the "License"); 5 // you may not use this file except in compliance with the License. 6 // You may obtain a copy of the License at 7 // 8 // http://www.apache.org/licenses/LICENSE-2.0 9 // 10 // Unless required by applicable law or agreed to in writing, software 11 // distributed under the License is distributed on an "AS IS" BASIS, 12 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 // See the License for the specific language governing permissions and 14 // limitations under the License. 15 // 16 17 #ifndef GRPC_SRC_CORE_LIB_SECURITY_CREDENTIALS_TOKEN_FETCHER_TOKEN_FETCHER_CREDENTIALS_H 18 #define GRPC_SRC_CORE_LIB_SECURITY_CREDENTIALS_TOKEN_FETCHER_TOKEN_FETCHER_CREDENTIALS_H 19 20 #include <grpc/event_engine/event_engine.h> 21 22 #include <atomic> 23 #include <memory> 24 #include <utility> 25 26 #include "absl/container/flat_hash_set.h" 27 #include "absl/functional/any_invocable.h" 28 #include "absl/status/statusor.h" 29 #include "absl/types/variant.h" 30 #include "src/core/lib/iomgr/polling_entity.h" 31 #include "src/core/lib/promise/arena_promise.h" 32 #include "src/core/lib/security/credentials/credentials.h" 33 #include "src/core/lib/transport/metadata.h" 34 #include "src/core/util/backoff.h" 35 #include "src/core/util/http_client/httpcli.h" 36 #include "src/core/util/orphanable.h" 37 #include "src/core/util/ref_counted.h" 38 #include "src/core/util/ref_counted_ptr.h" 39 #include "src/core/util/sync.h" 40 #include "src/core/util/time.h" 41 #include "src/core/util/useful.h" 42 43 namespace grpc_core { 44 45 // A base class for credentials that fetch tokens via an HTTP request. 46 // Subclasses must implement FetchToken(). 47 class TokenFetcherCredentials : public grpc_call_credentials { 48 public: 49 // Represents a token. 50 class Token : public RefCounted<Token> { 51 public: 52 Token(Slice token, Timestamp expiration); 53 54 // Returns the token's expiration time. ExpirationTime()55 Timestamp ExpirationTime() const { return expiration_; } 56 57 // Adds the token to the call's client initial metadata. 58 void AddTokenToClientInitialMetadata(ClientMetadata& metadata) const; 59 60 private: 61 Slice token_; 62 Timestamp expiration_; 63 }; 64 65 ~TokenFetcherCredentials() override; 66 67 void Orphaned() override; 68 69 ArenaPromise<absl::StatusOr<ClientMetadataHandle>> GetRequestMetadata( 70 ClientMetadataHandle initial_metadata, 71 const GetRequestMetadataArgs* args) override; 72 73 protected: 74 // Base class for fetch requests. 75 class FetchRequest : public InternallyRefCounted<FetchRequest> {}; 76 77 explicit TokenFetcherCredentials( 78 std::shared_ptr<grpc_event_engine::experimental::EventEngine> 79 event_engine = nullptr, 80 bool test_only_use_backoff_jitter = true); 81 82 // Fetches a token. The on_done callback will be invoked when complete. 83 virtual OrphanablePtr<FetchRequest> FetchToken( 84 Timestamp deadline, 85 absl::AnyInvocable<void(absl::StatusOr<RefCountedPtr<Token>>)> 86 on_done) = 0; 87 event_engine()88 grpc_event_engine::experimental::EventEngine& event_engine() const { 89 return *event_engine_; 90 } 91 pollent()92 grpc_polling_entity* pollent() { return &pollent_; } 93 94 private: 95 // A call that is waiting for a token fetch request to complete. 96 struct QueuedCall : public RefCounted<QueuedCall> { 97 std::atomic<bool> done{false}; 98 Waker waker; 99 grpc_polling_entity* pollent; 100 ClientMetadataHandle md; 101 absl::StatusOr<RefCountedPtr<Token>> result; 102 }; 103 104 class FetchState : public InternallyRefCounted<FetchState> { 105 public: 106 explicit FetchState(WeakRefCountedPtr<TokenFetcherCredentials> creds) 107 ABSL_EXCLUSIVE_LOCKS_REQUIRED(&TokenFetcherCredentials::mu_); 108 109 // Disabling thread safety annotations, since Orphan() is called 110 // by OrpahanablePtr<>, which does not have the right lock 111 // annotations. 112 void Orphan() override ABSL_NO_THREAD_SAFETY_ANALYSIS; 113 114 // Returns non-OK when we're in backoff. 115 absl::Status status() const; 116 117 RefCountedPtr<QueuedCall> QueueCall(ClientMetadataHandle initial_metadata) 118 ABSL_EXCLUSIVE_LOCKS_REQUIRED(&TokenFetcherCredentials::mu_); 119 120 private: 121 class BackoffTimer : public InternallyRefCounted<BackoffTimer> { 122 public: 123 BackoffTimer(RefCountedPtr<FetchState> fetch_state, absl::Status status) 124 ABSL_EXCLUSIVE_LOCKS_REQUIRED(&TokenFetcherCredentials::mu_); 125 126 // Disabling thread safety annotations, since Orphan() is called 127 // by OrpahanablePtr<>, which does not have the right lock 128 // annotations. 129 void Orphan() override ABSL_NO_THREAD_SAFETY_ANALYSIS; 130 status()131 absl::Status status() const { return status_; } 132 133 private: 134 void OnTimer(); 135 136 RefCountedPtr<FetchState> fetch_state_; 137 const absl::Status status_; 138 absl::optional<grpc_event_engine::experimental::EventEngine::TaskHandle> 139 timer_handle_ ABSL_GUARDED_BY(&TokenFetcherCredentials::mu_); 140 }; 141 142 struct Shutdown {}; 143 144 void StartFetchAttempt() 145 ABSL_EXCLUSIVE_LOCKS_REQUIRED(&TokenFetcherCredentials::mu_); 146 void TokenFetchComplete(absl::StatusOr<RefCountedPtr<Token>> token); 147 void ResumeQueuedCalls(absl::StatusOr<RefCountedPtr<Token>> token) 148 ABSL_EXCLUSIVE_LOCKS_REQUIRED(&TokenFetcherCredentials::mu_); 149 150 WeakRefCountedPtr<TokenFetcherCredentials> creds_; 151 // Pending token-fetch request or backoff timer, if any. 152 absl::variant<OrphanablePtr<FetchRequest>, OrphanablePtr<BackoffTimer>, 153 Shutdown> 154 state_ ABSL_GUARDED_BY(&TokenFetcherCredentials::mu_); 155 // Calls that are queued up waiting for the token. 156 absl::flat_hash_set<RefCountedPtr<QueuedCall>> queued_calls_ 157 ABSL_GUARDED_BY(&TokenFetcherCredentials::mu_); 158 // Backoff state. 159 BackOff backoff_ ABSL_GUARDED_BY(&TokenFetcherCredentials::mu_); 160 }; 161 cmp_impl(const grpc_call_credentials * other)162 int cmp_impl(const grpc_call_credentials* other) const override { 163 // TODO(yashykt): Check if we can do something better here 164 return QsortCompare(static_cast<const grpc_call_credentials*>(this), other); 165 } 166 167 std::shared_ptr<grpc_event_engine::experimental::EventEngine> event_engine_; 168 const bool test_only_use_backoff_jitter_; 169 170 Mutex mu_; 171 // Cached token, if any. 172 RefCountedPtr<Token> token_ ABSL_GUARDED_BY(&mu_); 173 // Fetch state, if any. 174 OrphanablePtr<FetchState> fetch_state_ ABSL_GUARDED_BY(&mu_); 175 176 grpc_polling_entity pollent_ ABSL_GUARDED_BY(&mu_); 177 }; 178 179 } // namespace grpc_core 180 181 #endif // GRPC_SRC_CORE_LIB_SECURITY_CREDENTIALS_TOKEN_FETCHER_TOKEN_FETCHER_CREDENTIALS_H 182