Repository: ray-project/ray
File: src/ray/rpc/authentication/authentication_token_loader.cc (C++)

Code before:
#include "ray/rpc/authentication/k8s_constants.h"
#include "ray/util/logging.h"

#ifdef _WIN32
#ifndef _WINDOWS_
#ifndef WIN32_LEAN_AND_MEAN  // Sorry for the inconvenience. Please include any related
                             // headers you need manually.
                             // (https://stackoverflow.com/a/8294669)
#define WIN32_LEAN_AND_MEAN  // Prevent inclusion of WinSock2.h
#endif
#include <Windows.h>  // Force inclusion of WinGDI here to resolve name conflict
#endif
#endif

namespace ray {
namespace rpc {

constexpr const char *kNoTokenErrorMessage =
    "Token authentication is enabled but Ray couldn't find an "
    "authentication token. "
    "Create a token file at ~/.ray/auth_token, "
    "or store the token in any file and set RAY_AUTH_TOKEN_PATH to point to it, "
    "or set the RAY_AUTH_TOKEN environment variable.";

constexpr int kRaySATokenDefaultTTLSeconds = 600;
constexpr int kRaySATokenExpirationBufferSeconds = 300;

std::optional<std::chrono::system_clock::time_point>
AuthenticationTokenLoader::GetTokenExpiration(const std::string &token) {
  std::vector<std::string> parts = absl::StrSplit(token, '.');
  if (parts.size() != 3) {
    RAY_LOG(WARNING) << "Invalid JWT token format.";
    return std::nullopt;
  }

  std::string payload;
  if (!absl::WebSafeBase64Unescape(parts[1], &payload)) {
    RAY_LOG(WARNING) << "Unable to base64 decode JWT token.";
    return std::nullopt;
  }

  try {
    auto json = nlohmann::json::parse(payload);
    if (json.contains("exp") && json["exp"].is_number()) {
      int64_t exp = json["exp"].get<int64_t>();
      return std::chrono::system_clock::from_time_t(exp);
    }
  } catch (...) {
    return std::nullopt;
  }


Review comment (performance):
just to confirm, but we dont expect jwt tokens to have 5min expiration time right (ie k8s tokens cannot be configured to rotate every 5min). because in that case we would never cache the token and keep reading from the the file every time.

Code after:
#include "ray/rpc/authentication/k8s_constants.h"
#include "ray/util/logging.h"

#ifdef _WIN32
#ifndef _WINDOWS_
#ifndef WIN32_LEAN_AND_MEAN  // Sorry for the inconvenience. Please include any related
                             // headers you need manually.
                             // (https://stackoverflow.com/a/8294669)
#define WIN32_LEAN_AND_MEAN  // Prevent inclusion of WinSock2.h
#endif
#include <Windows.h>  // Force inclusion of WinGDI here to resolve name conflict
#endif
#endif

namespace ray {
namespace rpc {

constexpr const char *kNoTokenErrorMessage =
    "Token authentication is enabled but Ray couldn't find an "
    "authentication token. "
    "Create a token file at ~/.ray/auth_token, "
    "or store the token in any file and set RAY_AUTH_TOKEN_PATH to point to it, "
    "or set the RAY_AUTH_TOKEN environment variable.";

std::optional<std::chrono::system_clock::time_point>
AuthenticationTokenLoader::GetJWTTokenExpiration(const std::string &token) {
  std::vector<std::string> parts = absl::StrSplit(token, '.');
  if (parts.size() != 3) {
    RAY_LOG(WARNING) << "Invalid JWT token format.";
    return std::nullopt;
  }

  std::string payload;
  if (!absl::WebSafeBase64Unescape(parts[1], &payload)) {
    RAY_LOG(WARNING) << "Unable to base64 decode JWT token.";
    return std::nullopt;
  }

  try {
    auto json = nlohmann::json::parse(payload);
    if (json.contains("exp") && json["exp"].is_number()) {
      int64_t exp = json["exp"].get<int64_t>();
      return std::chrono::system_clock::from_time_t(exp) -
             std::chrono::seconds(ray::rpc::k8s::kRaySATokenExpirationBufferSeconds);
    }
  } catch (...) {
    return std::nullopt;
  }

  return std::chrono::system_clock::now() +
         std::chrono::seconds(ray::rpc::k8s::kRaySATokenDefaultTTLSeconds);