From 4a2028c354a41845c7a7f145932ad808f136266e Mon Sep 17 00:00:00 2001 From: Martin McKeaveney Date: Thu, 23 Sep 2021 22:40:14 +0100 Subject: [PATCH] usageQuota middleware writing to couch --- packages/server/src/api/routes/application.js | 5 +- .../server/src/automations/steps/createRow.js | 4 +- packages/server/src/middleware/usageQuota.js | 25 ++-- packages/server/src/utilities/usageQuota.js | 110 +++++------------- .../server/src/utilities/usageQuota.old.js | 105 +++++++++++++++++ .../src/api/controllers/global/users.js | 24 ++++ 6 files changed, 177 insertions(+), 96 deletions(-) create mode 100644 packages/server/src/utilities/usageQuota.old.js diff --git a/packages/server/src/api/routes/application.js b/packages/server/src/api/routes/application.js index c1d39acbd5..ef4aacf708 100644 --- a/packages/server/src/api/routes/application.js +++ b/packages/server/src/api/routes/application.js @@ -2,11 +2,12 @@ const Router = require("@koa/router") const controller = require("../controllers/application") const authorized = require("../../middleware/authorized") const { BUILDER } = require("@budibase/auth/permissions") +const usage = require("../../middleware/usageQuota") const router = Router() router - .post("/api/applications", authorized(BUILDER), controller.create) + .post("/api/applications", authorized(BUILDER), usage, controller.create) .get("/api/applications/:appId/definition", controller.fetchAppDefinition) .get("/api/applications", controller.fetch) .get("/api/applications/:appId/appPackage", controller.fetchAppPackage) @@ -21,6 +22,6 @@ router authorized(BUILDER), controller.revertClient ) - .delete("/api/applications/:appId", authorized(BUILDER), controller.delete) + .delete("/api/applications/:appId", authorized(BUILDER), usage, controller.delete) module.exports = router diff --git a/packages/server/src/automations/steps/createRow.js b/packages/server/src/automations/steps/createRow.js index 9033004578..41e775b3de 100644 --- a/packages/server/src/automations/steps/createRow.js +++ b/packages/server/src/automations/steps/createRow.js @@ -60,7 +60,7 @@ exports.definition = { }, } -exports.run = async function ({ inputs, appId, apiKey, emitter }) { +exports.run = async function ({ inputs, appId, tenantId, emitter }) { if (inputs.row == null || inputs.row.tableId == null) { return { success: false, @@ -84,7 +84,7 @@ exports.run = async function ({ inputs, appId, apiKey, emitter }) { inputs.row ) if (env.USE_QUOTAS) { - await usage.update(apiKey, usage.Properties.ROW, 1) + await usage.update(tenantId, usage.Properties.ROW, 1) } await rowController.save(ctx) return { diff --git a/packages/server/src/middleware/usageQuota.js b/packages/server/src/middleware/usageQuota.js index 4647878721..4ad1092f6c 100644 --- a/packages/server/src/middleware/usageQuota.js +++ b/packages/server/src/middleware/usageQuota.js @@ -13,6 +13,7 @@ const DOMAIN_MAP = { upload: usageQuota.Properties.UPLOAD, views: usageQuota.Properties.VIEW, users: usageQuota.Properties.USER, + applications: usageQuota.Properties.APPS, // this will not be updated by endpoint calls // instead it will be updated by triggerInfo automationRuns: usageQuota.Properties.AUTOMATION, @@ -28,9 +29,9 @@ function getProperty(url) { module.exports = async (ctx, next) => { // if in development or a self hosted cloud usage quotas should not be executed - if (env.isDev() || env.SELF_HOSTED) { - return next() - } + // if (env.isDev() || env.SELF_HOSTED) { + // return next() + // } const db = new CouchDB(ctx.appId) let usage = METHOD_MAP[ctx.req.method] @@ -49,17 +50,17 @@ module.exports = async (ctx, next) => { } // update usage for uploads to be the total size - if (property === usageQuota.Properties.UPLOAD) { - const files = - ctx.request.files.file.length > 1 - ? Array.from(ctx.request.files.file) - : [ctx.request.files.file] - usage = files.map(file => file.size).reduce((total, size) => total + size) - } + // if (property === usageQuota.Properties.UPLOAD) { + // const files = + // ctx.request.files.file.length > 1 + // ? Array.from(ctx.request.files.file) + // : [ctx.request.files.file] + // usage = files.map(file => file.size).reduce((total, size) => total + size) + // } try { - await usageQuota.update(ctx.auth.apiKey, property, usage) + await usageQuota.update(ctx.user.tenantId, property, usage) return next() } catch (err) { - ctx.throw(403, err) + ctx.throw(400, err) } } diff --git a/packages/server/src/utilities/usageQuota.js b/packages/server/src/utilities/usageQuota.js index bfe71a4093..502fc4cad2 100644 --- a/packages/server/src/utilities/usageQuota.js +++ b/packages/server/src/utilities/usageQuota.js @@ -1,41 +1,6 @@ + const env = require("../environment") -const { apiKeyTable } = require("../db/dynamoClient") - -const DEFAULT_USAGE = { - rows: 0, - storage: 0, - views: 0, - automationRuns: 0, - users: 0, -} - -const DEFAULT_PLAN = { - rows: 1000, - // 1 GB - storage: 8589934592, - views: 10, - automationRuns: 100, - users: 10000, -} - -function buildUpdateParams(key, property, usage) { - return { - primary: key, - condition: - "attribute_exists(#quota) AND attribute_exists(#limits) AND #quota.#prop < #limits.#prop AND #quotaReset > :now", - expression: "ADD #quota.#prop :usage", - names: { - "#quota": "usageQuota", - "#prop": property, - "#limits": "usageLimits", - "#quotaReset": "quotaReset", - }, - values: { - ":usage": usage, - ":now": Date.now(), - }, - } -} +const { getGlobalDB } = require("@budibase/auth/tenancy") function getNewQuotaReset() { return Date.now() + 2592000000 @@ -47,59 +12,44 @@ exports.Properties = { VIEW: "views", USER: "users", AUTOMATION: "automationRuns", -} - -exports.getAPIKey = async appId => { - if (!env.USE_QUOTAS) { - return { apiKey: null } - } - return apiKeyTable.get({ primary: appId }) + APPS: "apps" } /** - * Given a specified API key this will add to the usage object for the specified property. - * @param {string} apiKey The API key which is to be updated. + * Given a specified tenantId this will add to the usage object for the specified property. + * @param {string} tenantId The tenant to update the usage quotas for. * @param {string} property The property which is to be added to (within the nested usageQuota object). * @param {number} usage The amount (this can be negative) to adjust the number by. * @returns {Promise} When this completes the API key will now be up to date - the quota period may have * also been reset after this call. */ -exports.update = async (apiKey, property, usage) => { - if (!env.USE_QUOTAS) { - return - } +exports.update = async (tenantId, property, usage) => { + // if (!env.USE_QUOTAS) { + // return + // } try { - await apiKeyTable.update(buildUpdateParams(apiKey, property, usage)) + const db = getGlobalDB() + const quota = await db.get("usage_quota") + // TODO: check if the quota needs reset + if (Date.now() >= quota.quotaReset) { + quota.quotaReset = getNewQuotaReset() + for (let prop of Object.keys(quota.usageQuota)) { + quota.usageQuota[prop] = 0 + } + } + + // increment the quota + quota.usageQuota[property] += usage + + if (quota.usageQuota[property] >= quota.usageLimits[property]) { + throw new Error(`You have exceeded your usage quota of ${quota.usageLimits[property]} ${property}.`) + } + + // update the usage quotas + await db.put(quota) } catch (err) { - // conditional check means the condition failed, need to check why - if (err.code === "ConditionalCheckFailedException") { - // get the API key so we can check it - const keyObj = await apiKeyTable.get({ primary: apiKey }) - // the usage quota or usage limits didn't exist - if (keyObj && (keyObj.usageQuota == null || keyObj.usageLimits == null)) { - keyObj.usageQuota = - keyObj.usageQuota == null ? DEFAULT_USAGE : keyObj.usageQuota - keyObj.usageLimits = - keyObj.usageLimits == null ? DEFAULT_PLAN : keyObj.usageLimits - keyObj.quotaReset = getNewQuotaReset() - await apiKeyTable.put({ item: keyObj }) - return - } - // we have in fact breached the reset period - else if (keyObj && keyObj.quotaReset <= Date.now()) { - // update the quota reset period and reset the values for all properties - keyObj.quotaReset = getNewQuotaReset() - for (let prop of Object.keys(keyObj.usageQuota)) { - if (prop === property) { - keyObj.usageQuota[prop] = usage > 0 ? usage : 0 - } else { - keyObj.usageQuota[prop] = 0 - } - } - await apiKeyTable.put({ item: keyObj }) - return - } - } + console.error(`Error updating usage quotas for ${property}`, err) throw err } + } diff --git a/packages/server/src/utilities/usageQuota.old.js b/packages/server/src/utilities/usageQuota.old.js new file mode 100644 index 0000000000..bfe71a4093 --- /dev/null +++ b/packages/server/src/utilities/usageQuota.old.js @@ -0,0 +1,105 @@ +const env = require("../environment") +const { apiKeyTable } = require("../db/dynamoClient") + +const DEFAULT_USAGE = { + rows: 0, + storage: 0, + views: 0, + automationRuns: 0, + users: 0, +} + +const DEFAULT_PLAN = { + rows: 1000, + // 1 GB + storage: 8589934592, + views: 10, + automationRuns: 100, + users: 10000, +} + +function buildUpdateParams(key, property, usage) { + return { + primary: key, + condition: + "attribute_exists(#quota) AND attribute_exists(#limits) AND #quota.#prop < #limits.#prop AND #quotaReset > :now", + expression: "ADD #quota.#prop :usage", + names: { + "#quota": "usageQuota", + "#prop": property, + "#limits": "usageLimits", + "#quotaReset": "quotaReset", + }, + values: { + ":usage": usage, + ":now": Date.now(), + }, + } +} + +function getNewQuotaReset() { + return Date.now() + 2592000000 +} + +exports.Properties = { + ROW: "rows", + UPLOAD: "storage", + VIEW: "views", + USER: "users", + AUTOMATION: "automationRuns", +} + +exports.getAPIKey = async appId => { + if (!env.USE_QUOTAS) { + return { apiKey: null } + } + return apiKeyTable.get({ primary: appId }) +} + +/** + * Given a specified API key this will add to the usage object for the specified property. + * @param {string} apiKey The API key which is to be updated. + * @param {string} property The property which is to be added to (within the nested usageQuota object). + * @param {number} usage The amount (this can be negative) to adjust the number by. + * @returns {Promise} When this completes the API key will now be up to date - the quota period may have + * also been reset after this call. + */ +exports.update = async (apiKey, property, usage) => { + if (!env.USE_QUOTAS) { + return + } + try { + await apiKeyTable.update(buildUpdateParams(apiKey, property, usage)) + } catch (err) { + // conditional check means the condition failed, need to check why + if (err.code === "ConditionalCheckFailedException") { + // get the API key so we can check it + const keyObj = await apiKeyTable.get({ primary: apiKey }) + // the usage quota or usage limits didn't exist + if (keyObj && (keyObj.usageQuota == null || keyObj.usageLimits == null)) { + keyObj.usageQuota = + keyObj.usageQuota == null ? DEFAULT_USAGE : keyObj.usageQuota + keyObj.usageLimits = + keyObj.usageLimits == null ? DEFAULT_PLAN : keyObj.usageLimits + keyObj.quotaReset = getNewQuotaReset() + await apiKeyTable.put({ item: keyObj }) + return + } + // we have in fact breached the reset period + else if (keyObj && keyObj.quotaReset <= Date.now()) { + // update the quota reset period and reset the values for all properties + keyObj.quotaReset = getNewQuotaReset() + for (let prop of Object.keys(keyObj.usageQuota)) { + if (prop === property) { + keyObj.usageQuota[prop] = usage > 0 ? usage : 0 + } else { + keyObj.usageQuota[prop] = 0 + } + } + await apiKeyTable.put({ item: keyObj }) + return + } + } + throw err + } +} diff --git a/packages/worker/src/api/controllers/global/users.js b/packages/worker/src/api/controllers/global/users.js index 1375240f34..b75c72290d 100644 --- a/packages/worker/src/api/controllers/global/users.js +++ b/packages/worker/src/api/controllers/global/users.js @@ -18,6 +18,7 @@ const { tryAddTenant, updateTenantId, } = require("@budibase/auth/tenancy") +const env = require("../../../environment") const PLATFORM_INFO_DB = StaticDatabases.PLATFORM_INFO.name @@ -138,6 +139,29 @@ exports.adminUser = async ctx => { include_docs: true, }) ) + + // write usage quotas for cloud + // if (!env.SELF_HOSTED) { + await db.post({ + _id: "usage_quota", + quotaReset: Date.now() + 2592000000, + usageQuota: { + automationRuns: 0, + rows: 0, + storage: 0, + apps: 0, + users: 0, + views: 0, + }, + usageLimits: { + automationRuns: 1000, + rows: 4000, + apps: 4, + // storage: 1000, + // users: 10 + }, + }) + // } if (response.rows.some(row => row.doc.admin)) { ctx.throw(