|
| 1 | +// Copyright 2025 The Ray Authors. |
| 2 | +// |
| 3 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +// you may not use this file except in compliance with the License. |
| 5 | +// You may obtain a copy of the License at |
| 6 | +// |
| 7 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +// |
| 9 | +// Unless required by applicable law or agreed to in writing, |
| 10 | +// software distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +// See the License for the specific language governing permissions |
| 13 | +// and limitations under the License. |
| 14 | + |
| 15 | +#include "ray/rpc/authentication/authentication_token_validator.h" |
| 16 | + |
| 17 | +#include "ray/rpc/authentication/authentication_mode.h" |
| 18 | +#include "ray/rpc/authentication/k8s_util.h" |
| 19 | +#include "ray/util/logging.h" |
| 20 | + |
| 21 | +namespace ray { |
| 22 | +namespace rpc { |
| 23 | + |
| 24 | +const std::chrono::minutes kCacheTTL(5); |
| 25 | + |
| 26 | +AuthenticationTokenValidator &AuthenticationTokenValidator::instance() { |
| 27 | + static AuthenticationTokenValidator instance; |
| 28 | + return instance; |
| 29 | +} |
| 30 | + |
| 31 | +bool AuthenticationTokenValidator::ValidateToken( |
| 32 | + const std::optional<AuthenticationToken> &expected_token, |
| 33 | + const AuthenticationToken &provided_token) { |
| 34 | + if (GetAuthenticationMode() == AuthenticationMode::TOKEN) { |
| 35 | + if (!expected_token.has_value() || expected_token->empty()) { |
| 36 | + return true; // No auth required on server side |
| 37 | + } |
| 38 | + |
| 39 | + return expected_token->Equals(provided_token); |
| 40 | + } else if (GetAuthenticationMode() == AuthenticationMode::K8S) { |
| 41 | + std::call_once(k8s::k8s_client_config_flag, k8s::InitK8sClientConfig); |
| 42 | + if (!k8s::k8s_client_initialized) { |
| 43 | + RAY_LOG(WARNING) << "Kubernetes client not initialized, K8s authentication failed."; |
| 44 | + return false; |
| 45 | + } |
| 46 | + |
| 47 | + // Check cache first. |
| 48 | + { |
| 49 | + std::lock_guard<std::mutex> lock(k8s_token_cache_mutex_); |
| 50 | + auto it = k8s_token_cache_.find(provided_token); |
| 51 | + if (it != k8s_token_cache_.end()) { |
| 52 | + if (std::chrono::steady_clock::now() < it->second.expiration) { |
| 53 | + RAY_LOG(DEBUG) << "K8s token found in cache and is valid."; |
| 54 | + return it->second.allowed; |
| 55 | + } else { |
| 56 | + RAY_LOG(DEBUG) << "K8s token in cache expired, removing from cache."; |
| 57 | + k8s_token_cache_.erase(it); |
| 58 | + } |
| 59 | + } |
| 60 | + } |
| 61 | + |
| 62 | + bool is_allowed = false; |
| 63 | + is_allowed = k8s::ValidateToken(provided_token); |
| 64 | + |
| 65 | + // Only cache validated tokens for now. We don't want to invalidate a token |
| 66 | + // due to unrelated errors from Kubernetes API server. This has the downside of |
| 67 | + // causing more load if an unauthenticated client continues to make calls. |
| 68 | + // TODO(andrewsykim): cache invalid tokens once k8s::ValidateToken can distinguish |
| 69 | + // between invalid token errors and server errors. |
| 70 | + if (is_allowed) { |
| 71 | + std::lock_guard<std::mutex> lock(k8s_token_cache_mutex_); |
| 72 | + k8s_token_cache_[provided_token] = {is_allowed, |
| 73 | + std::chrono::steady_clock::now() + kCacheTTL}; |
| 74 | + RAY_LOG(DEBUG) << "K8s token validated and saved to cached."; |
| 75 | + } |
| 76 | + |
| 77 | + return is_allowed; |
| 78 | + } |
| 79 | + |
| 80 | + RAY_LOG(DEBUG) << "Authentication mode is disabled, token considered valid."; |
| 81 | + return true; |
| 82 | +} |
| 83 | + |
| 84 | +} // namespace rpc |
| 85 | +} // namespace ray |
0 commit comments