add caching, cleanup bs

zvonand · zvonand · commit 9adf20d06399 · 2025-03-28T10:03:43.000Z
diff --git a/docs/en/operations/external-authenticators/tokens.md b/docs/en/operations/external-authenticators/tokens.md
@@ -24,6 +24,7 @@ To define an access token processor, add `access_token_processors` section to `c
         <gogoogle>
             <provider>Google</provider>
             <email_filter>^[A-Za-z0-9._%+-]+@example\.com$</email_filter>
+            <cache_lifetime>600</cache_lifetime>
         </gogoogle>
         <azuure>
             <provider>azure</provider>
@@ -41,10 +42,16 @@ Different providers have different sets of parameters.
 **Parameters**
 
 - `provider` -- name of identity provider. Mandatory, case-insensitive. Supported options: "Google", "Azure".
+- `cache_lifetime` --  maximum lifetime of cached token (in seconds). Optional, default: 3600.
 - `email_filter` -- Regex for validation of user emails. Optional parameter, only for Google IdP.
 - `client_id` -- Azure AD (Entra ID) client ID. Optional parameter, only for Azure IdP.
 - `tenant_id` -- Azure AD (Entra ID) tenant ID. Optional parameter, only for Azure IdP.
 
+### Tokens cache
+To reduce number of requests to IdP, tokens are cached internally for no longer then `cache_lifetime` seconds.
+If token expires sooner than `cache_lifetime`, then cache entry for this token will only be valid while token is valid.
+If token lifetime is longer than `cache_lifetime`, cache entry for this token will be valid for `cache_lifetime`. 
+
 ## IdP as External Authenticator {#idp-external-authenticator}
 
 Locally defined users can be authenticated with an access token. To allow this, `jwt` must be specified as user's authentication method. Example:
diff --git a/src/Access/AccessTokenProcessor.cpp b/src/Access/AccessTokenProcessor.cpp
@@ -9,7 +9,7 @@ namespace DB
 
 namespace
 {
-    /// The JSON reply from provider has only a few key-value pairs, so no need for SimdJSON/RapidJSON.
+    /// The JSON reply from provider has only a few key-value pairs, so no need for any advanced parsing.
     /// Reduce complexity by using picojson.
     picojson::object parseJSON(const String & json_string) {
         picojson::value jsonValue;
@@ -26,18 +26,20 @@ namespace
         return jsonValue.get<picojson::object>();
     }
 
-    std::string getValueByKey(const picojson::object & jsonObject, const std::string & key) {
+    template<typename ValueType = std::string>
+    ValueType getValueByKey(const picojson::object & jsonObject, const std::string & key) {
         auto it = jsonObject.find(key); // Find the key in the object
-        if (it == jsonObject.end()) {
+        if (it == jsonObject.end())
+        {
             throw std::runtime_error("Key not found: " + key);
         }
 
-        const picojson::value &value = it->second;
-        if (!value.is<std::string>()) {
-            throw std::runtime_error("Value for key '" + key + "' is not a string");
+        const picojson::value & value = it->second;
+        if (!value.is<ValueType>()) {
+            throw std::runtime_error("Value for key '" + key + "' has incorrect type.");
         }
 
-        return value.get<std::string>();
+        return value.get<ValueType>();
     }
 
     picojson::object getObjectFromURI(const Poco::URI & uri, const String & token = "")
@@ -96,9 +98,12 @@ std::unique_ptr<IAccessTokenProcessor> IAccessTokenProcessor::parseTokenProcesso
         String email_regex_str = config.hasProperty(prefix + ".email_filter") ? config.getString(
                 prefix + ".email_filter") : "";
 
+        UInt64 cache_lifetime = config.hasProperty(prefix + ".cache_lifetime") ? config.getUInt64(
+                prefix + ".cache_lifetime") : 3600;
+
         if (provider == "google")
         {
-            return std::make_unique<GoogleAccessTokenProcessor>(name, email_regex_str);
+            return std::make_unique<GoogleAccessTokenProcessor>(name, cache_lifetime, email_regex_str);
         }
         else if (provider == "azure")
         {
@@ -110,11 +115,9 @@ std::unique_ptr<IAccessTokenProcessor> IAccessTokenProcessor::parseTokenProcesso
                 throw Exception(ErrorCodes::INVALID_CONFIG_PARAMETER,
                                 "Could not parse access token processor {}: tenant_id must be specified", name);
 
-            String client_id_str = config.getString(prefix + ".client_id");
             String tenant_id_str = config.getString(prefix + ".tenant_id");
-            String client_secret_str  = config.hasProperty(prefix + ".client_secret") ? config.getString(prefix + ".client_secret") : "";
 
-            return std::make_unique<AzureAccessTokenProcessor>(name, email_regex_str, client_id_str, tenant_id_str, client_secret_str);
+            return std::make_unique<AzureAccessTokenProcessor>(name, cache_lifetime, email_regex_str, tenant_id_str);
         }
         else
             throw Exception(ErrorCodes::INVALID_CONFIG_PARAMETER,
@@ -132,10 +135,11 @@ bool GoogleAccessTokenProcessor::resolveAndValidate(const TokenCredentials & cre
 
     auto user_info = getUserInfo(token);
     String user_name = user_info["sub"];
+    bool has_email = user_info.contains("email");
 
     if (email_regex.ok())
     {
-        if (!user_info.contains("email"))
+        if (!has_email)
         {
             LOG_TRACE(getLogger("AccessTokenProcessor"), "{}: Failed to validate {} by e-mail", name, user_name);
             return false;
@@ -149,10 +153,54 @@ bool GoogleAccessTokenProcessor::resolveAndValidate(const TokenCredentials & cre
         }
 
     }
+
     /// Credentials are passed as const everywhere up the flow, so we have to comply,
     /// in this case const_cast looks acceptable.
     const_cast<TokenCredentials &>(credentials).setUserName(user_name);
-    const_cast<TokenCredentials &>(credentials).setGroups({});
+
+    auto token_info = getObjectFromURI(Poco::URI(token_info_uri), token);
+    if (token_info.contains("exp"))
+        const_cast<TokenCredentials &>(credentials).setExpiresAt(std::chrono::system_clock::from_time_t((getValueByKey<time_t>(token_info, "exp"))));
+
+    /// Groups info can only be retrieved if user email is known.
+    /// If no email found in user info, we skip this step and there are no external groups for the user.
+    if (has_email)
+    {
+        std::set<String> external_groups_names;
+        const Poco::URI get_groups_uri = Poco::URI("https://cloudidentity.googleapis.com/v1/groups/-/memberships:searchDirectGroups?query=member_key_id==" + user_info["email"] + "'");
+
+        try
+        {
+            auto groups_response = getObjectFromURI(get_groups_uri, token);
+
+            if (!groups_response.contains("memberships")) {
+                LOG_TRACE(getLogger("AccessTokenProcessor"),
+                          "{}: Failed to get Google groups: invalid content in response from server", name);
+                return true;
+            }
+
+            picojson::array groups_array = groups_response["memberships"].get<picojson::array>();
+
+            /// TODO: check for invalid JSON, LOG something meaningful
+            for (const auto & group: groups_array)
+            {
+                auto group_data = group.get<picojson::object>();
+                String group_name = getValueByKey(group_data["groupKey"].get<picojson::object>(), "id");
+                external_groups_names.insert(group_name);
+                LOG_TRACE(getLogger("AccessTokenProcessor"),
+                          "{}: User {}: new external group {}", name, user_name, group_name);
+            }
+
+            const_cast<TokenCredentials &>(credentials).setGroups(external_groups_names);
+        }
+        catch (const Exception & e)
+        {
+            /// Could not get groups info. Log it and skip it.
+            LOG_TRACE(getLogger("AccessTokenProcessor"),
+                      "{}: Failed to get Google groups, no external roles will be mapped. reason: {}", name, e.what());
+            return true;
+        }
+    }
 
     return true;
 }
@@ -177,8 +225,9 @@ std::unordered_map<String, String> GoogleAccessTokenProcessor::getUserInfo(const
 
 bool AzureAccessTokenProcessor::resolveAndValidate(const TokenCredentials & credentials)
 {
-    /// Token is a JWT in this case, but we cannot directly verify it against Azure AD JWKS. We will not trust any data in this token.
-    /// e.g. see here: https://stackoverflow.com/questions/60778634/failing-signature-validation-of-jwt-tokens-from-azure-ad
+    /// Token is a JWT in this case, but we cannot directly verify it against Azure AD JWKS.
+    /// We will not trust user data in this token except for 'exp' value to determine caching duration.
+    /// Explanation here: https://stackoverflow.com/questions/60778634/failing-signature-validation-of-jwt-tokens-from-azure-ad
     /// Let Azure validate it: only valid tokens will be accepted.
     /// Use GET https://graph.microsoft.com/oidc/userinfo to verify token and get sub at the same time
 
@@ -202,8 +251,49 @@ bool AzureAccessTokenProcessor::resolveAndValidate(const TokenCredentials & cred
         return false;
     }
 
-    /// TODO: do not store it in credentials.
-    const_cast<TokenCredentials &>(credentials).setGroups({});
+    try
+    {
+        const_cast<TokenCredentials &>(credentials).setExpiresAt(jwt::decode(token).get_expires_at());
+    }
+    catch (...) {
+        LOG_TRACE(getLogger("AccessTokenProcessor"),
+                  "{}: No expiration data found in a valid token, will use default cache lifetime", name);
+    }
+
+    std::set<String> external_groups_names;
+    const Poco::URI get_groups_uri = Poco::URI("https://graph.microsoft.com/v1.0/me/memberOf");
+
+    try
+    {
+        auto groups_response = getObjectFromURI(get_groups_uri, token);
+
+        if (!groups_response.contains("value")) {
+            LOG_TRACE(getLogger("AccessTokenProcessor"),
+                      "{}: Failed to get Azure groups: invalid content in response from server", name);
+            return true;
+        }
+
+        picojson::array groups_array = groups_response["value"].get<picojson::array>();
+
+        /// TODO: check for invalid JSON
+        for (const auto & group: groups_array)
+        {
+            auto group_data = group.get<picojson::object>();
+            String group_name = getValueByKey(group_data, "id");
+            external_groups_names.insert(group_name);
+            LOG_TRACE(getLogger("AccessTokenProcessor"),
+                      "{}: User {}: new external group {}", name, credentials.getUserName(), group_name);
+        }
+    }
+    catch (const Exception & e)
+    {
+        /// Could not get groups info. Log it and skip it.
+        LOG_TRACE(getLogger("AccessTokenProcessor"),
+                  "{}: Failed to get Azure groups, no external roles will be mapped. reason: {}", name, e.what());
+        return true;
+    }
+
+    const_cast<TokenCredentials &>(credentials).setGroups(external_groups_names);
 
     return true;
 }
diff --git a/src/Access/AccessTokenProcessor.h b/src/Access/AccessTokenProcessor.h
@@ -26,7 +26,12 @@ class GoogleAccessTokenProcessor;
 class IAccessTokenProcessor
 {
 public:
-    IAccessTokenProcessor(const String & name_, const String & email_regex_str) : name(name_), email_regex(email_regex_str)
+    IAccessTokenProcessor(const String & name_,
+                          const UInt64 cache_invalidation_interval_,
+                          const String & email_regex_str)
+                          : name(name_),
+                            cache_invalidation_interval(cache_invalidation_interval_),
+                            email_regex(email_regex_str)
     {
         if (!email_regex_str.empty())
         {
@@ -36,19 +41,12 @@ class IAccessTokenProcessor
         }
     }
 
-    String getName()
-    {
-        return name;
-    }
-
     virtual ~IAccessTokenProcessor() = default;
 
-    virtual bool resolveAndValidate(const TokenCredentials & credentials) = 0;
+    String getName() { return name; }
+    UInt64 getCacheInvalidationInterval() { return cache_invalidation_interval; }
 
-    virtual std::set<String> getGroups([[maybe_unused]] const TokenCredentials & credentials)
-    {
-        return {};
-    }
+    virtual bool resolveAndValidate(const TokenCredentials & credentials) = 0;
 
     static std::unique_ptr<DB::IAccessTokenProcessor> parseTokenProcessor(
         const Poco::Util::AbstractConfiguration & config,
@@ -57,14 +55,18 @@ class IAccessTokenProcessor
 
 protected:
     const String name;
+    const UInt64 cache_invalidation_interval;
     re2::RE2 email_regex;
 };
 
 
 class GoogleAccessTokenProcessor : public IAccessTokenProcessor
 {
 public:
-    GoogleAccessTokenProcessor(const String & name_, const String & email_regex_str) : IAccessTokenProcessor(name_, email_regex_str) {}
+    GoogleAccessTokenProcessor(const String & name_,
+                               const UInt64 cache_invalidation_interval_,
+                               const String & email_regex_str)
+                               : IAccessTokenProcessor(name_, cache_invalidation_interval_, email_regex_str) {}
 
     bool resolveAndValidate(const TokenCredentials & credentials) override;
 
@@ -80,24 +82,16 @@ class AzureAccessTokenProcessor : public IAccessTokenProcessor
 {
 public:
     AzureAccessTokenProcessor(const String & name_,
+                              const UInt64 cache_invalidation_interval_,
                               const String & email_regex_str,
-                              const String & client_id_,
-                              const String & tenant_id_,
-                              const String & client_secret_)
-                              : IAccessTokenProcessor(name_, email_regex_str),
-                                client_id(client_id_),
-                                tenant_id(tenant_id_),
-                                client_secret(client_secret_),
-                                jwks_uri_str("https://login.microsoftonline.com/" + tenant_id + "/discovery/v2.0/keys") {}
+                              const String & tenant_id_)
+                              : IAccessTokenProcessor(name_, cache_invalidation_interval_, email_regex_str),
+                                jwks_uri_str("https://login.microsoftonline.com/" + tenant_id_ + "/discovery/v2.0/keys") {}
 
     bool resolveAndValidate(const TokenCredentials & credentials) override;
 private:
     static const Poco::URI user_info_uri;
 
-    const String client_id;
-    const String tenant_id;
-    const String client_secret;
-
     const String jwks_uri_str;
 
     String validateTokenAndGetUsername(const String & token) const;
diff --git a/src/Access/Credentials.cpp b/src/Access/Credentials.cpp
@@ -101,6 +101,6 @@ const String & BasicCredentials::getPassword() const
 }
 
 /// Unless the token is validated, we will not use any data from it, including username.
-TokenCredentials::TokenCredentials(const String & token_) : Credentials(""), token(token_) {}
+TokenCredentials::TokenCredentials(const String & token_) : Credentials(""), token(token_), expires_at(std::chrono::system_clock::now() + std::chrono::hours(1)) {}
 
 }
diff --git a/src/Access/Credentials.h b/src/Access/Credentials.h
@@ -154,9 +154,18 @@ class TokenCredentials : public Credentials
     {
         groups = groups_;
     }
+    std::optional<std::chrono::system_clock::time_point> getExpiresAt() const
+    {
+        return expires_at;
+    }
+    void setExpiresAt(std::chrono::system_clock::time_point expires_at_)
+    {
+        expires_at = expires_at_;
+    }
 private:
     String token;
     std::set<String> groups;
+    std::optional<std::chrono::system_clock::time_point> expires_at;
 };
 
 }
diff --git a/src/Access/ExternalAuthenticators.cpp b/src/Access/ExternalAuthenticators.cpp
@@ -621,7 +621,6 @@ HTTPAuthClientParams ExternalAuthenticators::getHTTPAuthenticationParams(const S
     return it->second;
 }
 
-/// TODO: remove redundancy
 bool ExternalAuthenticators::resolveJWTCredentials(const TokenCredentials & credentials, bool throw_not_configured = true) const
 {
     std::lock_guard lock{mutex};
@@ -679,10 +678,48 @@ bool ExternalAuthenticators::checkAccessTokenCredentials(const TokenCredentials
     if (access_token_processors.empty())
         throw Exception(ErrorCodes::BAD_ARGUMENTS, "Access token authentication is not configured");
 
+    /// lookup token in local cache if not expired.
+    auto cached_entry_iter = access_token_cache.find(credentials.getToken());
+    if (cached_entry_iter != access_token_cache.end())
+    {
+        if (cached_entry_iter->second.expires_at <= std::chrono::system_clock::now())
+        {
+            LOG_TRACE(getLogger("AccessTokenAuthentication"), "Cache entry for user {} expired, removing", cached_entry_iter->second.user_name);
+            access_token_cache.erase(cached_entry_iter);
+        }
+        else
+        {
+            const auto & user_data = cached_entry_iter->second;
+            const_cast<TokenCredentials &>(credentials).setUserName(user_data.user_name);
+            const_cast<TokenCredentials &>(credentials).setGroups(user_data.external_roles);
+            LOG_TRACE(getLogger("AccessTokenAuthentication"), "Cache entry for user {} found, using it to authenticate", cached_entry_iter->second.user_name);
+            return true;
+        }
+    }
+
     for (const auto & it : access_token_processors)
     {
         if (it.second->resolveAndValidate(credentials))
         {
+            AccessTokenCacheEntry cache_entry;
+            cache_entry.user_name = credentials.getUserName();
+            cache_entry.external_roles = credentials.getGroups();
+
+            auto default_expiration_ts = std::chrono::system_clock::now()
+                                         + std::chrono::minutes(it.second->getCacheInvalidationInterval());
+
+            if (credentials.getExpiresAt().has_value())
+            {
+                if (credentials.getExpiresAt().value() < default_expiration_ts)
+                    cache_entry.expires_at = credentials.getExpiresAt().value();
+            }
+            else
+            {
+                cache_entry.expires_at = default_expiration_ts;
+            }
+            LOG_TRACE(getLogger("AccessTokenAuthentication"), "Cache entry for user {} added", cached_entry_iter->second.user_name);
+
+            access_token_cache[credentials.getToken()] = cache_entry;
             LOG_DEBUG(getLogger("AccessTokenAuthentication"), "Authenticated user {} with access token by {}", credentials.getUserName(), it.first);
             return true;
         }
diff --git a/src/Access/ExternalAuthenticators.h b/src/Access/ExternalAuthenticators.h
diff --git a/src/Access/TokenAccessStorage.cpp b/src/Access/TokenAccessStorage.cpp

Original file line number	Diff line number	Diff line change
`@@ -101,6 +101,6 @@ const String & BasicCredentials::getPassword() const`
`101`	`101`	`}`
`102`	`102`
`103`	`103`	`/// Unless the token is validated, we will not use any data from it, including username.`
`104`		`-TokenCredentials::TokenCredentials(const String & token_) : Credentials(""), token(token_) {}`
	`104`	`+TokenCredentials::TokenCredentials(const String & token_) : Credentials(""), token(token_), expires_at(std::chrono::system_clock::now() + std::chrono::hours(1)) {}`
`105`	`105`
`106`	`106`	`}`
Original file line number	Diff line number	Diff line change
`@@ -154,9 +154,18 @@ class TokenCredentials : public Credentials`
`154`	`154`	`{`
`155`	`155`	`groups = groups_;`
`156`	`156`	`}`
	`157`	`+ std::optional<std::chrono::system_clock::time_point> getExpiresAt() const`
	`158`	`+ {`
	`159`	`+ return expires_at;`
	`160`	`+ }`
	`161`	`+ void setExpiresAt(std::chrono::system_clock::time_point expires_at_)`
	`162`	`+ {`
	`163`	`+ expires_at = expires_at_;`
	`164`	`+ }`
`157`	`165`	`private:`
`158`	`166`	`String token;`
`159`	`167`	`std::set<String> groups;`
	`168`	`+ std::optional<std::chrono::system_clock::time_point> expires_at;`
`160`	`169`	`};`
`161`	`170`
`162`	`171`	`}`