From a470b72016f32c153261ede9abfdaa9893abe1f7 Mon Sep 17 00:00:00 2001 From: Dragory Date: Wed, 6 Mar 2019 23:31:24 +0200 Subject: [PATCH] Add attachmentStorage option --- package-lock.json | 1 - package.json | 1 + src/config.js | 10 ++- src/data/Thread.js | 18 ++-- src/data/attachments.js | 185 ++++++++++++++++++++++++++++----------- src/main.js | 2 - src/modules/reply.js | 4 - src/modules/webserver.js | 2 +- src/utils.js | 4 +- 9 files changed, 158 insertions(+), 69 deletions(-) diff --git a/package-lock.json b/package-lock.json index 6a551e3..a72305c 100644 --- a/package-lock.json +++ b/package-lock.json @@ -3277,7 +3277,6 @@ "version": "0.0.33", "resolved": "https://registry.npmjs.org/tmp/-/tmp-0.0.33.tgz", "integrity": "sha512-jRCJlojKnZ3addtTOjdIqoRuPEKBvNXcGYqzO6zWZX8KfKEpnGY5jfggJQ3EjKuu8D4bJRr0y+cYJFmYbImXGw==", - "dev": true, "requires": { "os-tmpdir": "~1.0.2" } diff --git a/package.json b/package.json index 3c3c133..a5300cd 100644 --- a/package.json +++ b/package.json @@ -19,6 +19,7 @@ "moment": "^2.21.0", "public-ip": "^2.0.1", "sqlite3": "^4.0.6", + "tmp": "0.0.33", "transliteration": "^1.6.2", "uuid": "^3.1.0" }, diff --git a/src/config.js b/src/config.js index d111c4b..4140007 100644 --- a/src/config.js +++ b/src/config.js @@ -80,6 +80,8 @@ const defaultConfig = { "relaySmallAttachmentsAsAttachments": false, "smallAttachmentLimit": 1024 * 1024 * 2, + "attachmentStorage": "local", + "attachmentStorageChannelId": null, "port": 8890, "url": null, @@ -130,7 +132,13 @@ for (const opt of required) { if (finalConfig.smallAttachmentLimit > 1024 * 1024 * 8) { finalConfig.smallAttachmentLimit = 1024 * 1024 * 8; - console.log('[WARN] smallAttachmentLimit capped at 8MB'); + console.warn('[WARN] smallAttachmentLimit capped at 8MB'); +} + +// Specific checks +if (finalConfig.attachmentStorage === 'discord' && ! finalConfig.attachmentStorageChannelId) { + console.error('Config option \'attachmentStorageChannelId\' is required with attachment storage \'discord\''); + process.exit(1); } // Make sure mainGuildId is internally always an array diff --git a/src/data/Thread.js b/src/data/Thread.js index a95b1a5..4a2badd 100644 --- a/src/data/Thread.js +++ b/src/data/Thread.js @@ -64,10 +64,18 @@ class Thread { if (replyAttachments.length > 0) { for (const attachment of replyAttachments) { - files.push(await attachments.attachmentToFile(attachment)); - const url = await attachments.getUrl(attachment.id, attachment.filename); + let savedAttachment; - logContent += `\n\n**Attachment:** ${url}`; + await Promise.all([ + attachments.attachmentToFile(attachment).then(file => { + files.push(file); + }), + attachments.saveAttachment(attachment).then(result => { + savedAttachment = result; + }) + ]); + + logContent += `\n\n**Attachment:** ${savedAttachment.url}`; } } @@ -131,10 +139,10 @@ class Thread { let attachmentFiles = []; for (const attachment of msg.attachments) { - await attachments.saveAttachment(attachment); + const savedAttachment = await attachments.saveAttachment(attachment); // Forward small attachments (<2MB) as attachments, just link to larger ones - const formatted = '\n\n' + await utils.formatAttachment(attachment); + const formatted = '\n\n' + await utils.formatAttachment(attachment, savedAttachment.url); logContent += formatted; // Logs always contain the link if (config.relaySmallAttachmentsAsAttachments && attachment.size <= 1024 * 1024 * 2) { diff --git a/src/data/attachments.js b/src/data/attachments.js index 5909ff7..e604ac6 100644 --- a/src/data/attachments.js +++ b/src/data/attachments.js @@ -1,57 +1,62 @@ const Eris = require('eris'); const fs = require('fs'); const https = require('https'); -const config = require('../config'); const {promisify} = require('util'); +const tmp = require('tmp'); +const config = require('../config'); +const utils = require('../utils'); const getUtils = () => require('../utils'); const access = promisify(fs.access); const readFile = promisify(fs.readFile); +const rename = promisify(fs.rename); -const attachmentDir = config.attachmentDir || `${__dirname}/../../attachments`; +const localAttachmentDir = config.attachmentDir || `${__dirname}/../../attachments`; const attachmentSavePromises = {}; -/** - * Returns the filesystem path for the given attachment id - * @param {String} attachmentId - * @returns {String} - */ -function getPath(attachmentId) { - return `${attachmentDir}/${attachmentId}`; +function getErrorResult(msg = null) { + return { + url: `Attachment could not be saved${msg ? ': ' + msg : ''}`, + failed: true + }; } /** * Attempts to download and save the given attachement * @param {Object} attachment * @param {Number=0} tries - * @returns {Promise} + * @returns {Promise<{ url: string }>} */ -async function saveAttachment(attachment) { - if (attachmentSavePromises[attachment.id]) { - return attachmentSavePromises[attachment.id]; - } +async function saveLocalAttachment(attachment) { + const targetPath = getLocalAttachmentPath(attachment.id); - const filepath = getPath(attachment.id); try { // If the file already exists, resolve immediately - await access(filepath); - return; + await access(targetPath); + const url = await getLocalAttachmentUrl(attachment.id, attachment.filename); + return { url }; } catch (e) {} - attachmentSavePromises[attachment.id] = saveAttachmentInner(attachment); - attachmentSavePromises[attachment.id] - .then(() => { - delete attachmentSavePromises[attachment.id]; - }, () => { - delete attachmentSavePromises[attachment.id]; - }); + // Download the attachment + const downloadResult = await downloadAttachment(attachment); - return attachmentSavePromises[attachment.id]; + // Move the temp file to the attachment folder + await rename(downloadResult.path, targetPath); + + // Resolve the attachment URL + const url = await getLocalAttachmentUrl(attachment.id, attachment.filename); + + return { url }; } -function saveAttachmentInner(attachment, tries = 0) { +/** + * @param {Object} attachment + * @param {Number} tries + * @returns {Promise<{ path: string, cleanup: function }>} + */ +function downloadAttachment(attachment, tries = 0) { return new Promise((resolve, reject) => { if (tries > 3) { console.error('Attachment download failed after 3 tries:', attachment); @@ -59,54 +64,130 @@ function saveAttachmentInner(attachment, tries = 0) { return; } - const filepath = getPath(attachment.id); - const writeStream = fs.createWriteStream(filepath); + tmp.file((err, filepath, fd, cleanupCallback) => { + const writeStream = fs.createWriteStream(filepath); - https.get(attachment.url, (res) => { - res.pipe(writeStream); - writeStream.on('finish', () => { - writeStream.end(); - resolve(); + https.get(attachment.url, (res) => { + res.pipe(writeStream); + writeStream.on('finish', () => { + writeStream.end(); + resolve({ + path: filepath, + cleanup: cleanupCallback + }); + }); + }).on('error', (err) => { + fs.unlink(filepath); + console.error('Error downloading attachment, retrying'); + resolve(downloadAttachment(attachment, tries++)); }); - }).on('error', (err) => { - fs.unlink(filepath); - console.error('Error downloading attachment, retrying'); - resolve(saveAttachmentInner(attachment, tries++)); }); }); } /** - * Attempts to download and save all attachments in the given message - * @param {Eris.Message} msg - * @returns {Promise} + * Returns the filesystem path for the given attachment id + * @param {String} attachmentId + * @returns {String} */ -function saveAttachmentsInMessage(msg) { - if (! msg.attachments || msg.attachments.length === 0) return Promise.resolve(); - return Promise.all(msg.attachments.map(saveAttachment)); +function getLocalAttachmentPath(attachmentId) { + return `${localAttachmentDir}/${attachmentId}`; } /** * Returns the self-hosted URL to the given attachment ID * @param {String} attachmentId * @param {String=null} desiredName Custom name for the attachment as a hint for the browser - * @returns {String} + * @returns {Promise} */ -function getUrl(attachmentId, desiredName = null) { +function getLocalAttachmentUrl(attachmentId, desiredName = null) { if (desiredName == null) desiredName = 'file.bin'; return getUtils().getSelfUrl(`attachments/${attachmentId}/${desiredName}`); } +/** + * @param {Object} attachment + * @returns {Promise<{ url: string }>} + */ +async function saveDiscordAttachment(attachment) { + if (attachment.size > 1024 * 1024 * 8) { + return getErrorResult('attachment too large (max 8MB)'); + } + + const attachmentChannelId = config.attachmentStorageChannelId; + const inboxGuild = utils.getInboxGuild(); + + if (! inboxGuild.channels.has(attachmentChannelId)) { + throw new Error('Attachment storage channel not found!'); + } + + const attachmentChannel = inboxGuild.channels.get(attachmentChannelId); + if (! (attachmentChannel instanceof Eris.TextChannel)) { + throw new Error('Attachment storage channel must be a text channel!'); + } + + const file = await attachmentToFile(attachment); + const savedAttachment = await createDiscordAttachmentMessage(attachmentChannel, file); + if (! savedAttachment) return getErrorResult(); + + return { url: savedAttachment.url }; +} + +async function createDiscordAttachmentMessage(channel, file, tries = 0) { + tries++; + + try { + const attachmentMessage = await channel.createMessage(undefined, file); + return attachmentMessage.attachments[0]; + } catch (e) { + if (tries > 3) { + console.error(`Attachment storage message could not be created after 3 tries: ${e.message}`); + return; + } + + return createDiscordAttachmentMessage(channel, file, tries); + } +} + +/** + * Turns the given attachment into a file object that can be sent forward as a new attachment + * @param {Object} attachment + * @returns {Promise<{file, name: string}>} + */ async function attachmentToFile(attachment) { - await saveAttachment(attachment); - const data = await readFile(getPath(attachment.id)); + const downloadResult = await downloadAttachment(attachment); + const data = await readFile(downloadResult.path); + downloadResult.cleanup(); return {file: data, name: attachment.filename}; } +/** + * Saves the given attachment based on the configured storage system + * @param {Object} attachment + * @returns {Promise<{ url: string }>} + */ +function saveAttachment(attachment) { + if (attachmentSavePromises[attachment.id]) { + return attachmentSavePromises[attachment.id]; + } + + if (config.attachmentStorage === 'local') { + attachmentSavePromises[attachment.id] = saveLocalAttachment(attachment); + } else if (config.attachmentStorage === 'discord') { + attachmentSavePromises[attachment.id] = saveDiscordAttachment(attachment); + } else { + throw new Error(`Unknown attachment storage option: ${config.attachmentStorage}`); + } + + attachmentSavePromises[attachment.id].then(() => { + delete attachmentSavePromises[attachment.id]; + }); + + return attachmentSavePromises[attachment.id]; +} + module.exports = { - getPath, - saveAttachment, - saveAttachmentsInMessage, - getUrl, - attachmentToFile + getLocalAttachmentPath, + attachmentToFile, + saveAttachment }; diff --git a/src/main.js b/src/main.js index 6975b9b..9c054a9 100644 --- a/src/main.js +++ b/src/main.js @@ -51,8 +51,6 @@ bot.on('messageCreate', async msg => { // AUTO-REPLY: If config.alwaysReply is enabled, send all chat messages in thread channels as replies if (! utils.isStaff(msg.member)) return; // Only staff are allowed to reply - if (msg.attachments.length) await attachments.saveAttachmentsInMessage(msg); - const replied = await thread.replyToUser(msg.member, msg.content.trim(), msg.attachments, config.alwaysReplyAnon || false); if (replied) msg.delete(); } else { diff --git a/src/modules/reply.js b/src/modules/reply.js index ac6931a..f041cc9 100644 --- a/src/modules/reply.js +++ b/src/modules/reply.js @@ -10,8 +10,6 @@ module.exports = bot => { if (! thread) return; const text = args.join(' ').trim(); - if (msg.attachments.length) await attachments.saveAttachmentsInMessage(msg); - const replied = await thread.replyToUser(msg.member, text, msg.attachments, false); if (replied) msg.delete(); }); @@ -23,8 +21,6 @@ module.exports = bot => { if (! thread) return; const text = args.join(' ').trim(); - if (msg.attachments.length) await attachments.saveAttachmentsInMessage(msg); - const replied = await thread.replyToUser(msg.member, text, msg.attachments, true); if (replied) msg.delete(); }); diff --git a/src/modules/webserver.js b/src/modules/webserver.js index fdad567..5d8b2c9 100644 --- a/src/modules/webserver.js +++ b/src/modules/webserver.js @@ -55,7 +55,7 @@ function serveAttachments(res, pathParts) { if (id.match(/^[0-9]+$/) === null) return notfound(res); if (desiredFilename.match(/^[0-9a-z._-]+$/i) === null) return notfound(res); - const attachmentPath = attachments.getPath(id); + const attachmentPath = attachments.getLocalAttachmentPath(id); fs.access(attachmentPath, (err) => { if (err) return notfound(res); diff --git a/src/utils.js b/src/utils.js index 5fd6a0b..69f3058 100644 --- a/src/utils.js +++ b/src/utils.js @@ -3,7 +3,6 @@ const bot = require('./bot'); const moment = require('moment'); const humanizeDuration = require('humanize-duration'); const publicIp = require('public-ip'); -const attachments = require('./data/attachments'); const config = require('./config'); class BotError extends Error {} @@ -123,11 +122,10 @@ function messageIsOnMainServer(msg) { * @param attachment * @returns {Promise} */ -async function formatAttachment(attachment) { +async function formatAttachment(attachment, attachmentUrl) { let filesize = attachment.size || 0; filesize /= 1024; - const attachmentUrl = await attachments.getUrl(attachment.id, attachment.filename); return `**Attachment:** ${attachment.filename} (${filesize.toFixed(1)}KB)\n${attachmentUrl}`; }