feat: support oauth (#214)

Author: shruti-mantri

.env.example

@@ -1,6 +1,12 @@
 PORT = 3000
 GCP_STDIO = false
 
+OAUTH_ENABLED = true
+GOOGLE_OAUTH_CLIENT_ID = YOUR_CLIENT_ID
+GOOGLE_OAUTH_CLIENT_SECRET = YOUR_CLIENT_SECRET
+GOOGLE_OAUTH_AUDIENCE = YOUR_AUDIENCE # generally, same as your client id
+GOOGLE_OAUTH_REDIRECT_URI = http://localhost:7777/oauth/callback
+
 OAUTH_PROTECTED_RESOURCE = http://localhost:${PORT}/mcp
 OAUTH_AUTHORIZATION_SERVER = http://localhost:${PORT}/auth/google
 OAUTH_AUTHORIZATION_ENDPOINT = https://accounts.google.com/o/oauth2/v2/auth

constants.js

@@ -5,3 +5,4 @@ export const SCOPES = {
 };
 export const BEARER_METHODS_SUPPORTED = ['header'];
 export const RESPONSE_TYPES_SUPPORTED = ['code'];
+export const GCLOUD_AUTH = 'gcloud_auth';

lib/clients.js

@@ -14,9 +14,14 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
+import { OAuth2Client } from 'google-auth-library';
+import { GCLOUD_AUTH } from '../constants.js';
+
+const CLIENT_ID = process.env.GOOGLE_OAUTH_CLIENT_ID;
+const CLIENT_SECRET = process.env.GOOGLE_OAUTH_CLIENT_SECRET;
+const REDIRECT_URI = process.env.GOOGLE_OAUTH_REDIRECT_URI;
 const AUTHORIZATION_HEADER = 'Authorization';
 const BEARER_PREFIX = 'Bearer';
-const GCLOUD_AUTH = 'gcloud_auth';
 
 function keyGenerator(projectId, accessToken) {
   return accessToken !== GCLOUD_AUTH ? projectId + accessToken : projectId;
@@ -32,6 +37,7 @@ const clients = {
   logging: new Map(),
   billing: new Map(),
   projects: new Map(),
+  oauth: new Map(),
 };
 
 function getAuthClient(accessToken) {
@@ -44,6 +50,47 @@ function getAuthClient(accessToken) {
   };
 }
 
+/**
+ * Wraps an OAuth2Client to be compatible with gRPC-based Google Cloud clients.
+ *
+ * Some Google Cloud SDK clients (like Cloud Run, Service Usage, etc.) use gRPC
+ * and expect the `getRequestHeaders()` method of the auth client to return a
+ * `Map` of headers. However, the standard `OAuth2Client` from `google-auth-library`
+ * returns a plain Javascript object (e.g., `{ Authorization: 'Bearer ...' }`).
+ *
+ * This wrapper uses a Proxy to intercept calls to `getRequestHeaders`. It calls
+ * the original method and converts the result from a plain object to a `Map`
+ * if necessary, ensuring compatibility with gRPC-based clients while maintaining
+ * the original behavior for other properties.
+ *
+ * @param {OAuth2Client} authClient - The original OAuth2Client instance.
+ * @returns {Proxy<OAuth2Client>} A proxy compatible with gRPC clients.
+ */
+function wrapForGrpc(authClient) {
+  return new Proxy(authClient, {
+    get(target, prop, receiver) {
+      if (prop === 'getRequestHeaders') {
+        return async (...args) => {
+          const headers = await target.getRequestHeaders(...args);
+          if (headers instanceof Map) {
+            return headers;
+          }
+          // Convert plain object (from OAuth2Client) to Map (expected by grpc-js)
+          const headerMap = new Map();
+          for (const [k, v] of Object.entries(headers)) {
+            headerMap.set(k, v);
+          }
+          return headerMap;
+        };
+      }
+      return Reflect.get(target, prop, receiver);
+    },
+  });
+}
+
+// Services that use HTTP/REST (or prefer Object headers) instead of gRPC (Map headers)
+const HTTP_SERVICES = ['storage', 'logging'];
+
 export async function getClient(
   service,
   key,
@@ -55,7 +102,14 @@ export async function getClient(
     const ClientClass = await loadClient();
     const finalOptions = { ...options };
     if (accessToken && accessToken !== GCLOUD_AUTH) {
-      finalOptions.authClient = getAuthClient(accessToken);
+      const oauthClient = await getOAuthClient(accessToken);
+      // Storage and Logging use HTTP and expect the native OAuth2Client (headers as Object).
+      // Other services (Run, ServiceUsage, etc.) use gRPC and expect headers as Map.
+      if (HTTP_SERVICES.includes(service)) {
+        finalOptions.authClient = oauthClient;
+      } else {
+        finalOptions.authClient = wrapForGrpc(oauthClient);
+      }
     }
     clients[service].set(key, new ClientClass(finalOptions));
   }
@@ -220,3 +274,19 @@ export async function getProjectsClient(accessToken = GCLOUD_AUTH) {
     accessToken
   );
 }
+
+/**
+ * Gets an OAuth2 Client for the specified access token.
+ * @param {string} accessToken - The access token.
+ * @returns {Promise<OAuth2Client>}
+ */
+export async function getOAuthClient(accessToken) {
+  // Use the access token itself as the key since it's unique per session/user
+  if (!clients.oauth.has(accessToken)) {
+    const client = new OAuth2Client(CLIENT_ID, CLIENT_SECRET, REDIRECT_URI);
+    client.setCredentials({ access_token: accessToken });
+    clients.oauth.set(accessToken, client);
+  }
+
+  return clients.oauth.get(accessToken);
+}

lib/cloud-api/auth.js

@@ -23,6 +23,7 @@ import { GoogleAuth } from 'google-auth-library';
  * @async
  * @returns {Promise<boolean>} A promise that resolves to true if GCP auth are found, and false otherwise.
  */
+//TODO: Rename ensureGCPCredentials to ensureADCCredentials
 export async function ensureGCPCredentials() {
   console.error('Checking for Google Cloud Application Default Credentials...');
   try {

lib/cloud-api/billing.js

@@ -30,8 +30,8 @@ import { getBillingClient } from '../clients.js';
  * @returns {Promise<Array<{name: string, displayName: string, open: boolean}>>} A promise that resolves to an array of billing account objects,
  * each with 'name', 'displayName', and 'open' status. Returns an empty array on error.
  */
-export async function listBillingAccounts() {
-  const client = await getBillingClient();
+export async function listBillingAccounts(accessToken) {
+  const client = await getBillingClient('global', accessToken);
   try {
     const [accounts] = await client.listBillingAccounts();
     if (!accounts || accounts.length === 0) {
@@ -59,9 +59,10 @@ export async function listBillingAccounts() {
  */
 export async function attachProjectToBillingAccount(
   projectId,
-  billingAccountName
+  billingAccountName,
+  accessToken
 ) {
-  const client = await getBillingClient();
+  const client = await getBillingClient(projectId, accessToken);
   const projectName = `projects/${projectId}`;
 
   if (!projectId) {
@@ -111,8 +112,8 @@ export async function attachProjectToBillingAccount(
  * @param {string} projectId - The ID of the project to check.
  * @returns {Promise<boolean>} A promise that resolves to true if billing is enabled, false otherwise.
  */
-export async function isBillingEnabled(projectId) {
-  const client = await getBillingClient();
+export async function isBillingEnabled(projectId, accessToken) {
+  const client = await getBillingClient(projectId, accessToken);
   const projectName = `projects/${projectId}`;
   try {
     // getProjectBillingInfo requires cloudbilling.googleapis.com API to check billing status
@@ -136,10 +137,14 @@ export async function isBillingEnabled(projectId) {
  * @param {string} projectId The project ID to check billing for.
  * @param {function} progressCallback A callback function for progress updates.
  */
-export async function ensureBillingEnabled(projectId, progressCallback) {
-  if (!(await isBillingEnabled(projectId))) {
+export async function ensureBillingEnabled(
+  projectId,
+  accessToken,
+  progressCallback
+) {
+  if (!(await isBillingEnabled(projectId, accessToken))) {
     // Billing is disabled, try to fix it.
-    const accounts = await listBillingAccounts();
+    const accounts = await listBillingAccounts(accessToken);
 
     if (accounts && accounts.length === 1 && accounts[0].open) {
       // Exactly one open account found, try to attach it.
@@ -151,7 +156,8 @@ export async function ensureBillingEnabled(projectId, progressCallback) {
 
       const attachmentResult = await attachProjectToBillingAccount(
         projectId,
-        account.name
+        account.name,
+        accessToken
       );
 
       if (!attachmentResult || !attachmentResult.billingEnabled) {

lib/cloud-api/build.js

@@ -53,21 +53,23 @@ export async function triggerCloudBuild(
   targetRepoName,
   targetImageUrl,
   hasDockerfile,
+  accessToken,
   progressCallback
 ) {
   const serviceName = targetImageUrl.split('/').pop().split(/[:@]/)[0];
   const parent = `projects/${projectId}/locations/${location}`;
   const serviceFullName = `${parent}/services/${serviceName}`;
-  const buildsClient = await getBuildsClient(projectId);
-  const runClient = await getRunClient(projectId);
+  const buildsClient = await getBuildsClient(projectId, accessToken);
+  const runClient = await getRunClient(projectId, accessToken);
   const serviceExists = await checkCloudRunServiceExists(
     projectId,
     location,
     serviceName,
+    accessToken,
     progressCallback
   );
-  const cloudBuildClient = await getCloudBuildClient(projectId);
-  const loggingClient = await getLoggingClient(projectId);
+  const cloudBuildClient = await getCloudBuildClient(projectId, accessToken);
+  const loggingClient = await getLoggingClient(projectId, accessToken);
   let buildSteps;
 
   const servicePatch = {

lib/cloud-api/helpers.js

@@ -100,8 +100,16 @@ async function checkAndEnableApi(
  * @returns {Promise<void>} A promise that resolves when the API is enabled.
  * @throws {Error} If the API fails to enable or if there's an issue checking its status.
  */
-export async function enableApiWithRetry(projectId, api, progressCallback) {
-  const serviceUsageClient = await getServiceUsageClient(projectId);
+export async function enableApiWithRetry(
+  projectId,
+  api,
+  accessToken,
+  progressCallback
+) {
+  const serviceUsageClient = await getServiceUsageClient(
+    projectId,
+    accessToken
+  );
   const serviceName = `projects/${projectId}/services/${api}`;
   try {
     await checkAndEnableApi(
@@ -148,20 +156,25 @@ export async function enableApiWithRetry(projectId, api, progressCallback) {
  * @throws {Error} If an API fails to enable or if there's an issue checking its status.
  * @returns {Promise<void>} A promise that resolves when all specified APIs are enabled.
  */
-export async function ensureApisEnabled(projectId, apis, progressCallback) {
+export async function ensureApisEnabled(
+  projectId,
+  apis,
+  accessToken,
+  progressCallback
+) {
   for (const api of PREREQUISITE_APIS) {
-    await enableApiWithRetry(projectId, api, progressCallback);
+    await enableApiWithRetry(projectId, api, accessToken, progressCallback);
   }
 
   // Before enabling other APIs, ensure billing is enabled.
-  await ensureBillingEnabled(projectId, progressCallback);
+  await ensureBillingEnabled(projectId, accessToken, progressCallback);
 
   const message = 'Checking and enabling required APIs...';
   console.log(message);
   if (progressCallback) progressCallback({ level: 'info', data: message });
 
   for (const api of apis) {
-    await enableApiWithRetry(projectId, api, progressCallback);
+    await enableApiWithRetry(projectId, api, accessToken, progressCallback);
   }
   const successMsg = 'All required APIs are enabled.';
   console.log(successMsg);

lib/cloud-api/projects.js

@@ -26,8 +26,8 @@ import { getProjectsClient } from '../clients.js';
  * @function listProjects
  * @returns {Promise<Array<{id: string}>>} A promise that resolves to an array of project objects, each with an 'id' property. Returns an empty array on error.
  */
-export async function listProjects() {
-  const client = await getProjectsClient();
+export async function listProjects(accessToken) {
+  const client = await getProjectsClient(accessToken);
   try {
     const [projects] = await client.searchProjects();
     return projects.map((project) => ({
@@ -71,8 +71,8 @@ export function generateProjectId() {
  * @param {string} [parent] - Optional. The resource name of the parent under which the project is to be created. e.g., "organizations/123" or "folders/456".
  * @returns {Promise<{projectId: string}|null>} A promise that resolves to an object containing the new project's ID.
  */
-export async function createProject(projectId, parent) {
-  const client = await getProjectsClient();
+export async function createProject(projectId, parent, accessToken) {
+  const client = await getProjectsClient(accessToken);
   let projectIdToUse = projectId;
 
   if (!projectIdToUse) {
@@ -112,10 +112,14 @@ export async function createProject(projectId, parent) {
  * @param {string} [parent] - Optional. The resource name of the parent under which the project is to be created. e.g., "organizations/123" or "folders/456".
  * @returns {Promise<{projectId: string, billingMessage: string}>} A promise that resolves to an object containing the project ID and a billing status message.
  */
-export async function createProjectAndAttachBilling(projectIdParam, parent) {
+export async function createProjectAndAttachBilling(
+  projectIdParam,
+  parent,
+  accessToken
+) {
   let newProject;
   try {
-    newProject = await createProject(projectIdParam, parent);
+    newProject = await createProject(projectIdParam, parent, accessToken);
   } catch (error) {
     throw new Error(`Failed to create project: ${error.message}`);
   }
@@ -128,7 +132,7 @@ export async function createProjectAndAttachBilling(projectIdParam, parent) {
   let billingMessage = `Project ${projectId} created successfully.`;
 
   try {
-    const billingAccounts = await listBillingAccounts();
+    const billingAccounts = await listBillingAccounts(accessToken);
     if (billingAccounts && billingAccounts.length > 0) {
       const firstBillingAccount = billingAccounts.find((acc) => acc.open); // Prefer an open account
       if (firstBillingAccount) {
@@ -137,7 +141,8 @@ export async function createProjectAndAttachBilling(projectIdParam, parent) {
         );
         const billingInfo = await attachProjectToBillingAccount(
           projectId,
-          firstBillingAccount.name
+          firstBillingAccount.name,
+          accessToken
         );
         if (billingInfo && billingInfo.billingEnabled) {
           billingMessage += ` It has been attached to billing account ${firstBillingAccount.displayName}.`;
@@ -171,8 +176,8 @@ export async function createProjectAndAttachBilling(projectIdParam, parent) {
  * @param {string} projectId - The ID of the project to delete.
  * @returns {Promise<void>} A promise that resolves when the delete operation is initiated.
  */
-export async function deleteProject(projectId) {
-  const client = await getProjectsClient();
+export async function deleteProject(projectId, accessToken) {
+  const client = await getProjectsClient(accessToken);
   try {
     console.log(`Attempting to delete project with ID: ${projectId}`);
     await client.deleteProject({ name: `projects/${projectId}` });

lib/cloud-api/registry.js

@@ -33,12 +33,16 @@ import { logAndProgress } from '../util/helpers.js';
  */
 export async function ensureArtifactRegistryRepoExists(
   projectId,
+  accessToken,
   location,
   repositoryId,
   format = 'DOCKER',
   progressCallback
 ) {
-  const artifactRegistryClient = await getArtifactRegistryClient(projectId);
+  const artifactRegistryClient = await getArtifactRegistryClient(
+    projectId,
+    accessToken
+  );
   const parent = `projects/${projectId}/locations/${location}`;
   const repoPath = artifactRegistryClient.repositoryPath(
     projectId,