feat: make middleware an extensible function and change name to axiomMiddleware

Author: gabrielelpidio

src/index.ts

@@ -1,6 +1,6 @@
 export { log, Logger, LogLevel, type LoggerConfig, type RequestReport } from './logger';
 export { EndpointType, throttle } from './shared';
-export { middleware, config } from './middleware';
+export { axiomMiddleware } from './middleware';
 export * from './platform/base';
 export * from './config';
 export { withAxiom, type AxiomRequest, withAxiomNextConfig, withAxiomRouteHandler } from './withAxiom';

src/middleware.ts

@@ -1,52 +1,55 @@
-import { NextResponse } from 'next/server';
-import type { NextFetchEvent, NextRequest } from 'next/server';
+import { NextMiddleware, NextResponse } from 'next/server';
 import { config as axiomConfig } from './config';
 import { EndpointType } from './shared';
 
 const webVitalsEndpoint = axiomConfig.getIngestURL(EndpointType.webVitals);
 const logsEndpoint = axiomConfig.getIngestURL(EndpointType.logs);
 
-const headers = {
-  authorization: 'Bearer ' + axiomConfig.token,
-  'Content-Type': 'application/json',
-};
-
-export async function middleware(request: NextRequest, event: NextFetchEvent): Promise<NextResponse<unknown> | void> {
-  // If the request is not for axiom, do nothing
-  // This is a safety check, as users may add a custom matcher
-  if (!request.nextUrl.pathname.startsWith('/_axiom')) return;
-
-  // Web vitals
-  if (request.nextUrl.pathname.startsWith('/_axiom/web-vitals')) {
-    // Forward the request to the axiom ingest endpoint
-    event.waitUntil(
-      fetch(webVitalsEndpoint, {
-        body: request.body,
-        method: 'POST',
-        headers,
-      }).catch(console.error)
-    );
-
-    // Return a 204 to the client
-    return new NextResponse(null, { status: 204 });
-  }
-
-  // Logs
-  if (request.nextUrl.pathname.startsWith('/_axiom/logs')) {
-    // Forward the request to the axiom ingest endpoint
-    event.waitUntil(
-      fetch(logsEndpoint, {
+export const axiomMiddleware =
+  (middleware?: NextMiddleware): NextMiddleware =>
+  async (request, event) => {
+    // If the request is not for axiom, do nothing
+    // This is a safety check, as users may add a custom matcher
+    if (!request.nextUrl.pathname.startsWith('/_axiom') || typeof axiomConfig.customEndpoint !== 'undefined') {
+      return middleware ? middleware(request, event) : undefined;
+    }
+
+    const headers = new Headers(request.headers);
+
+    // If we send the host and referrer headers, the request will fail
+    headers.delete('host');
+    headers.delete('referrer');
+
+    // Add the authorization header
+    headers.set('Authorization', 'Bearer ' + axiomConfig.token);
+    headers.set('Content-Type', 'application/json');
+
+    let endpoint: string | null = null;
+
+    if (request.nextUrl.pathname.startsWith('/_axiom/web-vitals')) {
+      endpoint = webVitalsEndpoint;
+    } else if (request.nextUrl.pathname.startsWith('/_axiom/logs')) {
+      endpoint = logsEndpoint;
+    }
+
+    // Web vitals
+    if (endpoint) {
+      const response = fetch(endpoint, {
         body: request.body,
         method: 'POST',
-        headers,
-      }).catch(console.error)
-    );
-
-    // Return a 204 to the client
-    return new NextResponse(null, { status: 204 });
-  }
-}
-
-export const config = {
-  matcher: '/_axiom/:path*',
-};
+        headers: headers,
+      }).catch(console.error);
+      if (typeof event.waitUntil !== 'undefined') {
+        // Forward the request to the axiom ingest endpoint
+        event.waitUntil(response);
+        // Return a 204 to the client because we are not waiting for the response
+        return new NextResponse(null, { status: 204 });
+      } else {
+        const res = await response;
+        if (res) {
+          return new NextResponse(res.body, { status: res.status });
+        }
+        return new NextResponse(null, { status: 500 });
+      }
+    }
+  };

tests/middleware.test.ts

@@ -1,6 +1,6 @@
 import { NextFetchEvent, NextRequest, NextResponse } from 'next/server';
 import { expect, test, vi, vitest } from 'vitest';
-import { middleware } from '../src/middleware';
+import { axiomMiddleware } from '../src/middleware';
 
 vi.hoisted(() => {
   // set axiom env vars before importing logger
@@ -12,7 +12,7 @@ test('middleware processes web vital requests', async () => {
   const context = {
     waitUntil: vitest.fn(),
   } as unknown as NextFetchEvent;
-  const response = (await middleware(
+  const response = (await axiomMiddleware()(
     new NextRequest('http://localhost/_axiom/web-vitals', {
       method: 'POST',
     }),
@@ -27,7 +27,7 @@ test('middleware processes logs requests', async () => {
   const context = {
     waitUntil: vitest.fn(),
   } as unknown as NextFetchEvent;
-  const response = (await middleware(
+  const response = (await axiomMiddleware()(
     new NextRequest('http://localhost/_axiom/logs', {
       method: 'POST',
     }),
@@ -42,7 +42,7 @@ test('middleware ignores non-axiom requests', async () => {
   const context = {
     waitUntil: vitest.fn(),
   } as unknown as NextFetchEvent;
-  const response = await middleware(
+  const response = await axiomMiddleware()(
     new NextRequest('http://localhost/', {
       method: 'GET',
     }),
@@ -56,7 +56,7 @@ test('all logs from the browser are sent to the server', async () => {
   const context = {
     waitUntil: vitest.fn(),
   } as unknown as NextFetchEvent;
-  const response = (await middleware(
+  const response = (await axiomMiddleware()(
     new NextRequest('http://localhost/_axiom/logs', {
       method: 'POST',
       body: JSON.stringify({ logs: ['hello, world!'] }),