From 965d6f4bc75ea520e863845533a397f5e2b43343 Mon Sep 17 00:00:00 2001 From: David Fan Date: Mon, 6 Jul 2026 17:56:25 +0000 Subject: [PATCH 1/3] Make GeneratorParams co-own its Model to fix use-after-free GeneratorParams stored only a bare 'const Config&' aliasing the model-owned Config, relying on an undocumented lifetime assumption. Destroying the model (OgaDestroyModel) before OgaCreateGenerator freed the model and its Config, so Generator::Generator's model.shared_from_this() read freed memory (heap-use-after-free, CWE-416). Add a std::shared_ptr member to GeneratorParams (declared before the Config reference so it is initialized first) and bind the Config through it. This mirrors Generator, which already co-owns the model. The benchmark-only GeneratorParams(const Config&) path leaves model_owner_ null by design. Add a C API regression test that destroys the model handle before creating the generator and asserts it succeeds without a use-after-free. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/generators.cpp | 3 ++- src/generators.h | 5 ++++- test/c_api_tests.cpp | 25 +++++++++++++++++++++++++ 3 files changed, 31 insertions(+), 2 deletions(-) diff --git a/src/generators.cpp b/src/generators.cpp index 30d4e9844e..975f711efc 100644 --- a/src/generators.cpp +++ b/src/generators.cpp @@ -290,7 +290,8 @@ GeneratorParams::GeneratorParams(const Config& config) } GeneratorParams::GeneratorParams(const Model& model) - : config{*model.config_.get()}, + : model_owner_{model.shared_from_this()}, + config{*model_owner_->config_.get()}, use_graph_capture{IsGraphCaptureEnabled(model.config_->model.decoder.session_options)}, use_multi_profile{IsMultiProfileEnabled(model.config_->model.decoder.session_options)}, p_device{model.p_device_scoring_} { diff --git a/src/generators.h b/src/generators.h index f70d599b9f..35c53e8ce0 100644 --- a/src/generators.h +++ b/src/generators.h @@ -73,7 +73,10 @@ struct GeneratorParams : std::enable_shared_from_this, LeakChec GeneratorParams(const Config& config); // This constructor is only used for internal generator benchmarks GeneratorParams(const Model& model); - const Config& config; // The model outlives the GeneratorParams + // Co-owns the model so the aliased Config below cannot be freed while this + // params object is alive. Null for the benchmark-only Config constructor. + std::shared_ptr model_owner_; + const Config& config; // Aliases model-owned Config; kept alive by model_owner_ Config::Search search{config.search}; // Copy of the search parameters from the config // Query the params to get the value set for a param diff --git a/test/c_api_tests.cpp b/test/c_api_tests.cpp index 90056e79ed..f305f69634 100644 --- a/test/c_api_tests.cpp +++ b/test/c_api_tests.cpp @@ -1182,6 +1182,31 @@ TEST(CAPITests, TopKTopPExceedsVocabSizeThrows) { } } +// Regression test: a GeneratorParams created from a model must keep that model +// alive, so destroying the model handle before creating the generator does not +// cause a use-after-free (GeneratorParams aliases the model-owned Config, and +// Generator::Generator calls model.shared_from_this()). +TEST(CAPITests, CreateGeneratorAfterDestroyModel) { + OgaModel* model = nullptr; + ASSERT_EQ(OgaCreateModel(PHI2_PATH, &model), nullptr); + ASSERT_NE(model, nullptr); + + OgaGeneratorParams* params = nullptr; + ASSERT_EQ(OgaCreateGeneratorParams(model, ¶ms), nullptr); + ASSERT_NE(params, nullptr); + + // Drop the external reference to the model. The params must keep the + // underlying model (and its Config) alive, so the handle stays usable. + OgaDestroyModel(model); + + OgaGenerator* generator = nullptr; + ASSERT_EQ(OgaCreateGenerator(model, params, &generator), nullptr); + ASSERT_NE(generator, nullptr); + + OgaDestroyGenerator(generator); + OgaDestroyGeneratorParams(params); +} + TEST(CAPITests, AdaptersTest) { #ifdef USE_CUDA using OutputType = Ort::Float16_t; From eb37db1bc003f7e5ae6bb497343e66dd3b9f8d1f Mon Sep 17 00:00:00 2001 From: David Fan Date: Tue, 7 Jul 2026 23:12:03 +0000 Subject: [PATCH 2/3] Address review: clarify post-destroy model handle comment in test Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- test/c_api_tests.cpp | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/test/c_api_tests.cpp b/test/c_api_tests.cpp index f305f69634..167077e6a1 100644 --- a/test/c_api_tests.cpp +++ b/test/c_api_tests.cpp @@ -1195,8 +1195,11 @@ TEST(CAPITests, CreateGeneratorAfterDestroyModel) { ASSERT_EQ(OgaCreateGeneratorParams(model, ¶ms), nullptr); ASSERT_NE(params, nullptr); - // Drop the external reference to the model. The params must keep the - // underlying model (and its Config) alive, so the handle stays usable. + // Drop the external reference to the model by destroying its handle. Because + // params co-owns the underlying Model (and its Config) via shared ownership, + // the object itself stays alive, so dereferencing the raw model pointer below + // remains valid. This does NOT imply the handle is generally usable after + // OgaDestroyModel; it is valid here only because another owner keeps it alive. OgaDestroyModel(model); OgaGenerator* generator = nullptr; From 339ef81e3f05a9603358d8200093c392ffa178c4 Mon Sep 17 00:00:00 2001 From: David Fan Date: Fri, 10 Jul 2026 03:08:05 +0000 Subject: [PATCH 3/3] Address review: rename model_owner_ to model_ Per reviewer feedback, rename the co-owning member from model_owner_ to model_. The aliased const Config& config member is retained because the benchmark-only GeneratorParams(const Config&) constructor has no Model (model_ is null there), so config cannot be universally derived from model_; keeping the alias also avoids churning the many params.config call sites. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/generators.cpp | 4 ++-- src/generators.h | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/generators.cpp b/src/generators.cpp index 975f711efc..53ab91cc92 100644 --- a/src/generators.cpp +++ b/src/generators.cpp @@ -290,8 +290,8 @@ GeneratorParams::GeneratorParams(const Config& config) } GeneratorParams::GeneratorParams(const Model& model) - : model_owner_{model.shared_from_this()}, - config{*model_owner_->config_.get()}, + : model_{model.shared_from_this()}, + config{*model_->config_.get()}, use_graph_capture{IsGraphCaptureEnabled(model.config_->model.decoder.session_options)}, use_multi_profile{IsMultiProfileEnabled(model.config_->model.decoder.session_options)}, p_device{model.p_device_scoring_} { diff --git a/src/generators.h b/src/generators.h index 35c53e8ce0..2bc6bd0292 100644 --- a/src/generators.h +++ b/src/generators.h @@ -75,8 +75,8 @@ struct GeneratorParams : std::enable_shared_from_this, LeakChec // Co-owns the model so the aliased Config below cannot be freed while this // params object is alive. Null for the benchmark-only Config constructor. - std::shared_ptr model_owner_; - const Config& config; // Aliases model-owned Config; kept alive by model_owner_ + std::shared_ptr model_; + const Config& config; // Aliases model-owned Config; kept alive by model_ Config::Search search{config.search}; // Copy of the search parameters from the config // Query the params to get the value set for a param