feat(agent-service): add opt-in access control and per-user isolation

agent-service/.env.example

@@ -18,4 +18,12 @@ LLM_ENDPOINT=http://localhost:9096
 # Texera backend services
 TEXERA_DASHBOARD_SERVICE_ENDPOINT=http://localhost:8080
 WORKFLOW_COMPILING_SERVICE_ENDPOINT=http://localhost:9090
-WORKFLOW_EXECUTION_SERVICE_ENDPOINT=http://localhost:8085
+WORKFLOW_EXECUTION_SERVICE_ENDPOINT=http://localhost:8085
+
+# Access control. When AGENT_AUTH_REQUIRED is true, every request must carry a
+# valid user JWT (Authorization: Bearer for HTTP, ?access-token= for the
+# WebSocket) and agents are isolated to their owning user. AUTH_JWT_SECRET is
+# the HS256 secret shared with the rest of Texera and is used to verify token
+# signatures. Defaults below preserve the previous permissive behavior.
+AGENT_AUTH_REQUIRED=false
+AUTH_JWT_SECRET=

agent-service/src/api/auth-api.test.ts

@@ -0,0 +1,159 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { afterEach, beforeEach, describe, expect, test } from "bun:test";
+import { createHmac } from "crypto";
+import {
+  createAuthHeaders,
+  extractUserFromToken,
+  getUidFromToken,
+  isAuthRequired,
+  validateToken,
+  verifyToken,
+} from "./auth-api";
+
+const SECRET = "unit-test-secret-key";
+
+function b64url(input: string | Buffer): string {
+  return Buffer.from(input).toString("base64").replace(/=/g, "").replace(/\+/g, "-").replace(/\//g, "_");
+}
+
+function signJwt(payload: Record<string, unknown>, opts: { secret?: string; alg?: string } = {}): string {
+  const header = b64url(JSON.stringify({ alg: opts.alg ?? "HS256", typ: "JWT" }));
+  const body = b64url(JSON.stringify(payload));
+  const sig = b64url(
+    createHmac("sha256", opts.secret ?? SECRET)
+      .update(`${header}.${body}`)
+      .digest()
+  );
+  return `${header}.${body}.${sig}`;
+}
+
+function futureExp(): number {
+  return Math.floor(Date.now() / 1000) + 3600;
+}
+
+const prevSecret = process.env.AUTH_JWT_SECRET;
+const prevRequired = process.env.AGENT_AUTH_REQUIRED;
+
+beforeEach(() => {
+  delete process.env.AUTH_JWT_SECRET;
+  delete process.env.AGENT_AUTH_REQUIRED;
+});
+
+afterEach(() => {
+  if (prevSecret === undefined) delete process.env.AUTH_JWT_SECRET;
+  else process.env.AUTH_JWT_SECRET = prevSecret;
+  if (prevRequired === undefined) delete process.env.AGENT_AUTH_REQUIRED;
+  else process.env.AGENT_AUTH_REQUIRED = prevRequired;
+});
+
+describe("extractUserFromToken / getUidFromToken", () => {
+  test("maps claims to UserInfo", () => {
+    const token = signJwt({ sub: "alice", userId: 9, email: "[email protected]", role: "ADMIN", exp: futureExp() });
+    expect(extractUserFromToken(token)).toEqual({ uid: 9, name: "alice", email: "[email protected]", role: "ADMIN" });
+    expect(getUidFromToken(token)).toBe(9);
+  });
+
+  test("getUidFromToken returns undefined for a malformed token", () => {
+    expect(getUidFromToken("not-a-jwt")).toBeUndefined();
+  });
+});
+
+describe("isAuthRequired", () => {
+  test("defaults to false", () => {
+    expect(isAuthRequired()).toBe(false);
+  });
+
+  test("is true for 'true' or '1'", () => {
+    process.env.AGENT_AUTH_REQUIRED = "true";
+    expect(isAuthRequired()).toBe(true);
+    process.env.AGENT_AUTH_REQUIRED = "1";
+    expect(isAuthRequired()).toBe(true);
+  });
+});
+
+describe("verifyToken", () => {
+  beforeEach(() => {
+    process.env.AUTH_JWT_SECRET = SECRET;
+  });
+
+  test("accepts a correctly signed, unexpired token", () => {
+    expect(verifyToken(signJwt({ sub: "u", userId: 1, exp: futureExp() }))).toBe(true);
+  });
+
+  test("rejects a token signed with the wrong secret", () => {
+    expect(verifyToken(signJwt({ sub: "u", userId: 1, exp: futureExp() }, { secret: "other" }))).toBe(false);
+  });
+
+  test("rejects an expired token", () => {
+    const exp = Math.floor(Date.now() / 1000) - 3600;
+    expect(verifyToken(signJwt({ sub: "u", userId: 1, exp }))).toBe(false);
+  });
+
+  test("rejects a token missing the subject claim", () => {
+    expect(verifyToken(signJwt({ userId: 1, exp: futureExp() }))).toBe(false);
+  });
+
+  test("rejects a non-HS256 algorithm", () => {
+    expect(verifyToken(signJwt({ sub: "u", userId: 1, exp: futureExp() }, { alg: "none" }))).toBe(false);
+  });
+
+  test("rejects a structurally invalid token", () => {
+    expect(verifyToken("a.b")).toBe(false);
+  });
+
+  test("returns false when no secret is configured", () => {
+    delete process.env.AUTH_JWT_SECRET;
+    expect(verifyToken(signJwt({ sub: "u", userId: 1, exp: futureExp() }))).toBe(false);
+  });
+});
+
+describe("validateToken", () => {
+  test("enforced mode requires a valid signature", () => {
+    process.env.AUTH_JWT_SECRET = SECRET;
+    process.env.AGENT_AUTH_REQUIRED = "true";
+    expect(validateToken(signJwt({ sub: "u", userId: 1, exp: futureExp() }))).toBe(true);
+    expect(validateToken(signJwt({ sub: "u", userId: 1, exp: futureExp() }, { secret: "x" }))).toBe(false);
+  });
+
+  test("permissive mode accepts any unexpired, decodable token", () => {
+    // No enforcement: a token signed with an arbitrary secret is still accepted
+    // as long as it is well-formed and unexpired (prior behavior).
+    expect(validateToken(signJwt({ sub: "u", userId: 1, exp: futureExp() }, { secret: "anything" }))).toBe(true);
+  });
+
+  test("permissive mode rejects an expired token", () => {
+    const exp = Math.floor(Date.now() / 1000) - 3600;
+    expect(validateToken(signJwt({ sub: "u", userId: 1, exp }))).toBe(false);
+  });
+
+  test("permissive mode rejects a malformed token", () => {
+    expect(validateToken("nonsense")).toBe(false);
+  });
+});
+
+describe("createAuthHeaders", () => {
+  test("builds a Bearer header with JSON content type", () => {
+    expect(createAuthHeaders("t.o.k")).toEqual({
+      Authorization: "Bearer t.o.k",
+      "Content-Type": "application/json",
+    });
+  });
+});

agent-service/src/api/auth-api.ts

@@ -17,22 +17,69 @@
  * under the License.
  */
 
+import { createHmac, timingSafeEqual } from "crypto";
 import type { UserInfo } from "../types/agent";
+import { createLogger } from "../logger";
 
 export type { UserInfo } from "../types/agent";
 
-function decodeJWT(token: string): any {
+const log = createLogger("Auth");
+
+// Matches the backend JwtAuth (org.apache.texera.auth): HMAC-SHA256 over
+// `${header}.${payload}` keyed by AUTH_JWT_SECRET (UTF-8), with a 30s clock skew.
+const JWT_ALGORITHM = "HS256";
+const CLOCK_SKEW_SECONDS = 30;
+
+// Auth settings are read from the environment at call time (not cached) so that
+// they can be toggled per request/per test without rebuilding the app.
+function getJwtSecret(): string {
+  return process.env.AUTH_JWT_SECRET ?? "";
+}
+
+/**
+ * Whether the agent service enforces authentication and per-user isolation.
+ *
+ * Opt-in via AGENT_AUTH_REQUIRED so the feature can be deployed and the client
+ * updated before enforcement is switched on. When disabled the service keeps
+ * its previous permissive behavior.
+ */
+export function isAuthRequired(): boolean {
+  const v = process.env.AGENT_AUTH_REQUIRED;
+  return v === "true" || v === "1";
+}
+
+function base64UrlToBuffer(segment: string): Buffer {
+  return Buffer.from(segment.replace(/-/g, "+").replace(/_/g, "/"), "base64");
+}
+
+function base64UrlEncode(buf: Buffer): string {
+  return buf.toString("base64").replace(/=/g, "").replace(/\+/g, "-").replace(/\//g, "_");
+}
+
+interface JwtParts {
+  header: any;
+  payload: any;
+}
+
+function decodeJwtParts(token: string): JwtParts {
+  const parts = token.split(".");
+  if (parts.length !== 3) {
+    throw new Error("Invalid JWT format");
+  }
   try {
-    const parts = token.split(".");
-    if (parts.length !== 3) {
-      throw new Error("Invalid JWT format");
-    }
-    return JSON.parse(Buffer.from(parts[1], "base64").toString("utf-8"));
+    return {
+      header: JSON.parse(base64UrlToBuffer(parts[0]).toString("utf-8")),
+      payload: JSON.parse(base64UrlToBuffer(parts[1]).toString("utf-8")),
+    };
   } catch (error) {
     throw new Error(`Failed to decode JWT: ${error}`);
   }
 }
 
+function decodeJWT(token: string): any {
+  return decodeJwtParts(token).payload;
+}
+
 export function extractUserFromToken(token: string): UserInfo {
   const payload = decodeJWT(token);
   return {
@@ -43,18 +90,68 @@ export function extractUserFromToken(token: string): UserInfo {
   };
 }
 
-function isTokenExpired(token: string): boolean {
+export function getUidFromToken(token: string): number | undefined {
   try {
-    const payload = decodeJWT(token);
-    if (!payload.exp) return false;
-    return Date.now() >= payload.exp * 1000;
+    const uid = decodeJWT(token).userId;
+    return typeof uid === "number" ? uid : undefined;
   } catch {
-    return true;
+    return undefined;
   }
 }
 
+function isExpired(payload: any): boolean {
+  if (typeof payload.exp !== "number") return true;
+  return Math.floor(Date.now() / 1000) > payload.exp + CLOCK_SKEW_SECONDS;
+}
+
+/**
+ * Cryptographically verify an HS256 JWT against AUTH_JWT_SECRET and check its
+ * expiry and required claims. Returns false (never throws) on any failure.
+ */
+export function verifyToken(token: string): boolean {
+  const secret = getJwtSecret();
+  if (!secret) {
+    log.warn("token verification requested but AUTH_JWT_SECRET is not set");
+    return false;
+  }
+
+  const parts = token.split(".");
+  if (parts.length !== 3) return false;
+
+  let parsed: JwtParts;
+  try {
+    parsed = decodeJwtParts(token);
+  } catch {
+    return false;
+  }
+
+  if (parsed.header?.alg !== JWT_ALGORITHM) return false;
+  // The backend requires both a subject and an expiration time.
+  if (!parsed.payload?.sub || typeof parsed.payload?.exp !== "number") return false;
+  if (isExpired(parsed.payload)) return false;
+
+  const expected = base64UrlEncode(createHmac("sha256", secret).update(`${parts[0]}.${parts[1]}`).digest());
+  const provided = parts[2];
+  if (expected.length !== provided.length) return false;
+  return timingSafeEqual(Buffer.from(expected), Buffer.from(provided));
+}
+
+/**
+ * Accept a token for use. When enforcement is on this requires a valid HS256
+ * signature and a live expiry; otherwise it falls back to an expiry-only check
+ * to preserve the service's prior behavior.
+ */
 export function validateToken(token: string): boolean {
-  return !isTokenExpired(token);
+  if (isAuthRequired()) {
+    return verifyToken(token);
+  }
+  try {
+    const payload = decodeJWT(token);
+    if (typeof payload.exp !== "number") return true;
+    return !isExpired(payload);
+  } catch {
+    return false;
+  }
 }
 
 export function createAuthHeaders(token: string): Record<string, string> {