Skip to content

Fix use-after-free by making GeneratorParams co-own its Model - #2270

Merged
jiafatom merged 3 commits into
mainfrom
fix/generatorparams-model-lifetime
Jul 13, 2026
Merged

Fix use-after-free by making GeneratorParams co-own its Model#2270
jiafatom merged 3 commits into
mainfrom
fix/generatorparams-model-lifetime

Conversation

@jiafatom

@jiafatom jiafatom commented Jul 6, 2026

Copy link
Copy Markdown
Contributor

Summary

Fixes a use-after-free (CWE-416) reachable through the public C ABI.

GeneratorParams stored only a bare const Config& aliasing the model-owned Config (src/generators.h), relying on the documented-but-unenforced assumption that "the model outlives the GeneratorParams". It did not take co-ownership of the model. The following ordinary C API sequence therefore triggers a UAF:

  1. OgaCreateModel — allocates the Model and its std::unique_ptr<Config> config_.
  2. OgaCreateGeneratorParams — binds config{*model.config_.get()}, a bare reference; the model's refcount is unchanged.
  3. OgaDestroyModelExternalRelease drops the last owner and frees the Model and its Config.
  4. OgaCreateGenerator(model, params)Generator::Generator runs model_{model.shared_from_this()} on the freed model → heap-use-after-free read.

Generator already co-owns the model via std::shared_ptr<const Model> model_; GeneratorParams did not follow the same discipline, leaving the window open.

Fix

  • src/generators.h: add std::shared_ptr<const Model> model_owner_ to GeneratorParams, declared before the config reference so it is initialized first.
  • src/generators.cpp: initialize model_owner_{model.shared_from_this()} and bind config{*model_owner_->config_.get()}. Now the model (and its Config) cannot be freed while a params handle is live, and a generator can never be constructed from a freed model. The benchmark-only GeneratorParams(const Config&) path leaves model_owner_ null by design.
  • test/c_api_tests.cpp: add CreateGeneratorAfterDestroyModel, reproducing the PoC ordering (destroy model handle, then create generator) and asserting success without a UAF.

Testing

Local build isn't available in my environment (root-owned build dir / needs the jiafa-dev container); relying on CI to build and run the C++ unit tests, ideally under ASan.

Copilot AI review requested due to automatic review settings July 6, 2026 17:56
@jiafatom
jiafatom requested a review from a team as a code owner July 6, 2026 17:56

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR fixes a use-after-free in the public C API by ensuring GeneratorParams co-owns the Model it was created from, preventing the model-owned Config from being freed while params are still alive.

Changes:

  • Add std::shared_ptr<const Model> model_owner_ to GeneratorParams to extend model lifetime.
  • Initialize model_owner_ via model.shared_from_this() and bind config through model_owner_.
  • Add a C API regression test covering the destroy-model-then-create-generator sequence.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File Description
test/c_api_tests.cpp Adds regression coverage for the C API lifetime/UAF scenario.
src/generators.h Extends GeneratorParams to co-own the model backing its aliased Config.
src/generators.cpp Initializes model_owner_ and binds config through it to prevent UAF.

Comment thread test/c_api_tests.cpp Outdated
@jiafatom
jiafatom force-pushed the fix/generatorparams-model-lifetime branch from 71a116a to 31a4c64 Compare July 7, 2026 23:06
Comment thread src/generators.h Outdated
jiafatom and others added 3 commits July 10, 2026 03:13
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<const Model> 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>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
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>
@jiafatom
jiafatom force-pushed the fix/generatorparams-model-lifetime branch from 8284fe9 to 339ef81 Compare July 10, 2026 03:14
@jiafatom
jiafatom merged commit 33bf39f into main Jul 13, 2026
62 checks passed
@jiafatom
jiafatom deleted the fix/generatorparams-model-lifetime branch July 13, 2026 17:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants