diff --git a/src/BaseAgent.ts b/src/BaseAgent.ts new file mode 100644 index 0000000..9003ea5 --- /dev/null +++ b/src/BaseAgent.ts @@ -0,0 +1,27 @@ +import { + Agent, + Dispatcher, +} from 'undici'; +import { AsyncLocalStorage } from 'node:async_hooks'; +import { FetchOpaque } from './FetchOpaqueInterceptor.js'; + +export interface BaseAgentOptions extends Agent.Options { + opaqueLocalStorage?: AsyncLocalStorage; +} + +export class BaseAgent extends Agent { + #opaqueLocalStorage?: AsyncLocalStorage; + + constructor(options: BaseAgentOptions) { + super(options); + this.#opaqueLocalStorage = options.opaqueLocalStorage; + } + + dispatch(options: Agent.DispatchOptions, handler: Dispatcher.DispatchHandler): boolean { + const opaque = this.#opaqueLocalStorage?.getStore(); + if (opaque) { + (handler as any).opaque = opaque; + } + return super.dispatch(options, handler); + } +} diff --git a/src/FetchOpaqueInterceptor.ts b/src/FetchOpaqueInterceptor.ts index dd741df..4d34e68 100644 --- a/src/FetchOpaqueInterceptor.ts +++ b/src/FetchOpaqueInterceptor.ts @@ -1,7 +1,6 @@ // const { AsyncLocalStorage } = require('node:async_hooks'); import { AsyncLocalStorage } from 'node:async_hooks'; import symbols from './symbols.js'; -import { Dispatcher } from 'undici'; // const RedirectHandler = require('../handler/redirect-handler') @@ -28,14 +27,3 @@ export interface FetchOpaque { export interface OpaqueInterceptorOptions { opaqueLocalStorage: AsyncLocalStorage; } - -export function fetchOpaqueInterceptor(opts: OpaqueInterceptorOptions) { - const opaqueLocalStorage = opts?.opaqueLocalStorage; - return (dispatch: Dispatcher['dispatch']): Dispatcher['dispatch'] => { - return function redirectInterceptor(opts: Dispatcher.DispatchOptions, handler: Dispatcher.DispatchHandler) { - const opaque = opaqueLocalStorage?.getStore(); - (handler as any).opaque = opaque; - return dispatch(opts, handler); - }; - }; -} diff --git a/src/HttpAgent.ts b/src/HttpAgent.ts index cb1200b..69befde 100644 --- a/src/HttpAgent.ts +++ b/src/HttpAgent.ts @@ -5,10 +5,11 @@ import { Dispatcher, buildConnector, } from 'undici'; +import { BaseAgent, BaseAgentOptions } from './BaseAgent.js'; export type CheckAddressFunction = (ip: string, family: number | string, hostname: string) => boolean; -export interface HttpAgentOptions extends Agent.Options { +export interface HttpAgentOptions extends BaseAgentOptions { lookup?: LookupFunction; checkAddress?: CheckAddressFunction; connect?: buildConnector.BuildOptions, @@ -31,7 +32,7 @@ class IllegalAddressError extends Error { } } -export class HttpAgent extends Agent { +export class HttpAgent extends BaseAgent { #checkAddress?: CheckAddressFunction; constructor(options: HttpAgentOptions) { diff --git a/src/fetch.ts b/src/fetch.ts index f47fd98..9df8655 100644 --- a/src/fetch.ts +++ b/src/fetch.ts @@ -8,6 +8,7 @@ import { Agent, getGlobalDispatcher, Pool, + Dispatcher, } from 'undici'; // eslint-disable-next-line @typescript-eslint/ban-ts-comment // @ts-ignore @@ -35,9 +36,10 @@ import { HttpMethod, RequestMeta, } from './Request.js'; -import { FetchOpaque, fetchOpaqueInterceptor } from './FetchOpaqueInterceptor.js'; +import { FetchOpaque } from './FetchOpaqueInterceptor.js'; import { RawResponseWithMeta, SocketInfo } from './Response.js'; import { IncomingHttpHeaders } from './IncomingHttpHeaders.js'; +import { BaseAgent, BaseAgentOptions } from './BaseAgent.js'; export interface UrllibRequestInit extends RequestInit { // default is true @@ -56,7 +58,7 @@ export type FetchResponseDiagnosticsMessage = { }; export class FetchFactory { - static #dispatcher: Agent; + static #dispatcher: Dispatcher.ComposedDispatcher; static #opaqueLocalStorage = new AsyncLocalStorage(); static getDispatcher() { @@ -68,17 +70,10 @@ export class FetchFactory { } static setClientOptions(clientOptions: ClientOptions) { - let dispatcherOption: Agent.Options = { - interceptors: { - Agent: [ - fetchOpaqueInterceptor({ - opaqueLocalStorage: FetchFactory.#opaqueLocalStorage, - }), - ], - Client: [], - }, + let dispatcherOption: BaseAgentOptions = { + opaqueLocalStorage: FetchFactory.#opaqueLocalStorage, }; - let dispatcherClazz: new (options: Agent.Options) => Agent = Agent; + let dispatcherClazz: new (options: BaseAgentOptions) => BaseAgent = BaseAgent; if (clientOptions?.lookup || clientOptions?.checkAddress) { dispatcherOption = { ...dispatcherOption, @@ -87,21 +82,21 @@ export class FetchFactory { connect: clientOptions.connect, allowH2: clientOptions.allowH2, } as HttpAgentOptions; - dispatcherClazz = HttpAgent as unknown as new (options: Agent.Options) => Agent; + dispatcherClazz = HttpAgent as unknown as new (options: BaseAgentOptions) => BaseAgent; } else if (clientOptions?.connect) { dispatcherOption = { ...dispatcherOption, connect: clientOptions.connect, allowH2: clientOptions.allowH2, } as HttpAgentOptions; - dispatcherClazz = Agent; + dispatcherClazz = BaseAgent; } else if (clientOptions?.allowH2) { // Support HTTP2 dispatcherOption = { ...dispatcherOption, allowH2: clientOptions.allowH2, } as HttpAgentOptions; - dispatcherClazz = Agent; + dispatcherClazz = BaseAgent; } FetchFactory.#dispatcher = new dispatcherClazz(dispatcherOption); initDiagnosticsChannel(); diff --git a/test/fetch.test.ts b/test/fetch.test.ts index 7474141..9429d95 100644 --- a/test/fetch.test.ts +++ b/test/fetch.test.ts @@ -46,6 +46,7 @@ describe('fetch.test.ts', () => { assert(requestDiagnosticsMessage!.request); assert(responseDiagnosticsMessage!.request); assert(responseDiagnosticsMessage!.response); + assert([ '127.0.0.1', '::1' ].includes(responseDiagnosticsMessage!.response.socket.localAddress)); assert(fetchDiagnosticsMessage!.fetch); assert(fetchResponseDiagnosticsMessage!.fetch);