| 1 | // This file is part of HFS - Copyright 2021-2023, Massimo Melina <a@rejetto.com> - License https://www.gnu.org/licenses/gpl-3.0.txt |
| 2 | |
| 3 | import { Readable } from 'stream' |
| 4 | import Koa from 'koa' |
| 5 | import { ThrottledStream, ThrottleGroup } from './ThrottledStream' |
| 6 | import { defineConfig } from './config' |
| 7 | import { isLocalHost } from './misc' |
| 8 | import { Connection, getConnection, updateConnection } from './connections' |
| 9 | import _ from 'lodash' |
| 10 | import events from './events' |
| 11 | import { storedMap } from './persistence' |
| 12 | |
| 13 | const mainThrottleGroup = new ThrottleGroup(Infinity) |
| 14 | |
| 15 | defineConfig('max_kbps', Infinity).sub(v => |
| 16 | mainThrottleGroup.updateLimit(v)) |
| 17 | |
| 18 | const ip2group: Record<string, { |
| 19 | count: number |
| 20 | group: ThrottleGroup |
| 21 | }> = {} |
| 22 | |
| 23 | const SymThrStr = Symbol('stream') |
| 24 | const SymTimeout = Symbol('timeout') |
| 25 | |
| 26 | const maxKbpsPerIp = defineConfig('max_kbps_per_ip', Infinity) |
| 27 | maxKbpsPerIp.sub(v => { |
| 28 | for (const [ip, {group}] of Object.entries(ip2group)) |
| 29 | if (ip) // empty-string = unlimited group |
| 30 | group.updateLimit(v) |
| 31 | }) |
| 32 | |
| 33 | export const throttler: Koa.Middleware = async (ctx, next) => { |
| 34 | await next() |
| 35 | let { body } = ctx |
| 36 | const downloadTotal: number = ctx.response.length |
| 37 | if (typeof body === 'string' || body && body instanceof Buffer) |
| 38 | ctx.body = body = Readable.from(body) |
| 39 | if (!body || !(body instanceof Readable)) |
| 40 | return |
| 41 | // we wrap the stream also for unlimited connections to get speed and other features |
| 42 | const noLimit = ctx.state.account?.ignore_limits || isLocalHost(ctx) |
| 43 | const ipGroup = ip2group[noLimit ? '' : ctx.ip] ||= { |
| 44 | count: 0, |
| 45 | group: new ThrottleGroup(noLimit ? Infinity : maxKbpsPerIp.get(), noLimit ? undefined : mainThrottleGroup), |
| 46 | } |
| 47 | const conn = getConnection(ctx) |
| 48 | if (!conn) throw 'assert throttler connection' |
| 49 | |
| 50 | const ts = conn[SymThrStr] = new ThrottledStream(ipGroup.group, conn[SymThrStr]) |
| 51 | const offset = ts.getBytesSent() |
| 52 | let closed = false |
| 53 | |
| 54 | const DELAY = 1000 |
| 55 | const update = _.debounce(() => { |
| 56 | const ts = conn[SymThrStr] as ThrottledStream |
| 57 | const outSpeed = roundSpeed(ts.getSpeed()) |
| 58 | const { state } = ctx |
| 59 | updateConnection(conn, { outSpeedKb: outSpeed, sent: conn.socket.bytesWritten }, |
| 60 | { opProgress: state.opTotal && ((state.opOffset || 0) + (ts.getBytesSent() - offset) / state.opTotal) }) |
| 61 | /* in case this stream stands still for a while (before the end), we'll have neither 'sent' or 'close' events, |
| 62 | * so who will take care to updateConnection? This artificial next-call will ensure just that */ |
| 63 | clearTimeout(conn[SymTimeout]) |
| 64 | if (outSpeed || !closed) |
| 65 | conn[SymTimeout] = setTimeout(update, DELAY) |
| 66 | }, DELAY, { leading: true, maxWait:DELAY }) |
| 67 | await totalSent.ready() |
| 68 | ts.on('sent', n => { |
| 69 | totalSent.set(x => x + n) |
| 70 | update() |
| 71 | }) |
| 72 | |
| 73 | ++ipGroup.count |
| 74 | ts.on('close', ()=> { |
| 75 | update.flush() |
| 76 | closed = true |
| 77 | if (--ipGroup.count) return // any left? |
| 78 | delete ip2group[ctx.ip] |
| 79 | }) |
| 80 | |
| 81 | ctx.state.originalStream = body |
| 82 | ctx.body = body.pipe(ts) |
| 83 | |
| 84 | if (downloadTotal !== undefined) // undefined will break SSE |
| 85 | ctx.response.length = downloadTotal // preserve this info |
| 86 | ts.once('close', () => // in case of compressed response, we offer calculation of real size |
| 87 | ctx.state.length = ts.getBytesSent() - offset) |
| 88 | } |
| 89 | |
| 90 | declare module "koa" { |
| 91 | interface DefaultState { |
| 92 | length?: number |
| 93 | originalStream?: Parameters<Koa.Middleware>[0]['body'] |
| 94 | } |
| 95 | } |
| 96 | |
| 97 | export function roundSpeed(n: number) { |
| 98 | return _.round(n, 1) || _.round(n, 3) // further precision if necessary |
| 99 | } |
| 100 | |
| 101 | export const totalSent = storedMap.singleSync<number>('totalSent', 0) |
| 102 | export const totalGot = storedMap.singleSync<number>('totalGot', 0) |
| 103 | export let totalOutSpeedKb = 0 |
| 104 | export let totalInSpeedKb = 0 |
| 105 | |
| 106 | let lastSent: number | undefined |
| 107 | let lastGot: number | undefined |
| 108 | let last = Date.now() |
| 109 | setInterval(() => { |
| 110 | const now = Date.now() |
| 111 | const past = now - last |
| 112 | last = now |
| 113 | { |
| 114 | const v = totalSent.get() |
| 115 | totalOutSpeedKb = roundSpeed((v - (lastSent ?? v)) / past) // lastSent is bytes, past is milliseconds, so the result is KB/s |
| 116 | lastSent = v |
| 117 | } |
| 118 | { |
| 119 | const v = totalGot.get() |
| 120 | totalInSpeedKb = roundSpeed((v - (lastGot ?? v)) / past) |
| 121 | lastGot = v |
| 122 | } |
| 123 | }, 1000) |
| 124 | |
| 125 | events.on('connection', (c: Connection) => { |
| 126 | const count = (data: Buffer | string) => totalGot.set(x => x + data.length) |
| 127 | c.socket.on('data', count) |
| 128 | c.socket.on('secure', s => s.on('data', count)) // secure sockets won't forward 'data' events to the plain ones |
| 129 | }) |