oidc-proxy new iam 2

ydb/library/testlib/service_mocks/session_service_mock.h

@@ -25,7 +25,7 @@ class TSessionServiceMock : public yandex::cloud::priv::oauth::v1::SessionServic
     THashSet<TString> AllowedAccessTokens;
 
     TSessionServiceMock() {
-        AuthorizationRequiredMessage.Setauthorize_url("https://auth.cloud.yandex.ru/oauth/authorize");
+        AuthorizationRequiredMessage.Setauthorize_url("https://auth.test.net/oauth/authorize");
     }
 
     grpc::Status Check(grpc::ServerContext*,

ydb/mvp/core/mvp_tokens.cpp

@@ -1,12 +1,15 @@
 #include <contrib/libs/jwt-cpp/include/jwt-cpp/jwt.h>
 #include <ydb/library/actors/http/http_proxy.h>
+#include <ydb/mvp/core/core_ydb.h>
 #include <ydb/public/api/grpc/ydb_auth_v1.grpc.pb.h>
 #include "mvp_tokens.h"
+#include <ydb/public/api/client/nc_private/iam/token_service.grpc.pb.h>
+#include <ydb/public/api/client/nc_private/iam/token_exchange_service.grpc.pb.h>
 
 namespace NMVP {
 
-TMvpTokenator* TMvpTokenator::CreateTokenator(const NMvp::TTokensConfig& tokensConfig,  const NActors::TActorId& httpProxy) {
-    return new TMvpTokenator(tokensConfig, httpProxy);
+TMvpTokenator* TMvpTokenator::CreateTokenator(const NMvp::TTokensConfig& tokensConfig,  const NActors::TActorId& httpProxy, const NMVP::EAuthProfile authProfile) {
+    return new TMvpTokenator(tokensConfig, httpProxy, authProfile);
 }
 
 TString TMvpTokenator::GetToken(const TString& name) {
@@ -21,8 +24,9 @@ TString TMvpTokenator::GetToken(const TString& name) {
     return token;
 }
 
-TMvpTokenator::TMvpTokenator(NMvp::TTokensConfig tokensConfig, const NActors::TActorId& httpProxy)
+TMvpTokenator::TMvpTokenator(NMvp::TTokensConfig tokensConfig, const NActors::TActorId& httpProxy, const NMVP::EAuthProfile authProfile)
     : HttpProxy(httpProxy)
+    , AuthProfile(authProfile)
 {
     if (tokensConfig.HasStaffApiUserTokenInfo()) {
         UpdateStaffApiUserToken(&tokensConfig.staffapiusertokeninfo());
@@ -93,7 +97,7 @@ void TMvpTokenator::Handle(TEvPrivate::TEvRefreshToken::TPtr event) {
     BLOG_ERROR("Token " << name << " not found");
 }
 
-void TMvpTokenator::Handle(TEvPrivate::TEvUpdateIamToken::TPtr event) {
+void TMvpTokenator::Handle(TEvPrivate::TEvUpdateIamTokenYandex::TPtr event) {
     TDuration refreshPeriod = SUCCESS_REFRESH_PERIOD;
     if (event->Get()->Status.Ok()) {
         BLOG_D("Updating token " << event->Get()->Name << " to " << event->Get()->Response.subject());
@@ -108,6 +112,21 @@ void TMvpTokenator::Handle(TEvPrivate::TEvUpdateIamToken::TPtr event) {
     RefreshQueue.push({TInstant::Now() + refreshPeriod, event->Get()->Name});
 }
 
+void TMvpTokenator::Handle(TEvPrivate::TEvUpdateIamTokenNebius::TPtr event) {
+    TDuration refreshPeriod = SUCCESS_REFRESH_PERIOD;
+    if (event->Get()->Status.Ok()) {
+        BLOG_D("Updating token " << event->Get()->Name << " to " << event->Get()->Subject);
+        {
+            auto guard = Guard(TokensLock);
+            Tokens[event->Get()->Name] = "Bearer " + std::move(event->Get()->Response.access_token());
+        }
+    } else {
+        BLOG_ERROR("Error refreshing token " << event->Get()->Name << ", status: " << event->Get()->Status.GRpcStatusCode << ", error: " << event->Get()->Status.Msg);
+        refreshPeriod = ERROR_REFRESH_PERIOD;
+    }
+    RefreshQueue.push({TInstant::Now() + refreshPeriod, event->Get()->Name});
+}
+
 void TMvpTokenator::Handle(TEvPrivate::TEvUpdateStaticCredentialsToken::TPtr event) {
     TDuration refreshPeriod = SUCCESS_REFRESH_PERIOD;
     if (event->Get()->Status.Ok() && event->Get()->Response.operation().has_result()) {
@@ -221,27 +240,54 @@ void TMvpTokenator::Handle(NHttp::TEvHttpProxy::TEvHttpIncomingResponse::TPtr ev
 
 void TMvpTokenator::UpdateJwtToken(const NMvp::TJwtInfo* jwtInfo) {
     auto now = std::chrono::system_clock::now();
-    auto expires_at = now + std::chrono::hours(1);
-    auto serviceAccountId = jwtInfo->accountid();
-    auto keyId = jwtInfo->keyid();
+    auto expiresAt = now + std::chrono::hours(1);
+    const auto& serviceAccountId = jwtInfo->accountid();
+    const auto& keyId = jwtInfo->keyid();
     std::set<std::string> audience;
     audience.insert(jwtInfo->audience());
-    auto algorithm = jwt::algorithm::ps256(jwtInfo->publickey(), jwtInfo->privatekey());
 
-    auto encoded_token = jwt::create()
-            .set_key_id(keyId)
-            .set_issuer(serviceAccountId)
-            .set_audience(audience)
-            .set_issued_at(now)
-            .set_expires_at(expires_at)
-            .sign(algorithm);
+    switch (AuthProfile) {
+        case NMVP::EAuthProfile::Yandex: {
+            auto algorithm = jwt::algorithm::ps256(jwtInfo->publickey(), jwtInfo->privatekey());
+            auto encodedToken = jwt::create()
+                    .set_key_id(keyId)
+                    .set_issuer(serviceAccountId)
+                    .set_audience(audience)
+                    .set_issued_at(now)
+                    .set_expires_at(expiresAt)
+                    .sign(algorithm);
+            yandex::cloud::priv::iam::v1::CreateIamTokenRequest request;
+            request.set_jwt(TString(encodedToken));
+            RequestCreateToken<yandex::cloud::priv::iam::v1::IamTokenService,
+                                yandex::cloud::priv::iam::v1::CreateIamTokenRequest,
+                                yandex::cloud::priv::iam::v1::CreateIamTokenResponse,
+                                TEvPrivate::TEvUpdateIamTokenYandex>(jwtInfo->name(), jwtInfo->endpoint(), request, &yandex::cloud::priv::iam::v1::IamTokenService::Stub::AsyncCreate);
 
-    yandex::cloud::priv::iam::v1::CreateIamTokenRequest request;
-    request.set_jwt(TString(encoded_token));
-    RequestCreateToken<yandex::cloud::priv::iam::v1::IamTokenService,
-                       yandex::cloud::priv::iam::v1::CreateIamTokenRequest,
-                       yandex::cloud::priv::iam::v1::CreateIamTokenResponse,
-                       TEvPrivate::TEvUpdateIamToken>(jwtInfo->name(), jwtInfo->endpoint(), request, &yandex::cloud::priv::iam::v1::IamTokenService::Stub::AsyncCreate);
+            break;
+        }
+        case NMVP::EAuthProfile::Nebius: {
+            auto algorithm = jwt::algorithm::rs256(jwtInfo->publickey(), jwtInfo->privatekey());
+            auto encodedToken = jwt::create()
+                    .set_key_id(keyId)
+                    .set_issuer(serviceAccountId)
+                    .set_subject(serviceAccountId)
+                    .set_issued_at(now)
+                    .set_expires_at(expiresAt)
+                    .sign(algorithm);
+            nebius::iam::v1::ExchangeTokenRequest request;
+            request.set_grant_type("urn:ietf:params:oauth:grant-type:token-exchange");
+            request.set_requested_token_type("urn:ietf:params:oauth:token-type:access_token");
+            request.set_subject_token_type("urn:ietf:params:oauth:token-type:jwt");
+            request.set_subject_token(TString(encodedToken));
+
+            RequestCreateToken<nebius::iam::v1::TokenExchangeService,
+                                nebius::iam::v1::ExchangeTokenRequest,
+                                nebius::iam::v1::CreateTokenResponse,
+                                TEvPrivate::TEvUpdateIamTokenNebius>(jwtInfo->name(), jwtInfo->endpoint(), request, &nebius::iam::v1::TokenExchangeService::Stub::AsyncExchange, serviceAccountId);
+
+            break;
+        }
+    }
 }
 
 void TMvpTokenator::UpdateOAuthToken(const NMvp::TOAuthInfo* oauthInfo) {
@@ -250,7 +296,7 @@ void TMvpTokenator::UpdateOAuthToken(const NMvp::TOAuthInfo* oauthInfo) {
     RequestCreateToken<yandex::cloud::priv::iam::v1::IamTokenService,
                        yandex::cloud::priv::iam::v1::CreateIamTokenRequest,
                        yandex::cloud::priv::iam::v1::CreateIamTokenResponse,
-                       TEvPrivate::TEvUpdateIamToken>(oauthInfo->name(), oauthInfo->endpoint(), request, &yandex::cloud::priv::iam::v1::IamTokenService::Stub::AsyncCreate);
+                       TEvPrivate::TEvUpdateIamTokenYandex>(oauthInfo->name(), oauthInfo->endpoint(), request, &yandex::cloud::priv::iam::v1::IamTokenService::Stub::AsyncCreate);
 }
 
 }

ydb/mvp/core/mvp_tokens.h

@@ -9,17 +9,30 @@
 #include <library/cpp/json/json_writer.h>
 #include <library/cpp/json/json_reader.h>
 #include <ydb/mvp/core/protos/mvp.pb.h>
+#include <ydb/mvp/core/core_ydb.h>
 #include <ydb/public/api/client/yc_private/iam/iam_token_service.grpc.pb.h>
+#include <ydb/public/api/client/nc_private/iam/token_service.grpc.pb.h>
+#include <ydb/public/api/client/nc_private/iam/token_exchange_service.grpc.pb.h>
 #include <ydb/public/api/protos/ydb_auth.pb.h>
 #include "grpc_log.h"
 
 namespace NMVP {
 
+enum EAuthProfile {
+    Yandex = 1,
+    Nebius = 2
+};
+
+const THashMap<TString, EAuthProfile> AuthProfileByName = {
+    { "yandex", EAuthProfile::Yandex },
+    { "nebius", EAuthProfile::Nebius }
+};
+
 class TMvpTokenator : public NActors::TActorBootstrapped<TMvpTokenator> {
 public:
     using TBase = NActors::TActorBootstrapped<TMvpTokenator>;
 
-    static TMvpTokenator* CreateTokenator(const NMvp::TTokensConfig& tokensConfig, const NActors::TActorId& httpProxy);
+    static TMvpTokenator* CreateTokenator(const NMvp::TTokensConfig& tokensConfig, const NActors::TActorId& httpProxy, const NMVP::EAuthProfile authProfile = NMVP::EAuthProfile::Yandex);
     TString GetToken(const TString& name);
 
 protected:
@@ -32,8 +45,9 @@ class TMvpTokenator : public NActors::TActorBootstrapped<TMvpTokenator> {
     struct TEvPrivate {
         enum EEv {
             EvRefreshToken = EventSpaceBegin(NActors::TEvents::ES_PRIVATE),
-            EvUpdateIamToken,
+            EvUpdateIamTokenYandex,
             EvUpdateStaticCredentialsToken,
+            EvUpdateIamTokenNebius,
             EvEnd
         };
 
@@ -50,32 +64,37 @@ class TMvpTokenator : public NActors::TActorBootstrapped<TMvpTokenator> {
         template <ui32 TEventType, typename TResponse>
         struct TEvUpdateToken : NActors::TEventLocal<TEvUpdateToken<TEventType, TResponse>, TEventType> {
             TString Name;
+            TString Subject;
             NYdbGrpc::TGrpcStatus Status;
             TResponse Response;
 
-            TEvUpdateToken(const TString& name, NYdbGrpc::TGrpcStatus&& status, TResponse&& response)
+            TEvUpdateToken(const TString& name, const TString& subject, NYdbGrpc::TGrpcStatus&& status, TResponse&& response)
                 : Name(name)
+                , Subject(subject)
                 , Status(status)
                 , Response(response)
             {}
         };
 
-        using TEvUpdateIamToken = TEvUpdateToken<EvUpdateIamToken, yandex::cloud::priv::iam::v1::CreateIamTokenResponse>;
+        using TEvUpdateIamTokenYandex = TEvUpdateToken<EvUpdateIamTokenYandex, yandex::cloud::priv::iam::v1::CreateIamTokenResponse>;
+        using TEvUpdateIamTokenNebius = TEvUpdateToken<EvUpdateIamTokenNebius, nebius::iam::v1::CreateTokenResponse>;
         using TEvUpdateStaticCredentialsToken = TEvUpdateToken<EvUpdateStaticCredentialsToken, Ydb::Auth::LoginResponse>;
     };
 
-    TMvpTokenator(NMvp::TTokensConfig tokensConfig, const NActors::TActorId& httpProxy);
+    TMvpTokenator(NMvp::TTokensConfig tokensConfig, const NActors::TActorId& httpProxy, const EAuthProfile authProfile);
     void Bootstrap();
     void HandlePeriodic();
     void Handle(TEvPrivate::TEvRefreshToken::TPtr event);
-    void Handle(TEvPrivate::TEvUpdateIamToken::TPtr event);
+    void Handle(TEvPrivate::TEvUpdateIamTokenYandex::TPtr event);
+    void Handle(TEvPrivate::TEvUpdateIamTokenNebius::TPtr event);
     void Handle(TEvPrivate::TEvUpdateStaticCredentialsToken::TPtr event);
     void Handle(NHttp::TEvHttpProxy::TEvHttpIncomingResponse::TPtr event);
 
     STATEFN(StateWork) {
         switch (ev->GetTypeRewrite()) {
             hFunc(TEvPrivate::TEvRefreshToken, Handle);
-            hFunc(TEvPrivate::TEvUpdateIamToken, Handle);
+            hFunc(TEvPrivate::TEvUpdateIamTokenYandex, Handle);
+            hFunc(TEvPrivate::TEvUpdateIamTokenNebius, Handle);
             hFunc(TEvPrivate::TEvUpdateStaticCredentialsToken, Handle);
             hFunc(NHttp::TEvHttpProxy::TEvHttpIncomingResponse, Handle);
             cFunc(NActors::TEvents::TSystem::Wakeup, HandlePeriodic);
@@ -119,6 +138,7 @@ class TMvpTokenator : public NActors::TActorBootstrapped<TMvpTokenator> {
     TTokenConfigs TokenConfigs;
     TSpinLock TokensLock;
     NActors::TActorId HttpProxy;
+    const NMVP::EAuthProfile AuthProfile;
     THashMap<NHttp::THttpRequest*, TString> HttpRequestNames;
 
     template <typename TGRpcService>
@@ -134,15 +154,15 @@ class TMvpTokenator : public NActors::TActorBootstrapped<TMvpTokenator> {
     }
 
     template <typename TGrpcService, typename TRequest, typename TResponse, typename TUpdateToken>
-    void RequestCreateToken(const TString& name, const TString& endpoint, TRequest& request, typename NYdbGrpc::TSimpleRequestProcessor<typename TGrpcService::Stub, TRequest, TResponse>::TAsyncRequest asyncRequest) {
+    void RequestCreateToken(const TString& name, const TString& endpoint, TRequest& request, typename NYdbGrpc::TSimpleRequestProcessor<typename TGrpcService::Stub, TRequest, TResponse>::TAsyncRequest asyncRequest, const TString& subject = "") {
         NActors::TActorId actorId = SelfId();
         NActors::TActorSystem* actorSystem = NActors::TActivationContext::ActorSystem();
         NYdbGrpc::TCallMeta meta;
         meta.Timeout = RPC_TIMEOUT;
         auto connection = CreateGRpcServiceConnection<TGrpcService>(endpoint);
         NYdbGrpc::TResponseCallback<TResponse> cb =
-            [actorId, actorSystem, name](NYdbGrpc::TGrpcStatus&& status, TResponse&& response) -> void {
-                actorSystem->Send(actorId, new TUpdateToken(name, std::move(status), std::move(response)));
+            [actorId, actorSystem, name, subject](NYdbGrpc::TGrpcStatus&& status, TResponse&& response) -> void {
+                actorSystem->Send(actorId, new TUpdateToken(name, subject, std::move(status), std::move(response)));
         };
         connection->DoRequest(request, std::move(cb), asyncRequest, meta);
     }

ydb/mvp/core/ya.make

@@ -42,6 +42,7 @@ PEERDIR(
     ydb/public/api/client/yc_private/ydb/v1
     ydb/public/api/client/yc_private/resourcemanager
     ydb/public/api/client/yc_private/iam
+    ydb/public/api/client/nc_private/iam
     contrib/libs/googleapis-common-protos
     contrib/libs/jwt-cpp
     contrib/libs/yaml-cpp

ydb/mvp/oidc_proxy/example-config.yaml

@@ -15,10 +15,12 @@ generic:
     token_file: path/to/ydb_token
 
 oidc:
+  client_id: "client-id"
   secret_name: "oidc-secret"
   session_service_endpoint: "grpcs://ss.private-api.cloud-preprod.yandex.net:8655"
   session_service_token_name: "ss-token"
   authorization_server_address: "https://auth-preprod.cloud.yandex.ru"
+  auth_profile: "yandex"
   allowed_proxy_hosts:
     - "*.yandex.net"
     - "*.yandex-team.ru"

ydb/mvp/oidc_proxy/mvp.cpp

@@ -74,7 +74,7 @@ int TMVP::Init() {
     ActorSystem.Register(NActors::CreateProcStatCollector(TDuration::Seconds(5), AppData.MetricRegistry = std::make_shared<NMonitoring::TMetricRegistry>()));
 
     BaseHttpProxyId = ActorSystem.Register(NHttp::CreateHttpProxy(AppData.MetricRegistry));
-    ActorSystem.Register(AppData.Tokenator = TMvpTokenator::CreateTokenator(TokensConfig, BaseHttpProxyId));
+    ActorSystem.Register(AppData.Tokenator = TMvpTokenator::CreateTokenator(TokensConfig, BaseHttpProxyId, OpenIdConnectSettings.AuthProfile));
 
     HttpProxyId = ActorSystem.Register(NHttp::CreateHttpCache(BaseHttpProxyId, GetCachePolicy));
 
@@ -225,9 +225,13 @@ void TMVP::TryGetOidcOptionsFromConfig(const YAML::Node& config) {
     }
 
     SecretName = oidc["secret_name"].as<std::string>("");
+    OpenIdConnectSettings.ClientId = oidc["client_id"].as<std::string>(OpenIdConnectSettings.DEFAULT_CLIENT_ID);
     OpenIdConnectSettings.SessionServiceEndpoint = oidc["session_service_endpoint"].as<std::string>("");
     OpenIdConnectSettings.SessionServiceTokenName = oidc["session_service_token_name"].as<std::string>("");
     OpenIdConnectSettings.AuthorizationServerAddress = oidc["authorization_server_address"].as<std::string>("");
+    OpenIdConnectSettings.AuthEndpoint = oidc["auth_endpoint"].as<std::string>(OpenIdConnectSettings.DEFAULT_AUTH_ENDPOINT);
+    OpenIdConnectSettings.TokenEndpoint = oidc["token_endpoint"].as<std::string>(OpenIdConnectSettings.DEFAULT_TOKEN_ENDPOINT);
+    OpenIdConnectSettings.ExchangeEndpoint = oidc["exchange_endpoint"].as<std::string>(OpenIdConnectSettings.DEFAULT_EXCHANGE_ENDPOINT);
     Cout << "Started processing allowed_proxy_hosts..." << Endl;
     for (const std::string& host : oidc["allowed_proxy_hosts"].as<std::vector<std::string>>()) {
         Cout << host << " added to allowed_proxy_hosts" << Endl;
@@ -283,6 +287,17 @@ void TMVP::TryGetGenericOptionsFromConfig(
             HttpsPort = server["https_port"].as<ui16>(0);
         }
     }
+
+    if (generic["auth_profile"]) {
+        auto name = generic["auth_profile"].as<std::string>("yandex");
+        auto it = AuthProfileByName.find(name);
+        if (it != AuthProfileByName.end()) {
+            OpenIdConnectSettings.AuthProfile = it->second;
+        } else {
+            ythrow yexception() << "Unknown auth profile: " << name;
+        }
+    }
+
 }
 
 THolder<NActors::TActorSystemSetup> TMVP::BuildActorSystemSetup(int argc, char** argv) {

ydb/mvp/oidc_proxy/oidc_client.cpp

@@ -1,13 +1,13 @@
 #include "oidc_client.h"
-#include "oidc_protected_page.h"
-#include "oidc_session_create.h"
+#include "oidc_protected_page_handler.h"
+#include "oidc_session_create_handler.h"
 
 void InitOIDC(NActors::TActorSystem& actorSystem,
               const NActors::TActorId& httpProxyId,
               const TOpenIdConnectSettings& settings) {
     actorSystem.Send(httpProxyId, new NHttp::TEvHttpProxy::TEvRegisterHandler(
                          "/auth/callback",
-                         actorSystem.Register(new NMVP::TSessionCreator(httpProxyId, settings))
+                         actorSystem.Register(new NMVP::TSessionCreateHandler(httpProxyId, settings))
                          )
                      );