✌️

packages/backend/src/core/HttpRequestService.ts

@@ -1,18 +1,31 @@
 import * as http from 'node:http';
 import * as https from 'node:https';
 import CacheableLookup from 'cacheable-lookup';
-import fetch from 'node-fetch';
 import { HttpProxyAgent, HttpsProxyAgent } from 'hpagent';
 import { Inject, Injectable } from '@nestjs/common';
 import { DI } from '@/di-symbols.js';
 import type { Config } from '@/config.js';
 import { StatusError } from '@/misc/status-error.js';
 import { bindThis } from '@/decorators.js';
-import type { Response } from 'node-fetch';
-import type { URL } from 'node:url';
+import * as undici from 'undici';
+import { LookupFunction } from 'node:net';
+import { TransformStream } from 'node:stream/web';
+import * as dns from 'node:dns';
+
+export type IpChecker = (ip: string) => boolean;
 
 @Injectable()
 export class HttpRequestService {
+	/**
+	 * Get http non-proxy agent (undici)
+	 */
+	private agent: undici.Agent;
+
+	/**
+	 * Get http proxy or non-proxy agent (undici)
+	 */
+	public proxiedAgent: undici.ProxyAgent | undici.Agent;
+
 	/**
 	 * Get http non-proxy agent
 	 */
@@ -33,30 +46,57 @@ export class HttpRequestService {
 	 */
 	public httpsAgent: https.Agent;
 
+	public readonly dnsCache: CacheableLookup;
+	public readonly clientDefaults: undici.Agent.Options;
+
 	constructor(
 		@Inject(DI.config)
 		private config: Config,
 	) {
-		const cache = new CacheableLookup({
+		this.dnsCache = new CacheableLookup({
 			maxTtl: 3600,	// 1hours
 			errorTtl: 30,	// 30secs
 			lookup: false,	// nativeのdns.lookupにfallbackしない
 		});
-		
+
+		this.clientDefaults = {
+			keepAliveTimeout: 4 * 1000,
+			keepAliveMaxTimeout: 10 * 60 * 1000,
+			keepAliveTimeoutThreshold: 1 * 1000,
+			strictContentLength: true,
+			connect: {
+				maxCachedSessions: 100, // TLSセッションのキャッシュ数 https://github.com/nodejs/undici/blob/v5.14.0/lib/core/connect.js#L80
+				lookup: this.dnsCache.lookup as LookupFunction, // https://github.com/nodejs/undici/blob/v5.14.0/lib/core/connect.js#L98
+			},
+		}
+
+		this.agent = new undici.Agent({
+			...this.clientDefaults,
+		});
+
 		this.http = new http.Agent({
 			keepAlive: true,
 			keepAliveMsecs: 30 * 1000,
-			lookup: cache.lookup,
+			lookup: this.dnsCache.lookup,
 		} as http.AgentOptions);
 		
 		this.https = new https.Agent({
 			keepAlive: true,
 			keepAliveMsecs: 30 * 1000,
-			lookup: cache.lookup,
+			lookup: this.dnsCache.lookup,
 		} as https.AgentOptions);
-		
+
 		const maxSockets = Math.max(256, config.deliverJobConcurrency ?? 128);
-		
+
+		this.proxiedAgent = config.proxy
+			? new undici.ProxyAgent({
+				...this.clientDefaults,
+				connections: maxSockets,
+
+				uri: config.proxy,
+			})
+			: this.agent;
+
 		this.httpAgent = config.proxy
 			? new HttpProxyAgent({
 				keepAlive: true,
@@ -86,19 +126,49 @@ export class HttpRequestService {
 	 * @param bypassProxy Allways bypass proxy
 	 */
 	@bindThis
-	public getAgentByUrl(url: URL, bypassProxy = false): http.Agent | https.Agent {
+	public getAgentByUrl(url: URL, bypassProxy = false): undici.Agent | undici.ProxyAgent {
+		if (bypassProxy || (this.config.proxyBypassHosts || []).includes(url.hostname)) {
+			return this.agent;
+		} else {
+			return this.proxiedAgent;
+		}
+	}
+
+	/**
+	 * Get agent by URL
+	 * @param url URL
+	 * @param bypassProxy Allways bypass proxy
+	 */
+	@bindThis
+	public getHttpAgentByUrl(url: URL, bypassProxy = false): http.Agent | https.Agent {
 		if (bypassProxy || (this.config.proxyBypassHosts || []).includes(url.hostname)) {
 			return url.protocol === 'http:' ? this.http : this.https;
 		} else {
 			return url.protocol === 'http:' ? this.httpAgent : this.httpsAgent;
 		}
 	}
+	/**
+	 * check ip
+	 */
+	@bindThis
+	public checkIp(url: URL, fn: IpChecker): Promise<boolean> {
+		const lookup = this.dnsCache.lookup as LookupFunction || dns.lookup;
+
+		return new Promise((resolve, reject) => {
+			lookup(url.hostname, {}, (err, ip) => {
+				if (err) {
+					resolve(false);
+				} else {
+					resolve(fn(ip));
+				}
+			});
+		});
+	}
 
 	@bindThis
-	public async getJson(url: string, accept = 'application/json, */*', timeout = 10000, headers?: Record<string, string>): Promise<unknown> {
-		const res = await this.getResponse({
+	public async getJson<T extends unknown>(url: string, accept = 'application/json, */*', timeout = 10000, headers?: Record<string, string>): Promise<T> {
+		const res = await this.fetch({
 			url,
-			method: 'GET',
 			headers: Object.assign({
 				'User-Agent': this.config.userAgent,
 				Accept: accept,
@@ -112,9 +182,8 @@ export class HttpRequestService {
 
 	@bindThis
 	public async getHtml(url: string, accept = 'text/html, */*', timeout = 10000, headers?: Record<string, string>): Promise<string> {
-		const res = await this.getResponse({
+		const res = await this.fetch({
 			url,
-			method: 'GET',
 			headers: Object.assign({
 				'User-Agent': this.config.userAgent,
 				Accept: accept,
@@ -126,14 +195,35 @@ export class HttpRequestService {
 	}
 
 	@bindThis
-	public async getResponse(args: {
+	public async fetch(args: {
 		url: string,
-		method: string,
+		method?: string,
 		body?: string,
-		headers: Record<string, string>,
+		headers?: Record<string, string>,
 		timeout?: number,
 		size?: number,
-	}): Promise<Response> {
+		redirect?: RequestRedirect | undefined,
+		dispatcher?: undici.Dispatcher,
+		ipCheckers?: {
+			type: 'black' | 'white',
+			fn: IpChecker,
+		}[],
+		noOkError?: boolean,
+	}): Promise<undici.Response> {
+		const url = new URL(args.url);
+
+		if (args.ipCheckers) {
+			for (const check of args.ipCheckers) {
+				const result = await this.checkIp(url, check.fn);
+				if (
+					(check.type === 'black' && result === true) ||
+					(check.type === 'white' && result === false)
+				) {
+					throw new StatusError('IP is not allowed', 403, 'IP is not allowed');
+				}
+			}
+		}
+
 		const timeout = args.timeout ?? 10 * 1000;
 
 		const controller = new AbortController();
@@ -141,20 +231,57 @@ export class HttpRequestService {
 			controller.abort();
 		}, timeout * 6);
 
-		const res = await fetch(args.url, {
-			method: args.method,
-			headers: args.headers,
-			body: args.body,
-			timeout,
-			size: args.size ?? 10 * 1024 * 1024,
-			agent: (url) => this.getAgentByUrl(url),
-			signal: controller.signal,
-		});
+		const res = await Promise.race([
+			undici.fetch(args.url, {
+				method: args.method ?? 'GET',
+				headers: args.headers,
+				body: args.body,
+				redirect: args.redirect,
+				dispatcher: args.dispatcher ?? this.getAgentByUrl(url),
+				keepalive: true,
+				signal: controller.signal,
+			}),
+			new Promise<null>((res) => setTimeout(() => res(null)))
+		]);
 
-		if (!res.ok) {
+		if (res == null) {
+			throw new StatusError(`Request Timeout`, 408, 'Request Timeout');
+		}
+
+		if (!res.ok && !args.noOkError) {
 			throw new StatusError(`${res.status} ${res.statusText}`, res.status, res.statusText);
 		}
 
-		return res;
+		return ({
+			...res,
+			body: this.fetchLimiter(res, args.size),
+		});
+	}
+
+	/**
+	 * Fetch body limiter
+	 * @param res undici.Response
+	 * @param size number of Max size (Bytes) (default: 10MiB)
+	 * @returns ReadableStream<Uint8Array> (provided by node:stream/web)
+	 */
+	@bindThis
+	private fetchLimiter(res: undici.Response, size: number = 10 * 1024 * 1024) {
+		if (res.body == null) return null;
+
+		let total = 0;
+		return res.body.pipeThrough(new TransformStream({
+			start() {},
+			transform(chunk, controller) {
+				// TypeScirptグローバルの定義はUnit8ArrayだがundiciはReadableStream<any>を渡してくるので一応変換
+				const uint8 = new Uint8Array(chunk);
+				total += uint8.length;
+				if (total > size) {
+					controller.error(new StatusError(`Payload Too Large`, 413, 'Payload Too Large'));
+				} else {
+					controller.enqueue(uint8);
+				}
+			},
+			flush() {},
+		}));
 	}
 }