diff --git a/src/index.ts b/src/index.ts index 2e7ce0b..3f05452 100644 --- a/src/index.ts +++ b/src/index.ts @@ -27,4 +27,4 @@ app.use(staticPlugin({ prefix: '/' })) app.listen(1337); console.log( `api is running at ${app.server?.hostname}:${app.server?.port}` -); \ No newline at end of file +); diff --git a/src/router/html.ts b/src/router/html.ts index ff5c8e2..2d53430 100644 --- a/src/router/html.ts +++ b/src/router/html.ts @@ -2,6 +2,11 @@ import { Elysia } from 'elysia'; import { m, eta, error } from '@/utils/html' import healthStatus from '@/utils/health'; import { checkIpRanges } from '@/utils/ranges'; +import { + buildRateLimitCookie, + createRateLimitCookieValue, + getRateLimitCookie, +} from '@/utils/rate-limit'; const app = new Elysia() app.get('/', async ({ set }) => { @@ -29,15 +34,20 @@ app.get('/save', async ({ query: { url }, set, headers, error }) => { websocket = process.env.ALTERNATIVE_WEBSOCKET! } const ytPatched = (process.env.YT_PATCHED || '').toLowerCase() === 'true' || process.env.YT_PATCHED === '1' + const rateLimitId = getRateLimitCookie(headers.cookie) || createRateLimitCookieValue() + if (!getRateLimitCookie(headers.cookie)) { + set.headers['Set-Cookie'] = buildRateLimitCookie(rateLimitId) + } set.headers['Content-Type'] = 'text/html; charset=utf-8' - return error(412, await m(eta.render('./save', { + return await m(eta.render('./save', { title: 'Save Video | PreserveTube', websocket, sitekey: process.env.SITEKEY, url, - ytPatched - }))) + ytPatched, + rateLimitId + })) }) app.get('/savechannel', async ({ query: { url }, set, headers, error }) => { @@ -46,22 +56,27 @@ app.get('/savechannel', async ({ query: { url }, set, headers, error }) => { const ranges = await checkIpRanges(headers['cf-connecting-ip'] || headers['x-forwarded-for'] || '') if (ranges.blocked) { set.headers['Content-Type'] = 'text/html; charset=utf-8' - return await m(eta.render('./blocked', { + return error(412, await m(eta.render('./blocked', { title: 'Blocked | PreserveTube' - })) + }))) } let websocket = process.env.WEBSOCKET if (healthStatus[process.env.METADATA!] != 'healthy') { websocket = process.env.ALTERNATIVE_WEBSOCKET! } + const rateLimitId = getRateLimitCookie(headers.cookie) || createRateLimitCookieValue() + if (!getRateLimitCookie(headers.cookie)) { + set.headers['Set-Cookie'] = buildRateLimitCookie(rateLimitId) + } set.headers['Content-Type'] = 'text/html; charset=utf-8' return await m(eta.render('./savechannel', { title: 'Save Channel | PreserveTube', websocket, sitekey: process.env.SITEKEY, - url + url, + rateLimitId })) }) diff --git a/src/router/search.ts b/src/router/search.ts index 18b27bb..56e786e 100644 --- a/src/router/search.ts +++ b/src/router/search.ts @@ -60,4 +60,4 @@ app.get('/search/channel', async ({ query: { url }, error, redirect }) => { }) app.onError(error) -export default app \ No newline at end of file +export default app diff --git a/src/router/websocket.ts b/src/router/websocket.ts index d08e395..d3c35d4 100644 --- a/src/router/websocket.ts +++ b/src/router/websocket.ts @@ -11,6 +11,7 @@ import { error } from '@/utils/html' import redis from '@/utils/redis'; import { parseSlop } from '@/utils/slop'; import { checkIpRanges } from '@/utils/ranges'; +import { getRateLimitSubjects } from '@/utils/rate-limit'; const app = new Elysia() const videoIds: Record = {} @@ -18,15 +19,19 @@ const videoIds: Record = {} const MB_LIMIT = 250 const saveKey = (videoId: string) => `save:${videoId}` -const checkMbLimit = async (hash: string, mb?: number): Promise => { - const key = `save-mb:${hash}` - const current = parseInt(await redis.get(key) || '0') - if (!mb) return current >= MB_LIMIT - if (current + mb > MB_LIMIT) return true +const checkMbLimit = async (subjects: string[], mb?: number): Promise => { + const keys = subjects.map(subject => `save-mb:${Bun.hash(subject).toString()}`) + const currentValues = await Promise.all(keys.map(key => redis.get(key))) + const currents = currentValues.map(value => parseInt(value || '0')) + + if (!mb) return currents.some(current => current >= MB_LIMIT) + if (currents.some(current => current + mb > MB_LIMIT)) return true const pipeline = redis.pipeline() - pipeline.incrby(key, mb) - pipeline.expire(key, 24 * 60 * 60) + for (const key of keys) { + pipeline.incrby(key, mb) + pipeline.expire(key, 24 * 60 * 60) + } await pipeline.exec() return false } @@ -89,7 +94,8 @@ const getRateLimitKey = (ip: string): string => { app.ws('/save', { query: t.Object({ - url: t.String() + url: t.String(), + rlid: t.Optional(t.String()) }), body: t.String(), open: async (ws) => { @@ -115,13 +121,16 @@ app.ws('/save', { ws.send(`DONE - ${process.env.FRONTEND}/watch?v=${videoId}`) ws.close() } else { - const hash = Bun.hash(getRateLimitKey(ws.data.headers['cf-connecting-ip'] || '0.0.0.0')) - const isLimited = await checkMbLimit(hash.toString()) + const subjects = getRateLimitSubjects( + getRateLimitKey(ws.data.headers['cf-connecting-ip'] || '0.0.0.0'), + ws.data.query.rlid + ) + const isLimited = await checkMbLimit(subjects) if (isLimited) { return sendError(ws, 'You have been ratelimited.
Is this an urgent archive? Please email me: admin@preservetube.com'); } - console.log(`saving (${hash}) - ${ws.data.path} - ${JSON.stringify(ws.data.query)}`) + console.log(`saving (${subjects.map(subject => Bun.hash(subject).toString()).join(',')}) - ${ws.data.path} - ${JSON.stringify(ws.data.query)}`) ws.send('DATA - This process is automatic. Your video will start archiving shortly.') ws.send('CAPTCHA - Solving a cryptographic challenge before downloading.') videoIds[ws.id] = videoId @@ -167,8 +176,11 @@ app.ws('/save', { } const mbsUsed = Math.ceil(downloadResult.size / (1024 * 1024)) - const hash = Bun.hash(getRateLimitKey(ws.data.headers['cf-connecting-ip'] || '0.0.0.0')) - const isMbLimited = await checkMbLimit(hash.toString(), mbsUsed) + const subjects = getRateLimitSubjects( + getRateLimitKey(ws.data.headers['cf-connecting-ip'] || '0.0.0.0'), + ws.data.query.rlid + ) + const isMbLimited = await checkMbLimit(subjects, mbsUsed) if (isMbLimited) { const file = fs.readdirSync('./videos/').find(f => f.includes(`${videoId}.`)) if (file) fs.unlinkSync('./videos/' + file) @@ -193,7 +205,8 @@ app.ws('/save', { app.ws('/savechannel', { query: t.Object({ - url: t.String() + url: t.String(), + rlid: t.Optional(t.String()) }), body: t.String(), open: async (ws) => { @@ -246,14 +259,17 @@ app.ws('/savechannel', { .executeTakeFirst() if (already) continue - const hash = Bun.hash(getRateLimitKey(ws.data.headers['cf-connecting-ip'] || '0.0.0.0')) - const isLimited = await checkMbLimit(hash.toString()) + const subjects = getRateLimitSubjects( + getRateLimitKey(ws.data.headers['cf-connecting-ip'] || '0.0.0.0'), + ws.data.query.rlid + ) + const isLimited = await checkMbLimit(subjects) if (isLimited) { sendError(ws, 'You have been ratelimited.
Is this an urgent archive? Please email me: admin@preservetube.com', false); break; } - console.log(`saving (${hash}) - ${ws.data.path} - ${video.video_id}`) + console.log(`saving (${subjects.map(subject => Bun.hash(subject).toString()).join(',')}) - ${ws.data.path} - ${video.video_id}`) const isSlop = await parseSlop(video.video_id, video.title.text, video.description_snippet?.text || '', channelId) @@ -271,7 +287,7 @@ app.ws('/savechannel', { const downloadResult = await downloadVideo(ws, video.video_id); if (!downloadResult.fail) { const mbsUsed = Math.ceil(downloadResult.size / (1024 * 1024)) - const isMbLimited = await checkMbLimit(hash.toString(), mbsUsed) + const isMbLimited = await checkMbLimit(subjects, mbsUsed) if (isMbLimited) { const file = fs.readdirSync('./videos/').find(f => f.includes(`${video.video_id}.`)) if (file) fs.unlinkSync('./videos/' + file) diff --git a/src/templates/save.eta b/src/templates/save.eta index 5f1ec86..71db2ab 100644 --- a/src/templates/save.eta +++ b/src/templates/save.eta @@ -45,6 +45,7 @@ const url = "<%= it.url %>" const websocket = "<%= it.websocket %>" const sitekey = "<%= it.sitekey %>" + const initialRateLimitId = "<%= it.rateLimitId %>" let ws = null const h1 = document.getElementById("title") const h3 = document.getElementById("bottom") @@ -53,8 +54,33 @@ const acceptInstructions = document.getElementById("accept-instructions") const declineInstructions = document.getElementById("decline-instructions") + function getRateLimitId() { + if (initialRateLimitId) { + return initialRateLimitId + } + + try { + const stored = localStorage.getItem("pt_rlid") + if (stored) { + return stored + } + } catch { + } + + return initialRateLimitId + } + + function persistRateLimitId(rateLimitId) { + try { + localStorage.setItem("pt_rlid", rateLimitId) + } catch { + } + } + function startSaveFlow() { - ws = new WebSocket(`${websocket}/save?url=${encodeURIComponent(url)}`) + const rateLimitId = getRateLimitId() + persistRateLimitId(rateLimitId) + ws = new WebSocket(`${websocket}/save?url=${encodeURIComponent(url)}&rlid=${encodeURIComponent(rateLimitId)}`) ws.onopen = function () { document.querySelector("footer").innerHTML = `page rendered by <%= hostname %>, connected via ${websocket.split('-')[1].split('.')[0]}` @@ -145,6 +171,7 @@ }, 2000) runFirstSaveGate() + persistRateLimitId(initialRateLimitId) <% } %> diff --git a/src/templates/savechannel.eta b/src/templates/savechannel.eta index 1f58e4f..8ea5c77 100644 --- a/src/templates/savechannel.eta +++ b/src/templates/savechannel.eta @@ -13,7 +13,27 @@ const url = "<%= it.url %>" const websocket = "<%= it.websocket %>" const sitekey = "<%= it.sitekey %>" - const ws = new WebSocket(`${websocket}/savechannel?url=${encodeURIComponent(url)}`) + const initialRateLimitId = "<%= it.rateLimitId %>" + function getRateLimitId() { + if (initialRateLimitId) { + return initialRateLimitId + } + + try { + const stored = localStorage.getItem("pt_rlid") + if (stored) { + return stored + } + } catch { + } + + return initialRateLimitId + } + try { + localStorage.setItem("pt_rlid", initialRateLimitId) + } catch { + } + const ws = new WebSocket(`${websocket}/savechannel?url=${encodeURIComponent(url)}&rlid=${encodeURIComponent(getRateLimitId())}`) const h1 = document.getElementById("title") const h3 = document.getElementById("bottom") const captcha = document.getElementById("captcha") @@ -58,4 +78,4 @@ .logs { text-align: center; } - \ No newline at end of file +