Skip to content

Commit

Permalink
fix: fix socket info in response (#555)
Browse files Browse the repository at this point in the history
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Introduced the `BaseAgent` class with enhanced asynchronous local
storage capabilities.
- Updated `HttpAgent` to extend `BaseAgent`, improving address
validation and error handling during dispatch operations.

- **Bug Fixes**
- Enhanced validation in the `dispatch` method of `HttpAgent` to ensure
hostname validity.

- **Tests**
- Added a new assertion in the fetch tests to verify the local address
of the response socket.

- **Chores**
- Removed the `fetchOpaqueInterceptor` function and associated imports
to streamline the codebase.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
  • Loading branch information
killagu authored Dec 4, 2024
1 parent 12c4e3c commit 629c7a3
Show file tree
Hide file tree
Showing 5 changed files with 41 additions and 29 deletions.
27 changes: 27 additions & 0 deletions src/BaseAgent.ts
Original file line number Diff line number Diff line change
@@ -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<FetchOpaque>;
}

export class BaseAgent extends Agent {
#opaqueLocalStorage?: AsyncLocalStorage<FetchOpaque>;

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);
}
}
12 changes: 0 additions & 12 deletions src/FetchOpaqueInterceptor.ts
Original file line number Diff line number Diff line change
@@ -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')

Expand All @@ -28,14 +27,3 @@ export interface FetchOpaque {
export interface OpaqueInterceptorOptions {
opaqueLocalStorage: AsyncLocalStorage<FetchOpaque>;
}

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);
};
};
}
5 changes: 3 additions & 2 deletions src/HttpAgent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -31,7 +32,7 @@ class IllegalAddressError extends Error {
}
}

export class HttpAgent extends Agent {
export class HttpAgent extends BaseAgent {
#checkAddress?: CheckAddressFunction;

constructor(options: HttpAgentOptions) {
Expand Down
25 changes: 10 additions & 15 deletions src/fetch.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import {
Agent,
getGlobalDispatcher,
Pool,
Dispatcher,
} from 'undici';
// eslint-disable-next-line @typescript-eslint/ban-ts-comment
// @ts-ignore
Expand Down Expand Up @@ -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
Expand All @@ -56,7 +58,7 @@ export type FetchResponseDiagnosticsMessage = {
};

export class FetchFactory {
static #dispatcher: Agent;
static #dispatcher: Dispatcher.ComposedDispatcher;
static #opaqueLocalStorage = new AsyncLocalStorage<FetchOpaque>();

static getDispatcher() {
Expand All @@ -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,
Expand All @@ -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();
Expand Down
1 change: 1 addition & 0 deletions test/fetch.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down

0 comments on commit 629c7a3

Please sign in to comment.