| 1 | import { AsyncLocalStorage } from 'node:async_hooks' |
| 2 | import { CallHandler, ExecutionContext, Injectable, NestInterceptor } from '@nestjs/common' |
| 3 | import { SSE_METADATA } from '@nestjs/common/constants' |
| 4 | import { Observable } from 'rxjs' |
| 5 | |
| 6 | export type PropagationHeaders = Record<string, string | string[] | undefined> |
| 7 | |
| 8 | interface Store { |
| 9 | headers: PropagationHeaders |
| 10 | } |
| 11 | |
| 12 | export const propagationContext = new AsyncLocalStorage<Store>() |
| 13 | |
| 14 | export function getRequestIdFromHeaders(headers: PropagationHeaders | undefined): string | undefined { |
| 15 | const requestId = headers?.['x-request-id'] |
| 16 | return Array.isArray(requestId) ? requestId[0] : requestId |
| 17 | } |
| 18 | |
| 19 | export function getCurrentRequestId(): string | undefined { |
| 20 | return getRequestIdFromHeaders(propagationContext.getStore()?.headers) |
| 21 | } |
| 22 | |
| 23 | export const COMMON_PROPAGATION_HEADERS = [ |
| 24 | 'x-request-id', |
| 25 | 'x-b3-traceid', |
| 26 | 'x-b3-spanid', |
| 27 | 'x-b3-parentspanid', |
| 28 | 'x-b3-sampled', |
| 29 | 'x-b3-flags', |
| 30 | 'x-ot-span-context', |
| 31 | 'grpc-trace-bin', |
| 32 | 'traceparent', |
| 33 | 'x-cloud-trace-context', |
| 34 | 'x-amzn-trace-id', |
| 35 | 'x-client-type', |
| 36 | 'x-client-version', |
| 37 | 'accept-language', |
| 38 | ] |
| 39 | |
| 40 | @Injectable() |
| 41 | export class PropagationInterceptor implements NestInterceptor { |
| 42 | public intercept(context: ExecutionContext, next: CallHandler): Observable<unknown> { |
| 43 | return propagationContext.run({ headers: this.getHeaders(context) }, () => next.handle()) |
| 44 | } |
| 45 | |
| 46 | private getHttpHeaders(context: ExecutionContext) { |
| 47 | const request = context.switchToHttp().getRequest() |
| 48 | const response = context.switchToHttp().getResponse() |
| 49 | |
| 50 | if (!Reflect.getMetadata(SSE_METADATA, context.getHandler())) |
| 51 | response.header('x-request-id', request.headers['x-request-id']) |
| 52 | |
| 53 | return request.headers |
| 54 | } |
| 55 | |
| 56 | private getWsHeaders(context: ExecutionContext) { |
| 57 | const host = context.switchToWs() |
| 58 | const socket = host.getClient() |
| 59 | const data = host.getData() |
| 60 | const headers = socket.handshake.headers |
| 61 | if (data?.headers) { |
| 62 | return { |
| 63 | ...headers, |
| 64 | ...data.headers, |
| 65 | } |
| 66 | } |
| 67 | return headers |
| 68 | } |
| 69 | |
| 70 | private getRpcHeaders(context: ExecutionContext) { |
| 71 | const rpcContext = context.switchToRpc().getContext() |
| 72 | const metadata = rpcContext.getMap ? rpcContext.getMap() : {} |
| 73 | return metadata |
| 74 | } |
| 75 | |
| 76 | private getHeaders(context: ExecutionContext) { |
| 77 | const type = context.getType() |
| 78 | if (type === 'http') |
| 79 | return this.getHttpHeaders(context) |
| 80 | |
| 81 | if (type === 'ws') |
| 82 | return this.getWsHeaders(context) |
| 83 | |
| 84 | if (type === 'rpc') |
| 85 | return this.getRpcHeaders(context) |
| 86 | |
| 87 | return {} |
| 88 | } |
| 89 | } |
| 90 |