fix: migrate withOAuthProtectedResource to defineMiddleware (#116)

* refactor: withOAuthProtectedResource

* fix: address PR feedback by @mandarini

Author: mandarini

e2e/supabase/functions/server-e2e/deno.json

@@ -5,7 +5,7 @@
     "@supabase/server/middleware/postgres-admin": "../_vendor/package/dist/middleware/postgres-admin/index.mjs",
     "@supabase/supabase-js": "npm:@supabase/supabase-js@2",
     "@supabase/supabase-js/cors": "npm:@supabase/supabase-js@2/cors",
-    "@supabase/middleware": "npm:@supabase/middleware@0.3.0",
+    "@supabase/middleware": "npm:@supabase/middleware@0.3.1",
     "jose": "npm:jose@6",
     "pg": "npm:pg@8"
   }

package.json

@@ -243,7 +243,7 @@
     "vitest": "^4.1.0"
   },
   "dependencies": {
-    "@supabase/middleware": "^0.3.0",
+    "@supabase/middleware": "^0.3.1",
     "jose": "^6.2.0"
   }
 }

pnpm-lock.yaml

@@ -22,11 +22,14 @@ describe('withOAuthProtectedResource - metadata route', () => {
     expect(body.bearer_methods_supported).toContain('header')
   })
 
-  it('ignores POST to /fn/oauth-protected-resource (passes through)', async () => {
+  it('passes POST to /fn/oauth-protected-resource through to the inner handler', async () => {
+    // Only GET matches the metadata route; every other method (and path) falls
+    // through to the inner handler rather than 404ing — see the path-routing
+    // describe block below.
     const res = await withOAuthProtectedResource(passthrough)(
       req('POST', '/my-fn/oauth-protected-resource'),
     )
-    expect(res.status).toBe(404)
+    expect(res.status).toBe(200)
   })
 })
 
