diff --git a/src/rest_utils.cpp b/src/rest_utils.cpp index b3fb2768b2..f96e9492d3 100644 --- a/src/rest_utils.cpp +++ b/src/rest_utils.cpp @@ -15,6 +15,7 @@ //***************************************************************************** #include "rest_utils.hpp" +#include #include #include #include @@ -46,6 +47,21 @@ enum : unsigned int { namespace ovms { +namespace { +// Integral types are always JSON-representable; only floating values can be NaN/Inf, +// which rapidjson cannot serialize (and which are invalid per the JSON spec). +template +inline bool isJsonRepresentable(T) { + return true; +} +inline bool isJsonRepresentable(float v) { + return std::isfinite(v); +} +inline bool isJsonRepresentable(double v) { + return std::isfinite(v); +} +} // namespace + static Status checkValField(size_t valFieldSize, size_t expectedElementsNumber) { if (valFieldSize != expectedElementsNumber) { std::stringstream ss; @@ -126,6 +142,9 @@ static void appendBinaryOutput(std::string& bytesOutputsBuffer, char* output, si appendBinaryOutput(bytesOutputsBuffer, (char*)tensor.contents().CONTENTS_FIELD().data(), expectedContentSize); \ } else { \ for (auto& number : tensor.contents().CONTENTS_FIELD()) { \ + if (!isJsonRepresentable(number)) \ + return Status(StatusCode::JSON_SERIALIZATION_ERROR, \ + "Output \"" + tensor.name() + "\" contains a non-finite (NaN/Inf) value that cannot be serialized to JSON"); \ writer.WRITER_TYPE(number); \ } \ } \ @@ -133,8 +152,13 @@ static void appendBinaryOutput(std::string& bytesOutputsBuffer, char* output, si if (binaryOutput) { \ appendBinaryOutput(bytesOutputsBuffer, (char*)response_proto.raw_output_contents(tensor_it).data(), expectedContentSize); \ } else { \ - for (size_t i = 0; i < response_proto.raw_output_contents(tensor_it).size(); i += sizeof(DATATYPE)) \ - writer.WRITER_TYPE(*(reinterpret_cast(response_proto.raw_output_contents(tensor_it).data() + i))); \ + for (size_t i = 0; i < response_proto.raw_output_contents(tensor_it).size(); i += sizeof(DATATYPE)) { \ + DATATYPE valToWrite = *(reinterpret_cast(response_proto.raw_output_contents(tensor_it).data() + i)); \ + if (!isJsonRepresentable(valToWrite)) \ + return Status(StatusCode::JSON_SERIALIZATION_ERROR, \ + "Output \"" + tensor.name() + "\" contains a non-finite (NaN/Inf) value that cannot be serialized to JSON"); \ + writer.WRITER_TYPE(valToWrite); \ + } \ } \ } diff --git a/src/test/rest_utils_test.cpp b/src/test/rest_utils_test.cpp index fe10425da9..62cec19da9 100644 --- a/src/test/rest_utils_test.cpp +++ b/src/test/rest_utils_test.cpp @@ -13,6 +13,7 @@ // See the License for the specific language governing permissions and // limitations under the License. //***************************************************************************** +#include #include #include @@ -103,6 +104,23 @@ TEST_F(KFSMakeJsonFromPredictResponseRawTest, Positive) { })"); } +TEST_F(KFSMakeJsonFromPredictResponseRawTest, NonFiniteInfFloatError) { + float bad[8] = {5.0f, std::numeric_limits::infinity(), -3.0f, 2.5f, 9.0f, 55.5f, -0.5f, -1.5f}; + proto.mutable_raw_output_contents(0)->assign(reinterpret_cast(bad), 8 * sizeof(float)); + EXPECT_EQ(makeJsonFromPredictResponse(proto, &json, inferenceHeaderContentLength), StatusCode::JSON_SERIALIZATION_ERROR); +} + +TEST_F(KFSMakeJsonFromPredictResponseRawTest, NonFiniteNaNFloatError) { + float bad[8] = {std::numeric_limits::quiet_NaN(), 10.0f, -3.0f, 2.5f, 9.0f, 55.5f, -0.5f, -1.5f}; + proto.mutable_raw_output_contents(0)->assign(reinterpret_cast(bad), 8 * sizeof(float)); + EXPECT_EQ(makeJsonFromPredictResponse(proto, &json, inferenceHeaderContentLength), StatusCode::JSON_SERIALIZATION_ERROR); +} + +TEST_F(KFSMakeJsonFromPredictResponseRawTest, IntOutputsUnaffectedByFiniteCheck) { + // sanity: integer outputs must still serialize fine (helper returns true for them) + ASSERT_EQ(makeJsonFromPredictResponse(proto, &json, inferenceHeaderContentLength), StatusCode::OK); +} + TEST_F(KFSMakeJsonFromPredictResponseRawTest, EmptyRawOutputContentsError) { proto.mutable_raw_output_contents()->Clear(); EXPECT_EQ(makeJsonFromPredictResponse(proto, &json, inferenceHeaderContentLength), StatusCode::REST_SERIALIZE_VAL_FIELD_INVALID_SIZE);