Files
ncnn/tests/test_modelbin.cpp
Dreladwyn 201145bdd9 fix crash reading past end of memory-mapped model file (#7021)
* fix crash reading past end of memory-mapped model file

DataReaderFromMemory used by mapped model loading was unbounded:
read()/reference() advanced past the mapped region, so a truncated
or corrupt .bin crashed with an access violation inside
ModelBinFromDataReader instead of failing cleanly.

Add a size-aware DataReaderFromMemory ctor that clamps reads and
fails references past the end, use it in the mapped branches of
Net::load_model, and propagate load_model()'s return value which
was previously dropped.

* fix bounded memory reader scans and add loading regressions

* f

* reuse bounded memory reader scan buffer

---------

Co-authored-by: Dr3ladwyn <Dr3ladwyn@users.noreply.github.com>
Co-authored-by: nihui <nihuini@tencent.com>
2026-10-12 00:01:46 +08:00

230 lines
6.2 KiB
C++

// Copyright 2026 Tencent
// SPDX-License-Identifier: BSD-3-Clause
#include <stdio.h>
#include "datareader.h"
#include "modelbin.h"
static int test_mat_from_bfloat16()
{
const unsigned short data[] = {0x3f80, 0xc020, 0x3f00};
ncnn::Mat m = ncnn::Mat::from_bfloat16(data, 3);
if (m.empty() || m.elemsize != 4u || m[0] != 1.f || m[1] != -2.5f || m[2] != 0.5f)
{
fprintf(stderr, "test_mat_from_bfloat16 failed\n");
return -1;
}
return 0;
}
static int test_modelbin_bfloat16()
{
const unsigned char model_data[] = {
0x83, 0x8b, 0x34, 0x01,
0x80, 0x3f, 0x20, 0xc0, 0x00, 0x3f, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x80, 0x40
};
const unsigned char* mem = model_data;
ncnn::DataReaderFromMemory dr(mem);
ncnn::ModelBinFromDataReader mb(dr);
ncnn::Mat m = mb.load(3, 0);
if (m.empty() || m.elemsize != 4u || m[0] != 1.f || m[1] != -2.5f || m[2] != 0.5f)
{
fprintf(stderr, "test_modelbin_bfloat16 failed\n");
return -1;
}
ncnn::Mat m2 = mb.load(1, 0);
if (m2.empty() || m2[0] != 4.f)
{
fprintf(stderr, "test_modelbin_bfloat16 alignment failed\n");
return -1;
}
return 0;
}
class DataReaderFromMemoryNoReference : public ncnn::DataReader
{
public:
DataReaderFromMemoryNoReference(const unsigned char* _data, size_t _size)
: data(_data), size(_size), offset(0)
{
}
virtual size_t read(void* buf, size_t read_size) const
{
if (offset >= size)
return 0;
size_t remain = size - offset;
if (read_size > remain)
read_size = remain;
memcpy(buf, data + offset, read_size);
offset += read_size;
return read_size;
}
private:
const unsigned char* data;
size_t size;
mutable size_t offset;
};
static int test_modelbin_int8(const ncnn::DataReader& dr, int w)
{
ncnn::ModelBinFromDataReader mb(dr);
ncnn::Mat m = mb.load(w, 3);
if (m.empty() || m.w != w || m.elemsize != 1u || m.elempack != 1)
{
fprintf(stderr, "test_modelbin_int8 shape failed w=%d\n", w);
return -1;
}
const signed char expected[] = {0, 1, -128, 127};
const signed char* p = m;
for (int i = 0; i < w; i++)
{
if (p[i] != expected[i])
{
fprintf(stderr, "test_modelbin_int8 value failed w=%d i=%d %d != %d\n", w, i, p[i], expected[i]);
return -1;
}
}
ncnn::Mat m2 = mb.load(1, 1);
if (m2.empty() || m2[0] != 4.f)
{
fprintf(stderr, "test_modelbin_int8 alignment failed w=%d\n", w);
return -1;
}
return 0;
}
static int test_modelbin_int8(int w)
{
unsigned char model_data[] = {
0x00, 0x01, 0x80, 0x7f,
0x00, 0x00, 0x80, 0x40
};
for (int i = w; i < 4; i++)
model_data[i] = 0;
// reference path
const unsigned char* mem = model_data;
ncnn::DataReaderFromMemory dr(mem);
int ret = test_modelbin_int8(dr, w);
if (ret != 0)
return ret;
// copy path
// DataReaderFromMemoryNoReference inherits DataReader::reference(),
// which returns 0, so ModelBinFromDataReader falls back to read().
DataReaderFromMemoryNoReference dr2(model_data, sizeof(model_data));
ret = test_modelbin_int8(dr2, w);
if (ret != 0)
return ret;
return 0;
}
static int test_modelbin_int8_short_read()
{
// three int8 elements require four bytes in the model because raw
// int8 data is padded to a 32-bit boundary
const unsigned char model_data[] = {
0x00, 0x01, 0x00
};
DataReaderFromMemoryNoReference dr(model_data, sizeof(model_data));
ncnn::ModelBinFromDataReader mb(dr);
ncnn::Mat m = mb.load(3, 3);
if (!m.empty())
{
fprintf(stderr, "test_modelbin_int8_short_read failed\n");
return -1;
}
return 0;
}
static int test_modelbin_int8()
{
return 0
|| test_modelbin_int8(1)
|| test_modelbin_int8(2)
|| test_modelbin_int8(3)
|| test_modelbin_int8(4)
|| test_modelbin_int8_short_read();
}
static int test_modelbin_bounded()
{
const unsigned char model_data[][16] = {
{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x80, 0x3f, 0x00, 0x00, 0x20, 0xc0, 0x00, 0x00, 0x00, 0x3f},
{0x47, 0x6b, 0x30, 0x01, 0x00, 0x3c, 0x00, 0xc1, 0x00, 0x38, 0x00, 0x00},
{0x83, 0x8b, 0x34, 0x01, 0x80, 0x3f, 0x20, 0xc0, 0x00, 0x3f, 0x00, 0x00},
{0x38, 0x4b, 0x0d, 0x00, 0x00, 0x01, 0x80, 0x00},
{0x00, 0x01, 0x80, 0x00},
{0x00, 0x00, 0x80, 0x3f, 0x00, 0x00, 0x20, 0xc0, 0x00, 0x00, 0x00, 0x3f}
};
const size_t sizes[] = {16, 12, 12, 8, 4, 12};
const int types[] = {0, 0, 0, 0, 3, 1};
for (size_t i = 0; i < sizeof(sizes) / sizeof(sizes[0]); i++)
{
// truncate at every byte, including format tags and alignment padding
for (size_t size = 0; size <= sizes[i]; size++)
{
const unsigned char* mem = model_data[i];
ncnn::DataReaderFromMemory dr(mem, size);
ncnn::ModelBinFromDataReader mb(dr);
ncnn::Mat m = mb.load(3, types[i]);
if (m.empty() != (size < sizes[i]) || mem != model_data[i] + size)
{
fprintf(stderr, "test_modelbin_bounded failed format=%zu size=%zu\n", i, size);
return -1;
}
if (size == sizes[i])
{
if (i == 3 || i == 4)
{
const signed char* p = m;
if (m.elemsize != 1u || p[0] != 0 || p[1] != 1 || p[2] != -128)
{
fprintf(stderr, "test_modelbin_bounded int8 values failed\n");
return -1;
}
}
else if (m.elemsize != 4u || m[0] != 1.f || m[1] != -2.5f || m[2] != 0.5f)
{
fprintf(stderr, "test_modelbin_bounded float values failed format=%zu\n", i);
return -1;
}
}
}
}
return 0;
}
int main()
{
return 0
|| test_mat_from_bfloat16()
|| test_modelbin_bfloat16()
|| test_modelbin_int8()
|| test_modelbin_bounded();
}