From de47a17be7ce541fcd5c8ee698c2ef4aa926b5fb Mon Sep 17 00:00:00 2001
From: tamaina <tamaina@hotmail.co.jp>
Date: Mon, 18 Mar 2019 13:20:49 +0900
Subject: [PATCH] Improve drive downloading (#4523)

* Improve drive file downloading

* fix name

* wtf crlf

* semicolon
---
 package.json                                               | 2 ++
 .../app/desktop/views/components/context-menu.menu.vue     | 2 +-
 src/client/app/desktop/views/components/drive.file.vue     | 4 +++-
 .../app/mobile/views/components/drive.file-detail.vue      | 7 ++++++-
 src/prelude/url.ts                                         | 4 ++++
 src/server/file/send-drive-file.ts                         | 7 ++++++-
 6 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/package.json b/package.json
index da436926f5..bb48d060cd 100644
--- a/package.json
+++ b/package.json
@@ -77,6 +77,7 @@
 		"@types/qrcode": "1.3.0",
 		"@types/ratelimiter": "2.1.28",
 		"@types/redis": "2.8.10",
+		"@types/rename": "1.0.1",
 		"@types/request": "2.48.1",
 		"@types/request-promise-native": "1.0.15",
 		"@types/request-stats": "3.0.0",
@@ -193,6 +194,7 @@
 		"recaptcha-promise": "0.1.3",
 		"reconnecting-websocket": "4.1.10",
 		"redis": "2.8.0",
+		"rename": "1.0.4",
 		"request": "2.88.0",
 		"request-promise-native": "1.0.7",
 		"request-stats": "3.0.0",
diff --git a/src/client/app/desktop/views/components/context-menu.menu.vue b/src/client/app/desktop/views/components/context-menu.menu.vue
index 1ae3c85d57..f2bb3bec23 100644
--- a/src/client/app/desktop/views/components/context-menu.menu.vue
+++ b/src/client/app/desktop/views/components/context-menu.menu.vue
@@ -6,7 +6,7 @@
 				<p @click="click(item)"><i v-if="item.icon" :class="$style.icon"><fa :icon="item.icon"/></i>{{ item.text }}</p>
 			</template>
 			<template v-else-if="item.type == 'link'">
-				<a :href="item.href" :target="item.target" @click="click(item)"><i v-if="item.icon" :class="$style.icon"><fa :icon="item.icon"/></i>{{ item.text }}</a>
+				<a :href="item.href" :target="item.target" @click="click(item)" :download="item.download"><i v-if="item.icon" :class="$style.icon"><fa :icon="item.icon"/></i>{{ item.text }}</a>
 			</template>
 			<template v-else-if="item.type == 'nest'">
 				<p><i v-if="item.icon" :class="$style.icon"><fa :icon="item.icon"/></i>{{ item.text }}...<span class="caret"><fa icon="caret-right"/></span></p>
diff --git a/src/client/app/desktop/views/components/drive.file.vue b/src/client/app/desktop/views/components/drive.file.vue
index fbd649e8f6..b9d202f555 100644
--- a/src/client/app/desktop/views/components/drive.file.vue
+++ b/src/client/app/desktop/views/components/drive.file.vue
@@ -38,6 +38,7 @@ import anime from 'animejs';
 import copyToClipboard from '../../../common/scripts/copy-to-clipboard';
 import updateAvatar from '../../api/update-avatar';
 import updateBanner from '../../api/update-banner';
+import { appendQuery } from '../../../../../prelude/url';
 
 export default Vue.extend({
 	i18n: i18n('desktop/views/components/drive.file.vue'),
@@ -88,9 +89,10 @@ export default Vue.extend({
 				action: this.copyUrl
 			}, {
 				type: 'link',
-				href: `${this.file.url}?download`,
+				href: appendQuery(this.file.url, 'download'),
 				text: this.$t('contextmenu.download'),
 				icon: 'download',
+				download: this.file.name
 			}, null, {
 				type: 'item',
 				text: this.$t('@.delete'),
diff --git a/src/client/app/mobile/views/components/drive.file-detail.vue b/src/client/app/mobile/views/components/drive.file-detail.vue
index 4d0a747fcb..d7432437e0 100644
--- a/src/client/app/mobile/views/components/drive.file-detail.vue
+++ b/src/client/app/mobile/views/components/drive.file-detail.vue
@@ -38,7 +38,7 @@
 	<div class="menu">
 		<div>
 			<ui-input readonly :value="file.url">URL</ui-input>
-			<ui-button link :href="`${file.url}?download`" :download="file.name"><fa icon="download"/> {{ $t('download') }}</ui-button>
+			<ui-button link :href="dlUrl" :download="file.name"><fa icon="download"/> {{ $t('download') }}</ui-button>
 			<ui-button @click="rename"><fa icon="pencil-alt"/> {{ $t('rename') }}</ui-button>
 			<ui-button @click="move"><fa :icon="['far', 'folder-open']"/> {{ $t('move') }}</ui-button>
 			<ui-button @click="toggleSensitive" v-if="file.isSensitive"><fa :icon="['far', 'eye']"/> {{ $t('unmark-as-sensitive') }}</ui-button>
@@ -61,6 +61,7 @@
 import Vue from 'vue';
 import i18n from '../../../i18n';
 import { gcd } from '../../../../../prelude/math';
+import { appendQuery } from '../../../../../prelude/url';
 
 export default Vue.extend({
 	i18n: i18n('mobile/views/components/drive.file-detail.vue'),
@@ -86,6 +87,10 @@ export default Vue.extend({
 			return this.file.properties.avgColor && this.file.properties.avgColor.length == 3 ? {
 				'background-color': `rgb(${ this.file.properties.avgColor.join(',') })`
 			} : {};
+		},
+
+		dlUrl(): string {
+			return appendQuery(this.file.url, 'download');
 		}
 	},
 
diff --git a/src/prelude/url.ts b/src/prelude/url.ts
index ff1012d4c1..a3613fc9b9 100644
--- a/src/prelude/url.ts
+++ b/src/prelude/url.ts
@@ -5,3 +5,7 @@ export function query(obj: {}): string {
 		.filter(([, v]) => Array.isArray(v) ? v.length : v !== undefined)
 		.reduce((a, [k, v]) => (a[k] = v, a), {} as Record<string, any>));
 }
+
+export function appendQuery(url: string, query: string): string {
+	return `${url}${/\?/.test(url) ? url.endsWith('?') ? '' : '&' : '?'}${query}`;
+}
diff --git a/src/server/file/send-drive-file.ts b/src/server/file/send-drive-file.ts
index 691d3bf848..c57648bb7a 100644
--- a/src/server/file/send-drive-file.ts
+++ b/src/server/file/send-drive-file.ts
@@ -1,6 +1,7 @@
 import * as Koa from 'koa';
 import * as send from 'koa-send';
 import * as mongodb from 'mongodb';
+import * as rename from 'rename';
 import DriveFile, { getDriveFileBucket } from '../../models/drive-file';
 import DriveFileThumbnail, { getDriveFileThumbnailBucket } from '../../models/drive-file-thumbnail';
 import DriveFileWebpublic, { getDriveFileWebpublicBucket } from '../../models/drive-file-webpublic';
@@ -62,10 +63,12 @@ export default async function(ctx: Koa.BaseContext) {
 
 		if (thumb != null) {
 			ctx.set('Content-Type', 'image/jpeg');
+			ctx.set('Content-Disposition', `filename="${rename(file.filename, { suffix: '-thumb', extname: '.jpeg' })}"`);
 			const bucket = await getDriveFileThumbnailBucket();
 			ctx.body = bucket.openDownloadStream(thumb._id);
 		} else {
 			if (file.contentType.startsWith('image/')) {
+				ctx.set('Content-Disposition', `filename="${file.filename}"`);
 				await sendRaw();
 			} else {
 				ctx.status = 404;
@@ -79,15 +82,17 @@ export default async function(ctx: Koa.BaseContext) {
 
 		if (web != null) {
 			ctx.set('Content-Type', file.contentType);
+			ctx.set('Content-Disposition', `filename="${rename(file.filename, { suffix: '-web' })}"`);
 
 			const bucket = await getDriveFileWebpublicBucket();
 			ctx.body = bucket.openDownloadStream(web._id);
 		} else {
+			ctx.set('Content-Disposition', `filename="${file.filename}"`);
 			await sendRaw();
 		}
 	} else {
 		if ('download' in ctx.query) {
-			ctx.set('Content-Disposition', 'attachment');
+			ctx.set('Content-Disposition', `attachment; filename="${file.filename}`);
 		}
 
 		await sendRaw();