Add support for GCP IAP JIT account provisioning (#330)

msukkari · web-flow · commit 9227b3caba7d · 2025-06-03T19:28:38.000-07:00
* initial gcp iap implementation

* gcp iap working

* add docs for gcp iap

* feedback

* changelog
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - Added copy button for filenames. [#328](https://github.com/sourcebot-dev/sourcebot/pull/328)
 - Added development docker compose file. [#328](https://github.com/sourcebot-dev/sourcebot/pull/328)
+- Added GCP IAP JIT provisioning. [#330](https://github.com/sourcebot-dev/sourcebot/pull/330)
 
 ### Fixed
 - Fixed issue with the symbol hover popover clipping at the top of the page. [#326](https://github.com/sourcebot-dev/sourcebot/pull/326)
diff --git a/docs/docs/configuration/auth/overview.mdx b/docs/docs/configuration/auth/overview.mdx
@@ -80,6 +80,16 @@ Optional environment variables:
 - `AUTH_EE_GOOGLE_CLIENT_ID`
 - `AUTH_EE_GOOGLE_CLIENT_SECRET`
 
+### GCP IAP
+---
+
+Custom provider built to enable automatic Sourcebot account registration/login when using GCP IAP.
+
+**Required environment variables**
+- `AUTH_EE_GCP_IAP_ENABLED`
+- `AUTH_EE_GCP_IAP_AUDIENCE`
+    - This can be found by selecting the ⋮ icon next to the IAP-enabled backend service and pressing `Get JWT audience code`
+
 ### Okta
 ---
 
diff --git a/docs/docs/configuration/environment-variables.mdx b/docs/docs/configuration/environment-variables.mdx
@@ -51,6 +51,8 @@ The following environment variables allow you to configure your Sourcebot deploy
 | `AUTH_EE_OKTA_CLIENT_ID` | `-` | <p>The client ID for Okta SSO authentication.</p> |
 | `AUTH_EE_OKTA_CLIENT_SECRET` | `-` | <p>The client secret for Okta SSO authentication.</p> |
 | `AUTH_EE_OKTA_ISSUER` | `-` | <p>The issuer URL for Okta SSO authentication.</p> |
+| `AUTH_EE_GCP_IAP_ENABLED` | `false` | <p>When enabled, allows Sourcebot to automatically register/login from a successful GCP IAP redirect</p> |
+| `AUTH_EE_GCP_IAP_AUDIENCE` | - | <p>The GCP IAP audience to use when verifying JWT tokens. Must be set to enable GCP IAP JIT provisioning</p> | 
 
 
 ### Review Agent Environment Variables
diff --git a/packages/web/package.json b/packages/web/package.json
@@ -114,6 +114,7 @@
     "embla-carousel-react": "^8.3.0",
     "escape-string-regexp": "^5.0.0",
     "fuse.js": "^7.0.0",
+    "google-auth-library": "^9.15.1",
     "graphql": "^16.9.0",
     "http-status-codes": "^2.3.0",
     "input-otp": "^1.4.2",
diff --git a/packages/web/src/app/[domain]/components/gcpIapAuth.tsx b/packages/web/src/app/[domain]/components/gcpIapAuth.tsx
@@ -0,0 +1,26 @@
+'use client';
+
+import { signIn } from "next-auth/react";
+import { useEffect } from "react";
+
+interface GcpIapAuthProps {
+    callbackUrl?: string;
+}
+
+export const GcpIapAuth = ({ callbackUrl }: GcpIapAuthProps) => {
+    useEffect(() => {
+        signIn("gcp-iap", {
+            redirectTo: callbackUrl ?? "/"
+        }).catch((error) => {
+            console.error("Error signing in with GCP IAP:", error);
+        });
+    }, [callbackUrl]);
+
+    return (
+        <div className="flex items-center justify-center min-h-screen">
+            <div className="text-center">
+                <p className="text-lg">Signing in with Google Cloud IAP...</p>
+            </div>
+        </div>
+    );
+}; 
diff --git a/packages/web/src/app/[domain]/layout.tsx b/packages/web/src/app/[domain]/layout.tsx
@@ -17,6 +17,7 @@ import { PendingApprovalCard } from "./components/pendingApproval";
 import { hasEntitlement } from "@/features/entitlements/server";
 import { getPublicAccessStatus } from "@/ee/features/publicAccess/publicAccess";
 import { env } from "@/env.mjs";
+import { GcpIapAuth } from "./components/gcpIapAuth";
 
 interface LayoutProps {
     children: React.ReactNode,
@@ -37,7 +38,12 @@ export default async function Layout({
     if (!publicAccessEnabled) {
         const session = await auth();
         if (!session) {
-            redirect('/login');
+            const ssoEntitlement = await hasEntitlement("sso");
+            if (ssoEntitlement && env.AUTH_EE_GCP_IAP_ENABLED && env.AUTH_EE_GCP_IAP_AUDIENCE) {
+                return <GcpIapAuth callbackUrl={`/${domain}`} />;
+            } else {
+                redirect('/login');
+            }
         }
 
         const membership = await prisma.userToOrg.findUnique({
diff --git a/packages/web/src/auth.ts b/packages/web/src/auth.ts
@@ -5,21 +5,17 @@ import EmailProvider from "next-auth/providers/nodemailer";
 import { PrismaAdapter } from "@auth/prisma-adapter"
 import { prisma } from "@/prisma";
 import { env } from "@/env.mjs";
-import { OrgRole, User } from '@sourcebot/db';
+import { User } from '@sourcebot/db';
 import 'next-auth/jwt';
 import type { Provider } from "next-auth/providers";
 import { verifyCredentialsRequestSchema } from './lib/schemas';
 import { createTransport } from 'nodemailer';
 import { render } from '@react-email/render';
 import MagicLinkEmail from './emails/magicLinkEmail';
-import { SINGLE_TENANT_ORG_DOMAIN, SINGLE_TENANT_ORG_ID } from './lib/constants';
 import bcrypt from 'bcryptjs';
-import { createAccountRequest } from './actions';
-import { getSSOProviders, handleJITProvisioning } from '@/ee/sso/sso';
+import { getSSOProviders } from '@/ee/sso/sso';
 import { hasEntitlement } from '@/features/entitlements/server';
-import { isServiceError } from './lib/utils';
-import { ServiceErrorException } from './lib/serviceError';
-import { createLogger } from "@sourcebot/logger";
+import { onCreateUser } from '@/lib/authUtils';
 
 export const runtime = 'nodejs';
 
@@ -37,8 +33,6 @@ declare module 'next-auth/jwt' {
     }
 }
 
-const logger = createLogger('web-auth');
-
 export const getProviders = () => {
     const providers: Provider[] = [];
 
@@ -134,91 +128,6 @@ export const getProviders = () => {
     return providers;
 }
 
-const onCreateUser = async ({ user }: { user: AuthJsUser }) => {
-    // In single-tenant mode, we assign the first user to sign
-    // up as the owner of the default org.
-    if (
-        env.SOURCEBOT_TENANCY_MODE === 'single'
-    ) {
-        const defaultOrg = await prisma.org.findUnique({
-            where: {
-                id: SINGLE_TENANT_ORG_ID,
-            },
-            include: {
-                members: {
-                    where: {
-                        role: {
-                            not: OrgRole.GUEST,
-                        }
-                    }
-                },
-            }
-        });
-
-        if (!defaultOrg) {
-            throw new Error("Default org not found on single tenant user creation");
-        }
-
-        // We can't use the getOrgMembers action here because we're not authed yet
-        const members = await prisma.userToOrg.findMany({
-            where: {
-                orgId: SINGLE_TENANT_ORG_ID,
-                role: {
-                    not: OrgRole.GUEST,
-                }
-            },
-        });
-
-        // Only the first user to sign up will be an owner of the default org.
-        const isFirstUser = members.length === 0;
-        if (isFirstUser) {
-            await prisma.$transaction(async (tx) => {
-                await tx.org.update({
-                    where: {
-                        id: SINGLE_TENANT_ORG_ID,
-                    },
-                    data: {
-                        members: {
-                            create: {
-                                role: OrgRole.OWNER,
-                                user: {
-                                    connect: {
-                                        id: user.id,
-                                    }
-                                }
-                            }
-                        }
-                    }
-                });
-
-                await tx.user.update({
-                    where: {
-                        id: user.id,
-                    },
-                    data: {
-                        pendingApproval: false,
-                    }
-                });
-            });
-        } else {
-            // TODO(auth): handle multi tenant case
-            if (env.AUTH_EE_ENABLE_JIT_PROVISIONING === 'true' && hasEntitlement("sso")) {
-                const res = await handleJITProvisioning(user.id!, SINGLE_TENANT_ORG_DOMAIN);
-                if (isServiceError(res)) {
-                    logger.error(`Failed to provision user ${user.id} for org ${SINGLE_TENANT_ORG_DOMAIN}: ${res.message}`);
-                    throw new ServiceErrorException(res);
-                }
-            } else {
-                const res = await createAccountRequest(user.id!, SINGLE_TENANT_ORG_DOMAIN);
-                if (isServiceError(res)) {
-                    logger.error(`Failed to provision user ${user.id} for org ${SINGLE_TENANT_ORG_DOMAIN}: ${res.message}`);
-                    throw new ServiceErrorException(res);
-                }
-            }
-        }
-    }
-}
-
 export const { handlers, signIn, signOut, auth } = NextAuth({
     secret: env.AUTH_SECRET,
     adapter: PrismaAdapter(prisma),
diff --git a/packages/web/src/ee/sso/sso.tsx b/packages/web/src/ee/sso/sso.tsx
@@ -12,7 +12,14 @@ import { OrgRole } from "@sourcebot/db";
 import { getSeats, SOURCEBOT_UNLIMITED_SEATS } from "@/features/entitlements/server";
 import { StatusCodes } from "http-status-codes";
 import { ErrorCode } from "@/lib/errorCodes";
+import { OAuth2Client } from "google-auth-library";
 import { sew } from "@/actions";
+import Credentials from "next-auth/providers/credentials";
+import type { User as AuthJsUser } from "next-auth";
+import { onCreateUser } from "@/lib/authUtils";
+import { createLogger } from "@sourcebot/logger";
+
+const logger = createLogger('web-sso');
 
 export const getSSOProviders = (): Provider[] => {
     const providers: Provider[] = [];
@@ -88,6 +95,82 @@ export const getSSOProviders = (): Provider[] => {
         }));
     }
 
+    if (env.AUTH_EE_GCP_IAP_ENABLED && env.AUTH_EE_GCP_IAP_AUDIENCE) {
+        providers.push(Credentials({
+            id: "gcp-iap",
+            name: "Google Cloud IAP",
+            credentials: {},
+            authorize: async (credentials, req) => {
+                try {
+                    const iapAssertion = req.headers?.get("x-goog-iap-jwt-assertion");
+                    if (!iapAssertion || typeof iapAssertion !== "string") {
+                        logger.warn("No IAP assertion found in headers");
+                        return null;
+                    }
+
+                    const oauth2Client = new OAuth2Client();
+                    
+                    const { pubkeys } = await oauth2Client.getIapPublicKeys();
+                    const ticket = await oauth2Client.verifySignedJwtWithCertsAsync(
+                        iapAssertion,
+                        pubkeys,
+                        env.AUTH_EE_GCP_IAP_AUDIENCE,
+                        ['https://cloud.google.com/iap']
+                    );
+
+                    const payload = ticket.getPayload();
+                    if (!payload) {
+                        logger.warn("Invalid IAP token payload");
+                        return null;
+                    }
+
+                    const email = payload.email;
+                    const name = payload.name || payload.email;
+                    const image = payload.picture;
+
+                    if (!email) {
+                        logger.warn("Missing email in IAP token");
+                        return null;
+                    }
+
+                    const existingUser = await prisma.user.findUnique({
+                        where: { email }
+                    });
+
+                    if (!existingUser) {
+                        const newUser = await prisma.user.create({
+                            data: {
+                                email,
+                                name,
+                                image,
+                            }
+                        });
+
+                        const authJsUser: AuthJsUser = {
+                            id: newUser.id,
+                            email: newUser.email,
+                            name: newUser.name,
+                            image: newUser.image,
+                        };
+
+                        await onCreateUser({ user: authJsUser });
+                        return authJsUser;
+                    } else {
+                        return {
+                            id: existingUser.id,
+                            email: existingUser.email,
+                            name: existingUser.name,
+                            image: existingUser.image,
+                        };
+                    }
+                } catch (error) {
+                    logger.error("Error verifying IAP token:", error);
+                    return null;
+                }
+            },
+        }));
+    }
+
     return providers;
 }
 
@@ -129,7 +212,7 @@ export const handleJITProvisioning = async (userId: string, domain: string): Pro
     });
 
     if (userToOrg) {
-        console.warn(`JIT provisioning skipped for user ${userId} since they're already a member of org ${domain}`);
+        logger.warn(`JIT provisioning skipped for user ${userId} since they're already a member of org ${domain}`);
         return true;
     }
 
diff --git a/packages/web/src/env.mjs b/packages/web/src/env.mjs
@@ -50,6 +50,9 @@ export const env = createEnv({
         AUTH_EE_MICROSOFT_ENTRA_ID_CLIENT_SECRET: z.string().optional(),
         AUTH_EE_MICROSOFT_ENTRA_ID_ISSUER: z.string().optional(),
 
+        AUTH_EE_GCP_IAP_ENABLED: booleanSchema.default('false'),
+        AUTH_EE_GCP_IAP_AUDIENCE: z.string().optional(),
+
         DATA_CACHE_DIR: z.string(),
 
         // Email
diff --git a/packages/web/src/lib/authUtils.ts b/packages/web/src/lib/authUtils.ts
diff --git a/yarn.lock b/yarn.lock
