fix: return to the chosen plan's checkout after sign-in

src/components/github-auth-button.tsx

@@ -1,12 +1,20 @@
 import { MarkGithubIcon } from '@primer/octicons-react'
 import { ArrowRight } from 'lucide-react'
 
-export function GitHubAuthButton({ label }: { label: string }) {
+export function GitHubAuthButton({
+  label,
+  next,
+}: {
+  label: string
+  // Same-site path to return to after sign-in; validated on the server.
+  next?: string
+}) {
+  const href = next
+    ? `/auth/github/start?next=${encodeURIComponent(next)}`
+    : '/auth/github/start'
+
   return (
-    <a
-      className="button-primary mt-7 w-full justify-center"
-      href="/auth/github/start"
-    >
+    <a className="button-primary mt-7 w-full justify-center" href={href}>
       <MarkGithubIcon size={18} />
       {label}
       <ArrowRight className="size-4" />

src/lib/safe-next.test.ts

@@ -0,0 +1,27 @@
+import { describe, expect, it } from 'vitest'
+
+import { DEFAULT_NEXT, safeNextPath } from './safe-next'
+
+describe('safeNextPath', () => {
+  it('keeps same-site paths with their query', () => {
+    expect(safeNextPath('/pricing?checkout=team')).toBe(
+      '/pricing?checkout=team',
+    )
+    expect(safeNextPath('/app')).toBe('/app')
+  })
+
+  it.each([
+    undefined,
+    null,
+    42,
+    '',
+    'pricing',
+    'https://evil.example/phish',
+    '//evil.example',
+    '/\\evil.example',
+    'javascript:alert(1)',
+    `/${'a'.repeat(600)}`,
+  ])('falls back to the dashboard for %j', (value) => {
+    expect(safeNextPath(value)).toBe(DEFAULT_NEXT)
+  })
+})

src/lib/safe-next.ts

@@ -0,0 +1,29 @@
+// Where to send someone after sign-in. Only same-site paths are allowed, so a
+// crafted `?next=` can never bounce users to another origin (open redirect).
+// Anything else falls back to the dashboard.
+
+export const DEFAULT_NEXT = '/app'
+
+export function safeNextPath(value: unknown): string {
+  if (typeof value !== 'string' || value.length > 512) return DEFAULT_NEXT
+  // Must be a root-relative path. `//host` and `/\host` are protocol-relative
+  // in browsers, so reject them before parsing.
+  if (
+    !value.startsWith('/') ||
+    value.startsWith('//') ||
+    value.startsWith('/\\')
+  ) {
+    return DEFAULT_NEXT
+  }
+
+  const base = 'https://jargons.invalid'
+  let url: URL
+  try {
+    url = new URL(value, base)
+  } catch {
+    return DEFAULT_NEXT
+  }
+  if (url.origin !== base) return DEFAULT_NEXT
+
+  return `${url.pathname}${url.search}${url.hash}`
+}

src/routes/auth.github.callback.tsx

@@ -1,6 +1,9 @@
 import { createFileRoute } from '@tanstack/react-router'
 
-import { completeGitHubAuthWithCode } from '../server/github-auth'
+import {
+  completeGitHubAuthWithCode,
+  consumeSignInNext,
+} from '../server/github-auth'
 
 export const Route = createFileRoute('/auth/github/callback')({
   server: {
@@ -19,7 +22,7 @@ export const Route = createFileRoute('/auth/github/callback')({
         try {
           await completeGitHubAuthWithCode(code)
 
-          return redirectTo(new URL('/app', url.origin))
+          return redirectTo(new URL(await consumeSignInNext(), url.origin))
         } catch (error) {
           console.error('GitHub auth callback failed', error)
 

src/routes/auth.github.start.tsx

@@ -1,17 +1,20 @@
 import { createFileRoute } from '@tanstack/react-router'
 
-import { getGitHubAuthorizeUrl } from '../server/github-auth'
+import { getGitHubAuthorizeUrl, setSignInNext } from '../server/github-auth'
 
 export const Route = createFileRoute('/auth/github/start')({
   server: {
     handlers: {
-      GET: async () =>
+      GET: async ({ request }) => {
+        await setSignInNext(new URL(request.url).searchParams.get('next'))
+
         // Not Response.redirect() — its immutable headers break if the
         // framework needs to append headers (e.g. Set-Cookie) afterwards.
-        new Response(null, {
+        return new Response(null, {
           status: 302,
           headers: { location: getGitHubAuthorizeUrl() },
-        }),
+        })
+      },
     },
   },
 })

src/routes/auth.sign-in.tsx

@@ -6,14 +6,15 @@ import { AuthShell } from '../components/auth-shell'
 import { GitHubAuthButton } from '../components/github-auth-button'
 
 export const Route = createFileRoute('/auth/sign-in')({
-  validateSearch: (search): { error?: string } => ({
+  validateSearch: (search): { error?: string; next?: string } => ({
     error: typeof search.error === 'string' ? search.error : undefined,
+    next: typeof search.next === 'string' ? search.next : undefined,
   }),
   component: SignInPage,
 })
 
 function SignInPage() {
-  const { error } = Route.useSearch()
+  const { error, next } = Route.useSearch()
 
   return (
     <AuthShell
@@ -32,7 +33,7 @@ function SignInPage() {
           Continue with the GitHub account connected to your Jargons workspace.
         </p>
 
-        <GitHubAuthButton label="Continue with GitHub" />
+        <GitHubAuthButton label="Continue with GitHub" next={next} />
 
         {error ? (
           <p className="mt-4 rounded-2xl border border-red-300/15 bg-red-300/[0.06] px-4 py-3 text-sm leading-6 text-red-200">
@@ -45,6 +46,7 @@ function SignInPage() {
           <Link
             className="text-amber-300 hover:text-amber-200"
             to="/auth/sign-up"
+            search={{ next }}
           >
             Create an account
           </Link>

src/routes/auth.sign-up.tsx

@@ -6,10 +6,15 @@ import { AuthShell } from '../components/auth-shell'
 import { GitHubAuthButton } from '../components/github-auth-button'
 
 export const Route = createFileRoute('/auth/sign-up')({
+  validateSearch: (search): { next?: string } => ({
+    next: typeof search.next === 'string' ? search.next : undefined,
+  }),
   component: SignUpPage,
 })
 
 function SignUpPage() {
+  const { next } = Route.useSearch()
+
   return (
     <AuthShell
       eyebrow="create workspace"
@@ -28,7 +33,7 @@ function SignUpPage() {
           repositories.
         </p>
 
-        <GitHubAuthButton label="Sign up with GitHub" />
+        <GitHubAuthButton label="Sign up with GitHub" next={next} />
 
         <div className="mt-5 grid gap-3">
           <SignupBenefit
@@ -53,6 +58,7 @@ function SignUpPage() {
           <Link
             className="text-amber-300 hover:text-amber-200"
             to="/auth/sign-in"
+            search={{ next }}
           >
             Sign in
           </Link>

src/routes/pricing.tsx

@@ -1,13 +1,22 @@
 import { Link, createFileRoute } from '@tanstack/react-router'
 import { ArrowRight, Check, LoaderCircle } from 'lucide-react'
-import { useState } from 'react'
+import { useEffect, useRef, useState } from 'react'
 
 import { OwlMark } from '../components/owl-mark'
-import { FREE_RUN_LIMIT, PAID_PLAN_IDS, PAID_PLANS } from '../lib/plans'
+import {
+  FREE_RUN_LIMIT,
+  isPaidPlan,
+  PAID_PLAN_IDS,
+  PAID_PLANS,
+} from '../lib/plans'
 import type { PaidPlan, Plan } from '../lib/plans'
 import { getBilling } from '../server/billing'
 
 export const Route = createFileRoute('/pricing')({
+  // `?checkout=<plan>` resumes a subscribe click that had to sign in first.
+  validateSearch: (search): { checkout?: PaidPlan } => ({
+    checkout: isPaidPlan(search.checkout) ? search.checkout : undefined,
+  }),
   loader: async () => {
     const billing = await getBilling()
     return {
@@ -18,14 +27,33 @@ export const Route = createFileRoute('/pricing')({
   component: PricingPage,
 })
 
+// Sign in, then come back here and go straight into `target`'s checkout.
+function signInThenCheckout(target: PaidPlan) {
+  const next = `/pricing?checkout=${target}`
+  window.location.assign(`/auth/sign-in?next=${encodeURIComponent(next)}`)
+}
+
 function PricingPage() {
   const { signedIn, plan } = Route.useLoaderData()
+  const { checkout } = Route.useSearch()
+  const navigate = Route.useNavigate()
   const [busy, setBusy] = useState<PaidPlan | null>(null)
   const [error, setError] = useState<PaidPlan | null>(null)
+  const resumed = useRef(false)
+
+  useEffect(() => {
+    if (!checkout || resumed.current) return
+    resumed.current = true
+    // Drop the param first so Back from the checkout doesn't restart it.
+    void navigate({ search: {}, replace: true })
+    if (signedIn && plan === 'free') {
+      void subscribe(checkout)
+    }
+  }, [checkout, signedIn, plan])
 
   async function subscribe(target: PaidPlan) {
     if (!signedIn) {
-      window.location.assign('/auth/sign-in')
+      signInThenCheckout(target)
       return
     }
     setBusy(target)
@@ -37,7 +65,7 @@ function PricingPage() {
         body: JSON.stringify({ plan: target }),
       })
       if (res.status === 401) {
-        window.location.assign('/auth/sign-in')
+        signInThenCheckout(target)
         return
       }
       const data = (await res.json().catch(() => ({}))) as { url?: string }

src/server/github-auth.ts

@@ -1,6 +1,7 @@
 import { createServerFn } from '@tanstack/react-start'
 
 import { loadDb } from '../db/load'
+import { DEFAULT_NEXT, safeNextPath } from '../lib/safe-next'
 import { getEnv, getOptionalEnv } from './env'
 
 type GitHubTokenResponse = {
@@ -23,6 +24,8 @@ const githubAuthorizeUrl = 'https://github.com/login/oauth/authorize'
 const githubTokenUrl = 'https://github.com/login/oauth/access_token'
 const githubUserUrl = 'https://api.github.com/user'
 const sessionCookieName = 'jargons_session'
+// Where to land after the GitHub round trip, set when sign-in starts.
+const signInNextCookieName = 'jargons_sign_in_next'
 
 export type CurrentUser = {
   id: string
@@ -270,6 +273,38 @@ export async function deleteSessionCookie() {
   deleteCookie(sessionCookieName, { path: '/' })
 }
 
+// Remember where to send the user once GitHub sign-in completes. Lax is enough:
+// the cookie only has to survive the top-level redirect back from GitHub.
+export async function setSignInNext(next: unknown): Promise<void> {
+  const path = safeNextPath(next)
+  const { setCookie, deleteCookie } =
+    await import('@tanstack/react-start/server')
+
+  if (path === DEFAULT_NEXT) {
+    deleteCookie(signInNextCookieName, { path: '/' })
+    return
+  }
+
+  setCookie(signInNextCookieName, path, {
+    maxAge: 60 * 10,
+    httpOnly: true,
+    path: '/',
+    sameSite: 'lax',
+    secure: getOptionalEnv('NODE_ENV', 'development') === 'production',
+  })
+}
+
+// Read and clear the post-sign-in destination (the dashboard when unset).
+export async function consumeSignInNext(): Promise<string> {
+  const { getCookie, deleteCookie } =
+    await import('@tanstack/react-start/server')
+  const path = safeNextPath(getCookie(signInNextCookieName))
+
+  deleteCookie(signInNextCookieName, { path: '/' })
+
+  return path
+}
+
 export async function exchangeGitHubCode(code: string) {
   const response = await fetch(githubTokenUrl, {
     method: 'POST',