diff --git a/backend/cpp/audio-cpp/grpc-server.cpp b/backend/cpp/audio-cpp/grpc-server.cpp index c582b823c..73e6ca1bd 100644 --- a/backend/cpp/audio-cpp/grpc-server.cpp +++ b/backend/cpp/audio-cpp/grpc-server.cpp @@ -62,16 +62,23 @@ std::atomic g_shutdown_requested{false}; // g_model_mu for its duration, so it has to work from a reference taken under // the lock and used after releasing it. Under a unique_ptr, a Free or a reload // arriving mid-request destroys the model out from under that reference. Every -// handler instead takes a counted reference through snapshot(), so Free drops -// the global's reference and whichever request finishes last destroys the model. +// handler instead takes a counted reference through snapshot_for(), so Free +// drops the global's reference and whichever request finishes last destroys the +// model. std::mutex g_model_mu; std::shared_ptr g_model; -// The only correct way for a handler to reach the model. Returns by value, so -// the caller owns a reference for as long as its local lives, and null when no -// model is loaded. Never return LoadedModel& from here: that is the shape that -// reintroduces the use-after-free. -std::shared_ptr snapshot() { +// The raw reference-taking primitive. Returns by value, so the caller owns a +// reference for as long as its local lives, and null when no model is loaded. +// Never return LoadedModel& from here: that is the shape that reintroduces the +// use-after-free. +// +// _unchecked because it performs NO request-level validation. Status is its +// only legitimate caller, because Status takes a HealthMessage, which carries +// no ModelIdentity to check. Every RPC that acts on a model must call +// snapshot_for instead; the name is deliberately unpleasant so that reaching +// for it is a visible decision rather than an omission. +std::shared_ptr snapshot_unchecked() { std::lock_guard lock(g_model_mu); return g_model; } @@ -130,7 +137,7 @@ int parse_device_index(const std::string &value) { // // The loaded identity is read off the LoadedModel rather than a separate // global, which is where this differs from llama-cpp. A handler holding the -// model through snapshot() is then necessarily judging against the identity +// model through snapshot_for() is then necessarily judging against the identity // THAT model was loaded with, and a concurrent reload cannot swap one without // the other. template @@ -153,6 +160,39 @@ GStatus check_model_identity(const audiocpp_backend::LoadedModel &model, "\", requested \"" + request->modelidentity() + "\""); } +// The ONE way an RPC handler should reach the model. Takes the counted +// reference, refuses when nothing is loaded, and runs the identity check, in +// that order. Returns null with `out` set to the status to return; on success +// returns the model and leaves `out` OK. +// +// This exists as a structural guarantee rather than a convenience. The identity +// check used to be two lines every handler had to remember, and nothing failed +// if a new handler forgot them: there is no C++ equivalent of +// pkg/grpc/model_identity_modalities_test.go, so the convention could only rot. +// Every handler already has to call something to obtain the model, so making +// the guarded call the shortest path means the default is correct and skipping +// it requires deliberately typing snapshot_unchecked. +// +// The order is forced: the identity being compared against belongs to the model +// this call is about to use, so it cannot precede taking the reference. And it +// must precede routing, or a mismatched request to a model that cannot serve +// the RPC leaks UNIMPLEMENTED instead of the NOT_FOUND the router matches on. +template +std::shared_ptr +snapshot_for(const Request *request, GStatus &out) { + auto model = snapshot_unchecked(); + if (model == nullptr) { + out = GStatus(grpc::StatusCode::FAILED_PRECONDITION, + "audio-cpp: no model is loaded; call LoadModel first"); + return nullptr; + } + out = check_model_identity(*model, request); + if (!out.ok()) { + return nullptr; + } + return model; +} + // Maps a thrown exception onto the gRPC status the client should see. GStatus to_status(const std::exception &err) { if (dynamic_cast(&err) != nullptr) { @@ -180,8 +220,11 @@ public: GStatus Status(ServerContext *, const backend::HealthMessage *, backend::StatusResponse *response) override { - response->set_state(snapshot() ? backend::StatusResponse::READY - : backend::StatusResponse::UNINITIALIZED); + // snapshot_unchecked is right here, and is the only place it is: + // HealthMessage carries no ModelIdentity, so there is nothing to check. + response->set_state(snapshot_unchecked() + ? backend::StatusResponse::READY + : backend::StatusResponse::UNINITIALIZED); return GStatus::OK; } @@ -257,7 +300,7 @@ public: GStatus Free(ServerContext *, const backend::HealthMessage *, backend::Result *result) override { // Drops the global's reference. Any handler still running holds its own - // from snapshot(), so the model is destroyed by whichever of them + // from snapshot_for(), so the model is destroyed by whichever of them // finishes last rather than underneath one of them. Destroying // LoadedModel releases its sessions first, then the model. std::shared_ptr released; @@ -285,20 +328,15 @@ public: GStatus VAD(ServerContext *, const backend::VADRequest *request, backend::VADResponse *response) override { try { - // snapshot(), and the result is held for the whole call. Free may + // snapshot_for, and the result is held for the whole call. Free may // arrive mid-request; it drops the global's reference only, so this - // one keeps the model alive until the handler returns. - const auto model = snapshot(); + // one keeps the model alive until the handler returns. It also + // performs the not-loaded and identity checks, so a stale route is + // refused before any work happens. + GStatus refusal = GStatus::OK; + const auto model = snapshot_for(request, refusal); if (model == nullptr) { - return GStatus(grpc::StatusCode::FAILED_PRECONDITION, - "audio-cpp: no model is loaded; call LoadModel first"); - } - // Before any work: a stale route must be refused, not served. - // Cannot precede snapshot(), because the identity being compared - // against belongs to the model this call is about to use. - const GStatus identity = check_model_identity(*model, request); - if (!identity.ok()) { - return identity; + return refusal; } audiocpp_backend::RequestShape shape; @@ -339,14 +377,10 @@ public: GStatus Diarize(ServerContext *, const backend::DiarizeRequest *request, backend::DiarizeResponse *response) override { try { - const auto model = snapshot(); + GStatus refusal = GStatus::OK; + const auto model = snapshot_for(request, refusal); if (model == nullptr) { - return GStatus(grpc::StatusCode::FAILED_PRECONDITION, - "audio-cpp: no model is loaded; call LoadModel first"); - } - const GStatus identity = check_model_identity(*model, request); - if (!identity.ok()) { - return identity; + return refusal; } audiocpp_backend::RequestShape shape; diff --git a/backend/cpp/audio-cpp/loaded_model.cpp b/backend/cpp/audio-cpp/loaded_model.cpp index ec5b0a7c3..9be67a17d 100644 --- a/backend/cpp/audio-cpp/loaded_model.cpp +++ b/backend/cpp/audio-cpp/loaded_model.cpp @@ -338,9 +338,11 @@ LoadedModel::LoadedModel(const std::string &resolved_path, } LoadedModel::Session LoadedModel::session_for(Rpc rpc, const RequestShape &shape, - const LaneEntry &lane) { + LaneEntry &lane) { // Proof of holding only. Nothing here reads it, and nothing should: its - // whole job is to make a caller that has not taken the lane fail to compile. + // whole job is to make a caller that has not taken the lane fail to + // compile. Non-const so it cannot bind to an inline acquire(), whose + // temporary would be released at the end of this call. (void)lane; const Route route = resolve_route(rpc, shape, capabilities_); if (!route.ok) { @@ -439,7 +441,7 @@ std::unique_ptr LoadedModel::acquire_owned(int requested_timeout_ms) engine::runtime::TaskResult run_offline(const LoadedModel::Session &session, const engine::runtime::TaskRequest &request, - const LaneEntry &lane) { + LaneEntry &lane) { // Proof of holding only, as in session_for. (void)lane; if (session.offline == nullptr) { diff --git a/backend/cpp/audio-cpp/loaded_model.h b/backend/cpp/audio-cpp/loaded_model.h index 9f2a4cf0c..d0711a38a 100644 --- a/backend/cpp/audio-cpp/loaded_model.h +++ b/backend/cpp/audio-cpp/loaded_model.h @@ -108,9 +108,26 @@ public: // session cache is an unsynchronised std::map and the sessions themselves // are not reentrant, so this must only be called with the lane held; the // lane admits one caller at a time, which is exactly the constraint the - // sessions impose. Pass the LaneEntry from acquire(). It is deliberately - // not a reference to `this`'s own lane member, because a caller can only - // produce a LaneEntry by having taken one. + // sessions impose. Pass the LaneEntry from acquire(). + // + // NON-CONST reference on purpose, and do not "tidy" it to const. A const + // reference binds to a temporary, which makes this compile: + // + // auto session = model->session_for(rpc, shape, model->acquire(0)); + // auto result = run_offline(session, task, model->acquire(0)); + // + // and each temporary dies at the end of its own full-expression, so the + // lane is released between the two calls. That is precisely the split this + // parameter exists to prevent, and it is the form a future caller is most + // likely to reach for because it reads as tidy. Requiring an lvalue forces + // a named entry whose scope spans both calls. + // + // What it proves is bounded, so do not over-trust it: it proves A lane was + // taken, not THIS model's lane. A caller determined to defeat it can + // construct an entry on an unrelated InferenceLane and pass that. It + // therefore catches the two mistakes that actually happen, forgetting the + // lane entirely and taking it after routing, and does not catch lane + // identity. // // STATE CONTRACT, and it is the CALLER'S to honour. Sessions are cached per // (task, mode), so a streaming session is normally the same warm object the @@ -131,8 +148,7 @@ public: // // Offline sessions need no such care: their interface has no reset and // run() takes a whole request. - Session session_for(Rpc rpc, const RequestShape &shape, - const LaneEntry &lane); + Session session_for(Rpc rpc, const RequestShape &shape, LaneEntry &lane); // Takes the inference lane, or throws LaneUnavailable. Serializes runs // against this model. `requested_timeout_ms` is a per-request wait hint @@ -200,7 +216,10 @@ private: // `lane` is a PROOF OF HOLDING, unused at runtime, for the same reason // session_for takes one: the session is not reentrant and prepare() mutates it, // so running without the lane is a data race. Making it a parameter turns that -// into a compile error instead of a comment. +// into a compile error instead of a comment. Non-const for the same reason as +// session_for's: a const reference would bind to `model.acquire(0)` written +// inline, and that temporary dies at the end of this call, releasing the lane +// before the caller's next one. // // Throws CapabilityError when the session is not an offline one. That should be // unreachable through session_for, which already refuses a non-offline session @@ -208,6 +227,6 @@ private: // dereference. engine::runtime::TaskResult run_offline(const LoadedModel::Session &session, const engine::runtime::TaskRequest &request, - const LaneEntry &lane); + LaneEntry &lane); } // namespace audiocpp_backend