Skip to content

Commit 3bff165

Browse files
Paola3stefaniaPaola3stefania
committed
fixes
1 parent 290d1f6 commit 3bff165

7 files changed

Lines changed: 206 additions & 37 deletions

File tree

‎packages/agent-auth/src/client.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,11 @@ export const agentAuthClient = () => {
1414
"/agent/revoke": "POST",
1515
"/agent/rotate-key": "POST",
1616
"/agent/cleanup": "POST",
17+
"/agent/request-scope": "POST",
18+
"/agent/approve-scope": "POST",
19+
"/agent/workgroup/create": "POST",
20+
"/agent/workgroup/update": "POST",
21+
"/agent/workgroup/delete": "POST",
1722
},
1823
$ERROR_CODES: AGENT_AUTH_ERROR_CODES,
1924
} satisfies BetterAuthClientPlugin;

‎packages/agent-auth/src/index.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -279,6 +279,10 @@ export const agentAuth = (options?: AgentAuthOptions) => {
279279
requestScope: routes.requestScope,
280280
scopeRequestStatus: routes.scopeRequestStatus,
281281
approveScope: routes.approveScope,
282+
createWorkgroup: routes.createWorkgroup,
283+
listWorkgroups: routes.listWorkgroups,
284+
updateWorkgroup: routes.updateWorkgroup,
285+
deleteWorkgroup: routes.deleteWorkgroup,
282286
},
283287
rateLimit: buildRateLimits(options?.rateLimit),
284288
schema,

‎packages/agent-auth/src/mcp-tools.ts‎

Lines changed: 75 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -583,26 +583,37 @@ export function createAgentMCPTools(
583583
},
584584
);
585585

