cloud/blockstore/libs/ydbstats/ydbauth.cpp - main - githubmirror/nbs

cloud/blockstore/libs/ydbstats/ydbauth.cpp
188 строк5 KB
Kirill Pleshivtsev
Introduce UnsafeExtractValue for workaround broken semantics of TFuture::ExtractValue() (#5318)
25 фев 2026, 18:36
25 фев 2026, 18:365ab7a4d
#include "ydbauth.h"

#include "config.h"
#include "ydbstats.h"

#include <cloud/blockstore/libs/kikimr/events.h>

#include <cloud/storage/core/libs/common/error.h>
#include <cloud/storage/core/libs/common/future_helper.h>
#include <cloud/storage/core/libs/common/scheduler.h>
#include <cloud/storage/core/libs/diagnostics/logging.h>
#include <cloud/storage/core/libs/iam/iface/client.h>

#include <contrib/ydb/public/sdk/cpp/client/ydb_driver/driver.h>
#include <contrib/ydb/public/sdk/cpp/client/ydb_scheme/scheme.h>
#include <contrib/ydb/public/sdk/cpp/client/ydb_table/table.h>

#include <util/generic/hash.h>
#include <util/generic/hash_set.h>
#include <util/generic/maybe.h>
#include <util/generic/string.h>
#include <util/generic/vector.h>
#include <util/stream/file.h>
#include <util/string/builder.h>
#include <util/system/rwlock.h>

namespace NCloud::NBlockStore::NYdbStats {

using namespace NIamClient;

using namespace NThreading;

using namespace NYdb;

namespace {

////////////////////////////////////////////////////////////////////////////////

class TYdbTokenProvider final
    : public ICredentialsProvider
    , public std::enable_shared_from_this<TYdbTokenProvider>
{
private:
    const ISchedulerPtr Scheduler;
    const IIamTokenClientPtr Client;
    const TDuration TokenRefreshTimeBeforeExpiration;

    mutable TRWMutex Mutex;
    mutable TStringType Token;
    mutable TInstant ExpiresAt;

public:
    TYdbTokenProvider(
            ISchedulerPtr scheduler,
            IIamTokenClientPtr client,
            TDuration tokenRefreshTimeBeforeExpiration,
            TTokenInfo initialToken)
        : Scheduler(std::move(scheduler))
        , Client(std::move(client))
        , TokenRefreshTimeBeforeExpiration(
              tokenRefreshTimeBeforeExpiration)
        , Token(std::move(initialToken.Token))
        , ExpiresAt(initialToken.ExpiresAt)
    {}

    TStringType GetAuthInfo() const override
    {
        TReadGuard guard(Mutex);

        return Token;
    }

    bool IsValid() const override
    {
        TReadGuard guard(Mutex);
        return !Token.empty();
    }

    void Init()
    {
        TInstant expiresAt;
        {
            TReadGuard guard(Mutex);
            expiresAt = ExpiresAt;
        }

        ScheduleRefreshToken(expiresAt);
    }

private:
    void ScheduleRefreshToken(TInstant expiresAt) const
    {
        auto deadline =
            Max(expiresAt - TokenRefreshTimeBeforeExpiration, Now());
        Scheduler->Schedule(
            deadline,
            [weak = weak_from_this()]
            {
                if (auto self = weak.lock()) {
                    self->RefreshToken();
                }
            });
    }

    void RefreshToken() const
    {
        Client->GetTokenAsync().Subscribe(
            [weak = weak_from_this()](
                const TFuture<TResultOrError<TTokenInfo>>& future)
            {
                auto self = weak.lock();
                if (!self) {
                    return;
                }

                TInstant expiresAt;

                {
                    TWriteGuard guard(self->Mutex);
                    auto result = UnsafeExtractValue(future);
                    if (!HasError(result)) {
                        auto tokenInfo = result.ExtractResult();
                        self->Token = std::move(tokenInfo.Token);
                        self->ExpiresAt = tokenInfo.ExpiresAt;
                    }

                    expiresAt = self->ExpiresAt;
                }

                self->ScheduleRefreshToken(expiresAt);
            });
    }
};

using TYdbTokenProviderPtr = std::shared_ptr<TYdbTokenProvider>;

////////////////////////////////////////////////////////////////////////////////

class TIamCredentialsProviderFactory final
    : public ICredentialsProviderFactory
{
private:
    const TDuration IamTokenRefreshTimeBeforeExpiration;
    const TTokenInfo InitialTokenInfo;
    const ISchedulerPtr Scheduler;
    const IIamTokenClientPtr Client;

public:
    TIamCredentialsProviderFactory(
            TDuration iamTokenRefreshTimeBeforeExpiration,
            TTokenInfo initialTokenInfo,
            ISchedulerPtr scheduler,
            IIamTokenClientPtr client)
        : IamTokenRefreshTimeBeforeExpiration(
              iamTokenRefreshTimeBeforeExpiration)
        , InitialTokenInfo(std::move(initialTokenInfo))
        , Scheduler(std::move(scheduler))
        , Client(std::move(client))
    {}

    TCredentialsProviderPtr CreateProvider() const
    {
        auto res = make_shared<TYdbTokenProvider>(
            Scheduler,
            Client,
            IamTokenRefreshTimeBeforeExpiration,
            InitialTokenInfo);
        res->Init();
        return res;
    }
};

}  // namespace

TCredentialsProviderFactoryPtr CreateIamCredentialsProviderFactory(
    TDuration iamTokenRefreshTimeBeforeExpiration,
    TTokenInfo initialTokenInfo,
    ISchedulerPtr scheduler,
    IIamTokenClientPtr client)
{
    return std::make_shared<TIamCredentialsProviderFactory>(
        iamTokenRefreshTimeBeforeExpiration,
        std::move(initialTokenInfo),
        std::move(scheduler),
        std::move(client));
}

}   // namespace NCloud::NBlockStore::NYdbStats