Merge branch 'fetch' into serve-stream

packages/backend/src/core/DownloadService.ts

@@ -8,11 +8,12 @@ import got, * as Got from 'got';
 import chalk from 'chalk';
 import { DI } from '@/di-symbols.js';
 import type { Config } from '@/config.js';
-import { HttpRequestService } from '@/core/HttpRequestService.js';
+import { HttpRequestService, UndiciFetcher } from '@/core/HttpRequestService.js';
 import { createTemp } from '@/misc/create-temp.js';
 import { StatusError } from '@/misc/status-error.js';
 import { LoggerService } from '@/core/LoggerService.js';
 import type Logger from '@/logger.js';
+import { buildConnector } from 'undici';
 
 const pipeline = util.promisify(stream.pipeline);
 import { bindThis } from '@/decorators.js';
@@ -20,6 +21,7 @@ import { bindThis } from '@/decorators.js';
 @Injectable()
 export class DownloadService {
 	private logger: Logger;
+	private undiciFetcher: UndiciFetcher;
 
 	constructor(
 		@Inject(DI.config)
@@ -29,69 +31,58 @@ export class DownloadService {
 		private loggerService: LoggerService,
 	) {
 		this.logger = this.loggerService.getLogger('download');
+
+		this.undiciFetcher = new UndiciFetcher(this.httpRequestService.getStandardUndiciFetcherOption(
+			{
+				connect: process.env.NODE_ENV === 'development' ?
+					this.httpRequestService.clientDefaults.connect
+					:
+					this.httpRequestService.getConnectorWithIpCheck(
+						buildConnector({
+							...this.httpRequestService.clientDefaults.connect,
+						}),
+						(ip) => !this.isPrivateIp(ip)
+					),
+				bodyTimeout: 30 * 1000,
+			},
+			{
+				connect: this.httpRequestService.clientDefaults.connect,
+			}
+		), this.logger);
 	}
 
 	@bindThis
-	public gotUrl(url: string): Got.Request {
+	public fetchUrl(url: string): any {
 		this.logger.info(`Downloading ${chalk.cyan(url)} ...`);
 	
 		const timeout = 30 * 1000;
 		const operationTimeout = 60 * 1000;
 		const maxSize = this.config.maxFileSize ?? 262144000;
-	
-		const req = got.stream(url, {
-			headers: {
-				'User-Agent': this.config.userAgent,
-			},
-			timeout: {
-				lookup: timeout,
-				connect: timeout,
-				secureConnect: timeout,
-				socket: timeout,	// read timeout
-				response: timeout,
-				send: timeout,
-				request: operationTimeout,	// whole operation timeout
-			},
-			agent: {
-				http: this.httpRequestService.httpAgent,
-				https: this.httpRequestService.httpsAgent,
-			},
-			http2: false,	// default
-			retry: {
-				limit: 0,
-			},
-		}).on('response', (res: Got.Response) => {
-			if ((process.env.NODE_ENV === 'production' || process.env.NODE_ENV === 'test') && !this.config.proxy && res.ip) {
-				if (this.isPrivateIp(res.ip)) {
-					this.logger.warn(`Blocked address: ${res.ip}`);
-					req.destroy();
-				}
-			}
-	
-			const contentLength = res.headers['content-length'];
-			if (contentLength != null) {
-				const size = Number(contentLength);
-				if (size > maxSize) {
-					this.logger.warn(`maxSize exceeded (${size} > ${maxSize}) on response`);
-					req.destroy();
-				}
-			}
-		}).on('downloadProgress', (progress: Got.Progress) => {
-			if (progress.transferred > maxSize) {
-				this.logger.warn(`maxSize exceeded (${progress.transferred} > ${maxSize}) on downloadProgress`);
-				req.destroy();
-			}
-		});
 
-		return req;
+		const response = await this.undiciFetcher.fetch(
+			url,
+			{
+				method: 'GET',
+				headers: {
+					'User-Agent': this.config.userAgent,
+				},
+			}
+		);
+
+		if (response.body === null) {
+			throw new StatusError('No body', 400, 'No body');
+		}
+	
+		this.logger.succ(`Download finished: ${chalk.cyan(url)}`);
+
+		return response;
 	}
 
 	@bindThis
-	public async pipeRequestToFile(req: Got.Request, path: string): Promise<void> {
-		const copied = req.pipe(new stream.PassThrough());
+	public async pipeRequestToFile(response: any, path: string): Promise<void> {
 		try {
 			this.logger.info(`Saving File to ${chalk.cyanBright(path)} from downloading ...`);
-			await pipeline(copied, fs.createWriteStream(path));
+			await pipeline(stream.Readable.fromWeb(response.body), fs.createWriteStream(path));
 		} catch (e) {
 			if (e instanceof Got.HTTPError) {
 				throw new StatusError(`${e.response.statusCode} ${e.response.statusMessage}`, e.response.statusCode, e.response.statusMessage);
@@ -125,7 +116,7 @@ export class DownloadService {
 			cleanup();
 		}
 	}
-	
+
 	@bindThis
 	private isPrivateIp(ip: string): boolean {
 		for (const net of this.config.allowedPrivateNetworks ?? []) {
@@ -135,6 +126,6 @@ export class DownloadService {
 			}
 		}
 
-		return PrivateIp(ip);
+		return PrivateIp(ip) ?? false;
 	}
 }