Tune download (#2)

* s2-2

* allowedPrivateNetworks

* style

* Proxyの間にあると誤解しそうなのでconfigの記述順を変更

* Fix error handler

1 files changed, 49 insertions, 18 deletions

diff --git a/src/misc/download-url.ts b/src/misc/download-url.ts
index 43e061c715..463fb555bb 100644
--- a/src/misc/download-url.ts
+++ b/src/misc/download-url.ts
@@ -1,13 +1,13 @@
 import * as fs from 'fs';
 import * as stream from 'stream';
 import * as util from 'util';
-import { URL } from 'url';
-import fetch from 'node-fetch';
-import { getAgentByUrl } from './fetch';
-import { AbortController } from 'abort-controller';
+import got, * as Got from 'got';
+import { httpAgent, httpsAgent } from './fetch';
 import config from '@/config/index';
 import * as chalk from 'chalk';
 import Logger from '@/services/logger';
+import * as IPCIDR from 'ip-cidr';
+const PrivateIp = require('private-ip');
 
 const pipeline = util.promisify(stream.pipeline);
 
@@ -15,26 +15,57 @@ export async function downloadUrl(url: string, path: string) {
 	const logger = new Logger('download');
 
 	logger.info(`Downloading ${chalk.cyan(url)} ...`);
-	const controller = new AbortController();
-	setTimeout(() => {
-		controller.abort();
-	}, 60 * 1000);
 
-	const response = await fetch(new URL(url).href, {
+	const timeout = 30 * 1000;
+	const operationTimeout = 60 * 1000;
+
+	const req = got.stream(url, {
 		headers: {
 			'User-Agent': config.userAgent
 		},
-		timeout: 10 * 1000,
-		signal: controller.signal,
-		agent: getAgentByUrl,
+		timeout: {
+			lookup: timeout,
+			connect: timeout,
+			secureConnect: timeout,
+			socket: timeout,	// read timeout
+			response: timeout,
+			send: timeout,
+			request: operationTimeout,	// whole operation timeout
+		},
+		agent: {
+			http: httpAgent,
+			https: httpsAgent,
+		},
+		retry: 0,
+	}).on('response', (res: Got.Response) => {
+		if ((process.env.NODE_ENV === 'production' || process.env.NODE_ENV === 'test') && !config.proxy && res.ip) {
+			if (isPrivateIp(res.ip)) {
+				logger.warn(`Blocked address: ${res.ip}`);
+				req.destroy();
+			}
+		}
+	}).on('error', (e: any) => {
+		if (e.name === 'HTTPError') {
+			const statusCode = e.response?.statusCode;
+			const statusMessage = e.response?.statusMessage;
+			e.name = `StatusError`;
+			e.statusCode = statusCode;
+			e.message = `${statusCode} ${statusMessage}`;
+		}
 	});
 
-	if (!response.ok) {
-		logger.error(`Got ${response.status} (${url})`);
-		throw response.status;
-	}
-
-	await pipeline(response.body, fs.createWriteStream(path));
+	await pipeline(req, fs.createWriteStream(path));
 
 	logger.succ(`Download finished: ${chalk.cyan(url)}`);
 }
+
+function isPrivateIp(ip: string) {
+	for (const net of config.allowedPrivateNetworks || []) {
+		const cidr = new IPCIDR(net);
+		if (cidr.contains(ip)) {
+			return false;
+		}
+	}
+
+	return PrivateIp(ip);
+}