diff --git a/shared/api/speech_features.hpp b/shared/api/speech_features.hpp index 5bb09edc8..fafa586dd 100644 --- a/shared/api/speech_features.hpp +++ b/shared/api/speech_features.hpp @@ -194,7 +194,9 @@ class LogMel { } OrtxStatus Compute(const ortc::Tensor& stft_norm, ortc::Tensor& logmel) { - assert(stft_norm.Shape().size() == 3 && stft_norm.Shape()[0] == 1); + if (stft_norm.Shape().size() != 3 || stft_norm.Shape()[0] != 1) { + return {kOrtxErrorInvalidArgument, "[LogMel]: Input STFT tensor must have shape [1, freq, time]."}; + } const std::vector& stft_shape = stft_norm.Shape(); const int64_t stft_freq = stft_shape[1]; // freq bins (e.g., 257) @@ -202,6 +204,12 @@ class LogMel { const int64_t mel_freq = mel_filters_.nr(); // n_mel const int64_t mel_input_freq = mel_filters_.nc(); // expected input freq bins + if (stft_time < 2 || stft_freq < 1) { + return {kOrtxErrorInvalidArgument, + "[LogMel]: STFT output is degenerate (too few time steps or frequency bins). " + "The input audio data may be too short or malformed."}; + } + // Remove last frequency bin from STFT (to mimic Python stft[:, :, :-1]) const int64_t mag_time = stft_time - 1; diff --git a/test/pp_api_test/test_decode_audio.cc b/test/pp_api_test/test_decode_audio.cc index 9c6421104..4bddb3586 100644 --- a/test/pp_api_test/test_decode_audio.cc +++ b/test/pp_api_test/test_decode_audio.cc @@ -403,3 +403,49 @@ TEST(DecodeAudioTest, BatchDecodeFailFast) { EXPECT_EQ(results[0], nullptr) << "fail-fast must clear already-produced results"; EXPECT_EQ(results[1], nullptr); } + +// Regression test: decoding a known-malformed audio buffer must fail cleanly (no crash) +// and must not populate the output result slot on error. +TEST(DecodeAudioTest, MalformedAudioDoesNotCrash) { + // 17 bytes of malformed data that previously triggered a heap-buffer-overflow + // in the LogMel feature extractor via a degenerate spectrogram. + const uint8_t crash_data[] = { + 0xff, 0xff, 0x07, 0xfa, 0xe6, 0xe6, 0xe6, 0xe6, + 0xe6, 0xe6, 0xe6, 0xe6, 0xe6, 0xe6, 0xe6, 0xe6, 0xe6}; + + const void* data_ptrs[1] = {crash_data}; + int64_t sizes[1] = {static_cast(sizeof(crash_data))}; + + ort_extensions::OrtxObjectPtr raw_audios; + ASSERT_EQ(OrtxCreateRawAudios(raw_audios.ToBeAssigned(), data_ptrs, sizes, 1), kOrtxOK); + + // Attempt to decode — should fail gracefully, not crash. + OrtxTensorResult* result_ptr = nullptr; + extError_t err = OrtxDecodeAudios(raw_audios.get(), 0, /*stereo_to_mono=*/1, &result_ptr, 1); + ort_extensions::OrtxObjectPtr holder(result_ptr); // cleanup if populated + EXPECT_NE(err, kOrtxOK); + EXPECT_EQ(result_ptr, nullptr); +} + +// Verify that the LogMel feature extractor rejects malformed audio that produces +// a degenerate spectrogram, rather than crashing with an OOB read. +TEST(ExtractorTest, MalformedAudioLogMelDoesNotCrash) { + // Use a valid but extremely short WAV so decoding succeeds and LogMel sees a degenerate STFT. + auto wav = MakeSineWav(440.0f, 1.0f / 16000.0f, 16000); + + const void* data_ptrs[1] = {wav.data()}; + int64_t sizes[1] = {static_cast(wav.size())}; + + ort_extensions::OrtxObjectPtr raw_audios; + ASSERT_EQ(OrtxCreateRawAudios(raw_audios.ToBeAssigned(), data_ptrs, sizes, 1), kOrtxOK); + + // Use the Whisper feature extraction config. + ort_extensions::OrtxObjectPtr feature_extractor( + OrtxCreateSpeechFeatureExtractor, "data/whisper/feature_extraction.json"); + ASSERT_EQ(feature_extractor.Code(), kOrtxOK) << OrtxGetLastErrorMessage(); + + ort_extensions::OrtxObjectPtr result; + extError_t err = OrtxSpeechLogMel(feature_extractor.get(), raw_audios.get(), result.ToBeAssigned()); + // Should return an error, not crash. + EXPECT_NE(err, kOrtxOK); +}