From 3d18c0c0ba9bb09b2c5bdf77a6a3dee41ce24058 Mon Sep 17 00:00:00 2001
From: Menci <huanghaorui301@gmail.com>
Date: Thu, 30 May 2019 20:40:41 +0800
Subject: [PATCH] Add option for download testdata though different host

---
 config-example.json |  3 ++-
 modules/api.js      |  6 ++----
 modules/api_v2.js   | 28 ++++++++++++++++++++++++++++
 modules/problem.js  | 20 ++++++++++++++++++--
 utility.js          |  6 ++++++
 5 files changed, 56 insertions(+), 7 deletions(-)

diff --git a/config-example.json b/config-example.json
index 6447968..17bbd3a 100644
--- a/config-example.json
+++ b/config-example.json
@@ -125,5 +125,6 @@
   },
   "no_cdn": false,
   "submissions_page_fast_pagination": false,
-  "username_regex": "^[a-zA-Z0-9\\-\\_]+$"
+  "username_regex": "^[a-zA-Z0-9\\-\\_]+$",
+  "site_for_download": ""
 }
diff --git a/modules/api.js b/modules/api.js
index 6abaaf4..8731053 100644
--- a/modules/api.js
+++ b/modules/api.js
@@ -42,8 +42,7 @@ app.post('/api/forget', async (req, res) => {
       expiresIn: '12h'
     });
 
-    const currentProto = req.get("X-Forwarded-Proto") || req.protocol;
-    const vurl = currentProto + '://' + req.get('host') + syzoj.utils.makeUrl(['api', 'forget_confirm'], { token: token });
+    const vurl = syzoj.utils.getCurrentLocation(req, true) + syzoj.utils.makeUrl(['api', 'forget_confirm'], { token: token });
     try {
       await Email.send(user.email,
         `${user.username} 的 ${syzoj.config.title} 密码重置邮件`,
@@ -93,8 +92,7 @@ app.post('/api/sign_up', async (req, res) => {
         expiresIn: '2d'
       });
 
-      const currentProto = req.get("X-Forwarded-Proto") || req.protocol;
-      const vurl = currentProto + '://' + req.get('host') + syzoj.utils.makeUrl(['api', 'sign_up_confirm'], { token: token });
+      const vurl = syzoj.utils.getCurrentLocation(req, true) + syzoj.utils.makeUrl(['api', 'sign_up_confirm'], { token: token });
       try {
         await Email.send(req.body.email,
           `${req.body.username} 的 ${syzoj.config.title} 注册验证邮件`,
diff --git a/modules/api_v2.js b/modules/api_v2.js
index 2e88299..e5c6a26 100644
--- a/modules/api_v2.js
+++ b/modules/api_v2.js
@@ -1,3 +1,5 @@
+const jwt = require('jsonwebtoken');
+
 app.get('/api/v2/search/users/:keyword*?', async (req, res) => {
   try {
     let User = syzoj.model('user');
@@ -104,3 +106,29 @@ app.apiRouter.post('/api/v2/markdown', async (req, res) => {
     res.send(e);
   }
 });
+
+function verifyJWT(token) {
+  try {
+    jwt.verify(token, syzoj.config.session_secret);
+    return true;
+  } catch (e) {
+    return false;
+  }
+}
+
+app.apiRouter.get('/api/v2/download/:token', async (req, res) => {
+  try {
+    const token = req.params.token, data = jwt.decode(token);
+    if (!data) throw new ErrorMessage("无效的令牌。");
+    if (verifyJWT(token)) {
+      res.download(data.filename, data.sendName);
+    } else {
+      res.redirect(data.originUrl);
+    }
+  } catch (e) {
+    syzoj.log(e);
+    res.render('error', {
+      err: e
+    });
+  }
+})
diff --git a/modules/problem.js b/modules/problem.js
index 751b89f..b7ae601 100644
--- a/modules/problem.js
+++ b/modules/problem.js
@@ -7,6 +7,7 @@ let Article = syzoj.model('article');
 
 const randomstring = require('randomstring');
 const fs = require('fs-extra');
+const jwt = require('jsonwebtoken');
 
 let Judger = syzoj.lib('judger');
 let CodeFormatter = syzoj.lib('code_formatter');
@@ -816,6 +817,20 @@ app.post('/problem/:id/testdata/delete/:filename', async (req, res) => {
   }
 });
 
+function downloadOrRedirect(req, res, filename, sendName) {
+  if (syzoj.config.site_for_download) {
+    res.redirect(syzoj.config.site_for_download + syzoj.utils.makeUrl(['api', 'v2', 'download', jwt.sign({
+      filename: filename,
+      sendName: sendName,
+      originUrl: syzoj.utils.getCurrentLocation(req)
+    }, syzoj.config.session_secret, {
+      expiresIn: '2m'
+    })]));
+  } else {
+    res.download(filename, sendName);
+  }
+}
+
 app.get('/problem/:id/testdata/download/:filename?', async (req, res) => {
   try {
     let id = parseInt(req.params.id);
@@ -833,7 +848,8 @@ app.get('/problem/:id/testdata/download/:filename?', async (req, res) => {
     let path = require('path');
     let filename = req.params.filename ? path.join(problem.getTestdataPath(), req.params.filename) : (problem.getTestdataArchivePath());
     if (!await syzoj.utils.isFile(filename)) throw new ErrorMessage('文件不存在。');
-    res.download(filename, path.basename(filename));
+
+    downloadOrRedirect(req, res, filename, path.basename(filename));
   } catch (e) {
     syzoj.log(e);
     res.status(404);
@@ -866,7 +882,7 @@ app.get('/problem/:id/download/additional_file', async (req, res) => {
 
     if (!problem.additional_file) throw new ErrorMessage('无附加文件。');
 
-    res.download(problem.additional_file.getPath(), `additional_file_${id}.zip`);
+    downloadOrRedirect(req, res, problem.additional_file.getPath(), `additional_file_${id}.zip`);
   } catch (e) {
     syzoj.log(e);
     res.status(404);
diff --git a/utility.js b/utility.js
index 4af19ed..f8a0903 100644
--- a/utility.js
+++ b/utility.js
@@ -317,5 +317,11 @@ module.exports = {
       }
       attempt();
     });
+  },
+  getCurrentLocation(req, hostOnly) {
+    const currentProto = req.get("X-Forwarded-Proto") || req.protocol,
+          host = currentProto + '://' + req.get('host');
+    if (hostOnly) return host;
+    else return host + req.originalUrl;
   }
 };