Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 53 additions & 1 deletion packages/cli/src/acp/acpClient.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -208,7 +208,16 @@
});

expect(response.protocolVersion).toBe(acp.PROTOCOL_VERSION);
expect(response.authMethods).toHaveLength(3);
expect(response.authMethods).toHaveLength(4);
const gatewayAuth = response.authMethods?.find(
(m) => m.id === AuthType.GATEWAY,
);
expect(gatewayAuth?._meta).toEqual({
gateway: {
protocol: 'google',
restartRequired: 'false',
},
});
const geminiAuth = response.authMethods?.find(
(m) => m.id === AuthType.USE_GEMINI,
);
Expand All @@ -228,6 +237,8 @@
expect(mockConfig.refreshAuth).toHaveBeenCalledWith(
AuthType.LOGIN_WITH_GOOGLE,
undefined,
undefined,
undefined,
);
expect(mockSettings.setValue).toHaveBeenCalledWith(
SettingScope.User,
Expand All @@ -247,6 +258,8 @@
expect(mockConfig.refreshAuth).toHaveBeenCalledWith(
AuthType.USE_GEMINI,
'test-api-key',
undefined,
undefined,
);
expect(mockSettings.setValue).toHaveBeenCalledWith(
SettingScope.User,
Expand All @@ -255,6 +268,45 @@
);
});

it('should authenticate correctly with gateway method', async () => {
await agent.authenticate({
methodId: AuthType.GATEWAY,
_meta: {
gateway: {
baseUrl: 'https://example.com',
headers: { Authorization: 'Bearer token' },
},
},
} as unknown as acp.AuthenticateRequest);

expect(mockConfig.refreshAuth).toHaveBeenCalledWith(
AuthType.GATEWAY,
undefined,
'https://example.com',
{ Authorization: 'Bearer token' },
);
expect(mockSettings.setValue).toHaveBeenCalledWith(
SettingScope.User,
'security.auth.selectedType',
AuthType.GATEWAY,
);
});

it('should throw acp.RequestError when gateway payload is malformed', async () => {
await expect(
agent.authenticate({
methodId: AuthType.GATEWAY,
_meta: {
gateway: {
// Invalid baseUrl
baseUrl: 123,
headers: { Authorization: 'Bearer token' },
},
},
} as unknown as acp.AuthenticateRequest),
).rejects.toThrow(/Malformed gateway payload/);
});

it('should create a new session', async () => {
vi.useFakeTimers();
mockConfig.getContentGeneratorConfig = vi.fn().mockReturnValue({
Expand Down Expand Up @@ -333,8 +385,8 @@
name: expect.stringContaining('Auto'),
}),
expect.objectContaining({
modelId: 'gemini-3.1-pro-preview',

Check warning on line 388 in packages/cli/src/acp/acpClient.test.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-3.1". Please make sure this change is appropriate to submit.
name: 'gemini-3.1-pro-preview',

Check warning on line 389 in packages/cli/src/acp/acpClient.test.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-3.1". Please make sure this change is appropriate to submit.
}),
]),
);
Expand Down
60 changes: 57 additions & 3 deletions packages/cli/src/acp/acpClient.ts
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,8 @@
private sessions: Map<string, Session> = new Map();
private clientCapabilities: acp.ClientCapabilities | undefined;
private apiKey: string | undefined;
private baseUrl: string | undefined;
private customHeaders: Record<string, string> | undefined;

constructor(
private config: Config,
Expand Down Expand Up @@ -131,6 +133,17 @@
name: 'Vertex AI',
description: 'Use an API key with Vertex AI GenAI API',
},
{
id: AuthType.GATEWAY,
name: 'AI API Gateway',
description: 'Use a custom AI API Gateway',
_meta: {
gateway: {
protocol: 'google',
restartRequired: 'false',
},
},
},
];

await this.config.initialize();
Expand Down Expand Up @@ -179,7 +192,38 @@
if (apiKey) {
this.apiKey = apiKey;
}
await this.config.refreshAuth(method, apiKey ?? this.apiKey);

// Extract gateway details if present
const gatewaySchema = z.object({
baseUrl: z.string().optional(),
headers: z.record(z.string()).optional(),
});

let baseUrl: string | undefined;
let headers: Record<string, string> | undefined;

if (meta?.['gateway']) {
const result = gatewaySchema.safeParse(meta['gateway']);
if (result.success) {
baseUrl = result.data.baseUrl;
headers = result.data.headers;
} else {
throw new acp.RequestError(
-32602,
`Malformed gateway payload: ${result.error.message}`,
);
}
}

this.baseUrl = baseUrl;
this.customHeaders = headers;
Comment on lines 192 to +219

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.

security-high high

The authenticate method updates the agent's internal state (this.apiKey, this.baseUrl, and this.customHeaders) before successfully completing the authentication process via this.config.refreshAuth. If refreshAuth fails (e.g., due to an invalid API key or an unreachable gateway), the agent is left with these untrusted configuration values in its memory.

These values are subsequently used in newSession and other methods. This creates a vulnerability where a malicious client can send a failed authenticate request to "poison" the agent's configuration. For example, an attacker could set a malicious baseUrl in a failed authentication attempt. If the user later starts a new session using a standard authentication method (like USE_GEMINI), the agent will use the malicious baseUrl from the previous failed attempt, effectively redirecting the user's AI requests and potentially leaking their real API key and chat history to the attacker's gateway.

To remediate this, update the internal state only after this.config.refreshAuth has successfully completed.

