commit 3bb4ee67577f68107b11272216a9fd9c6cbb28a5
parent 183d774937c14e9b3f36ba298c296111747744aa
Author: Randy Concepcion <rconcepcion@mozilla.com>
Date: Wed, 19 Nov 2025 16:04:05 +0000
Bug 1997986 - Add ML security layer initially as a pass-through component for LLM calls and tool dispatch. r=tarek,ai-ondevice-reviewers
- Moved JSActor pattern from security-specific JSActor back to parent
MLEngine JSActor
- Move validation logic for both request/response from child to parent
to avoid IPC round trip (improve performance)
Differential Revision: https://phabricator.services.mozilla.com/D271495
Diffstat:
6 files changed, 173 insertions(+), 9 deletions(-)
diff --git a/toolkit/components/ml/actors/MLEngineParent.sys.mjs b/toolkit/components/ml/actors/MLEngineParent.sys.mjs
@@ -1040,6 +1040,30 @@ export class MLEngine {
}
/**
+ * Validates an inference request before sending to child process.
+ *
+ * @param {object} request - The request to validate
+ * @returns {object|null} The validated request, or null if blocked
+ * @private
+ */
+ #validateRequest(request) {
+ lazy.console.debug("[MLSecurity] Validating request:", request);
+ return request;
+ }
+
+ /**
+ * Validates an inference response after receiving from child process.
+ *
+ * @param {object} response - The response to validate
+ * @returns {object|null} The validated response, or null if blocked
+ * @private
+ */
+ #validateResponse(response) {
+ lazy.console.debug("[MLSecurity] Validating response:", response);
+ return response;
+ }
+
+ /**
* Observes shutdown events from the child process.
*
* When the inference process is shutdown, we want to set the port to null and throw an error.
@@ -1308,12 +1332,19 @@ export class MLEngine {
});
}
if (response) {
- const totalTime =
- response.metrics.tokenizingTime + response.metrics.inferenceTime;
- Glean.firefoxAiRuntime.runInferenceSuccess[
- this.getGleanLabel()
- ].accumulateSingleSample(totalTime);
- request.resolve(response);
+ // Validate response before returning to caller
+ const validatedResponse = this.#validateResponse(response);
+ if (!validatedResponse) {
+ request.reject(new Error("Response failed security validation"));
+ } else {
+ const totalTime =
+ validatedResponse.metrics.tokenizingTime +
+ validatedResponse.metrics.inferenceTime;
+ Glean.firefoxAiRuntime.runInferenceSuccess[
+ this.getGleanLabel()
+ ].accumulateSingleSample(totalTime);
+ request.resolve(validatedResponse);
+ }
} else {
request.reject(error);
}
@@ -1477,6 +1508,12 @@ export class MLEngine {
throw new Error("Port does not exist");
}
+ // Validate request before sending to child process
+ const validatedRequest = this.#validateRequest(request);
+ if (!validatedRequest) {
+ throw new Error("Request failed security validation");
+ }
+
const resourcesPromise = this.getInferenceResources();
const beforeRun = ChromeUtils.now();
@@ -1484,7 +1521,7 @@ export class MLEngine {
{
type: "EnginePort:Run",
requestId,
- request,
+ request: validatedRequest,
engineRunOptions: { enableInferenceProgress: false },
},
transferables
@@ -1572,12 +1609,18 @@ export class MLEngine {
throw new Error("The port is null");
}
+ // Validate request before sending to child process
+ const validatedRequest = this.#validateRequest(request);
+ if (!validatedRequest) {
+ throw new Error("Request failed security validation");
+ }
+
// Send the request to the engine via postMessage with optional transferables
this.#port.postMessage(
{
type: "EnginePort:Run",
requestId,
- request,
+ request: validatedRequest,
engineRunOptions: { enableInferenceProgress: true },
},
transferables
diff --git a/toolkit/components/ml/moz.build b/toolkit/components/ml/moz.build
@@ -12,11 +12,15 @@ JAR_MANIFESTS += ["jar.mn"]
with Files("**"):
BUG_COMPONENT = ("Core", "Machine Learning: On Device")
-DIRS += ["actors"]
+DIRS += [
+ "actors",
+]
if CONFIG["MOZ_WIDGET_TOOLKIT"] != "android":
DIRS += ["backends/llama"]
+XPCSHELL_TESTS_MANIFESTS += ["tests/xpcshell/xpcshell.toml"]
+
BROWSER_CHROME_MANIFESTS += [
"tests/browser/browser.toml",
"tests/browser/perftest.toml",
diff --git a/toolkit/components/ml/tests/browser/browser.toml b/toolkit/components/ml/tests/browser/browser.toml
@@ -22,6 +22,8 @@ support-files = [
["browser_ml_engine_rs_hub.js"]
+["browser_ml_engine_security.js"]
+
["browser_ml_native.js"]
skip-if = [
"os == 'android'",
diff --git a/toolkit/components/ml/tests/browser/browser_ml_engine_security.js b/toolkit/components/ml/tests/browser/browser_ml_engine_security.js
@@ -0,0 +1,67 @@
+/* Any copyright is dedicated to the Public Domain.
+ http://creativecommons.org/publicdomain/zero/1.0/ */
+
+"use strict";
+
+/**
+ * Browser integration tests for ML security layer scaffolding.
+ * These tests verify that the security layer infrastructure is correctly
+ * integrated into MLEngine.
+ */
+
+/**
+ * Test that MLEngine can be instantiated with security layer integrated.
+ */
+add_task(async function test_mlengine_instantiation() {
+ info("Testing MLEngine instantiation with security layer");
+
+ await SpecialPowers.pushPrefEnv({
+ set: [["browser.ml.enable", true]],
+ });
+
+ const tab = await BrowserTestUtils.openNewForegroundTab(
+ gBrowser,
+ "data:text/html,<meta charset=utf-8>MLEngine Test"
+ );
+
+ try {
+ const { MLEngine } = ChromeUtils.importESModule(
+ "resource://gre/actors/MLEngineParent.sys.mjs"
+ );
+
+ // Create engine instance
+ const engine = new MLEngine({
+ mlEngineParent: {},
+ pipelineOptions: {
+ engineId: "browser-test-instantiation",
+ featureId: "test-feature",
+ taskName: "test-task",
+ },
+ notificationsCallback: null,
+ });
+
+ Assert.ok(engine, "MLEngine instantiates successfully");
+ Assert.equal(
+ engine.engineId,
+ "browser-test-instantiation",
+ "Engine ID is set correctly"
+ );
+ Assert.equal(
+ engine.engineStatus,
+ "uninitialized",
+ "Initial status is uninitialized"
+ );
+
+ // Verify engine is tracked
+ const retrieved = MLEngine.getInstance("browser-test-instantiation");
+ Assert.equal(retrieved, engine, "Engine is tracked in instances map");
+
+ // Clean up
+ await MLEngine.removeInstance("browser-test-instantiation", false, false);
+
+ info("MLEngine instantiation works with security layer");
+ } finally {
+ BrowserTestUtils.removeTab(tab);
+ await SpecialPowers.popPrefEnv();
+ }
+});
diff --git a/toolkit/components/ml/tests/xpcshell/test_ml_engine_security.js b/toolkit/components/ml/tests/xpcshell/test_ml_engine_security.js
@@ -0,0 +1,43 @@
+/* Any copyright is dedicated to the Public Domain.
+ http://creativecommons.org/publicdomain/zero/1.0/ */
+
+"use strict";
+
+const { MLEngine } = ChromeUtils.importESModule(
+ "resource://gre/actors/MLEngineParent.sys.mjs"
+);
+
+/**
+ * Test that MLEngine properly manages request lifecycle.
+ * This ensures security layer integration doesn't break request tracking.
+ */
+add_task(async function test_request_lifecycle_management() {
+ info("Testing request lifecycle management");
+
+ const engine = new MLEngine({
+ mlEngineParent: {},
+ pipelineOptions: { engineId: "test-request-lifecycle" },
+ notificationsCallback: null,
+ });
+
+ // Verify engine starts with no pending requests
+ // The #requests map is private, but we can test behavior
+
+ // Without a port, run() should throw before creating a request
+ let errorThrown = false;
+ try {
+ await engine.run({ args: ["test"] });
+ } catch (error) {
+ errorThrown = true;
+ Assert.ok(
+ error.message.includes("Port does not exist"),
+ "Should fail on port check"
+ );
+ }
+
+ Assert.ok(errorThrown, "run() without port should throw");
+
+ await MLEngine.removeInstance("test-request-lifecycle", false, false);
+
+ info("Request lifecycle management works correctly");
+});
diff --git a/toolkit/components/ml/tests/xpcshell/xpcshell.toml b/toolkit/components/ml/tests/xpcshell/xpcshell.toml
@@ -0,0 +1,5 @@
+[DEFAULT]
+head = ""
+skip-if = ["os == 'android'"]
+
+["test_ml_engine_security.js"]
