Merge pull request #92 from imbue-ai/preston/google-oauth

Add `latchkey prepare` command and PKCE for Google OAuth

Author: hynek-urban

src/cliCommands.ts

@@ -46,6 +46,8 @@ import {
   LoginCancelledError,
   LoginFailedError,
   NoCurlCredentialsNotSupportedError,
+  PrepareInputInvalidError,
+  PrepareNotSupportedError,
   Service,
 } from './services/index.js';
 import {
@@ -77,6 +79,7 @@ import {
   authList,
   authBrowser,
   authBrowserPrepare,
+  prepareService,
   UnknownServiceError,
   BrowserNotConfiguredError,
   PreparationRequiredError,
@@ -734,6 +737,50 @@ export function registerCommands(program: Command, deps: CliDependencies): void
       }
     });
 
+  authCommand
+    .command('prepare')
+    .description(
+      "Register a service's client details (e.g. an OAuth client id/secret) from a JSON payload, for use during login."
+    )
+    .argument('<service_name>', 'Name of the service to prepare')
+    .argument(
+      '<json>',
+      'Service-specific registration JSON, e.g. \'{"clientId":"...","clientSecret":"..."}\''
+    )
+    .addHelpText(
+      'after',
+      `\nExample:\n  $ latchkey auth prepare google-gmail '{"clientId":"<id>","clientSecret":"<secret>"}'`
+    )
+    .action(async (serviceName: string, json: string) => {
+      if (deps.config.gatewayUrl !== null) {
+        await forwardToGateway(deps, {
+          command: 'auth prepare',
+          params: { serviceName, json },
+        });
+        deps.log(`Done`);
+        return;
+      }
+      try {
+        const encryptedStorage = await createEncryptedStorageFromConfig(deps.config);
+        const apiCredentialStore = new ApiCredentialStore(
+          deps.config.credentialStorePath,
+          encryptedStorage
+        );
+        prepareService(deps.registry, apiCredentialStore, serviceName, json);
+        deps.log(`Done`);
+      } catch (error) {
+        if (
+          error instanceof UnknownServiceError ||
+          error instanceof PrepareNotSupportedError ||
+          error instanceof PrepareInputInvalidError
+        ) {
+          deps.errorLog(`Error: ${error.message}`);
+          deps.exit(1);
+        }
+        throw error;
+      }
+    });
+
   program
     .command('curl')
     .description('Run curl with API credential injection.')

src/gateway/latchkeyEndpoint.ts

@@ -17,6 +17,7 @@ import {
   authList,
   authBrowser,
   authBrowserPrepare,
+  prepareService,
   UnknownServiceError,
   BrowserNotConfiguredError,
   PreparationRequiredError,
@@ -26,7 +27,12 @@ import {
   BrowserFlowsNotSupportedError,
   GraphicalEnvironmentNotFoundError,
 } from '../playwrightUtils.js';
-import { LoginCancelledError, LoginFailedError } from '../services/index.js';
+import {
+  LoginCancelledError,
+  LoginFailedError,
+  PrepareInputInvalidError,
+  PrepareNotSupportedError,
+} from '../services/index.js';
 
 const serviceNameParamsMessage = "missing required argument 'service_name'";
 const serviceNameParams = z.object(
@@ -68,12 +74,20 @@ const AuthBrowserPrepareRequestSchema = z.object({
   params: serviceNameParams,
 });
 
+const PrepareRequestSchema = z.object({
+  command: z.literal('auth prepare'),
+  params: serviceNameParams.extend({
+    json: z.string(),
+  }),
+});
+
 export const LatchkeyRequestSchema = z.discriminatedUnion('command', [
   ServicesListRequestSchema,
   ServicesInfoRequestSchema,
   AuthListRequestSchema,
   AuthBrowserRequestSchema,
   AuthBrowserPrepareRequestSchema,
+  PrepareRequestSchema,
 ]);
 
 export type LatchkeyRequest = z.infer<typeof LatchkeyRequestSchema>;
@@ -87,6 +101,8 @@ const KNOWN_ERROR_CLASSES: readonly (abstract new (...args: never[]) => Error)[]
   PreparationRequiredError,
   LoginCancelledError,
   LoginFailedError,
+  PrepareNotSupportedError,
+  PrepareInputInvalidError,
 ];
 
 function isKnownError(error: unknown): error is Error {
@@ -158,6 +174,14 @@ async function dispatch(
         deps.config,
         parsed.params.serviceName
       );
+
+    case 'auth prepare':
+      return prepareService(
+        deps.registry,
+        apiCredentialStore,
+        parsed.params.serviceName,
+        parsed.params.json
+      );
   }
 }
 

