eirtube/eirtubeMods/ratelimiting.js

const {render} = require("pinski/plugins")
const {getToken} = require("../utils/getuser")
const constants = require("../utils/constants")
const db = require("../utils/db")
const crypto = require("crypto")

let users = {}
let buckets = {}

class User {
	constructor(req) {
		this.uuid = crypto.randomUUID()
		this.token = User.getToken(req)
		this.ip = User.getIP(req)

		users[this.token] = this
		users[this.ip] = this
	}

	static getToken(req) {
		return getToken(req)
	}

	static getIP(req) {
		return req.headers["x-forwarded-for"] || req.socket.remoteAddress || null
	}
}

class Bucket {
	constructor() {
		this.left = constants.server_setup.ratelimiting.max_bucket_size,
		this.lastModified = Date.now()
	}

	drain() {
		// Refill bucket
		this.left += Math.floor(((Date.now() - this.lastModified) / 1000) / constants.server_setup.ratelimiting.bucket_refill_rate_seconds)
		this.left = Math.min(this.left, constants.server_setup.ratelimiting.max_bucket_size)

		if (this.left > 0) {
			this.left--
			this.lastModified = Date.now()
			return true
		}
		return false
	}
}

function authorized(req, doRateLimit, responseHeaders) {
	if (constants.server_setup.ratelimiting.enabled && doRateLimit) {
		let bucket = module.exports.getBucket(req)
		if (!bucket.drain())
			return { success: false, message: "ratelimited", timeLeftSeconds: ((bucket.lastModified + (1000 * constants.server_setup.ratelimiting.bucket_refill_rate_seconds)) - Date.now()) / 1000 }
	}

	const token = getToken(req, responseHeaders)
	let out = db.prepare("SELECT * FROM SeenTokens WHERE token = ?").pluck().all(token)
	return { success: Object.keys(req.headers).some((header) => req.headers[header] == "same-origin" || (header == "referer" && req.headers[header].indexOf("/watch?v=") > -1)) || out.length > 0, message: "auth" }
}

// Clear users and buckets that have been inactive
setInterval(() => {
	for (const k of Object.keys(buckets)) {
		const b = buckets[k]

		if ((Date.now() - b.lastModified) < 1000 * 60 * 30)
			continue

		for (const k2 of Object.keys(users)) {
			const u = users[k2]
			if (u.uuid == k)
				delete users[k2]
		}
		delete buckets[k]
	} 
}, 1000 * 60 * 30)

module.exports = {
	getUser: function(req) {
		const token = User.getToken(req)
		if (token && users[token]) {
			delete users[User.getIP(req)] // If a user has tokens enabled, we don't want their IP address to block other users on the same network
			return users[token]
		}
		const ip = User.getIP(req)
		if (users[ip])
			return users[ip]
		return new User(req)
	},
	getBucket: function(req) {
		const user = module.exports.getUser(req)
		const bucket = buckets[user.uuid]
		if (bucket)
			return bucket
		buckets[user.uuid] = new Bucket()
		return buckets[user.uuid]
	},
	remove: function(req) {
		const token = User.getToken(req)
		if (token && users[token]) {
			const uuid = users[token].uuid
			delete users[token]
			delete users[User.getIP(req)]
			delete buckets[uuid]
			return
		}
		const ip = User.getIP(req)
		if (users[ip]) {
			const uuid = users[ip].uuid
			delete users[ip]
			delete buckets[uuid]
		}
	},

	authorize: function(req, doRateLimit, responseHeaders) {
		const output = authorized(req, doRateLimit, responseHeaders)
		if (output.success == false) {
			if (output.message == "ratelimited")
				return render(429, "pug/errors/local-rate-limited.pug", { timeLeftSeconds: output.timeLeftSeconds }, {
					...responseHeaders,
					"Retry-After": Math.ceil(output.timeLeftSeconds)
				})
			else if (output.message == "auth")
				// return render(403, "pug/errors/message-error.pug", { error: new Error("Access denied") }, responseHeaders)
				return {
					statusCode: 403,
					headers: responseHeaders,
					contentType: "application/json",
					content: {
						error: "Access denied"
					}
				}
		}
		return true
	}
}