基于时间戳的文件下载secret

hydro/handler/fs.js

@@ -0,0 +1,18 @@
+const file = require('../model/file');
+const { Route, Handler } = require('../service/server');
+
+class FileDownloadHandler extends Handler {
+    async get({ id, secret, name }) {
+        this.response.attachment(name || id);
+        this.response.body = await file.get(id, secret);
+    }
+}
+
+async function apply() {
+    Route('/fs/:id/:secret', module.exports.FileDownloadHandler);
+}
+
+global.Hydro.handler.file = module.exports = {
+    FileDownloadHandler,
+    apply,
+};

hydro/handler/problem.js

@@ -1,6 +1,7 @@
 const fs = require('fs');
 const paginate = require('../lib/paginate');
 const validator = require('../lib/validator');
+const file = require('../model/file');
 const problem = require('../model/problem');
 const record = require('../model/record');
 const user = require('../model/user');
@@ -10,7 +11,6 @@ const bus = require('../service/bus');
 const {
     Route, Connection, Handler, ConnectionHandler,
 } = require('../service/server');
-const gridfs = require('../service/gridfs');
 const {
     NoProblemError, ProblemDataNotFoundError, BadRequestError,
     SolutionNotFoundError,
@@ -269,8 +269,8 @@ class ProblemDataUploadHandler extends ProblemManageHandler {
 
     async get() {
         if (this.pdoc.data && typeof this.pdoc.data === 'object') {
-            const files = await gridfs.find({ _id: this.pdoc.data }).toArray();
+            const f = await file.getMeta(this.pdoc.data);
-            this.md5 = files[0].md5;
+            this.md5 = f.md5;
         }
         this.response.body.md5 = this.md5;
     }
@@ -280,8 +280,8 @@ class ProblemDataUploadHandler extends ProblemManageHandler {
         const r = fs.createReadStream(this.request.files.file.path);
         await problem.setTestdata(this.pdoc._id, r);
         if (this.pdoc.data && typeof this.pdoc.data === 'object') {
-            const files = await gridfs.find({ _id: this.pdoc.data }).toArray();
+            const f = await file.getMeta(this.pdoc.data);
-            this.md5 = files[0].md5;
+            this.md5 = f.md5;
         }
         this.response.body.md5 = this.md5;
     }
@@ -291,9 +291,8 @@ class ProblemDataDownloadHandler extends ProblemDetailHandler {
     async get({ pid }) {
         if (this.uid !== this.pdoc.owner) this.checkPerm([PERM_READ_PROBLEM_DATA, PERM_JUDGE]);
         if (!this.pdoc.data) throw new ProblemDataNotFoundError(pid);
-        else if (typeof this.pdoc.data === 'string') [, this.ctx.setRedirect] = this.pdoc.data.split('from:');
+        else if (typeof this.pdoc.data === 'string') [, this.response.redirect] = this.pdoc.data.split('from:');
-        this.response.attachment(`${this.pdoc.title}.zip`);
+        this.response.redirect = file.url(this.pdoc.data, this.pdoc.title);
-        this.response.body = gridfs.openDownloadStream(this.pdoc.data);
     }
 }
 
@@ -306,12 +305,14 @@ class ProblemSolutionHandler extends ProblemDetailHandler {
             page,
             await system.get('SOLUTION_PER_PAGE'),
         );
-        const uids = [this.pdoc.owner]; const
+        const uids = [this.pdoc.owner];
-            docids = [];
+        const docids = [];
         for (const psdoc of psdocs) {
             docids.push(psdoc._id);
             uids.push(psdoc.owner);
-            if (psdoc.reply.length) { for (const psrdoc of psdoc.reply) uids.push(psrdoc.owner); }
+            if (psdoc.reply.length) {
+                for (const psrdoc of psdoc.reply) uids.push(psrdoc.owner);
+            }
         }
         const udict = await user.getList(uids);
         const pssdict = solution.getListStatus(docids, this.uid);

hydro/loader.js

@@ -240,7 +240,7 @@ async function load() {
     const builtinModel = [
         'blacklist', 'builtin', 'contest', 'discussion', 'message',
         'opcount', 'problem', 'record', 'setting', 'solution',
-        'token', 'training', 'user',
+        'token', 'training', 'user', 'file',
     ];
     for (const i of builtinModel) {
         const m = require(`./model/${i}`);
@@ -255,6 +255,7 @@ async function load() {
     const builtinHandler = [
         'home', 'problem', 'record', 'judge', 'user',
         'contest', 'training', 'discussion', 'manage', 'import',
+        'fs',
     ];
     for (const i of builtinHandler) require(`./handler/${i}`);
     await model();

hydro/model/blacklist.js

@@ -1,4 +1,4 @@
-const db = require('../service/db.js');
+const db = require('../service/db');
 
 const coll = db.collection('blacklist');
 

hydro/model/file.js

@@ -0,0 +1,69 @@
+const { ForbiddenError } = require('../error');
+const db = require('../service/db');
+const fs = require('../service/gridfs');
+const hash = require('../lib/hash.hydro');
+
+const coll = db.collection('file');
+
+async function add(stream, filename, meta = {}) {
+    const file = await coll.insertOne(meta);
+    const w = fs.openUploadStreamWithId(file.insertedId, filename);
+    await new Promise((resolve, reject) => {
+        stream.on('end', resolve);
+        stream.on('error', reject);
+        stream.pipe(w);
+    });
+    const c = await fs.find({ _id: file.insertedId }).toArray();
+    await coll.updateOne({ _id: file.insertedId }, { $set: { md5: c[0].md5, size: c[0].size } });
+    return file.insertedId;
+}
+
+function del(_id) {
+    return Promise.all([
+        coll.deleteOne({ _id }),
+        fs.delete(_id),
+    ]);
+}
+
+async function inc(_id) {
+    const file = await coll.findOne({ _id });
+    file.count++;
+}
+
+async function dec(_id) {
+    const file = await coll.findOneAndUpdate(
+        { _id },
+        { $inc: { count: -1 } },
+        { returnOriginal: false },
+    );
+    if (!file.count) await del(_id);
+    return file.count;
+}
+
+async function get(_id, secret = null) {
+    const file = await coll.findOne({ _id });
+    if (typeof secret !== 'undefined') {
+        const timestamp = Math.floor(Number(new Date())).toString();
+        if (!(
+            hash.check(file.secret, timestamp, secret)
+            || hash.check(file.secret, timestamp - 1, secret))) {
+            throw new ForbiddenError();
+        }
+    }
+    return fs.openDownloadStream(_id);
+}
+
+function getMeta(_id) {
+    return coll.findOne({ _id });
+}
+
+async function url(_id, name) {
+    const file = await coll.findOne({ _id });
+    const secret = hash.hash(file.secret, Math.floor(Number(new Date())).toString());
+    if (name) return `/fs/${_id}/${secret}?name=${Buffer.from(name).toString('base64')}`;
+    return `/fs/${_id}/${secret}`;
+}
+
+global.Hydro.model.blacklist = module.exports = {
+    add, get, del, inc, dec, getMeta, url,
+};

hydro/model/problem.js

@@ -1,9 +1,9 @@
 const { ObjectID } = require('bson');
 const { STATUS_ACCEPTED } = require('./builtin').STATUS;
+const file = require('./file');
 const { ProblemNotFoundError } = require('../error');
 const validator = require('../lib/validator');
-const db = require('../service/db.js');
+const db = require('../service/db');
-const gridfs = require('../service/gridfs');
 
 const coll = db.collection('problem');
 const collStatus = db.collection('problem.status');
@@ -155,20 +155,15 @@ async function updateStatus(pid, uid, rid, status) {
 }
 async function setTestdata(_id, readStream) {
     const pdoc = await getById(_id);
-    const f = gridfs.openUploadStream('data.zip');
+    const id = await file.add(readStream, 'data.zip');
-    await new Promise((resolve, reject) => {
+    if (pdoc.data && typeof pdoc.data === 'object') file.dec(this.pdoc.data);
-        readStream.pipe(f);
+    return await edit(_id, { data: id }); // eslint-disable-line no-return-await
-        f.once('finish', resolve);
-        f.once('error', reject);
-    });
-    if (pdoc.data && typeof pdoc.data === 'object') gridfs.delete(this.pdoc.data);
-    return await edit(_id, { data: f.id }); // eslint-disable-line no-return-await
 }
 
 async function getData(pid) {
     const pdoc = await get(pid);
     if (!pdoc.data) return null;
-    return gridfs.openDownloadStream(pdoc.data);
+    return file.get(pdoc.data);
 }
 
 global.Hydro.model.problem = module.exports = {

module/migrate-vijos/script.js

@@ -252,6 +252,13 @@ const tasks = {
     },
     'fs.files': async (docs) => docs,
     'fs.chunks': async (docs) => docs,
+    file: async (docs) => docs.map((doc) => ({
+        _id: doc._id,
+        count: doc.metadata.link,
+        secret: doc.metadata.secret,
+        size: doc.length,
+        md5: doc.md5,
+    })),
 };
 
 const cursor = {
@@ -268,6 +275,7 @@ const cursor = {
     record: (s) => s.collection('record').find(),
     'fs.files': (s) => s.collection('fs.files').find(),
     'fs.chunks': (s) => s.collection('fs.chunks').find(),
+    file: (s) => s.collection('fs.files').find(),
 };
 
 async function task(name, src, report) {
@@ -314,7 +322,7 @@ async function migrateVijos({
     }
     const d = [
         'problem', 'problem.status', 'solution', 'discussion', 'discussion.reply',
-        'contest', 'training', 'training.status', 'record',
+        'contest', 'training', 'training.status', 'record', 'file',
     ];
     for (const i of d) {
         await dst.collection(i).deleteMany();
@@ -322,6 +330,7 @@ async function migrateVijos({
     const t = [
         'user', 'problem', 'problem.status', 'solution', 'discussion',
         'discussion.reply', 'contest', 'training', 'training.status', 'record',
+        'file',
     ];
     for (const i of t) await task(i, src, report);
 }