diff options
| author | robot-piglet <[email protected]> | 2025-05-23 18:38:49 +0300 |
|---|---|---|
| committer | robot-piglet <[email protected]> | 2025-05-23 18:50:15 +0300 |
| commit | bed215f88992e53d74c041c0908555bbd70033ab (patch) | |
| tree | 848218907ce0ccf2d13d3a42caa5a04e08b19b3e | |
| parent | 969b824500454fea4fe524572add55633471e5cf (diff) | |
Intermediate changes
commit_hash:2267d618cccbb0a15c30048335abce1d7350ce07
| -rw-r--r-- | yt/yt/client/driver/shuffle_commands.cpp | 20 |
1 files changed, 20 insertions, 0 deletions
diff --git a/yt/yt/client/driver/shuffle_commands.cpp b/yt/yt/client/driver/shuffle_commands.cpp index 31a4d93eb7b..691870e399d 100644 --- a/yt/yt/client/driver/shuffle_commands.cpp +++ b/yt/yt/client/driver/shuffle_commands.cpp @@ -7,6 +7,7 @@ #include <yt/yt/client/formats/config.h> #include <yt/yt/client/signature/signature.h> +#include <yt/yt/client/signature/validator.h> #include <yt/yt/client/table_client/adapters.h> #include <yt/yt/client/table_client/table_output.h> @@ -14,6 +15,7 @@ namespace NYT::NDriver { +using namespace NApi; using namespace NConcurrency; using namespace NFormats; using namespace NTableClient; @@ -80,6 +82,15 @@ void TReadShuffleDataCommand::DoExecute(ICommandContextPtr context) { auto client = context->GetClient(); + const auto& signatureValidator = context->GetDriver()->GetSignatureValidator(); + auto validationSuccessful = WaitFor(signatureValidator->Validate(SignedShuffleHandle.Underlying())) + .ValueOrThrow(); + if (!validationSuccessful) { + auto shuffleHandle = ConvertTo<TShuffleHandlePtr>(TYsonStringBuf(SignedShuffleHandle.Underlying()->Payload())); + THROW_ERROR_EXCEPTION("Signature validation failed for shuffle handle") + << TErrorAttribute("shuffle_handle", shuffleHandle); + } + std::optional<std::pair<int, int>> writerIndexRange; if (WriterIndexBegin.has_value()) { writerIndexRange = std::pair(*WriterIndexBegin, *WriterIndexEnd); @@ -140,6 +151,15 @@ void TWriteShuffleDataCommand::DoExecute(ICommandContextPtr context) { auto client = context->GetClient(); + const auto& signatureValidator = context->GetDriver()->GetSignatureValidator(); + auto validationSuccessful = WaitFor(signatureValidator->Validate(SignedShuffleHandle.Underlying())) + .ValueOrThrow(); + if (!validationSuccessful) { + auto shuffleHandle = ConvertTo<TShuffleHandlePtr>(TYsonStringBuf(SignedShuffleHandle.Underlying()->Payload())); + THROW_ERROR_EXCEPTION("Signature validation failed for shuffle handle") + << TErrorAttribute("shuffle_handle", shuffleHandle); + } + Options.OverwriteExistingWriterData = OverwriteExistingWriterData; auto writer = WaitFor(context->GetClient()->CreateShuffleWriter( |