@@ -67,12 +70,15 @@ describe('withOAuthProtectedResource - method pass-through', () => {
 })
 
 describe('withOAuthProtectedResource - path routing', () => {
-  it('returns 404 for unrecognized sub-paths', async () => {
-    // /my-fn/something is not a registered route under the my-fn function
+  it('passes unrecognized sub-paths through to the inner handler (deliberate: AI-995)', async () => {
+    // Was a blanket 404 under the old hand-written closure — an accidental
+    // side effect of being a standalone wrapper, not a deliberate contract.
+    // The defineMiddleware conversion passes through instead, since that's
+    // what fits the composition model: routing is the inner handler's job.
     const res = await withOAuthProtectedResource(passthrough)(
       req('POST', '/my-fn/something'),
     )
-    expect(res.status).toBe(404)
+    expect(res.status).toBe(200)
   })
 
   it('infers function name from first path segment', async () => {
@@ -203,14 +209,14 @@ describe('unauthorizedResponse', () => {
 })
 
 describe('withOAuthProtectedResource - platform argument', () => {
-  it('forwards the platform second argument to the inner handler', async () => {
+  it('no longer forwards the raw platform argument — the inner handler receives ctx instead', async () => {
     let seen: unknown
     const handler = async (_req: Request, platformArg?: unknown) => {
       seen = platformArg
       return new Response('ok')
     }
     const env = { MY_BINDING: 'value' }
     await withOAuthProtectedResource(handler)(req('POST', '/my-fn'), env)
-    expect(seen).toBe(env)
+    expect(seen).not.toBe(env)
   })
 })

src/oauth-protected-resource/with-oauth-protected-resource.test.ts

@@ -1,6 +1,15 @@
+import { defineMiddleware } from '@supabase/middleware'
+import type { Middleware } from '@supabase/middleware'
+
 import { resourceMetadataResponse } from './responses.js'
 import { getResourceMetadataUrl, inferFunctionName } from './url.js'
 
+/** Shape contributed at `ctx.oauthProtectedResource`. */
+export interface OAuthProtectedResourceContribution {
+  /** Absolute URL of this resource's OAuth Protected Resource Metadata document (RFC 9728). */
+  resourceMetadataUrl: string
+}
+
 /**
  * Wraps a request handler with OAuth 2.1 Protected Resource behavior (RFC 9728)
  * for Supabase Edge Functions.
@@ -9,11 +18,13 @@ import { getResourceMetadataUrl, inferFunctionName } from './url.js'
  *   (with permissive CORS, including the `OPTIONS` preflight, so browser-based clients can read it)
  * - Enriches a `401` from the inner handler with `WWW-Authenticate: Bearer resource_metadata="..."`,
  *   unless the handler already set a `WWW-Authenticate` header (its value wins)
- * - Returns `404` for any other path (Edge Functions are single-endpoint - the inner handler owns `/{fn}` only)
+ * - Passes any other path through to the inner handler unchanged (composition,
+ *   not routing, decides what happens to it)
  *
- * The returned handler's optional second parameter is the host's platform
- * argument (a Workers `env`, a Deno `ServeHandlerInfo`) and is forwarded to
- * the inner handler unchanged — required for `withSupabase` to capture it.
+ * Contributes `ctx.oauthProtectedResource` (the resolved metadata URL) to the
+ * downstream context. When nested under `withSupabase`, this key is present at
+ * runtime but not yet reflected in the handler's `SupabaseContext` type — see
+ * `withSupabase`'s type note.
  *
  * @category Middleware
  *
@@ -32,61 +43,75 @@ import { getResourceMetadataUrl, inferFunctionName } from './url.js'
  * )
  * ```
  */
-export function withOAuthProtectedResource(
-  handler: (req: Request, platformArg?: unknown) => Promise<Response>,
-): (req: Request, platformArg?: unknown) => Promise<Response> {
-  return async (req: Request, platformArg?: unknown): Promise<Response> => {
-    const url = new URL(req.url)
-    const fn = inferFunctionName(req)
-    if (!fn) return new Response('Not Found', { status: 404 })
-    const basePath = `/${fn}`
+export const withOAuthProtectedResource: Middleware<
+  'oauthProtectedResource',
+  undefined,
+  Record<never, never>,
+  OAuthProtectedResourceContribution
+> = defineMiddleware<
+  'oauthProtectedResource',
+  undefined,
+  Record<never, never>,
+  OAuthProtectedResourceContribution
+>({
+  key: 'oauthProtectedResource',
+  run: () =>
+    async function* (req) {
+      const url = new URL(req.url)
+      const fn = inferFunctionName(req)
+      const metadataPath = fn ? `/${fn}/oauth-protected-resource` : undefined
 
-    // RFC 9728 — OAuth Protected Resource Metadata
-    if (
-      req.method === 'GET' &&
-      url.pathname === `${basePath}/oauth-protected-resource`
-    ) {
-      return resourceMetadataResponse(req)
-    }
+      // RFC 9728 — OAuth Protected Resource Metadata
+      if (
+        metadataPath &&
+        req.method === 'GET' &&
+        url.pathname === metadataPath
+      ) {
+        return resourceMetadataResponse(req)
+      }
 
-    // CORS preflight for the metadata route — browser-based clients (e.g.
-    // MCP Inspector) fetch the discovery document cross-origin.
-    if (
-      req.method === 'OPTIONS' &&
-      url.pathname === `${basePath}/oauth-protected-resource`
-    ) {
-      return new Response(null, {
-        status: 204,
-        headers: {
-          'Access-Control-Allow-Origin': '*',
-          'Access-Control-Allow-Methods': 'GET, OPTIONS',
-          'Access-Control-Allow-Headers': 'content-type, mcp-protocol-version',
-        },
-      })
-    }
+      // CORS preflight for the metadata route — browser-based clients (e.g.
+      // MCP Inspector) fetch the discovery document cross-origin.
+      if (
+        metadataPath &&
+        req.method === 'OPTIONS' &&
+        url.pathname === metadataPath
+      ) {
+        return new Response(null, {
+          status: 204,
+          headers: {
+            'Access-Control-Allow-Origin': '*',
+            'Access-Control-Allow-Methods': 'GET, OPTIONS',
+            'Access-Control-Allow-Headers':
+              'content-type, mcp-protocol-version',
+          },
+        })
+      }
 
-    if (url.pathname !== basePath) {
-      return new Response('Not Found', { status: 404 })
-    }
+      const resourceMetadataUrl = getResourceMetadataUrl(req)
+      const response = yield {
+        oauthProtectedResource: { resourceMetadataUrl },
+      }
 
-    const response = await handler(req, platformArg)
+      // Enrich a 401 with WWW-Authenticate so clients can discover the auth
+      // server — unless the handler already set one (its value wins, e.g. an
+      // RFC 6750 error or a custom resource_metadata override).
+      if (
+        response.status === 401 &&
+        !response.headers.has('WWW-Authenticate')
+      ) {
+        const headers = new Headers(response.headers)
+        headers.set(
+          'WWW-Authenticate',
+          `Bearer resource_metadata="${resourceMetadataUrl}"`,
+        )
+        return new Response(response.body, {
+          status: 401,
+          statusText: response.statusText,
+          headers,
+        })
+      }
 
-    // Enrich a 401 with WWW-Authenticate so clients can discover the auth
-    // server — unless the handler already set one (its value wins, e.g. an
-    // RFC 6750 error or a custom resource_metadata override).
-    if (response.status === 401 && !response.headers.has('WWW-Authenticate')) {
-      const headers = new Headers(response.headers)
-      headers.set(
-        'WWW-Authenticate',
-        `Bearer resource_metadata="${getResourceMetadataUrl(req)}"`,
-      )
-      return new Response(response.body, {
-        status: 401,
-        statusText: response.statusText,
-        headers,
-      })
-    }
-
-    return response
-  }
-}
+      return response
+    },
+})

src/oauth-protected-resource/with-oauth-protected-resource.ts

@@ -3,6 +3,7 @@ import { defineMiddleware, getEnv } from '@supabase/middleware'
 
 import { _resetAllowDeprecationWarned } from './core/utils/deprecation.js'
 import { EnvError } from './errors.js'
+import { withOAuthProtectedResource } from './oauth-protected-resource/with-oauth-protected-resource.js'
 import { withSupabase } from './with-supabase.js'
 
 const baseEnv = {
@@ -353,6 +354,40 @@ describe('withSupabase', () => {
     })
   })
 
+  describe('nested under an upstream middleware', () => {
+    // Nested under another entry, withSupabase must spread the upstream context
+    // rather than reseed — otherwise it drops upstream ctx keys and clobbers the
+    // platform env the entry captured (silently breaking getEnv on Workers).
+    it('preserves upstream ctx keys and the platform env captured by the entry', async () => {
+      let seenMetadataUrl: string | undefined
+      let seenBinding: string | undefined
+
+      const composed = withOAuthProtectedResource(
+        withSupabase({ auth: 'none', env: baseEnv }, async (_req, ctx) => {
+          // Present at runtime but not on the SupabaseContext type yet, so cast.
+          const upstream = ctx as {
+            oauthProtectedResource?: { resourceMetadataUrl: string }
+          }
+          seenMetadataUrl = upstream.oauthProtectedResource?.resourceMetadataUrl
+          seenBinding = getEnv('NESTED_TEST_BINDING')
+          return Response.json({ ok: true })
+        }),
+      )
+
+      // Workers-style entry invocation: fetch(request, env). withOAuthProtected-
+      // Resource is the entry, so it seeds the context with this env.
+      const res = await composed(new Request('http://localhost/my-fn'), {
+        NESTED_TEST_BINDING: 'from-platform',
+      })
+
+      expect(res.status).toBe(200)
+      // Upstream contribution survived withSupabase's context construction.
+      expect(seenMetadataUrl).toContain('/my-fn/oauth-protected-resource')
+      // Platform env captured by the entry was not clobbered by a reseed.
+      expect(seenBinding).toBe('from-platform')
+    })
+  })
+
   describe('client construction errors', () => {
     it('maps client-construction EnvError to a 500 JSON response', async () => {
       const handler = withSupabase(

src/with-supabase.test.ts

@@ -4,7 +4,7 @@ import { AuthError, CreateSupabaseClientError, EnvError } from './errors.js'
 import { withSupabaseAdminClient } from './middleware/admin-client/index.js'
 import { withSupabaseClient } from './middleware/client/index.js'
 import type { SupabaseContext, WithSupabaseConfig } from './types.js'
-import { seedContext } from '@supabase/middleware'
+import { isContext, seedContext } from '@supabase/middleware'
 import type { Entry } from '@supabase/middleware'
 
 type AnyEntry = Entry<string, object, unknown>
@@ -170,15 +170,15 @@ export function withSupabase<Database = unknown>(
 
     let response: Response
     try {
-      // seedContext() stamps the engine's context marker so middleware entries
-      // recognise this as an upstream context, and captures the host's second
-      // fetch argument (a Workers `env`, a Deno `ServeHandlerInfo`) as the
-      // platform env behind the engine's importable getEnv — without the
-      // forward, Workers bindings would be invisible to middleware. The
-      // verified auth identity is seeded alongside it; the client middleware
-      // read `authMode` / `authKeyName` to mirror the verified credentials.
+      // As the entry point, `platformArg` is the host env — seed a context from
+      // it (captured behind getEnv). Nested under another middleware, it's an
+      // already-seeded context: reuse it, or reseeding would clobber the platform
+      // env and drop upstream ctx keys.
+      const baseContext = isContext(platformArg)
+        ? platformArg
+        : seedContext(platformArg)
       response = await composed(req, {
-        ...seedContext(platformArg),
+        ...baseContext,
         userClaims: auth.userClaims,
         jwtClaims: auth.jwtClaims,
         authMode: auth.authMode,