References
  1. Security-sensitive settings should not use a merge strategy that allows less-trusted configuration scopes (like a workspace) to completely override more-trusted scopes (like global user settings). This comment highlights a similar issue where untrusted input from a failed authentication attempt can 'poison' security-sensitive internal state, effectively overriding it.


await this.config.refreshAuth(
method,
apiKey ?? this.apiKey,
baseUrl,
headers,
);
} catch (e) {
throw new acp.RequestError(-32000, getAcpErrorMessage(e));
}
Expand Down Expand Up @@ -209,7 +253,12 @@
let isAuthenticated = false;
let authErrorMessage = '';
try {
await config.refreshAuth(authType, this.apiKey);
await config.refreshAuth(
authType,
this.apiKey,
this.baseUrl,
this.customHeaders,
);
isAuthenticated = true;

// Extra validation for Gemini API key
Expand Down Expand Up @@ -371,7 +420,12 @@
// This satisfies the security requirement to verify the user before executing
// potentially unsafe server definitions.
try {
await config.refreshAuth(selectedAuthType, this.apiKey);
await config.refreshAuth(
selectedAuthType,
this.apiKey,
this.baseUrl,
this.customHeaders,
);
} catch (e) {
debugLogger.error(`Authentication failed: ${e}`);
throw acp.RequestError.authRequired();
Expand Down Expand Up @@ -1545,7 +1599,7 @@
value: PREVIEW_GEMINI_MODEL_AUTO,
title: getDisplayString(PREVIEW_GEMINI_MODEL_AUTO),
description: useGemini31
? 'Let Gemini CLI decide the best model for the task: gemini-3.1-pro, gemini-3-flash'

Check warning on line 1602 in packages/cli/src/acp/acpClient.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-3.1". Please make sure this change is appropriate to submit.
: 'Let Gemini CLI decide the best model for the task: gemini-3-pro, gemini-3-flash',
});
}
Expand Down
2 changes: 2 additions & 0 deletions packages/core/src/config/config.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -500,6 +500,8 @@
config,
authType,
undefined,
undefined,
undefined,
);
// Verify that contentGeneratorConfig is updated
expect(config.getContentGeneratorConfig()).toEqual(mockContentConfig);
Expand Down Expand Up @@ -2487,7 +2489,7 @@
mockCodeAssistServer.retrieveUserQuota.mockResolvedValue({
buckets: [
{
modelId: 'gemini-3.1-pro-preview',

Check warning on line 2492 in packages/core/src/config/config.test.ts

View workflow job for this annotation

GitHub Actions / Lint

Found sensitive keyword "gemini-3.1". Please make sure this change is appropriate to submit.
remainingAmount: '100',
remainingFraction: 1.0,
},
Expand Down
9 changes: 8 additions & 1 deletion packages/core/src/config/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1190,7 +1190,12 @@ export class Config implements McpContext {
return this.contentGenerator;
}

async refreshAuth(authMethod: AuthType, apiKey?: string) {
async refreshAuth(
authMethod: AuthType,
apiKey?: string,
baseUrl?: string,
customHeaders?: Record<string, string>,
) {
// Reset availability service when switching auth
this.modelAvailabilityService.reset();

Expand All @@ -1217,6 +1222,8 @@ export class Config implements McpContext {
this,
authMethod,
apiKey,
baseUrl,
customHeaders,
);
this.contentGenerator = await createContentGenerator(
newContentGeneratorConfig,
Expand Down
22 changes: 20 additions & 2 deletions packages/core/src/core/contentGenerator.ts
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ export enum AuthType {
USE_VERTEX_AI = 'vertex-ai',
LEGACY_CLOUD_SHELL = 'cloud-shell',
COMPUTE_ADC = 'compute-default-credentials',
GATEWAY = 'gateway',
}

/**
Expand Down Expand Up @@ -93,12 +94,16 @@ export type ContentGeneratorConfig = {
vertexai?: boolean;
authType?: AuthType;
proxy?: string;
baseUrl?: string;
customHeaders?: Record<string, string>;
};

export async function createContentGeneratorConfig(
config: Config,
authType: AuthType | undefined,
apiKey?: string,
baseUrl?: string,
customHeaders?: Record<string, string>,
): Promise<ContentGeneratorConfig> {
const geminiApiKey =
apiKey ||
Expand All @@ -115,6 +120,8 @@ export async function createContentGeneratorConfig(
const contentGeneratorConfig: ContentGeneratorConfig = {
authType,
proxy: config?.getProxy(),
baseUrl,
customHeaders,
};

// If we are using Google auth or we are in Cloud Shell, there is nothing else to validate for now
Expand Down Expand Up @@ -203,9 +210,13 @@ export async function createContentGenerator(

if (
config.authType === AuthType.USE_GEMINI ||
config.authType === AuthType.USE_VERTEX_AI
config.authType === AuthType.USE_VERTEX_AI ||
config.authType === AuthType.GATEWAY
) {
let headers: Record<string, string> = { ...baseHeaders };
if (config.customHeaders) {
headers = { ...headers, ...config.customHeaders };
}
if (gcConfig?.getUsageStatisticsEnabled()) {
const installationManager = new InstallationManager();
const installationId = installationManager.getInstallationId();
Expand All @@ -214,7 +225,14 @@ export async function createContentGenerator(
'x-gemini-api-privileged-user-id': `${installationId}`,
};
}
const httpOptions = { headers };
const httpOptions: {
baseUrl?: string;
headers: Record<string, string>;
} = { headers };

if (config.baseUrl) {
httpOptions.baseUrl = config.baseUrl;
}

const googleGenAI = new GoogleGenAI({
apiKey: config.apiKey === '' ? undefined : config.apiKey,
Expand Down
Loading