diff --git a/src/nnue/network.cpp b/src/nnue/network.cpp index 20149c73ceb..57a422c5725 100644 --- a/src/nnue/network.cpp +++ b/src/nnue/network.cpp @@ -63,6 +63,8 @@ namespace Stockfish::Eval::NNUE { namespace Detail { +constexpr u32 MaxDescriptionSize = 1024 * 1024; + // Read evaluation function parameters template bool read_parameters(std::istream& stream, T& reference) { @@ -306,20 +308,29 @@ bool Network::read_header(std::istream& stream, u32* hashValue, std::string* des version = read_little_endian(stream); *hashValue = read_little_endian(stream); size = read_little_endian(stream); - if (!stream || version != Version) + if (!stream || version != Version || size > Detail::MaxDescriptionSize) return false; + desc->resize(size); - stream.read(&(*desc)[0], size); + + if (size) + stream.read(desc->data(), size); return !stream.fail(); } // Write network header bool Network::write_header(std::ostream& stream, u32 hashValue, const std::string& desc) const { + if (desc.size() > Detail::MaxDescriptionSize) + return false; + write_little_endian(stream, Version); write_little_endian(stream, hashValue); write_little_endian(stream, u32(desc.size())); - stream.write(&desc[0], desc.size()); + + if (!desc.empty()) + stream.write(desc.data(), desc.size()); + return !stream.fail(); }