diff --git a/packages/hydrojudge/src/hosts/hydro.ts b/packages/hydrojudge/src/hosts/hydro.ts index ba1032c39..9bcad2147 100644 --- a/packages/hydrojudge/src/hosts/hydro.ts +++ b/packages/hydrojudge/src/hosts/hydro.ts @@ -3,7 +3,7 @@ import path from 'path'; import PQueue from 'p-queue'; import superagent from 'superagent'; import WebSocket from 'ws'; -import { fs } from '@hydrooj/utils'; +import { fs, pipeRequest } from '@hydrooj/utils'; import { LangConfig } from '@hydrooj/utils/lib/lang'; import * as sysinfo from '@hydrooj/utils/lib/sysinfo'; import type { JudgeResultBody } from 'hydrooj'; @@ -91,30 +91,15 @@ export default class Hydro implements Session { files: filenames, }); if (!res.body.links) throw new FormatError('problem not exist'); - const that = this; - // eslint-disable-next-line no-inner-declarations - async function download(name: string) { - if (name.includes('/')) await fs.ensureDir(path.join(filePath, name.split('/')[0])); - const w = fs.createWriteStream(path.join(filePath, name)); - that.get(res.body.links[name]).pipe(w); - await new Promise((resolve, reject) => { - const timeout = setTimeout(() => reject(new Error(`DownloadTimeout(${name}, 60s)`)), 60000); - w.on('error', (e) => { - clearTimeout(timeout); - reject(new Error(`DownloadFail(${name}): ${e.message}`)); - }); - w.on('finish', () => { - clearTimeout(timeout); - resolve(null); - }); - }); - } const tasks = []; const queue = new PQueue({ concurrency: 10 }); for (const name in res.body.links) { - tasks.push(queue.add(() => download(name))); + tasks.push(queue.add(async () => { + if (name.includes('/')) await fs.ensureDir(path.join(filePath, name.split('/')[0])); + const w = fs.createWriteStream(path.join(filePath, name)); + await pipeRequest(this.get(res.body.links[name]), w, 60000, name); + })); } - queue.start(); await Promise.all(tasks); await fs.writeFile(path.join(filePath, 'etags'), JSON.stringify(version)); } @@ -126,12 +111,7 @@ export default class Hydro implements Session { name = name.split('#')[0]; const res = await this.post('judge/files', { id: name }); const target = path.join(getConfig('tmp_dir'), name.replace(/\//g, '_')); - const w = fs.createWriteStream(target); - this.get(res.body.url).pipe(w); - await new Promise((resolve, reject) => { - w.on('finish', resolve); - w.on('error', (e) => reject(new Error(`DownloadFail(${name}): ${e.message}`))); - }); + await pipeRequest(this.get(res.body.url), fs.createWriteStream(target), 60000, name); return target; } diff --git a/packages/hydrojudge/src/sandbox/client.ts b/packages/hydrojudge/src/sandbox/client.ts index 1dc63c789..a5579a836 100644 --- a/packages/hydrojudge/src/sandbox/client.ts +++ b/packages/hydrojudge/src/sandbox/client.ts @@ -1,5 +1,6 @@ import fs from 'fs'; import superagent from 'superagent'; +import { pipeRequest } from '@hydrooj/utils'; import { getConfig } from '../config'; import { SandboxRequest, SandboxResult, SandboxVersion } from './interface'; @@ -12,11 +13,7 @@ const client = new Proxy({ getFile(fileId: string, dest?: string): Promise { if (dest) { const w = fs.createWriteStream(dest); - superagent.get(`${url}/file/${fileId}`).pipe(w); - return new Promise((resolve, reject) => { - w.on('finish', () => resolve(null)); - w.on('error', reject); - }); + return pipeRequest(superagent.get(`${url}/file/${fileId}`), w, 60000, fileId) as any; } return superagent.get(`${url}/file/${fileId}`).responseType('arraybuffer').then((res) => res.body); }, diff --git a/packages/hydrooj/src/handler/judge.ts b/packages/hydrooj/src/handler/judge.ts index 47bac0dc7..9f82a99a7 100644 --- a/packages/hydrooj/src/handler/judge.ts +++ b/packages/hydrooj/src/handler/judge.ts @@ -154,11 +154,13 @@ export async function end(body: Partial) { } export class JudgeFilesDownloadHandler extends Handler { + noCheckPermView = true; + notUsage = true; + async get() { this.response.body = 'ok'; } - noCheckPermView = true; @post('id', Types.String, true) @post('files', Types.Set, true) @post('pid', Types.UnsignedInt, true) @@ -207,6 +209,8 @@ export async function processJudgeFileCallback(rid: ObjectId, filename: string, } export class JudgeFileUpdateHandler extends Handler { + notUsage = true; + @post('rid', Types.ObjectId) @post('name', Types.Filename) async post(domainId: string, rid: ObjectId, filename: string) { diff --git a/packages/hydrooj/src/handler/misc.ts b/packages/hydrooj/src/handler/misc.ts index 937265c87..aff58df8d 100644 --- a/packages/hydrooj/src/handler/misc.ts +++ b/packages/hydrooj/src/handler/misc.ts @@ -104,6 +104,7 @@ export class FSDownloadHandler extends Handler { export class StorageHandler extends Handler { noCheckPermView = true; + notUsage = true; @param('target', Types.Name) @param('filename', Types.Filename, true) diff --git a/packages/hydrooj/src/service/server.ts b/packages/hydrooj/src/service/server.ts index 8ec9fc3d0..c821771dd 100644 --- a/packages/hydrooj/src/service/server.ts +++ b/packages/hydrooj/src/service/server.ts @@ -178,6 +178,7 @@ export class HandlerCommon { export class Handler extends HandlerCommon { loginMethods: any; noCheckPermView = false; + notUsage = false; allowCors = false; __param: Record[]>; @@ -207,7 +208,7 @@ export class Handler extends HandlerCommon { this.context.pendingError = new CsrfTokenError(); } } - if (!argv.options.benchmark) await this.limitRate('global', 5, 100); + if (!argv.options.benchmark && !this.notUsage) await this.limitRate('global', 5, 100); if (!this.noCheckPermView && !this.user.hasPriv(PRIV.PRIV_VIEW_ALL_DOMAIN)) this.checkPerm(PERM.PERM_VIEW); this.loginMethods = Object.keys(global.Hydro.module.oauth) .map((key) => ({ diff --git a/packages/utils/lib/utils.ts b/packages/utils/lib/utils.ts index 0f8895931..f4cecfe06 100644 --- a/packages/utils/lib/utils.ts +++ b/packages/utils/lib/utils.ts @@ -8,7 +8,7 @@ import fs from 'fs-extra'; import moment, { isMoment, Moment } from 'moment-timezone'; import { ObjectId } from 'mongodb'; import Logger from 'reggol'; -import type { Request } from 'superagent'; +import type * as superagent from 'superagent'; export * as yaml from 'js-yaml'; export * as fs from 'fs-extra'; @@ -82,10 +82,10 @@ export function isClass(obj: any, strict = false): obj is new (...args: any) => return false; } -function isSuperagentRequest(t: NodeJS.ReadableStream | Request): t is Request { +function isSuperagentRequest(t: NodeJS.ReadableStream | superagent.Request): t is superagent.Request { return 'req' in t; } -export function streamToBuffer(input: NodeJS.ReadableStream | Request, maxSize = 0): Promise { +export function streamToBuffer(input: NodeJS.ReadableStream | superagent.Request, maxSize = 0): Promise { let stream: NodeJS.ReadableStream; if (isSuperagentRequest(input)) { const s = new PassThrough(); @@ -323,4 +323,26 @@ export async function extractZip(zip: AdmZip, dest: string, overwrite = false, s } } +export async function pipeRequest(req: superagent.Request, w: fs.WriteStream, timeout?: number, name?: string) { + try { + await req.buffer(false).timeout({ + response: Math.min(10000, timeout), + deadline: timeout, + }).parse((resp, cb) => { + if (resp.statusCode !== 200) throw new Error(`${resp.statusCode}`); + else { + resp.pipe(w); + resp.on('end', () => { + cb(null, undefined); + }); + resp.on('error', (err) => { + cb(err, undefined); + }); + } + }); + } catch (e) { + throw new Error(`Download${e.errno === 'ETIMEDOUT' ? 'Timedout' : 'Error'}(${name ? `${name}, ` : ''}${e.message})`); + } +} + export * from '@hydrooj/utils/lib/common';