From 8823046e4bdac4ff29613acb1d0be460c8bb22ca Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Fri, 13 Oct 2023 09:59:10 -0700 Subject: [PATCH] Add check to avoid doing illegal memory access from an invalid iterator from std::prev() PiperOrigin-RevId: 573248334 --- .../framework/tool/switch_demux_calculator.cc | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/mediapipe/framework/tool/switch_demux_calculator.cc b/mediapipe/framework/tool/switch_demux_calculator.cc index c066d470a..a326621db 100644 --- a/mediapipe/framework/tool/switch_demux_calculator.cc +++ b/mediapipe/framework/tool/switch_demux_calculator.cc @@ -13,6 +13,7 @@ // limitations under the License. #include +#include #include #include #include @@ -223,8 +224,8 @@ absl::Status SwitchDemuxCalculator::RecordPackets(CalculatorContext* cc) { // Returns the channel index for a Timestamp. int SwitchDemuxCalculator::ChannelIndex(Timestamp timestamp) { - auto it = std::prev(channel_history_.upper_bound(timestamp)); - return it->second; + auto it = channel_history_.upper_bound(timestamp); + return it == channel_history_.begin() ? -1 : std::prev(it)->second; } // Dispatches all queued input packets with known channels. @@ -237,10 +238,12 @@ absl::Status SwitchDemuxCalculator::SendActivePackets(CalculatorContext* cc) { auto& queue = input_queue_[input_id]; while (!queue.empty() && queue.front().Timestamp() <= channel_settled) { int channel_index = ChannelIndex(queue.front().Timestamp()); - std::string output_tag = tool::ChannelTag(tag, channel_index); - auto output_id = cc->Outputs().GetId(output_tag, index); - if (output_id.IsValid()) { - cc->Outputs().Get(output_id).AddPacket(queue.front()); + if (channel_index != -1) { + std::string output_tag = tool::ChannelTag(tag, channel_index); + auto output_id = cc->Outputs().GetId(output_tag, index); + if (output_id.IsValid()) { + cc->Outputs().Get(output_id).AddPacket(queue.front()); + } } queue.pop(); }