Skip to content

Commit

Permalink
fix(astro): Isolate request instrumentation middleware (#9709)
Browse files Browse the repository at this point in the history 
Add scope isolation to our sever-side request instrumentation,
similarly to the request handler in
Sveltekit.
  • Loading branch information
@Lms24
Lms24 authored Dec 4, 2023
1 parent abbc57f commit bfc63cb
Show file tree
Hide file tree
Showing 3 changed files with 167 additions and 109 deletions.
201 changes: 113 additions & 88 deletions packages/astro/src/server/middleware.ts
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,11 @@
import { captureException, configureScope, getCurrentHub, startSpan } from '@sentry/node';
import {
captureException,
configureScope,
continueTrace,
getCurrentHub,
runWithAsyncContext,
startSpan,
} from '@sentry/node';
import type { Hub, Span } from '@sentry/types';
import {
addNonEnumerableProperty,
Expand Down Expand Up @@ -56,107 +63,125 @@ type AstroLocalsWithSentry = Record<string, unknown> & {
__sentry_wrapped__?: boolean;
};

export const handleRequest: (options?: MiddlewareOptions) => MiddlewareResponseHandler = (
options = { trackClientIp: false, trackHeaders: false },
) => {
export const handleRequest: (options?: MiddlewareOptions) => MiddlewareResponseHandler = options => {
const handlerOptions = { trackClientIp: false, trackHeaders: false, ...options };

return async (ctx, next) => {
// Make sure we don't accidentally double wrap (e.g. user added middleware and integration auto added it)
const locals = ctx.locals as AstroLocalsWithSentry;
if (locals && locals.__sentry_wrapped__) {
return next();
// if there is an active span, we know that this handle call is nested and hence
// we don't create a new domain for it. If we created one, nested server calls would
// create new transactions instead of adding a child span to the currently active span.
if (getCurrentHub().getScope().getSpan()) {
return instrumentRequest(ctx, next, handlerOptions);
}
addNonEnumerableProperty(locals, '__sentry_wrapped__', true);
return runWithAsyncContext(() => {
return instrumentRequest(ctx, next, handlerOptions);
});
};
};

const method = ctx.request.method;
const headers = ctx.request.headers;
async function instrumentRequest(
ctx: Parameters<MiddlewareResponseHandler>[0],
next: Parameters<MiddlewareResponseHandler>[1],
options: MiddlewareOptions,
): Promise<Response> {
// Make sure we don't accidentally double wrap (e.g. user added middleware and integration auto added it)
const locals = ctx.locals as AstroLocalsWithSentry;
if (locals && locals.__sentry_wrapped__) {
return next();
}
addNonEnumerableProperty(locals, '__sentry_wrapped__', true);

const { dynamicSamplingContext, traceparentData, propagationContext } = tracingContextFromHeaders(
headers.get('sentry-trace') || undefined,
headers.get('baggage'),
);
const { method, headers } = ctx.request;

const allHeaders: Record<string, string> = {};
const traceCtx = continueTrace({
sentryTrace: headers.get('sentry-trace') || undefined,
baggage: headers.get('baggage'),
});

const allHeaders: Record<string, string> = {};

if (options.trackHeaders) {
headers.forEach((value, key) => {
allHeaders[key] = value;
});
}

if (options.trackClientIp) {
configureScope(scope => {
scope.setPropagationContext(propagationContext);

if (options.trackClientIp) {
scope.setUser({ ip_address: ctx.clientAddress });
}
scope.setUser({ ip_address: ctx.clientAddress });
});
}

try {
// storing res in a variable instead of directly returning is necessary to
// invoke the catch block if next() throws
const res = await startSpan(
{
name: `${method} ${interpolateRouteFromUrlAndParams(ctx.url.pathname, ctx.params)}`,
op: 'http.server',
origin: 'auto.http.astro',
status: 'ok',
...traceparentData,
metadata: {
source: 'route',
dynamicSamplingContext: traceparentData && !dynamicSamplingContext ? {} : dynamicSamplingContext,
},
data: {
method,
url: stripUrlQueryAndFragment(ctx.url.href),
...(ctx.url.search && { 'http.query': ctx.url.search }),
...(ctx.url.hash && { 'http.fragment': ctx.url.hash }),
...(options.trackHeaders && { headers: allHeaders }),
},
try {
// storing res in a variable instead of directly returning is necessary to
// invoke the catch block if next() throws
const res = await startSpan(
{
...traceCtx,
name: `${method} ${interpolateRouteFromUrlAndParams(ctx.url.pathname, ctx.params)}`,
op: 'http.server',
origin: 'auto.http.astro',
status: 'ok',
metadata: {
...traceCtx?.metadata,
source: 'route',
},
async span => {
const originalResponse = await next();

if (span && originalResponse.status) {
span.setHttpStatus(originalResponse.status);
}

const hub = getCurrentHub();
const client = hub.getClient();
const contentType = originalResponse.headers.get('content-type');

const isPageloadRequest = contentType && contentType.startsWith('text/html');
if (!isPageloadRequest || !client) {
return originalResponse;
}

// Type case necessary b/c the body's ReadableStream type doesn't include
// the async iterator that is actually available in Node
// We later on use the async iterator to read the body chunks
// see https://github.com/microsoft/TypeScript/issues/39051
const originalBody = originalResponse.body as NodeJS.ReadableStream | null;
if (!originalBody) {
return originalResponse;
}

const newResponseStream = new ReadableStream({
start: async controller => {
for await (const chunk of originalBody) {
const html = typeof chunk === 'string' ? chunk : new TextDecoder().decode(chunk);
const modifiedHtml = addMetaTagToHead(html, hub, span);
controller.enqueue(new TextEncoder().encode(modifiedHtml));
}
controller.close();
},
});

return new Response(newResponseStream, originalResponse);
data: {
method,
url: stripUrlQueryAndFragment(ctx.url.href),
...(ctx.url.search && { 'http.query': ctx.url.search }),
...(ctx.url.hash && { 'http.fragment': ctx.url.hash }),
...(options.trackHeaders && { headers: allHeaders }),
},
);
return res;
} catch (e) {
sendErrorToSentry(e);
throw e;
}
// TODO: flush if serveless (first extract function)
};
};
},
async span => {
const originalResponse = await next();

if (span && originalResponse.status) {
span.setHttpStatus(originalResponse.status);
}

const hub = getCurrentHub();
const client = hub.getClient();
const contentType = originalResponse.headers.get('content-type');

const isPageloadRequest = contentType && contentType.startsWith('text/html');
if (!isPageloadRequest || !client) {
return originalResponse;
}

// Type case necessary b/c the body's ReadableStream type doesn't include
// the async iterator that is actually available in Node
// We later on use the async iterator to read the body chunks
// see https://github.com/microsoft/TypeScript/issues/39051
const originalBody = originalResponse.body as NodeJS.ReadableStream | null;
if (!originalBody) {
return originalResponse;
}

const decoder = new TextDecoder();

const newResponseStream = new ReadableStream({
start: async controller => {
for await (const chunk of originalBody) {
const html = typeof chunk === 'string' ? chunk : decoder.decode(chunk);
const modifiedHtml = addMetaTagToHead(html, hub, span);
controller.enqueue(new TextEncoder().encode(modifiedHtml));
}
controller.close();
},
});

return new Response(newResponseStream, originalResponse);
},
);
return res;
} catch (e) {
sendErrorToSentry(e);
throw e;
}
// TODO: flush if serverless (first extract function)
}

/**
* This function optimistically assumes that the HTML coming in chunks will not be split
Expand Down
74 changes: 54 additions & 20 deletions packages/astro/test/server/middleware.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,18 @@ vi.mock('../../src/server/meta', () => ({

describe('sentryMiddleware', () => {
const startSpanSpy = vi.spyOn(SentryNode, 'startSpan');

const getSpanMock = vi.fn(() => {});
// @ts-expect-error only returning a partial hub here
vi.spyOn(SentryNode, 'getCurrentHub').mockImplementation(() => {
return {
getScope: () => ({
getSpan: getSpanMock,
}),
getClient: () => ({}),
};
});

const nextResult = Promise.resolve(new Response(null, { status: 200, headers: new Headers() }));

afterEach(() => {
Expand Down Expand Up @@ -86,10 +98,6 @@ describe('sentryMiddleware', () => {
});

it('attaches tracing headers', async () => {
const scope = { setUser: vi.fn(), setPropagationContext: vi.fn() };
// @ts-expect-error, only passing a partial Scope object
const configureScopeSpy = vi.spyOn(SentryNode, 'configureScope').mockImplementation(cb => cb(scope));

const middleware = handleRequest();
const ctx = {
request: {
Expand All @@ -108,17 +116,6 @@ describe('sentryMiddleware', () => {
// @ts-expect-error, a partial ctx object is fine here
await middleware(ctx, next);

expect(configureScopeSpy).toHaveBeenCalledTimes(1);
expect(scope.setPropagationContext).toHaveBeenCalledWith({
dsc: {
release: '1.0.0',
},
parentSpanId: '1234567890123456',
sampled: true,
spanId: expect.any(String),
traceId: '12345678901234567890123456789012',
});

expect(startSpanSpy).toHaveBeenCalledWith(
expect.objectContaining({
metadata: {
Expand Down Expand Up @@ -174,11 +171,6 @@ describe('sentryMiddleware', () => {
});

it('injects tracing <meta> tags into the HTML of a pageload response', async () => {
vi.spyOn(SentryNode, 'getCurrentHub').mockImplementation(() => ({
// @ts-expect-error this is fine
getClient: () => ({}),
}));

const middleware = handleRequest();

const ctx = {
Expand Down Expand Up @@ -261,6 +253,48 @@ describe('sentryMiddleware', () => {

expect(html).toBe(originalHtml);
});

describe('async context isolation', () => {
const runWithAsyncContextSpy = vi.spyOn(SentryNode, 'runWithAsyncContext');
afterEach(() => {
vi.clearAllMocks();
runWithAsyncContextSpy.mockRestore();
});

it('starts a new async context if no span is active', async () => {
getSpanMock.mockReturnValueOnce(undefined);
const handler = handleRequest();
const ctx = {};
const next = vi.fn();

try {
// @ts-expect-error, a partial ctx object is fine here
await handler(ctx, next);
} catch {
// this is fine, it's not required to pass in this test
}

expect(runWithAsyncContextSpy).toHaveBeenCalledTimes(1);
});

it("doesn't start a new async context if a span is active", async () => {
// @ts-expect-error, a empty span is fine here
getSpanMock.mockReturnValueOnce({});

const handler = handleRequest();
const ctx = {};
const next = vi.fn();

try {
// @ts-expect-error, a partial ctx object is fine here
await handler(ctx, next);
} catch {
// this is fine, it's not required to pass in this test
}

expect(runWithAsyncContextSpy).not.toHaveBeenCalled();
});
});
});

describe('interpolateRouteFromUrlAndParams', () => {
Expand Down
1 change: 0 additions & 1 deletion packages/core/src/tracing/trace.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -226,7 +226,6 @@ export function continueTrace<V>(
* These values can be obtained from incoming request headers,
* or in the browser from `<meta name="sentry-trace">` and `<meta name="baggage">` HTML tags.
*
* It also takes an optional `request` option, which if provided will also be added to the scope & transaction metadata.
* The callback receives a transactionContext that may be used for `startTransaction` or `startSpan`.
*/
export function continueTrace<V>(
Expand Down

0 comments on commit bfc63cb

Please sign in to comment.