src/services/core/base.ts

@@ -3,6 +3,7 @@
  */
 
 import type { Browser, BrowserContext, Page, Response } from 'playwright';
+import type { z, ZodTypeAny } from 'zod';
 import {
   ApiCredentialStatus,
   ApiCredentials,
@@ -38,6 +39,58 @@ export class LoginFailedError extends Error {
   }
 }
 
+/**
+ * Thrown when `latchkey auth prepare` is run for a service that does not declare a
+ * prepare schema (the base default — services opt in by setting one).
+ */
+export class PrepareNotSupportedError extends Error {
+  constructor(serviceName: string) {
+    super(
+      `Service '${serviceName}' does not support 'latchkey auth prepare'. ` +
+        `Use 'latchkey services info ${serviceName}' to see how to authenticate.`
+    );
+    this.name = 'PrepareNotSupportedError';
+  }
+}
+
+/**
+ * Thrown when the JSON passed to `latchkey auth prepare` is malformed or does not
+ * match the service's prepare schema. The whole command is rejected and
+ * nothing is stored.
+ */
+export class PrepareInputInvalidError extends Error {
+  constructor(serviceName: string, detail: string) {
+    super(`Invalid prepare input for '${serviceName}': ${detail}`);
+    this.name = 'PrepareInputInvalidError';
+  }
+}
+
+/**
+ * Validate a parsed JSON value against a service's prepare schema and build the
+ * resulting credentials. Centralizes validation so each service's
+ * `prepareFromJson` only expresses its schema and build step. Throws
+ * `PrepareInputInvalidError` (with the failing fields) on any schema mismatch;
+ * nothing is built unless the input fully validates.
+ */
+export function buildPreparedCredentials<Schema extends ZodTypeAny>(
+  serviceName: string,
+  schema: Schema,
+  parsedJson: unknown,
+  build: (validatedInput: z.infer<Schema>) => ApiCredentials
+): ApiCredentials {
+  const result = schema.safeParse(parsedJson);
+  if (!result.success) {
+    const detail = result.error.issues
+      .map((issue) => {
+        const path = issue.path.join('.');
+        return path ? `${path}: ${issue.message}` : issue.message;
+      })
+      .join('; ');
+    throw new PrepareInputInvalidError(serviceName, detail);
+  }
+  return build(result.data as z.infer<Schema>);
+}
+
 export function isBrowserClosedError(error: Error): boolean {
   const message = error.message.toLowerCase();
   return (
@@ -50,6 +103,22 @@ export function isBrowserClosedError(error: Error): boolean {
   );
 }
 
+/**
+ * Detects the Playwright/CDP error raised when a response body can no longer be
+ * retrieved (`Network.getResponseBody` reports "No resource with given
+ * identifier found"). This happens for responses that retain no readable body —
+ * redirects, evicted or cached resources, or bodies fetched after the page has
+ * navigated onward. Callers that read response bodies opportunistically should
+ * treat this as inconclusive rather than fatal.
+ */
+export function isResponseBodyUnavailableError(error: Error): boolean {
+  const message = error.message.toLowerCase();
+  return (
+    message.includes('no resource with given identifier') ||
+    message.includes('network.getresponsebody')
+  );
+}
+
 export function isTimeoutError(error: Error): boolean {
   return error.name === 'TimeoutError';
 }
@@ -126,6 +195,18 @@ export abstract class Service {
     throw new NoCurlCredentialsNotSupportedError(this.name);
   }
 
+  /**
+   * Build credentials from a parsed JSON payload for `latchkey auth prepare`.
+   *
+   * Optional, like `getSession`/`refreshCredentials`: services opt in by
+   * implementing it (typically via `buildPreparedCredentials` with a Zod
+   * schema). When a service does not implement it, prepare is "not supported"
+   * — the default that lets every service stay closed until it declares a
+   * schema. Implementations validate `parsedJson` and throw
+   * `PrepareInputInvalidError` on mismatch.
+   */
+  prepareFromJson?(parsedJson: unknown): ApiCredentials;
+
   /**
    * Get a new session for the login flow.
    * Services that don't support browser login should not implement this method.

src/services/google/base.ts

@@ -20,15 +20,19 @@ import {
 } from '../../playwrightUtils.js';
 import {
   exchangeCodeForTokens,
+  generateCodeChallenge,
+  generateCodeVerifier,
   refreshAccessToken,
   startOAuthCallbackServer,
 } from '../../oauthUtils.js';
 import {
   Service,
   BrowserFollowupServiceSession,
+  buildPreparedCredentials,
   LoginFailedError,
   LoginCancelledError,
   isBrowserClosedError,
+  isResponseBodyUnavailableError,
   isTimeoutError,
 } from '../core/base.js';
 import type { EncryptedStorage } from '../../encryptedStorage.js';
@@ -573,9 +577,14 @@ function checkGoogleLoginResponse(
         .catch((error: unknown) => {
           // The response body can become unreadable if the page/context
           // closes while it's still being read (e.g. the automation
-          // navigates onward, or the user closes the browser). Treat that
-          // specific race as inconclusive; let any other error propagate.
-          if (error instanceof Error && isBrowserClosedError(error)) {
+          // navigates onward, or the user closes the browser), or if the
+          // response simply retains no readable body (redirects, cached or
+          // evicted resources). Login detection here is best-effort, so treat
+          // those cases as inconclusive; let any other error propagate.
+          if (
+            error instanceof Error &&
+            (isBrowserClosedError(error) || isResponseBodyUnavailableError(error))
+          ) {
             return;
           }
           throw error;
@@ -829,13 +838,21 @@ class GoogleServiceSession extends BrowserFollowupServiceSession {
       );
       const redirectUri = `http://localhost:${port.toString()}/oauth2callback`;
 
+      // PKCE (RFC 7636): bind the authorization code to a one-time verifier so a
+      // stolen code cannot be redeemed without it. We keep sending the client
+      // secret too so this is confidential-client + PKCE, defense-in-depth.
+      const codeVerifier = generateCodeVerifier();
+      const codeChallenge = generateCodeChallenge(codeVerifier);
+
       const authUrl = new URL('https://accounts.google.com/o/oauth2/v2/auth');
       authUrl.searchParams.set('client_id', clientId);
       authUrl.searchParams.set('redirect_uri', redirectUri);
       authUrl.searchParams.set('response_type', 'code');
       authUrl.searchParams.set('scope', allScopes.join(' '));
       authUrl.searchParams.set('access_type', 'offline');
       authUrl.searchParams.set('prompt', 'consent');
+      authUrl.searchParams.set('code_challenge', codeChallenge);
+      authUrl.searchParams.set('code_challenge_method', 'S256');
 
       await page.goto(authUrl.toString());
 
@@ -845,7 +862,8 @@ class GoogleServiceSession extends BrowserFollowupServiceSession {
         code,
         clientId,
         clientSecret,
-        redirectUri
+        redirectUri,
+        codeVerifier
       );
       const accessTokenExpiresAt = new Date(Date.now() + tokens.expires_in * 1000).toISOString();
 
@@ -866,6 +884,20 @@ class GoogleServiceSession extends BrowserFollowupServiceSession {
   }
 }
 
+/**
+ * JSON accepted by `latchkey auth prepare <google-service>`: the OAuth client
+ * credentials to use for that service. `.strict()` rejects unknown keys so
+ * typos are reported instead of silently ignored.
+ */
+export const GooglePrepareInputSchema = z
+  .object({
+    clientId: z.string().min(1),
+    clientSecret: z.string().min(1),
+  })
+  .strict();
+
+export type GooglePrepareInput = z.infer<typeof GooglePrepareInputSchema>;
+
 /**
  * Abstract base class for individual Google API services.
  *
@@ -878,6 +910,19 @@ export abstract class GoogleService extends Service {
 
   protected abstract readonly config: GoogleServiceConfig;
 
+  /**
+   * Google services accept an OAuth client's id/secret prepared
+   * in advance via `latchkey auth prepare`, stored as token-less OAuth credentials until login.
+   */
+  override prepareFromJson(parsedJson: unknown): ApiCredentials {
+    return buildPreparedCredentials(
+      this.name,
+      GooglePrepareInputSchema,
+      parsedJson,
+      ({ clientId, clientSecret }) => new OAuthCredentials(clientId, clientSecret)
+    );
+  }
+
   setCredentialsExample(serviceName: string): string {
     return `latchkey auth set ${serviceName} -H "Authorization: Bearer <token>"`;
   }

src/services/index.ts

@@ -12,6 +12,8 @@ export {
   LoginCancelledError,
   LoginFailedError,
   NoCurlCredentialsNotSupportedError,
+  PrepareNotSupportedError,
+  PrepareInputInvalidError,
 } from './core/base.js';
 export { RegisteredService } from './core/registered.js';