judge: allow unlimited access & fix file corruption when errored (#726)

Co-authored-by: undefined <i@undefined.moe>

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);
+        await pipeRequest(this.get(res.body.url), fs.createWriteStream(target), 60000, name);
-        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}`)));
-        });
         return target;
     }
 

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<Buffer> {
         if (dest) {
             const w = fs.createWriteStream(dest);
-            superagent.get(`${url}/file/${fileId}`).pipe(w);
+            return pipeRequest(superagent.get(`${url}/file/${fileId}`), w, 60000, fileId) as any;
-            return new Promise((resolve, reject) => {
-                w.on('finish', () => resolve(null));
-                w.on('error', reject);
-            });
         }
         return superagent.get(`${url}/file/${fileId}`).responseType('arraybuffer').then((res) => res.body);
     },

packages/hydrooj/src/handler/judge.ts

@@ -154,11 +154,13 @@ export async function end(body: Partial<JudgeResultBody>) {
 }
 
 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) {

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)

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<string, decorators.ParamOption<any>[]>;
 
@@ -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) => ({

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<Buffer> {
+export function streamToBuffer(input: NodeJS.ReadableStream | superagent.Request, maxSize = 0): Promise<Buffer> {
     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';