586-
if (tokenRes.ok) {
587-
const tokenData = (await tokenRes.json()) as {
588-
access_token: string;
586+
const resText = await tokenRes.text();
587+
let resJson: Record<string, unknown>;
588+
try {
589+
resJson = JSON.parse(resText);
590+
} catch {
591+
if (storage.removePendingFlow) await storage.removePendingFlow(url);
592+
return {
593+
content: [
594+
{
595+
type: "text" as const,
596+
text: `Device auth returned non-JSON response: ${resText.slice(0, 200)}`,
597+
},
598+
],
589599
};
590-
accessToken = tokenData.access_token;
600+
}
601+
602+
if (tokenRes.ok) {
603+
accessToken = resJson.access_token as string;
591604
break;
592605
}
593606

594-
const errorData = (await tokenRes.json()) as {
595-
error: string;
596-
};
607+
const error = resJson.error as string | undefined;
597608

598-
if (errorData.error === "authorization_pending") {
609+
if (error === "authorization_pending") {
599610
continue;
600611
}
601-
if (errorData.error === "slow_down") {
612+
if (error === "slow_down") {
602613
await new Promise((resolve) => setTimeout(resolve, pollInterval));
603614
continue;
604615
}
605-
if (errorData.error === "access_denied") {
616+
if (error === "access_denied") {
606617
if (storage.removePendingFlow) await storage.removePendingFlow(url);
607618
return {
608619
content: [
@@ -613,7 +624,7 @@ export function createAgentMCPTools(
613624
],
614625
};
615626
}
616-
if (errorData.error === "expired_token") {
627+
if (error === "expired_token") {
617628
if (storage.removePendingFlow) await storage.removePendingFlow(url);
618629
return {
619630
content: [
@@ -630,7 +641,7 @@ export function createAgentMCPTools(
630641
content: [
631642
{
632643
type: "text" as const,
633-
text: `Device auth failed: ${errorData.error}`,
644+
text: `Device auth failed: ${error}`,
634645
},
635646
],
636647
};
@@ -875,15 +886,40 @@ export function createAgentMCPTools(
875886
};
876887
}
877888

889+
const fullUrl = reqPath.startsWith("http")
890+
? reqPath
891+
: `${connection.appUrl}${reqPath}`;
892+
893+
// Prevent SSRF: only allow requests to the app the agent connected to
894+
try {
895+
const target = new URL(fullUrl);
896+
const allowed = new URL(connection.appUrl);
897+
if (target.origin !== allowed.origin) {
898+
return {
899+
content: [
900+
{
901+
type: "text" as const,
902+
text: `Blocked: request origin ${target.origin} does not match connected app ${allowed.origin}.`,
903+
},
904+
],
905+
};
906+
}
907+
} catch {
908+
return {
909+
content: [
910+
{
911+
type: "text" as const,
912+
text: `Invalid URL: ${fullUrl}`,
913+
},
914+
],
915+
};
916+
}
917+
878918
const jwt = await signAgentJWT({
879919
agentId,
880920
privateKey: connection.keypair.privateKey,
881921
});
882922

883-
const fullUrl = reqPath.startsWith("http")
884-
? reqPath
885-
: `${connection.appUrl}${reqPath}`;
886-
887923
const headers: Record<string, string> = {
888924
Authorization: `Bearer ${jwt}`,
889925
};
@@ -960,29 +996,40 @@ export function createAgentMCPTools(
960996
},
961997
);
962998

963-
if (tokenRes.ok) {
964-
const tokenData = (await tokenRes.json()) as {
965-
access_token: string;
999+
const resText = await tokenRes.text();
1000+
let resJson: Record<string, unknown>;
1001+
try {
1002+
resJson = JSON.parse(resText);
1003+
} catch {
1004+
if (storage.removePendingFlow) await storage.removePendingFlow(url);
1005+
return {
1006+
content: [
1007+
{
1008+
type: "text" as const,
1009+
text: `Device auth returned non-JSON response: ${resText.slice(0, 200)}`,
1010+
},
1011+
],
9661012
};
967-
accessToken = tokenData.access_token;
1013+
}
1014+
1015+
if (tokenRes.ok) {
1016+
accessToken = resJson.access_token as string;
9681017
break;
9691018
}
9701019

971-
const errorData = (await tokenRes.json()) as {
972-
error: string;
973-
};
1020+
const error = resJson.error as string | undefined;
9741021

975-
if (errorData.error === "authorization_pending") {
1022+
if (error === "authorization_pending") {
9761023
await new Promise((resolve) => setTimeout(resolve, pollInterval));
9771024
continue;
9781025
}
979-
if (errorData.error === "slow_down") {
1026+
if (error === "slow_down") {
9801027
await new Promise((resolve) =>
9811028
setTimeout(resolve, pollInterval * 2),
9821029
);
9831030
continue;
9841031
}
985-
if (errorData.error === "access_denied") {
1032+
if (error === "access_denied") {
9861033
if (storage.removePendingFlow) await storage.removePendingFlow(url);
9871034
return {
9881035
content: [
@@ -993,7 +1040,7 @@ export function createAgentMCPTools(
9931040
],
9941041
};
9951042
}
996-
if (errorData.error === "expired_token") {
1043+
if (error === "expired_token") {
9971044
if (storage.removePendingFlow) await storage.removePendingFlow(url);
9981045
return {
9991046
content: [
@@ -1010,7 +1057,7 @@ export function createAgentMCPTools(
10101057
content: [
10111058
{
10121059
type: "text" as const,
1013-
text: `Device auth failed: ${errorData.error}`,
1060+
text: `Device auth failed: ${error}`,
10141061
},
10151062
],
10161063
};

‎packages/agent-auth/src/routes/index.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,9 @@ export function createAgentRoutes(opts: ResolvedAgentAuthOptions) {
2222
createAgent: createAgent(opts),
2323
listAgents: listAgents(),
2424
getAgent: getAgent(),
25-
updateAgent: updateAgent(),
25+
updateAgent: updateAgent(opts),
2626
revokeAgent: revokeAgent(),
27-
rotateKey: rotateKey(),
27+
rotateKey: rotateKey(opts),
2828
getAgentSession: getAgentSession(),
2929
cleanupAgents: cleanupAgents(),
3030
requestScope: requestScope(),

‎packages/agent-auth/src/routes/rotate-key.ts‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,11 @@ import { APIError } from "@better-auth/core/error";
33
import { getSessionFromCtx } from "better-auth/api";
44
import * as z from "zod";
55
import { AGENT_AUTH_ERROR_CODES as ERROR_CODES } from "../error-codes";
6-
import type { Agent } from "../types";
6+
import type { Agent, ResolvedAgentAuthOptions } from "../types";
77

88
const AGENT_TABLE = "agent";
99

10-
export function rotateKey() {
10+
export function rotateKey(opts: ResolvedAgentAuthOptions) {
1111
return createAuthEndpoint(
1212
"/agent/rotate-key",
1313
{
@@ -50,6 +50,15 @@ export function rotateKey() {
5050
throw APIError.from("BAD_REQUEST", ERROR_CODES.INVALID_PUBLIC_KEY);
5151
}
5252

53+
const kty = publicKey.kty as string;
54+
const crv = (publicKey.crv as string) ?? null;
55+
const keyAlg = crv ?? kty;
56+
if (!opts.allowedKeyAlgorithms.includes(keyAlg)) {
57+
throw new APIError("BAD_REQUEST", {
58+
message: `Key algorithm "${keyAlg}" is not allowed. Accepted: ${opts.allowedKeyAlgorithms.join(", ")}`,
59+
});
60+
}
61+
5362
const kid = (publicKey.kid as string) ?? null;
5463

5564
await ctx.context.adapter.update<Agent>({

‎packages/agent-auth/src/routes/update-agent.ts‎

Lines changed: 41 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import { APIError } from "@better-auth/core/error";
33
import { getSessionFromCtx } from "better-auth/api";
44
import * as z from "zod";
55
import { AGENT_AUTH_ERROR_CODES as ERROR_CODES } from "../error-codes";
6-
import type { Agent } from "../types";
6+
import type { Agent, ResolvedAgentAuthOptions } from "../types";
77

88
const AGENT_TABLE = "agent";
99

@@ -20,7 +20,7 @@ const updateAgentBodySchema = z.object({
2020
.optional(),
2121
});
2222

23-
export function updateAgent() {
23+
export function updateAgent(opts: ResolvedAgentAuthOptions) {
2424
return createAuthEndpoint(
2525
"/agent/update",
2626
{
@@ -52,13 +52,50 @@ export function updateAgent() {
5252
throw APIError.from("NOT_FOUND", ERROR_CODES.AGENT_NOT_FOUND);
5353
}
5454

55+
let resolvedScopes = scopes;
56+
const resolvedRole = role;
57+
58+
if (role !== undefined && opts.roles?.[role]) {
59+
resolvedScopes = resolvedScopes ?? opts.roles[role];
60+
}
61+
62+
if (resolvedScopes && resolvedScopes.length > 0 && opts.validateScopes) {
63+
if (typeof opts.validateScopes === "function") {
64+
const valid = await opts.validateScopes(resolvedScopes);
65+
if (!valid) {
66+
throw APIError.from("BAD_REQUEST", ERROR_CODES.UNKNOWN_SCOPES);
67+
}
68+
} else {
69+
const knownScopes = new Set(Object.values(opts.roles ?? {}).flat());
70+
const invalid = resolvedScopes.filter(
71+
(s: string) => !knownScopes.has(s),
72+
);
73+
if (invalid.length > 0) {
74+
throw new APIError("BAD_REQUEST", {
75+
message: `${ERROR_CODES.UNKNOWN_SCOPES} Unrecognized: ${invalid.join(", ")}.`,
76+
});
77+
}
78+
}
79+
}
80+
81+
if (
82+
resolvedRole !== undefined &&
83+
opts.roles &&
84+
!opts.roles[resolvedRole]
85+
) {
86+
throw new APIError("BAD_REQUEST", {
87+
message: `Unknown role "${resolvedRole}". Known roles: ${Object.keys(opts.roles).join(", ")}.`,
88+
});
89+
}
90+
5591
const updates: Record<string, string | Date | null> = {
5692
updatedAt: new Date(),
5793
};
5894

5995
if (name !== undefined) updates.name = name;
60-
if (scopes !== undefined) updates.scopes = JSON.stringify(scopes);
61-
if (role !== undefined) updates.role = role;
96+
if (resolvedScopes !== undefined)
97+
updates.scopes = JSON.stringify(resolvedScopes);
98+
if (resolvedRole !== undefined) updates.role = resolvedRole;
6299
if (metadata !== undefined) updates.metadata = JSON.stringify(metadata);
63100

64101
const updated = await ctx.context.adapter.update<Agent>({

0 commit comments

Comments
 (0)